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

Zs/put overwrite policy

1 files changed, 87 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())