add separate PreCache step for GcReferenceChecker (#578)

- Improvement: GCv2: Use separate PreCache step to improve concurrency when checking references
- Improvement: GCv2: Improved verbose logging
- Improvement: GCv2: Sort chunks to read by block/offset when finding references
- Improvement: GCv2: Exit as soon as no more unreferenced items are left

4 files changed, 71 insertions, 6 deletions

diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index 03c7f4b95..cc727787f 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -1274,6 +1274,17 @@ BlockStore::GetBlockPath(const std::filesystem::path& BlocksBasePath, const uint
 	return Path.ToPath();
 }
 
+Ref<BlockStoreFile>
+BlockStore::GetBlockFile(uint32_t BlockIndex)
+{
+	RwLock::SharedLockScope _(m_InsertLock);
+	if (auto It = m_ChunkBlocks.find(BlockIndex); It != m_ChunkBlocks.end())
+	{
+		return It->second;
+	}
+	return {};
+}
+
 #if ZEN_WITH_TESTS
 
 TEST_CASE("blockstore.blockstoredisklocation")
diff --git a/src/zenstore/gc.cpp b/src/zenstore/gc.cpp
index e2ab34d1e..2660c2643 100644
--- a/src/zenstore/gc.cpp
+++ b/src/zenstore/gc.cpp
@@ -381,6 +381,7 @@ WriteReferencerStats(CbObjectWriter& Writer, const GcReferencerStats& Stats, boo
 	Writer.EndObject();
 
 	Writer << "CreateReferenceCheckers" << ToTimeSpan(Stats.CreateReferenceCheckersMS);
+	Writer << "PreCacheState" << ToTimeSpan(Stats.PreCacheStateMS);
 	Writer << "LockState" << ToTimeSpan(Stats.LockStateMS);
 	Writer << "Elapsed" << ToTimeSpan(Stats.ElapsedMS);
 };
@@ -449,6 +450,7 @@ WriteGCResult(CbObjectWriter& Writer, const GcResult& Result, bool HumanReadable
 
 	Writer << "RemoveExpiredData" << ToTimeSpan(Result.RemoveExpiredDataMS);
 	Writer << "CreateReferenceCheckers" << ToTimeSpan(Result.CreateReferenceCheckersMS);
+	Writer << "PreCacheState" << ToTimeSpan(Result.PreCacheStateMS);
 	Writer << "LockState" << ToTimeSpan(Result.LockStateMS);
 
 	Writer << "CreateReferencePruners" << ToTimeSpan(Result.CreateReferencePrunersMS);
@@ -507,8 +509,8 @@ Add(GcStats& Sum, const GcStats& Sub)
 void
 Sum(GcReferencerStats& Stat)
 {
-	Stat.ElapsedMS =
-		Stat.RemoveExpiredDataStats.ElapsedMS + Stat.CompactStoreStats.ElapsedMS + Stat.CreateReferenceCheckersMS + Stat.LockStateMS;
+	Stat.ElapsedMS = Stat.RemoveExpiredDataStats.ElapsedMS + Stat.CompactStoreStats.ElapsedMS + Stat.CreateReferenceCheckersMS +
+					 Stat.PreCacheStateMS + Stat.LockStateMS;
 }
 
 void
@@ -518,6 +520,7 @@ Add(GcReferencerStats& Sum, const GcReferencerStats& Sub)
 	Add(Sum.CompactStoreStats, Sub.CompactStoreStats);
 
 	Sum.CreateReferenceCheckersMS += Sub.CreateReferenceCheckersMS;
+	Sum.PreCacheStateMS += Sub.PreCacheStateMS;
 	Sum.LockStateMS += Sub.LockStateMS;
 
 	Sum.ElapsedMS += Sub.ElapsedMS;
@@ -802,10 +805,54 @@ GcManager::CollectGarbage(const GcSettings& Settings)
 					}
 				}
 
-				ZEN_INFO("GCV2: Locking state for {} reference checkers", ReferenceCheckers.size());
 				{
-					SCOPED_TIMER(uint64_t ElapsedMS = Timer.GetElapsedTimeMs(); Result.WriteBlockMS = std::chrono::milliseconds(ElapsedMS);
-								 ZEN_INFO("GCV2: Writes blocked for {}", NiceTimeSpanMs(ElapsedMS)));
+					ZEN_INFO("GCV2: Precaching state for {} reference checkers", ReferenceCheckers.size());
+					if (!ReferenceCheckers.empty())
+					{
+						if (CheckGCCancel())
+						{
+							return Sum(Result, true);
+						}
+						ZEN_TRACE_CPU("GcV2::PreCache");
+
+						Latch WorkLeft(1);
+
+						{
+							SCOPED_TIMER(Result.PreCacheStateMS = std::chrono::milliseconds(Timer.GetElapsedTimeMs());
+										 if (Ctx.Settings.Verbose) {
+											 ZEN_INFO("GCV2: Precached state using {} reference checkers in {}",
+													  ReferenceCheckers.size(),
+													  NiceTimeSpanMs(Result.PreCacheStateMS.count()));
+										 });
+							for (auto& It : ReferenceCheckers)
+							{
+								if (CheckGCCancel())
+								{
+									WorkLeft.CountDown();
+									WorkLeft.Wait();
+									return Sum(Result, true);
+								}
+
+								GcReferenceChecker*						   Checker = It.first.get();
+								size_t									   Index   = It.second;
+								std::pair<std::string, GcReferencerStats>& Stats   = Result.ReferencerStats[Index];
+								WorkLeft.AddCount(1);
+								ThreadPool.ScheduleWork([&Ctx, Checker, Index, &Stats, &WorkLeft]() {
+									auto _ = MakeGuard([&WorkLeft]() { WorkLeft.CountDown(); });
+									SCOPED_TIMER(Stats.second.PreCacheStateMS = std::chrono::milliseconds(Timer.GetElapsedTimeMs()););
+									Checker->PreCache(Ctx);
+								});
+							}
+							WorkLeft.CountDown();
+							WorkLeft.Wait();
+						}
+					}
+				}
+
+				SCOPED_TIMER(uint64_t ElapsedMS = Timer.GetElapsedTimeMs(); Result.WriteBlockMS = std::chrono::milliseconds(ElapsedMS);
+							 ZEN_INFO("GCV2: Writes blocked for {}", NiceTimeSpanMs(ElapsedMS)));
+				{
+					ZEN_INFO("GCV2: Locking state for {} reference checkers", ReferenceCheckers.size());
 					if (!ReferenceCheckers.empty())
 					{
 						if (CheckGCCancel())
@@ -849,7 +896,8 @@ GcManager::CollectGarbage(const GcSettings& Settings)
 							WorkLeft.Wait();
 						}
 					}
-
+				}
+				{
 					ZEN_INFO("GCV2: Removing unreferenced data for {} reference pruners", ReferencePruners.size());
 					{
 						const auto GetUnusedReferences = [&ReferenceCheckers, &Ctx](std::span<IoHash> References) -> std::vector<IoHash> {
diff --git a/src/zenstore/include/zenstore/blockstore.h b/src/zenstore/include/zenstore/blockstore.h
index 82e1c71c6..dcd4b5e87 100644
--- a/src/zenstore/include/zenstore/blockstore.h
+++ b/src/zenstore/include/zenstore/blockstore.h
@@ -180,6 +180,8 @@ public:
 
 	inline uint64_t TotalSize() const { return m_TotalSize.load(std::memory_order::relaxed); }
 
+	Ref<BlockStoreFile> GetBlockFile(uint32_t BlockIndex);
+
 private:
 	uint32_t GetFreeBlockIndex(uint32_t StartProbeIndex, RwLock::ExclusiveLockScope&, std::filesystem::path& OutBlockPath) const;
 
diff --git a/src/zenstore/include/zenstore/gc.h b/src/zenstore/include/zenstore/gc.h
index 7a6249970..698b0d4e8 100644
--- a/src/zenstore/include/zenstore/gc.h
+++ b/src/zenstore/include/zenstore/gc.h
@@ -87,6 +87,7 @@ struct GcReferencerStats
 	GcCompactStoreStats CompactStoreStats;
 
 	std::chrono::milliseconds CreateReferenceCheckersMS = {};
+	std::chrono::milliseconds PreCacheStateMS			= {};
 	std::chrono::milliseconds LockStateMS				= {};
 	std::chrono::milliseconds ElapsedMS					= {};
 };
@@ -112,6 +113,7 @@ struct GcResult
 	// Wall times, not sum of each
 	std::chrono::milliseconds RemoveExpiredDataMS		= {};
 	std::chrono::milliseconds CreateReferenceCheckersMS = {};
+	std::chrono::milliseconds PreCacheStateMS			= {};
 	std::chrono::milliseconds LockStateMS				= {};
 
 	std::chrono::milliseconds CreateReferencePrunersMS = {};
@@ -171,6 +173,8 @@ public:
 	// Destructor should unlock what was locked in LockState
 	virtual ~GcReferenceChecker() = default;
 
+	virtual void PreCache(GcCtx& Ctx) = 0;
+
 	// Lock the state and make sure no references changes, usually a read-lock is taken until the destruction
 	// of the instance. Called once before any calls to RemoveUsedReferencesFromSet
 	// The implementation should be as fast as possible as LockState is part of a stop the world (from changes)