Make sure ZenCacheMemoryLayer handles dropped buckets correctly (just like ZenCacheDiskLayer)

1 files changed, 87 insertions, 32 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 9ebec1e76..d4b8bff39 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -404,14 +404,14 @@ ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCa
 {
 	RwLock::SharedLockScope _(m_Lock);
 
-	auto it = m_Buckets.find(std::string(InBucket));
+	auto It = m_Buckets.find(std::string(InBucket));
 
-	if (it == m_Buckets.end())
+	if (It == m_Buckets.end())
 	{
 		return false;
 	}
 
-	CacheBucket* Bucket = &it->second;
+	CacheBucket* Bucket = It->second.get();
 
 	_.ReleaseNow();
 
@@ -425,14 +425,15 @@ ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCa
 void
 ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const ZenCacheValue& Value)
 {
-	CacheBucket* Bucket = nullptr;
+	const auto	 BucketName = std::string(InBucket);
+	CacheBucket* Bucket		= nullptr;
 
 	{
 		RwLock::SharedLockScope _(m_Lock);
 
 		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
 		{
-			Bucket = &It->second;
+			Bucket = It->second.get();
 		}
 	}
 
@@ -444,11 +445,12 @@ ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const
 
 		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
 		{
-			Bucket = &It->second;
+			Bucket = It->second.get();
 		}
 		else
 		{
-			Bucket = &m_Buckets[std::string(InBucket)];
+			auto InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>());
+			Bucket			  = InsertResult.first->second.get();
 		}
 	}
 
@@ -458,11 +460,37 @@ ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const
 }
 
 bool
-ZenCacheMemoryLayer::DropBucket(std::string_view Bucket)
+ZenCacheMemoryLayer::DropBucket(std::string_view InBucket)
 {
 	RwLock::ExclusiveLockScope _(m_Lock);
 
-	return !!m_Buckets.erase(std::string(Bucket));
+	auto It = m_Buckets.find(std::string(InBucket));
+
+	if (It != m_Buckets.end())
+	{
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It);
+		Bucket.Drop();
+		return true;
+	}
+	return false;
+}
+
+void
+ZenCacheMemoryLayer::Drop()
+{
+	RwLock::ExclusiveLockScope				  _(m_Lock);
+	std::vector<std::unique_ptr<CacheBucket>> Buckets;
+	Buckets.reserve(m_Buckets.size());
+	while (!m_Buckets.empty())
+	{
+		const auto&	 It		= m_Buckets.begin();
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It->first);
+		Bucket.Drop();
+	}
 }
 
 void
@@ -472,7 +500,7 @@ ZenCacheMemoryLayer::Scrub(ScrubContext& Ctx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.Scrub(Ctx);
+		Kv.second->Scrub(Ctx);
 	}
 }
 
@@ -486,7 +514,7 @@ ZenCacheMemoryLayer::GatherAccessTimes(zen::access_tracking::AccessTimes& Access
 	for (auto& Kv : m_Buckets)
 	{
 		std::vector<KeyAccessTime>& Bucket = AccessTimes.Buckets[Kv.first];
-		Kv.second.GatherAccessTimes(Bucket);
+		Kv.second->GatherAccessTimes(Bucket);
 	}
 }
 
@@ -505,7 +533,7 @@ ZenCacheMemoryLayer::TotalSize() const
 
 	for (auto& Kv : m_Buckets)
 	{
-		TotalSize += Kv.second.TotalSize();
+		TotalSize += Kv.second->TotalSize();
 	}
 
 	return TotalSize;
@@ -570,6 +598,13 @@ ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue
 	m_TotalSize.fetch_add(Value.Value.GetSize(), std::memory_order::relaxed);
 }
 
+void
+ZenCacheMemoryLayer::CacheBucket::Drop()
+{
+	RwLock::ExclusiveLockScope _(m_BucketLock);
+	m_CacheMap.clear();
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 ZenCacheDiskLayer::CacheBucket::CacheBucket(std::string BucketName) : m_BucketName(std::move(BucketName)), m_BucketId(Oid::Zero)
@@ -1207,7 +1242,7 @@ DeleteBucketFromDisk(const std::filesystem::path& BucketDir, std::string_view Bu
 	} while (true);
 }
 
-void
+bool
 ZenCacheDiskLayer::CacheBucket::Drop()
 {
 	RwLock::ExclusiveLockScope _(m_IndexLock);
@@ -1221,9 +1256,10 @@ ZenCacheDiskLayer::CacheBucket::Drop()
 	m_BlockStore.Close();
 	m_SlogFile.Close();
 
-	DeleteBucketFromDisk(m_BucketDir, m_BucketName);
+	bool Deleted = DeleteBucketFromDisk(m_BucketDir, m_BucketName);
 
 	m_Index.clear();
+	return Deleted;
 }
 
 void
@@ -1948,11 +1984,11 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCach
 	{
 		RwLock::SharedLockScope _(m_Lock);
 
-		auto it = m_Buckets.find(BucketName);
+		auto It = m_Buckets.find(BucketName);
 
-		if (it != m_Buckets.end())
+		if (It != m_Buckets.end())
 		{
-			Bucket = it->second.get();
+			Bucket = It->second.get();
 		}
 	}
 
@@ -1962,9 +1998,9 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCach
 
 		RwLock::ExclusiveLockScope _(m_Lock);
 
-		if (auto it = m_Buckets.find(BucketName); it != m_Buckets.end())
+		if (auto It = m_Buckets.find(BucketName); It != m_Buckets.end())
 		{
-			Bucket = it->second.get();
+			Bucket = It->second.get();
 		}
 		else
 		{
@@ -1995,11 +2031,11 @@ ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const Z
 	{
 		RwLock::SharedLockScope _(m_Lock);
 
-		auto it = m_Buckets.find(BucketName);
+		auto It = m_Buckets.find(BucketName);
 
-		if (it != m_Buckets.end())
+		if (It != m_Buckets.end())
 		{
-			Bucket = it->second.get();
+			Bucket = It->second.get();
 		}
 	}
 
@@ -2009,9 +2045,9 @@ ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const Z
 
 		RwLock::ExclusiveLockScope _(m_Lock);
 
-		if (auto it = m_Buckets.find(BucketName); it != m_Buckets.end())
+		if (auto It = m_Buckets.find(BucketName); It != m_Buckets.end())
 		{
-			Bucket = it->second.get();
+			Bucket = It->second.get();
 		}
 		else
 		{
@@ -2072,17 +2108,15 @@ ZenCacheDiskLayer::DropBucket(std::string_view InBucket)
 {
 	RwLock::ExclusiveLockScope _(m_Lock);
 
-	auto it = m_Buckets.find(std::string(InBucket));
+	auto It = m_Buckets.find(std::string(InBucket));
 
-	if (it != m_Buckets.end())
+	if (It != m_Buckets.end())
 	{
-		CacheBucket& Bucket = *it->second;
-		m_DroppedBuckets.push_back(std::move(it->second));
-		m_Buckets.erase(it);
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It);
 
-		Bucket.Drop();
-
-		return true;
+		return Bucket.Drop();
 	}
 
 	// Make sure we remove the folder even if we don't know about the bucket
@@ -2091,6 +2125,27 @@ ZenCacheDiskLayer::DropBucket(std::string_view InBucket)
 	return DeleteBucketFromDisk(BucketPath, InBucket);
 }
 
+bool
+ZenCacheDiskLayer::Drop()
+{
+	RwLock::ExclusiveLockScope _(m_Lock);
+
+	std::vector<std::unique_ptr<CacheBucket>> Buckets;
+	Buckets.reserve(m_Buckets.size());
+	while (!m_Buckets.empty())
+	{
+		const auto&	 It		= m_Buckets.begin();
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It->first);
+		if (!Bucket.Drop())
+		{
+			return false;
+		}
+	}
+	return true;
+}
+
 void
 ZenCacheDiskLayer::Flush()
 {