Merge pull request #106 from EpicGames/de/safer-delete-cache-bucket

Safer delete cache bucket

1 files changed, 168 insertions, 71 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 3189a14cc..9ebec1e76 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -572,7 +572,7 @@ ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue
 
 //////////////////////////////////////////////////////////////////////////
 
-ZenCacheDiskLayer::CacheBucket::CacheBucket(std::string BucketName) : m_BucketName(std::move(BucketName))
+ZenCacheDiskLayer::CacheBucket::CacheBucket(std::string BucketName) : m_BucketName(std::move(BucketName)), m_BucketId(Oid::Zero)
 {
 }
 
@@ -581,19 +581,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)
 {
 	using namespace std::literals;
@@ -611,7 +598,10 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bo
 	if (Manifest)
 	{
 		m_BucketId = Manifest["BucketId"].AsObjectId();
-		m_IsOk	   = m_BucketId != Oid::Zero;
+		if (m_BucketId == Oid::Zero)
+		{
+			return false;
+		}
 	}
 	else if (AllowCreate)
 	{
@@ -625,7 +615,7 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bo
 	}
 	else
 	{
-		return;
+		return false;
 	}
 
 	OpenLog(BucketDir, IsNew);
@@ -641,7 +631,7 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bo
 		}
 	}
 
-	m_IsOk = true;
+	return true;
 }
 
 void
@@ -1158,11 +1148,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 +1169,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,12 +1176,54 @@ 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);
+
+	std::vector<std::unique_ptr<RwLock::ExclusiveLockScope>> ShardLocks;
+	ShardLocks.reserve(256);
+	for (RwLock& Lock : m_ShardedLocks)
+	{
+		ShardLocks.push_back(std::make_unique<RwLock::ExclusiveLockScope>(Lock));
+	}
 	m_BlockStore.Close();
 	m_SlogFile.Close();
-	DeleteDirectories(m_BucketDir);
+
+	DeleteBucketFromDisk(m_BucketDir, m_BucketName);
+
+	m_Index.clear();
 }
 
 void
@@ -1703,7 +1725,8 @@ ZenCacheDiskLayer::CollectGarbage(GcContext& GcCtx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.CollectGarbage(GcCtx);
+		CacheBucket& Bucket = *Kv.second;
+		Bucket.CollectGarbage(GcCtx);
 	}
 }
 
@@ -1716,7 +1739,7 @@ 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;
 			Bucket.UpdateAccessTimes(Kv.second);
 		}
 	}
@@ -1929,7 +1952,7 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCach
 
 		if (it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
 		}
 	}
 
@@ -1941,22 +1964,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->OpenOrCreate(BucketPath))
+			{
+				m_Buckets.erase(BucketName);
+				return false;
+			}
 		}
 	}
 
 	ZEN_ASSERT(Bucket != nullptr);
-
 	return Bucket->Get(HashKey, OutValue);
 }
 
@@ -1973,7 +1999,7 @@ ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const Z
 
 		if (it != m_Buckets.end())
 		{
-			Bucket = &it->second;
+			Bucket = it->second.get();
 		}
 	}
 
@@ -1985,26 +2011,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->OpenOrCreate(BucketPath))
+			{
+				m_Buckets.erase(BucketName);
+				return;
+			}
 		}
 	}
 
 	ZEN_ASSERT(Bucket != nullptr);
 
-	if (Bucket->IsOk())
-	{
-		Bucket->Put(HashKey, Value);
-	}
+	Bucket->Put(HashKey, Value);
 }
 
 void
@@ -2023,26 +2050,20 @@ ZenCacheDiskLayer::DiscoverBuckets()
 		// New bucket needs to be created
 		if (auto It = m_Buckets.find(BucketName); It != m_Buckets.end())
 		{
+			continue;
 		}
-		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);
-
-			if (Bucket.IsOk())
-			{
-				ZEN_INFO("Discovered bucket '{}'", BucketName);
-			}
-			else
-			{
-				ZEN_WARN("Found directory '{}' in our base directory '{}' but it is not a valid bucket", BucketName, m_RootDir);
+		if (!Bucket.OpenOrCreate(BucketPath, /* AllowCreate */ false))
+		{
+			ZEN_WARN("Found directory '{}' in our base directory '{}' but it is not a valid bucket", BucketName, m_RootDir);
 
-				m_Buckets.erase(InsertResult.first);
-			}
+			m_Buckets.erase(InsertResult.first);
+			continue;
 		}
+		ZEN_INFO("Discovered bucket '{}'", BucketName);
 	}
 }
 
@@ -2055,19 +2076,19 @@ 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);
 
+		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
@@ -2080,7 +2101,8 @@ ZenCacheDiskLayer::Flush()
 		Buckets.reserve(m_Buckets.size());
 		for (auto& Kv : m_Buckets)
 		{
-			Buckets.push_back(&Kv.second);
+			CacheBucket* Bucket = Kv.second.get();
+			Buckets.push_back(Bucket);
 		}
 	}
 
@@ -2097,7 +2119,8 @@ ZenCacheDiskLayer::Scrub(ScrubContext& Ctx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.Scrub(Ctx);
+		CacheBucket& Bucket = *Kv.second;
+		Bucket.Scrub(Ctx);
 	}
 }
 
@@ -2108,7 +2131,8 @@ ZenCacheDiskLayer::GatherReferences(GcContext& GcCtx)
 
 	for (auto& Kv : m_Buckets)
 	{
-		Kv.second.GatherReferences(GcCtx);
+		CacheBucket& Bucket = *Kv.second;
+		Bucket.GatherReferences(GcCtx);
 	}
 }
 
@@ -2120,7 +2144,7 @@ ZenCacheDiskLayer::TotalSize() const
 
 	for (auto& Kv : m_Buckets)
 	{
-		TotalSize += Kv.second.TotalSize();
+		TotalSize += Kv.second->TotalSize();
 	}
 
 	return TotalSize;
@@ -2274,7 +2298,7 @@ ZenCacheStore::GetNamespace(std::string_view Namespace)
 void
 ZenCacheStore::IterateNamespaces(const std::function<void(std::string_view Namespace, ZenCacheNamespace& Store)>& Callback) const
 {
-	std::vector<std::pair<std::string, ZenCacheNamespace&> > Namespaces;
+	std::vector<std::pair<std::string, ZenCacheNamespace&>> Namespaces;
 	{
 		RwLock::SharedLockScope _(m_NamespacesLock);
 		Namespaces.reserve(m_Namespaces.size());
@@ -3118,6 +3142,79 @@ TEST_CASE("z$.namespaces")
 	}
 }
 
+TEST_CASE("z$.drop.bucket")
+{
+	using namespace testutils;
+
+	const auto CreateCacheValue = [](size_t Size) -> CbObject {
+		std::vector<uint8_t> Buf;
+		Buf.resize(Size);
+
+		CbObjectWriter Writer;
+		Writer.AddBinary("Binary"sv, Buf.data(), Buf.size());
+		return Writer.Save();
+	};
+
+	ScopedTemporaryDirectory TempDir;
+	CreateDirectories(TempDir.Path());
+
+	IoHash Key1;
+	IoHash Key2;
+
+	auto PutValue =
+		[&CreateCacheValue](ZenCacheStore& Zcs, std::string_view Namespace, std::string_view Bucket, size_t KeyIndex, size_t Size) {
+			// Create a cache record
+			IoHash	 Key		= CreateKey(KeyIndex);
+			CbObject CacheValue = CreateCacheValue(Size);
+
+			IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
+			Buffer.SetContentType(ZenContentType::kCbObject);
+
+			ZenCacheValue PutValue = {.Value = Buffer};
+			Zcs.Put(Namespace, Bucket, Key, PutValue);
+			return Key;
+		};
+	auto GetValue = [](ZenCacheStore& Zcs, std::string_view Namespace, std::string_view Bucket, const IoHash& Key) {
+		ZenCacheValue GetValue;
+		Zcs.Get(Namespace, Bucket, Key, GetValue);
+		return GetValue;
+	};
+	WorkerThreadPool Workers(1);
+	{
+		CasGc		  Gc;
+		ZenCacheStore Zcs(Gc, {.BasePath = TempDir.Path() / "cache", .AllowAutomaticCreationOfNamespaces = true});
+		const auto	  Bucket	= "teardrinker"sv;
+		const auto	  Namespace = "mynamespace"sv;
+
+		Key1 = PutValue(Zcs, Namespace, Bucket, 42, 4096);
+		Key2 = PutValue(Zcs, Namespace, Bucket, 43, 2048);
+
+		ZenCacheValue Value1 = GetValue(Zcs, Namespace, Bucket, Key1);
+		CHECK(Value1.Value);
+
+		std::atomic_bool WorkComplete = false;
+		Workers.ScheduleWork([&]() {
+			zen::Sleep(100);
+			Value1.Value = IoBuffer{};
+			Value1		 = GetValue(Zcs, Namespace, Bucket, Key1);
+			CHECK(!Value1.Value);
+			WorkComplete = true;
+		});
+		// On Windows, DropBucket() will be blocked as long as we hold a reference to a buffer in the bucket
+		// Our DropBucket execution blocks any incoming request from completing until we are done with the drop
+		Zcs.DropBucket(Namespace, Bucket);
+		while (!WorkComplete)
+		{
+			zen::Sleep(1);
+		}
+		CHECK(!Value1.Value);
+
+		// Entire bucket should be dropped, but doing a request should will re-create the namespace but it must still be empty
+		ZenCacheValue Value2 = GetValue(Zcs, Namespace, Bucket, Key2);
+		CHECK(!Value2.Value);
+	}
+}
+
 TEST_CASE("z$.blocked.disklayer.put")
 {
 	ScopedTemporaryDirectory TempDir;