Merge branch 'main' into zs/web-ui-blank-import-name-fix

2 files changed, 255 insertions, 13 deletions

diff --git a/src/zenstore/buildstore/buildstore.cpp b/src/zenstore/buildstore/buildstore.cpp
index d6d727aa9..cf518c06f 100644
--- a/src/zenstore/buildstore/buildstore.cpp
+++ b/src/zenstore/buildstore/buildstore.cpp
@@ -193,10 +193,12 @@ BuildStore::BuildStore(const BuildStoreConfig& Config, GcManager& Gc)
 
 		m_Gc.AddGcReferencer(*this);
 		m_Gc.AddGcReferenceLocker(*this);
+		m_Gc.AddGcStorage(this);
 	}
 	catch (const std::exception& Ex)
 	{
 		ZEN_ERROR("Failed to initialize build store. Reason: '{}'", Ex.what());
+		m_Gc.RemoveGcStorage(this);
 		m_Gc.RemoveGcReferenceLocker(*this);
 		m_Gc.RemoveGcReferencer(*this);
 	}
@@ -207,6 +209,7 @@ BuildStore::~BuildStore()
 	try
 	{
 		ZEN_TRACE_CPU("BuildStore::~BuildStore");
+		m_Gc.RemoveGcStorage(this);
 		m_Gc.RemoveGcReferenceLocker(*this);
 		m_Gc.RemoveGcReferencer(*this);
 		Flush();
@@ -552,6 +555,44 @@ BuildStore::Flush()
 	}
 }
 
+BuildStore::StorageStats
+BuildStore::GetStorageStats() const
+{
+	StorageStats Result;
+	{
+		RwLock::SharedLockScope _(m_Lock);
+		Result.EntryCount = m_BlobLookup.size();
+
+		for (auto LookupIt : m_BlobLookup)
+		{
+			const BlobIndex	 ReadBlobIndex = LookupIt.second;
+			const BlobEntry& ReadBlobEntry = m_BlobEntries[ReadBlobIndex];
+			if (ReadBlobEntry.Payload)
+			{
+				const PayloadEntry& Payload = m_PayloadEntries[ReadBlobEntry.Payload];
+				uint64_t			Size	= Payload.GetSize();
+				if ((Payload.GetFlags() & PayloadEntry::kStandalone) != 0)
+				{
+					Result.LargeBlobCount++;
+					Result.LargeBlobBytes += Size;
+				}
+				else
+				{
+					Result.SmallBlobCount++;
+					Result.SmallBlobBytes += Size;
+				}
+			}
+			if (ReadBlobEntry.Metadata)
+			{
+				const MetadataEntry& Metadata = m_MetadataEntries[ReadBlobEntry.Metadata];
+				Result.MetadataCount++;
+				Result.MetadataByteCount += Metadata.Location.Size;
+			}
+		}
+	}
+	return Result;
+}
+
 #if ZEN_WITH_TESTS
 std::optional<AccessTime>
 BuildStore::GetLastAccessTime(const IoHash& Key) const
@@ -1273,24 +1314,103 @@ BuildStore::RemoveExpiredData(GcCtx& Ctx, GcStats& Stats)
 		}
 	});
 
-	const GcClock::Tick ExpireTicks = Ctx.Settings.BuildStoreExpireTime.time_since_epoch().count();
+	const GcClock::Tick					   ExpireTicks = Ctx.Settings.BuildStoreExpireTime.time_since_epoch().count();
+	std::vector<IoHash>					   ExpiredBlobs;
+	tsl::robin_set<IoHash, IoHash::Hasher> SizeDroppedBlobs;
 
-	std::vector<IoHash> ExpiredBlobs;
 	{
-		RwLock::SharedLockScope __(m_Lock);
-		for (const auto& It : m_BlobLookup)
+		struct SizeInfo
 		{
-			const BlobIndex	 ReadBlobIndex = It.second;
-			const BlobEntry& ReadBlobEntry = m_BlobEntries[ReadBlobIndex];
+			const IoHash Key;
+			uint32_t	 SecondsSinceEpoch = 0;
+			uint64_t	 BlobSize		   = 0;
+		};
+
+		bool		   DiskSizeExceeded = false;
+		const uint64_t CurrentDiskSize =
+			m_LargeBlobStore.StorageSize().DiskSize + m_SmallBlobStore.StorageSize().DiskSize + m_MetadataBlockStore.TotalSize();
+		if (CurrentDiskSize > m_Config.MaxDiskSpaceLimit)
+		{
+			DiskSizeExceeded = true;
+		}
+
+		uint64_t ExpiredDataSize = 0;
+
+		std::vector<SizeInfo> NonExpiredBlobSizeInfos;
+
+		{
+			RwLock::SharedLockScope __(m_Lock);
+			if (DiskSizeExceeded)
+			{
+				NonExpiredBlobSizeInfos.reserve(m_BlobLookup.size());
+			}
+			for (const auto& It : m_BlobLookup)
+			{
+				const BlobIndex	 ReadBlobIndex = It.second;
+				const BlobEntry& ReadBlobEntry = m_BlobEntries[ReadBlobIndex];
+				uint64_t		 Size		   = 0;
+				if (ReadBlobEntry.Payload)
+				{
+					const PayloadEntry& Payload = m_PayloadEntries[ReadBlobEntry.Payload];
+					Size += Payload.GetSize();
+				}
+				if (ReadBlobEntry.Metadata)
+				{
+					const MetadataEntry& Metadata = m_MetadataEntries[ReadBlobEntry.Metadata];
+					Size += Metadata.Location.Size;
+				}
+
+				const GcClock::Tick AccessTick = ReadBlobEntry.LastAccessTime;
+				if (AccessTick < ExpireTicks)
+				{
+					ExpiredBlobs.push_back(It.first);
+					ExpiredDataSize += ExpiredDataSize;
+				}
+				else if (DiskSizeExceeded)
+				{
+					NonExpiredBlobSizeInfos.emplace_back(SizeInfo{.Key				 = It.first,
+																  .SecondsSinceEpoch = ReadBlobEntry.LastAccessTime.GetSecondsSinceEpoch(),
+																  .BlobSize			 = Size});
+				}
+			}
+			Stats.CheckedCount += m_BlobLookup.size();
+			Stats.FoundCount += ExpiredBlobs.size();
+		}
 
-			const GcClock::Tick AccessTick = ReadBlobEntry.LastAccessTime;
-			if (AccessTick < ExpireTicks)
+		if (DiskSizeExceeded)
+		{
+			const uint64_t NewSizeLimit =
+				m_Config.MaxDiskSpaceLimit -
+				(m_Config.MaxDiskSpaceLimit >> 4);	// Remove a bit more than just below the limit so we have some space to grow
+			if ((CurrentDiskSize - ExpiredDataSize) > NewSizeLimit)
 			{
-				ExpiredBlobs.push_back(It.first);
+				std::vector<size_t> NonExpiredOrder;
+				NonExpiredOrder.resize(NonExpiredBlobSizeInfos.size());
+				for (size_t Index = 0; Index < NonExpiredOrder.size(); Index++)
+				{
+					NonExpiredOrder[Index] = Index;
+				}
+				std::sort(NonExpiredOrder.begin(), NonExpiredOrder.end(), [&NonExpiredBlobSizeInfos](const size_t Lhs, const size_t Rhs) {
+					const SizeInfo& LhsInfo = NonExpiredBlobSizeInfos[Lhs];
+					const SizeInfo& RhsInfo = NonExpiredBlobSizeInfos[Rhs];
+					return LhsInfo.SecondsSinceEpoch < RhsInfo.SecondsSinceEpoch;
+				});
+
+				auto It = NonExpiredOrder.begin();
+				while (It != NonExpiredOrder.end())
+				{
+					const SizeInfo& Info = NonExpiredBlobSizeInfos[*It];
+					if ((CurrentDiskSize - ExpiredDataSize) < NewSizeLimit)
+					{
+						break;
+					}
+					ExpiredDataSize += Info.BlobSize;
+					ExpiredBlobs.push_back(Info.Key);
+					SizeDroppedBlobs.insert(Info.Key);
+					It++;
+				}
 			}
 		}
-		Stats.CheckedCount += m_BlobLookup.size();
-		Stats.FoundCount += ExpiredBlobs.size();
 	}
 
 	std::vector<IoHash> RemovedBlobs;
@@ -1318,7 +1438,7 @@ BuildStore::RemoveExpiredData(GcCtx& Ctx, GcStats& Stats)
 
 					const GcClock::Tick AccessTick = ReadBlobEntry.LastAccessTime;
 
-					if (AccessTick < ExpireTicks)
+					if (SizeDroppedBlobs.contains(ExpiredBlob) || (AccessTick < ExpireTicks))
 					{
 						if (ReadBlobEntry.Payload)
 						{
@@ -1388,6 +1508,21 @@ BuildStore::LockState(GcCtx& Ctx)
 	return Locks;
 }
 
+void
+BuildStore::ScrubStorage(ScrubContext& ScrubCtx)
+{
+	ZEN_UNUSED(ScrubCtx);
+	// TODO
+}
+
+GcStorageSize
+BuildStore::StorageSize() const
+{
+	GcStorageSize Result;
+	Result.DiskSize = m_MetadataBlockStore.TotalSize();
+	return Result;
+}
+
 /*
 		___________              __
 		\__    ___/___   _______/  |_  ______
@@ -1731,6 +1866,94 @@ TEST_CASE("BuildStore.GC")
 	}
 }
 
+TEST_CASE("BuildStore.SizeLimit")
+{
+	using namespace blockstore::testing;
+
+	ScopedTemporaryDirectory _;
+
+	BuildStoreConfig Config = {.MaxDiskSpaceLimit = 1024u * 1024u};
+	Config.RootDirectory	= _.Path() / "build_store";
+
+	std::vector<IoHash>	  CompressedBlobsHashes;
+	std::vector<IoBuffer> BlobMetaPayloads;
+	{
+		GcManager  Gc;
+		BuildStore Store(Config, Gc);
+		for (size_t I = 0; I < 64; I++)
+		{
+			IoBuffer		 Blob = CreateSemiRandomBlob(65537 + I * 7);
+			CompressedBuffer CompressedBlob =
+				CompressedBuffer::Compress(SharedBuffer(std::move(Blob)), OodleCompressor::Mermaid, OodleCompressionLevel::None);
+			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
+			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
+			Payload.SetContentType(ZenContentType::kCompressedBinary);
+			Store.PutBlob(CompressedBlobsHashes.back(), Payload);
+		}
+		for (const IoHash& BlobHash : CompressedBlobsHashes)
+		{
+			BlobMetaPayloads.push_back(MakeMetaData(BlobHash, {{"blobHash", fmt::format("{}", BlobHash)}}));
+			BlobMetaPayloads.back().SetContentType(ZenContentType::kCbObject);
+		}
+		Store.PutMetadatas(CompressedBlobsHashes, BlobMetaPayloads);
+
+		{
+			for (size_t I = 0; I < 64; I++)
+			{
+				const IoHash& Key		 = CompressedBlobsHashes[I];
+				GcClock::Tick AccessTick = (GcClock::Now() + std::chrono::minutes(I)).time_since_epoch().count();
+
+				Store.SetLastAccessTime(Key, AccessTime(AccessTick));
+			}
+		}
+	}
+	{
+		GcManager  Gc;
+		BuildStore Store(Config, Gc);
+
+		{
+			GcResult Result = Gc.CollectGarbage(GcSettings{.BuildStoreExpireTime = GcClock::Now() - std::chrono::hours(1),
+														   .CollectSmallObjects	 = true,
+														   .IsDeleteMode		 = true,
+														   .Verbose				 = true});
+
+			uint32_t DeletedBlobs = 0;
+
+			CHECK(!Result.WasCancelled);
+			for (const IoHash& BlobHash : CompressedBlobsHashes)
+			{
+				IoBuffer Blob = Store.GetBlob(BlobHash);
+				if (!Blob)
+				{
+					DeletedBlobs++;
+				}
+				else
+				{
+					IoBuffer DecompressedBlob = CompressedBuffer::FromCompressedNoValidate(std::move(Blob)).Decompress().AsIoBuffer();
+					CHECK(DecompressedBlob);
+					CHECK(IoHash::HashBuffer(DecompressedBlob) == BlobHash);
+				}
+			}
+			CHECK(DeletedBlobs == 50);
+
+			std::vector<IoBuffer> MetadataPayloads = Store.GetMetadatas(CompressedBlobsHashes, nullptr);
+			CHECK(MetadataPayloads.size() == BlobMetaPayloads.size());
+			for (size_t I = 0; I < MetadataPayloads.size(); I++)
+			{
+				const IoBuffer& MetadataPayload = MetadataPayloads[I];
+				if (I < DeletedBlobs)
+				{
+					CHECK(!MetadataPayload);
+				}
+				else
+				{
+					CHECK(IoHash::HashBuffer(MetadataPayload) == IoHash::HashBuffer(BlobMetaPayloads[I]));
+				}
+			}
+		}
+	}
+}
+
 void
 buildstore_forcelink()
 {
diff --git a/src/zenstore/include/zenstore/buildstore/buildstore.h b/src/zenstore/include/zenstore/buildstore/buildstore.h
index d88e682de..adf48dc26 100644
--- a/src/zenstore/include/zenstore/buildstore/buildstore.h
+++ b/src/zenstore/include/zenstore/buildstore/buildstore.h
@@ -24,9 +24,10 @@ struct BuildStoreConfig
 	uint32_t			  SmallBlobBlockStoreAlignement		   = 16;
 	uint32_t			  MetadataBlockStoreMaxBlockSize	   = 64 * 1024 * 1024;
 	uint32_t			  MetadataBlockStoreAlignement		   = 8;
+	uint64_t			  MaxDiskSpaceLimit					   = 1u * 1024u * 1024u * 1024u * 1024u;  // 1TB
 };
 
-class BuildStore : public GcReferencer, public GcReferenceLocker  //, public GcStorage
+class BuildStore : public GcReferencer, public GcReferenceLocker, public GcStorage
 {
 public:
 	explicit BuildStore(const BuildStoreConfig& Config, GcManager& Gc);
@@ -48,10 +49,24 @@ public:
 
 	void Flush();
 
+	struct StorageStats
+	{
+		uint64_t EntryCount		   = 0;
+		uint64_t LargeBlobCount	   = 0;
+		uint64_t LargeBlobBytes	   = 0;
+		uint64_t SmallBlobCount	   = 0;
+		uint64_t SmallBlobBytes	   = 0;
+		uint64_t MetadataCount	   = 0;
+		uint64_t MetadataByteCount = 0;
+	};
+
+	StorageStats GetStorageStats() const;
+
 #if ZEN_WITH_TESTS
 	std::optional<AccessTime> GetLastAccessTime(const IoHash& Key) const;
 	bool					  SetLastAccessTime(const IoHash& Key, const AccessTime& Time);
 #endif	// ZEN_WITH_TESTS
+
 private:
 	LoggerRef Log() { return m_Log; }
 
@@ -71,6 +86,10 @@ private:
 	//////// GcReferenceLocker
 	virtual std::vector<RwLock::SharedLockScope> LockState(GcCtx& Ctx) override;
 
+	//////// GcStorage
+	virtual void		  ScrubStorage(ScrubContext& ScrubCtx) override;
+	virtual GcStorageSize StorageSize() const override;
+
 #pragma pack(push)
 #pragma pack(1)
 	struct PayloadEntry