Use rename/delete and keep pointer for dropped buckets

2 files changed, 129 insertions, 77 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 1db99280e..be9e408b4 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -580,19 +580,6 @@ ZenCacheDiskLayer::CacheBucket::~CacheBucket()
 {
 }
 
-bool
-ZenCacheDiskLayer::CacheBucket::Delete(std::filesystem::path BucketDir)
-{
-	if (std::filesystem::exists(BucketDir))
-	{
-		DeleteDirectories(BucketDir);
-
-		return true;
-	}
-
-	return false;
-}
-
 void
 ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bool AllowCreate)
 {
@@ -1158,11 +1145,6 @@ ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(const DiskLocation& Loc,
 bool
 ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
 {
-	if (!m_IsOk)
-	{
-		return false;
-	}
-
 	RwLock::SharedLockScope _(m_IndexLock);
 	auto					It = m_Index.find(HashKey);
 	if (It == m_Index.end())
@@ -1184,11 +1166,6 @@ ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutVal
 void
 ZenCacheDiskLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
 {
-	if (!m_IsOk)
-	{
-		return;
-	}
-
 	if (Value.Value.Size() >= m_LargeObjectThreshold)
 	{
 		return PutStandaloneCacheValue(HashKey, Value);
@@ -1196,10 +1173,43 @@ ZenCacheDiskLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue&
 	PutInlineCacheValue(HashKey, Value);
 }
 
+static bool
+DeleteBucketFromDisk(const std::filesystem::path& BucketDir, std::string_view BucketName)
+{
+	int dropIndex = 0;
+	do
+	{
+		if (!std::filesystem::exists(BucketDir))
+		{
+			return false;
+		}
+
+		std::string			  DroppedBucketName = fmt::format("[dropped]{}({})", BucketName, dropIndex);
+		std::filesystem::path DroppedBucketPath = BucketDir.parent_path() / DroppedBucketName;
+		if (std::filesystem::exists(DroppedBucketPath))
+		{
+			dropIndex++;
+			continue;
+		}
+
+		std::error_code Ec;
+		std::filesystem::rename(BucketDir, DroppedBucketPath, Ec);
+		if (!Ec)
+		{
+			DeleteDirectories(DroppedBucketPath);
+			return true;
+		}
+		// TODO: Do we need to bail at some point?
+		zen::Sleep(100);
+	} while (true);
+}
+
 void
 ZenCacheDiskLayer::CacheBucket::Drop()
 {
-	RwLock::ExclusiveLockScope								 _(m_IndexLock);
+	RwLock::ExclusiveLockScope _(m_IndexLock);
+	m_IsOk = false;
+
 	std::vector<std::unique_ptr<RwLock::ExclusiveLockScope>> ShardLocks;
 	ShardLocks.reserve(256);
 	for (RwLock& Lock : m_ShardedLocks)
@@ -1208,8 +1218,10 @@ ZenCacheDiskLayer::CacheBucket::Drop()
 	}
 	m_BlockStore.Close();
 	m_SlogFile.Close();
+
+	DeleteBucketFromDisk(m_BucketDir, m_BucketName);
+
 	m_Index.clear();
-	DeleteDirectories(m_BucketDir);
 }
 
 void
@@ -1711,7 +1723,12 @@ ZenCacheDiskLayer::CollectGarbage(GcContext& GcCtx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.CollectGarbage(GcCtx);
+		CacheBucket& Bucket = *Kv.second;
+		if (!Bucket.IsOk())
+		{
+			continue;
+		}
+		Bucket.CollectGarbage(GcCtx);
 	}
 }
 
@@ -1724,7 +1741,11 @@ ZenCacheDiskLayer::UpdateAccessTimes(const zen::access_tracking::AccessTimes& Ac
 	{
 		if (auto It = m_Buckets.find(Kv.first); It != m_Buckets.end())
 		{
-			CacheBucket& Bucket = It->second;
+			CacheBucket& Bucket = *It->second;
+			if (!Bucket.IsOk())
+			{
+				continue;
+			}
 			Bucket.UpdateAccessTimes(Kv.second);
 		}
 	}
@@ -1937,7 +1958,11 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCach
 
 		if (it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
+			if (!Bucket->IsOk())
+			{
+				return false;
+			}
 		}
 	}
 
@@ -1949,22 +1974,25 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCach
 
 		if (auto it = m_Buckets.find(BucketName); it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
 		}
 		else
 		{
-			auto It = m_Buckets.try_emplace(BucketName, BucketName);
-			Bucket	= &It.first->second;
+			auto InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>(BucketName));
+			Bucket			  = InsertResult.first->second.get();
 
 			std::filesystem::path BucketPath = m_RootDir;
 			BucketPath /= BucketName;
 
 			Bucket->OpenOrCreate(BucketPath);
+			if (!Bucket->IsOk())
+			{
+				return false;
+			}
 		}
 	}
 
 	ZEN_ASSERT(Bucket != nullptr);
-
 	return Bucket->Get(HashKey, OutValue);
 }
 
@@ -1981,7 +2009,11 @@ ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const Z
 
 		if (it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
+			if (!Bucket->IsOk())
+			{
+				return;
+			}
 		}
 	}
 
@@ -1993,26 +2025,27 @@ ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const Z
 
 		if (auto it = m_Buckets.find(BucketName); it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
 		}
 		else
 		{
-			auto It = m_Buckets.try_emplace(BucketName, BucketName);
-			Bucket	= &It.first->second;
+			auto InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>(BucketName));
+			Bucket			  = InsertResult.first->second.get();
 
 			std::filesystem::path BucketPath = m_RootDir;
 			BucketPath /= BucketName;
 
 			Bucket->OpenOrCreate(BucketPath);
+			if (!Bucket->IsOk())
+			{
+				return;
+			}
 		}
 	}
 
 	ZEN_ASSERT(Bucket != nullptr);
 
-	if (Bucket->IsOk())
-	{
-		Bucket->Put(HashKey, Value);
-	}
+	Bucket->Put(HashKey, Value);
 }
 
 void
@@ -2034,9 +2067,8 @@ ZenCacheDiskLayer::DiscoverBuckets()
 		}
 		else
 		{
-			auto InsertResult = m_Buckets.try_emplace(BucketName, BucketName);
-
-			CacheBucket& Bucket = InsertResult.first->second;
+			auto		 InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>(BucketName));
+			CacheBucket& Bucket		  = *InsertResult.first->second;
 
 			Bucket.OpenOrCreate(BucketPath, /* AllowCreate */ false);
 
@@ -2063,19 +2095,23 @@ ZenCacheDiskLayer::DropBucket(std::string_view InBucket)
 
 	if (it != m_Buckets.end())
 	{
-		CacheBucket* Bucket = &it->second;
-
-		Bucket->Drop();
-
+		CacheBucket& Bucket = *it->second;
+		m_DroppedBuckets.push_back(std::move(it->second));
 		m_Buckets.erase(it);
+		if (!Bucket.IsOk())
+		{
+			return false;
+		}
+
+		Bucket.Drop();
 
 		return true;
 	}
 
+	// Make sure we remove the folder even if we don't know about the bucket
 	std::filesystem::path BucketPath = m_RootDir;
 	BucketPath /= std::string(InBucket);
-
-	return CacheBucket::Delete(BucketPath);
+	return DeleteBucketFromDisk(BucketPath, InBucket);
 }
 
 void
@@ -2088,7 +2124,12 @@ ZenCacheDiskLayer::Flush()
 		Buckets.reserve(m_Buckets.size());
 		for (auto& Kv : m_Buckets)
 		{
-			Buckets.push_back(&Kv.second);
+			CacheBucket* Bucket = Kv.second.get();
+			if (!Bucket->IsOk())
+			{
+				continue;
+			}
+			Buckets.push_back(Bucket);
 		}
 	}
 
@@ -2105,7 +2146,12 @@ ZenCacheDiskLayer::Scrub(ScrubContext& Ctx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.Scrub(Ctx);
+		CacheBucket& Bucket = *Kv.second;
+		if (!Bucket.IsOk())
+		{
+			continue;
+		}
+		Bucket.Scrub(Ctx);
 	}
 }
 
@@ -2116,7 +2162,12 @@ ZenCacheDiskLayer::GatherReferences(GcContext& GcCtx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.GatherReferences(GcCtx);
+		CacheBucket& Bucket = *Kv.second;
+		if (!Bucket.IsOk())
+		{
+			continue;
+		}
+		Bucket.GatherReferences(GcCtx);
 	}
 }
 
@@ -2128,7 +2179,7 @@ ZenCacheDiskLayer::TotalSize() const
 
 	for (auto& Kv : m_Buckets)
 	{
-		TotalSize += Kv.second.TotalSize();
+		TotalSize += Kv.second->TotalSize();
 	}
 
 	return TotalSize;
diff --git a/zenserver/cache/structuredcachestore.h b/zenserver/cache/structuredcachestore.h
index 13ba95f1b..bae15231b 100644
--- a/zenserver/cache/structuredcachestore.h
+++ b/zenserver/cache/structuredcachestore.h
@@ -233,17 +233,15 @@ private:
 		CacheBucket(std::string BucketName);
 		~CacheBucket();
 
-		void		OpenOrCreate(std::filesystem::path BucketDir, bool AllowCreate = true);
-		static bool Delete(std::filesystem::path BucketDir);
-		bool		Get(const IoHash& HashKey, ZenCacheValue& OutValue);
-		void		Put(const IoHash& HashKey, const ZenCacheValue& Value);
-		void		Drop();
-		void		Flush();
-		void		SaveManifest();
-		void		Scrub(ScrubContext& Ctx);
-		void		GatherReferences(GcContext& GcCtx);
-		void		CollectGarbage(GcContext& GcCtx);
-		void		UpdateAccessTimes(const std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes);
+		void OpenOrCreate(std::filesystem::path BucketDir, bool AllowCreate = true);
+		bool Get(const IoHash& HashKey, ZenCacheValue& OutValue);
+		void Put(const IoHash& HashKey, const ZenCacheValue& Value);
+		void Drop();
+		void Flush();
+		void Scrub(ScrubContext& Ctx);
+		void GatherReferences(GcContext& GcCtx);
+		void CollectGarbage(GcContext& GcCtx);
+		void UpdateAccessTimes(const std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes);
 
 		inline bool		IsOk() const { return m_IsOk; }
 		inline uint64_t TotalSize() const { return m_TotalSize.load(std::memory_order::relaxed); }
@@ -292,16 +290,18 @@ private:
 
 		std::atomic_uint64_t m_TotalSize{};
 
-		void	 BuildPath(PathBuilderBase& Path, const IoHash& HashKey);
-		void	 PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
-		bool	 GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey, ZenCacheValue& OutValue);
-		void	 PutInlineCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
-		bool	 GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue);
-		void	 MakeIndexSnapshot();
-		uint64_t ReadIndexFile();
-		uint64_t ReadLog(uint64_t LogPosition);
-		uint64_t MigrateLegacyData(bool CleanSource);
-		void	 OpenLog(const std::filesystem::path& BucketDir, const bool IsNew);
+		void		BuildPath(PathBuilderBase& Path, const IoHash& HashKey);
+		void		PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
+		bool		GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey, ZenCacheValue& OutValue);
+		void		PutInlineCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
+		bool		GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue);
+		void		MakeIndexSnapshot();
+		uint64_t	ReadIndexFile();
+		uint64_t	ReadLog(uint64_t LogPosition);
+		uint64_t	MigrateLegacyData(bool CleanSource);
+		void		OpenLog(const std::filesystem::path& BucketDir, const bool IsNew);
+		static bool Delete(std::filesystem::path BucketDir);
+		void		SaveManifest();
 
 		// These locks are here to avoid contention on file creation, therefore it's sufficient
 		// that we take the same lock for the same hash
@@ -314,9 +314,10 @@ private:
 		inline RwLock& LockForHash(const IoHash& Hash) { return m_ShardedLocks[Hash.Hash[19]]; }
 	};
 
-	std::filesystem::path						 m_RootDir;
-	mutable RwLock								 m_Lock;
-	std::unordered_map<std::string, CacheBucket> m_Buckets;	 // TODO: make this case insensitive
+	std::filesystem::path										  m_RootDir;
+	mutable RwLock												  m_Lock;
+	std::unordered_map<std::string, std::unique_ptr<CacheBucket>> m_Buckets;  // TODO: make this case insensitive
+	std::vector<std::unique_ptr<CacheBucket>>					  m_DroppedBuckets;
 
 	ZenCacheDiskLayer(const ZenCacheDiskLayer&) = delete;
 	ZenCacheDiskLayer& operator=(const ZenCacheDiskLayer&) = delete;