Merge pull request #434 from ue-foundation/zs/put-overwrite-policy

Zs/put overwrite policy

4 files changed, 98 insertions, 29 deletions

diff --git a/src/zenserver/cache/httpstructuredcache.cpp b/src/zenserver/cache/httpstructuredcache.cpp
index bb0c55618..acb6053d9 100644
--- a/src/zenserver/cache/httpstructuredcache.cpp
+++ b/src/zenserver/cache/httpstructuredcache.cpp
@@ -997,8 +997,14 @@ HttpStructuredCacheService::HandleGetCacheRecord(HttpServerRequest& Request, con
 
 				if (Success && StoreLocal)
 				{
-					m_CacheStore.Put(RequestContext, Ref.Namespace, Ref.BucketSegment, Ref.HashKey, ClientResultValue, {}, nullptr);
-					m_CacheStats.WriteCount++;
+					const bool				 Overwrite = !EnumHasAllFlags(PolicyFromUrl, CachePolicy::QueryLocal);
+					ZenCacheStore::PutResult PutResult =
+						m_CacheStore
+							.Put(RequestContext, Ref.Namespace, Ref.BucketSegment, Ref.HashKey, ClientResultValue, {}, Overwrite, nullptr);
+					if (PutResult.Status == zen::PutStatus::Success)
+					{
+						m_CacheStats.WriteCount++;
+					}
 				}
 			}
 			else if (AcceptType == ZenContentType::kCbPackage)
@@ -1083,30 +1089,35 @@ HttpStructuredCacheService::HandleGetCacheRecord(HttpServerRequest& Request, con
 
 						if (StoreLocal)
 						{
-							m_CacheStore.Put(RequestContext,
-											 Ref.Namespace,
-											 Ref.BucketSegment,
-											 Ref.HashKey,
-											 CacheValue,
-											 ReferencedAttachments,
-											 nullptr);
-							m_CacheStats.WriteCount++;
-
-							if (!WriteAttachmentBuffers.empty())
+							const bool				 Overwrite = !EnumHasAllFlags(PolicyFromUrl, CachePolicy::QueryLocal);
+							ZenCacheStore::PutResult PutResult = m_CacheStore.Put(RequestContext,
+																				  Ref.Namespace,
+																				  Ref.BucketSegment,
+																				  Ref.HashKey,
+																				  CacheValue,
+																				  ReferencedAttachments,
+																				  Overwrite,
+																				  nullptr);
+							if (PutResult.Status == zen::PutStatus::Success)
 							{
-								std::vector<CidStore::InsertResult> InsertResults =
-									m_CidStore.AddChunks(WriteAttachmentBuffers, WriteRawHashes);
-								for (const CidStore::InsertResult& Result : InsertResults)
+								m_CacheStats.WriteCount++;
+
+								if (!WriteAttachmentBuffers.empty())
 								{
-									if (Result.New)
+									std::vector<CidStore::InsertResult> InsertResults =
+										m_CidStore.AddChunks(WriteAttachmentBuffers, WriteRawHashes);
+									for (const CidStore::InsertResult& Result : InsertResults)
 									{
-										Count.New++;
+										if (Result.New)
+										{
+											Count.New++;
+										}
 									}
 								}
-							}
 
-							WriteAttachmentBuffers = {};
-							WriteRawHashes		   = {};
+								WriteAttachmentBuffers = {};
+								WriteRawHashes		   = {};
+							}
 						}
 
 						BinaryWriter MemStream;
@@ -1197,6 +1208,24 @@ HttpStructuredCacheService::HandlePutCacheRecord(HttpServerRequest& Request, con
 		return Request.WriteResponse(HttpResponseCode::InsufficientStorage);
 	}
 
+	auto WriteFailureResponse = [&Request](const ZenCacheStore::PutResult& PutResult) {
+		ZEN_UNUSED(PutResult);
+
+		HttpResponseCode ResponseCode = HttpResponseCode::InternalServerError;
+		switch (PutResult.Status)
+		{
+			case zen::PutStatus::Conflict:
+				ResponseCode = HttpResponseCode::Conflict;
+				break;
+			case zen::PutStatus::Invalid:
+				ResponseCode = HttpResponseCode::BadRequest;
+				break;
+		}
+
+		return PutResult.Message.empty() ? Request.WriteResponse(ResponseCode)
+										 : Request.WriteResponse(ResponseCode, zen::HttpContentType::kText, PutResult.Message);
+	};
+
 	const HttpContentType ContentType = Request.RequestContentType();
 
 	Body.SetContentType(ContentType);
@@ -1225,13 +1254,20 @@ HttpStructuredCacheService::HandlePutCacheRecord(HttpServerRequest& Request, con
 		{
 			RawHash = IoHash::HashBuffer(SharedBuffer(Body));
 		}
-		m_CacheStore.Put(RequestContext,
-						 Ref.Namespace,
-						 Ref.BucketSegment,
-						 Ref.HashKey,
-						 {.Value = Body, .RawSize = RawSize, .RawHash = RawHash},
-						 {},
-						 nullptr);
+		const bool Overwrite = !EnumHasAllFlags(PolicyFromUrl, CachePolicy::QueryLocal);
+		// TODO: Propagation for rejected PUTs
+		ZenCacheStore::PutResult PutResult = m_CacheStore.Put(RequestContext,
+															  Ref.Namespace,
+															  Ref.BucketSegment,
+															  Ref.HashKey,
+															  {.Value = Body, .RawSize = RawSize, .RawHash = RawHash},
+															  {},
+															  Overwrite,
+															  nullptr);
+		if (PutResult.Status != zen::PutStatus::Success)
+		{
+			return WriteFailureResponse(PutResult);
+		}
 		m_CacheStats.WriteCount++;
 
 		if (HasUpstream && EnumHasAllFlags(PolicyFromUrl, CachePolicy::StoreRemote))
@@ -1280,7 +1316,21 @@ HttpStructuredCacheService::HandlePutCacheRecord(HttpServerRequest& Request, con
 			TotalCount++;
 		});
 
-		m_CacheStore.Put(RequestContext, Ref.Namespace, Ref.BucketSegment, Ref.HashKey, {.Value = Body}, ReferencedAttachments, nullptr);
+		const bool Overwrite = !EnumHasAllFlags(PolicyFromUrl, CachePolicy::QueryLocal);
+
+		// TODO: Propagation for rejected PUTs
+		ZenCacheStore::PutResult PutResult = m_CacheStore.Put(RequestContext,
+															  Ref.Namespace,
+															  Ref.BucketSegment,
+															  Ref.HashKey,
+															  {.Value = Body},
+															  ReferencedAttachments,
+															  Overwrite,
+															  nullptr);
+		if (PutResult.Status != zen::PutStatus::Success)
+		{
+			return WriteFailureResponse(PutResult);
+		}
 		m_CacheStats.WriteCount++;
 
 		ZEN_DEBUG("PUTCACHERECORD - '{}/{}/{}' {} '{}' attachments '{}/{}' (valid/total) in {}",
@@ -1373,10 +1423,18 @@ HttpStructuredCacheService::HandlePutCacheRecord(HttpServerRequest& Request, con
 			return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Invalid attachment(s)"sv);
 		}
 
+		const bool Overwrite = !EnumHasAllFlags(Policy, CachePolicy::QueryLocal);
+
 		ZenCacheValue CacheValue;
 		CacheValue.Value = CacheRecord.GetBuffer().AsIoBuffer();
 		CacheValue.Value.SetContentType(ZenContentType::kCbObject);
-		m_CacheStore.Put(RequestContext, Ref.Namespace, Ref.BucketSegment, Ref.HashKey, CacheValue, ReferencedAttachments);
+		// TODO: Propagation for rejected PUTs
+		ZenCacheStore::PutResult PutResult =
+			m_CacheStore.Put(RequestContext, Ref.Namespace, Ref.BucketSegment, Ref.HashKey, CacheValue, ReferencedAttachments, Overwrite);
+		if (PutResult.Status != zen::PutStatus::Success)
+		{
+			return WriteFailureResponse(PutResult);
+		}
 		m_CacheStats.WriteCount++;
 
 		if (!WriteAttachmentBuffers.empty())
diff --git a/src/zenserver/config.cpp b/src/zenserver/config.cpp
index 1f9ae5fb6..d53bedad0 100644
--- a/src/zenserver/config.cpp
+++ b/src/zenserver/config.cpp
@@ -518,6 +518,7 @@ ParseConfigFile(const std::filesystem::path& Path,
 	LuaOptions.AddOption("cache.enable"sv, ServerOptions.StructuredCacheConfig.Enabled);
 	LuaOptions.AddOption("cache.writelog"sv, ServerOptions.StructuredCacheConfig.WriteLogEnabled, "cache-write-log"sv);
 	LuaOptions.AddOption("cache.accesslog"sv, ServerOptions.StructuredCacheConfig.AccessLogEnabled, "cache-access-log"sv);
+	LuaOptions.AddOption("cache.limitoverwrites"sv, ServerOptions.StructuredCacheConfig.LimitOverwrites, "cache-limit-overwrites"sv);
 
 	LuaOptions.AddOption("cache.memlayer.sizethreshold"sv,
 						 ServerOptions.StructuredCacheConfig.BucketConfig.MemCacheSizeThreshold,
@@ -1060,6 +1061,13 @@ ParseCliOptions(int argc, char* argv[], ZenServerOptions& ServerOptions)
 					   cxxopts::value<bool>(ServerOptions.StructuredCacheConfig.AccessLogEnabled)->default_value("false"),
 					   "");
 
+	options.add_option("cache",
+					   "",
+					   "cache-limit-overwrites",
+					   "Whether to require policy flag pattern before allowing overwrites",
+					   cxxopts::value<bool>(ServerOptions.StructuredCacheConfig.LimitOverwrites)->default_value("false"),
+					   "");
+
 	options.add_option(
 		"cache",
 		"",
diff --git a/src/zenserver/config.h b/src/zenserver/config.h
index 9753e3ae2..7d90aa4c1 100644
--- a/src/zenserver/config.h
+++ b/src/zenserver/config.h
@@ -132,6 +132,7 @@ struct ZenStructuredCacheConfig
 	bool																Enabled			 = true;
 	bool																WriteLogEnabled	 = false;
 	bool																AccessLogEnabled = false;
+	bool																LimitOverwrites	 = false;
 	std::vector<std::pair<std::string, ZenStructuredCacheBucketConfig>> PerBucketConfigs;
 	ZenStructuredCacheBucketConfig										BucketConfig;
 	uint64_t															MemTargetFootprintBytes = 512 * 1024 * 1024;
diff --git a/src/zenserver/zenserver.cpp b/src/zenserver/zenserver.cpp
index 54da51d77..6c9f2ed21 100644
--- a/src/zenserver/zenserver.cpp
+++ b/src/zenserver/zenserver.cpp
@@ -552,6 +552,7 @@ ZenServer::InitializeStructuredCache(const ZenServerOptions& ServerOptions)
 	Config.AllowAutomaticCreationOfNamespaces = true;
 	Config.Logging							  = {.EnableWriteLog  = ServerOptions.StructuredCacheConfig.WriteLogEnabled,
 												 .EnableAccessLog = ServerOptions.StructuredCacheConfig.AccessLogEnabled};
+
 	for (const auto& It : ServerOptions.StructuredCacheConfig.PerBucketConfigs)
 	{
 		const std::string&					   BucketName	   = It.first;
@@ -569,6 +570,7 @@ ZenServer::InitializeStructuredCache(const ZenServerOptions& ServerOptions)
 		ServerOptions.StructuredCacheConfig.BucketConfig.MemCacheSizeThreshold,
 	Config.NamespaceConfig.DiskLayerConfig.BucketConfig.LargeObjectThreshold =
 		ServerOptions.StructuredCacheConfig.BucketConfig.LargeObjectThreshold,
+	Config.NamespaceConfig.DiskLayerConfig.BucketConfig.LimitOverwrites = ServerOptions.StructuredCacheConfig.LimitOverwrites;
 	Config.NamespaceConfig.DiskLayerConfig.MemCacheTargetFootprintBytes = ServerOptions.StructuredCacheConfig.MemTargetFootprintBytes;
 	Config.NamespaceConfig.DiskLayerConfig.MemCacheTrimIntervalSeconds	= ServerOptions.StructuredCacheConfig.MemTrimIntervalSeconds;
 	Config.NamespaceConfig.DiskLayerConfig.MemCacheMaxAgeSeconds		= ServerOptions.StructuredCacheConfig.MemMaxAgeSeconds;