Merge pull request #92 from EpicGames/de/bucket-standalone-temp-file-cleanupv1.0.1.6

Make sure CacheBucket::PutStandaloneCacheValue cleans up the temp file

1 files changed, 106 insertions, 5 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 05c80c5bf..ce55b24b6 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -1725,6 +1725,8 @@ ZenCacheDiskLayer::UpdateAccessTimes(const zen::access_tracking::AccessTimes& Ac
 void
 ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value)
 {
+	uint64_t NewFileSize = Value.Value.Size();
+
 	TemporaryFile DataFile;
 
 	std::error_code Ec;
@@ -1734,12 +1736,28 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 		throw std::system_error(Ec, fmt::format("Failed to open temporary file for put in '{}'", m_BucketDir));
 	}
 
+	bool CleanUpTempFile = false;
+	auto __				 = MakeGuard([&] {
+		 if (CleanUpTempFile)
+		 {
+			 std::error_code Ec;
+			 std::filesystem::remove(DataFile.GetPath(), Ec);
+			 if (Ec)
+			 {
+				 ZEN_WARN("Failed to clean up temporary file '{}' for put in '{}', reason '{}'",
+						  DataFile.GetPath(),
+						  m_BucketDir,
+						  Ec.message());
+			 }
+		 }
+	 });
+
 	DataFile.WriteAll(Value.Value, Ec);
 	if (Ec)
 	{
 		throw std::system_error(Ec,
 								fmt::format("Failed to write payload ({} bytes) to temporary file '{}' for put in '{}'",
-											NiceBytes(Value.Value.Size()),
+											NiceBytes(NewFileSize),
 											DataFile.GetPath().string(),
 											m_BucketDir));
 	}
@@ -1748,7 +1766,7 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 	BuildPath(DataFilePath, HashKey);
 	std::filesystem::path FsPath{DataFilePath.ToPath()};
 
-	// We retry to open the file since it can be held open for read.
+	// We retry to move the file since it can be held open for read.
 	// This happens if the server processes a Get request for the file or
 	// if we are busy sending the file upstream
 	int RetryCount = 3;
@@ -1757,7 +1775,30 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 		Ec.clear();
 		{
 			RwLock::ExclusiveLockScope ValueLock(LockForHash(HashKey));
+
 			DataFile.MoveTemporaryIntoPlace(FsPath, Ec);
+
+			// Once we have called MoveTemporaryIntoPlace automatic clean up the temp file
+			// will be disabled as the file handle has already been closed
+			CleanUpTempFile = Ec ? true : false;
+
+			if (Ec)
+			{
+				std::error_code ExistingEc;
+				uint64_t		OldFileSize = std::filesystem::file_size(FsPath, ExistingEc);
+				if (!ExistingEc && (OldFileSize == NewFileSize))
+				{
+					ZEN_INFO(
+						"Failed to move temporary file '{}' to '{}'. Target file has same size, assuming concurrent write of same value, "
+						"move "
+						"failed with reason '{}'",
+						DataFile.GetPath(),
+						FsPath.string(),
+						m_BucketDir,
+						Ec.message());
+					return;
+				}
+			}
 		}
 
 		if (!Ec)
@@ -1773,9 +1814,10 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 				EntryFlags |= DiskLocation::kCompressed;
 			}
 
-			DiskLocation Loc(Value.Value.Size(), EntryFlags);
+			DiskLocation Loc(NewFileSize, EntryFlags);
 			IndexEntry	 Entry = IndexEntry(Loc, GcClock::TickCount());
 
+			uint64_t				   OldFileSize = 0;
 			RwLock::ExclusiveLockScope _(m_IndexLock);
 			if (auto It = m_Index.find(HashKey); It == m_Index.end())
 			{
@@ -1785,11 +1827,19 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 			else
 			{
 				// TODO: should check if write is idempotent and bail out if it is?
-				It.value() = Entry;
+				OldFileSize = It.value().Location.Size();
+				It.value()	= Entry;
 			}
 
 			m_SlogFile.Append({.Key = HashKey, .Location = Loc});
-			m_TotalSize.fetch_add(Loc.Size(), std::memory_order::relaxed);
+			if (OldFileSize <= NewFileSize)
+			{
+				m_TotalSize.fetch_add(NewFileSize - OldFileSize, std::memory_order::relaxed);
+			}
+			else
+			{
+				m_TotalSize.fetch_sub(OldFileSize - NewFileSize, std::memory_order::relaxed);
+			}
 			return;
 		}
 
@@ -2971,6 +3021,57 @@ TEST_CASE("z$.threadedinsert")	// * doctest::skip(true))
 	}
 }
 
+TEST_CASE("z$.blocked.disklayer.put")
+{
+	ScopedTemporaryDirectory TempDir;
+
+	GcStorageSize CacheSize;
+
+	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();
+	};
+
+	CasGc			  Gc;
+	ZenCacheNamespace Zcs(Gc, TempDir.Path() / "cache");
+
+	CbObject CacheValue = CreateCacheValue(64 * 1024 + 64);
+
+	IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
+	Buffer.SetContentType(ZenContentType::kCbObject);
+
+	size_t Key	   = Buffer.Size();
+	IoHash HashKey = IoHash::HashBuffer(&Key, sizeof(uint32_t));
+	Zcs.Put("test_bucket", HashKey, {.Value = Buffer});
+
+	ZenCacheValue BufferGet;
+	CHECK(Zcs.Get("test_bucket", HashKey, BufferGet));
+
+	// Overwriting with a value of same size should go fine
+	Zcs.Put("test_bucket", HashKey, {.Value = Buffer});
+
+	CbObject CacheValue2 = CreateCacheValue(64 * 1024 + 64 + 1);
+	IoBuffer Buffer2	 = CacheValue2.GetBuffer().AsIoBuffer();
+	Buffer2.SetContentType(ZenContentType::kCbObject);
+#	if ZEN_PLATFORM_WINDOWS
+	// On Windows platform, overwriting with different size while we have
+	// it open for read should throw exception if file is held open
+	CHECK_THROWS(Zcs.Put("test_bucket", HashKey, {.Value = Buffer2}));
+#	else
+	// Other platforms should handle overwrite just fine
+	Zcs.Put("test_bucket", HashKey, {.Value = Buffer2});
+#	endif
+
+	BufferGet = ZenCacheValue{};
+
+	// Read access has been removed, we should now be able to overwrite it
+	Zcs.Put("test_bucket", HashKey, {.Value = Buffer2});
+}
+
 #endif
 
 void