Handle PartialOnError cache policy when using z$ rest endpoint.

1 files changed, 71 insertions, 63 deletions

diff --git a/zenserver/cache/structuredcache.cpp b/zenserver/cache/structuredcache.cpp
index 4397e3723..875dbf67b 100644
--- a/zenserver/cache/structuredcache.cpp
+++ b/zenserver/cache/structuredcache.cpp
@@ -205,6 +205,7 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 	const ZenContentType AcceptType		 = Request.AcceptContentType();
 	const bool			 SkipData		 = (Policy & CachePolicy::SkipData) == CachePolicy::SkipData;
 	const bool			 SkipAttachments = (Policy & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
+	const bool			 PartialOnError	 = (Policy & CachePolicy::PartialOnError) == CachePolicy::PartialOnError;
 	const bool			 QueryUpstream	 = m_UpstreamCache && ((Policy & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);
 
 	bool		  Success = false;
@@ -214,43 +215,45 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 	{
 		Success = true;
 
-		if (!SkipData && AcceptType == ZenContentType::kCbPackage)
+		if (AcceptType == ZenContentType::kCbPackage)
 		{
-			CbPackage	 Package;
-			CbObjectView CacheRecord(LocalCacheValue.Value.Data());
-			uint32_t	 AttachmentCount = 0;
-			uint32_t	 ValidCount		 = 0;
+			CbPackage Package;
+			uint32_t  MissingCount = 0;
 
-			if (!SkipAttachments)
-			{
-				CacheRecord.IterateAttachments([this, &Package, &AttachmentCount, &ValidCount](CbFieldView AttachmentHash) {
+			CbObjectView CacheRecord(LocalCacheValue.Value.Data());
+			CacheRecord.IterateAttachments([this, SkipAttachments, &MissingCount, &Package](CbFieldView AttachmentHash) {
+				if (SkipAttachments && MissingCount == 0)
+				{
+					if (!m_CidStore.ContainsChunk(AttachmentHash.AsHash()))
+					{
+						MissingCount++;
+					}
+				}
+				else
+				{
 					if (IoBuffer Chunk = m_CidStore.FindChunkByCid(AttachmentHash.AsHash()))
 					{
 						Package.AddAttachment(CbAttachment(CompressedBuffer::FromCompressed(SharedBuffer(Chunk))));
-						ValidCount++;
 					}
-					AttachmentCount++;
-				});
-
-				if (ValidCount != AttachmentCount)
-				{
-					// Success = false;
-					ZEN_WARN("GET - '{}/{}' '{}' is partial, found '{}' of '{}' attachments",
-							 Ref.BucketSegment,
-							 Ref.HashKey,
-							 ToString(AcceptType),
-							 ValidCount,
-							 AttachmentCount);
+					else
+					{
+						MissingCount++;
+					}
 				}
-			}
+			});
+
+			Success = MissingCount == 0 || PartialOnError;
 
-			Package.SetObject(LoadCompactBinaryObject(LocalCacheValue.Value));
+			if (Success)
+			{
+				Package.SetObject(LoadCompactBinaryObject(LocalCacheValue.Value));
 
-			BinaryWriter MemStream;
-			Package.Save(MemStream);
+				BinaryWriter MemStream;
+				Package.Save(MemStream);
 
-			LocalCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
-			LocalCacheValue.Value.SetContentType(HttpContentType::kCbPackage);
+				LocalCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
+				LocalCacheValue.Value.SetContentType(HttpContentType::kCbPackage);
+			}
 		}
 	}
 
@@ -283,7 +286,7 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 	// Issue upstream query asynchronously in order to keep requests flowing without
 	// hogging I/O servicing threads with blocking work
 
-	Request.WriteResponseAsync([this, AcceptType, SkipData, SkipAttachments, Ref](HttpServerRequest& AsyncRequest) {
+	Request.WriteResponseAsync([this, AcceptType, SkipData, SkipAttachments, PartialOnError, Ref](HttpServerRequest& AsyncRequest) {
 		bool		  Success = false;
 		ZenCacheValue UpstreamCacheValue;
 
@@ -292,9 +295,9 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 		if (GetUpstreamCacheResult UpstreamResult = m_UpstreamCache->GetCacheRecord({Ref.BucketSegment, Ref.HashKey}, AcceptType);
 			UpstreamResult.Success)
 		{
-			Success					 = true;
-			UpstreamCacheValue.Value = UpstreamResult.Value;
+			Success = true;
 
+			UpstreamCacheValue.Value = UpstreamResult.Value;
 			UpstreamCacheValue.Value.SetContentType(AcceptType);
 
 			if (AcceptType == ZenContentType::kBinary || AcceptType == ZenContentType::kCbObject)
@@ -302,19 +305,7 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 				if (AcceptType == ZenContentType::kCbObject)
 				{
 					const CbValidateError ValidationResult = ValidateCompactBinary(UpstreamResult.Value, CbValidateMode::All);
-
-					if (ValidationResult == CbValidateError::None)
-					{
-						CbObjectView CacheRecord(UpstreamResult.Value.Data());
-
-						CbObjectWriter IndexData;
-						IndexData.BeginArray("references");
-						CacheRecord.IterateAttachments([&](CbFieldView Attachment) { IndexData.AddHash(Attachment.AsHash()); });
-						IndexData.EndArray();
-
-						UpstreamCacheValue.IndexData = IndexData.Save();
-					}
-					else
+					if (ValidationResult != CbValidateError::None)
 					{
 						Success = false;
 						ZEN_WARN("Get - '{}/{}' '{}' FAILED, invalid compact binary object from upstream",
@@ -334,43 +325,60 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 				CbPackage Package;
 				if (Package.TryLoad(UpstreamCacheValue.Value))
 				{
-					CbObject CacheRecord	 = Package.GetObject();
-					uint32_t AttachmentCount = 0;
-					uint32_t ValidCount		 = 0;
+					CbObject		CacheRecord = Package.GetObject();
+					AttachmentCount Count;
 
-					CacheRecord.IterateAttachments([this, &Package, &Ref, &AttachmentCount, &ValidCount](CbFieldView AttachmentHash) {
-						if (const CbAttachment* Attachment = Package.FindAttachment(AttachmentHash.AsHash()))
+					CacheRecord.IterateAttachments([this, &Package, &Ref, &Count, SkipAttachments](CbFieldView HashView) {
+						if (const CbAttachment* Attachment = Package.FindAttachment(HashView.AsHash()))
 						{
-							if (CompressedBuffer Chunk = Attachment->AsCompressedBinary())
+							if (CompressedBuffer Compressed = Attachment->AsCompressedBinary())
 							{
-								m_CidStore.AddChunk(Chunk);
-								ValidCount++;
+								auto InsertResult = m_CidStore.AddChunk(Compressed);
+								if (InsertResult.New)
+								{
+									Count.New++;
+								}
+								Count.Valid++;
 							}
 							else
 							{
-								ZEN_WARN("Get - '{}/{}' '{}' FAILED, upstream attachment not compressed",
+								ZEN_WARN("Uncompressed payload '{}' from upstream cache record '{}/{}'",
+										 HashView.AsHash(),
 										 Ref.BucketSegment,
-										 Ref.HashKey,
-										 ToString(ZenContentType::kCbPackage));
+										 Ref.HashKey);
+								Count.Invalid++;
+							}
+						}
+						else if (IoBuffer Chunk = m_CidStore.FindChunkByCid(HashView.AsHash()))
+						{
+							if (!SkipAttachments)
+							{
+								Package.AddAttachment(CbAttachment(CompressedBuffer::FromCompressed(SharedBuffer(Chunk))));
 							}
+							Count.Valid++;
 						}
-						AttachmentCount++;
+						Count.Total++;
 					});
 
-					if (ValidCount == AttachmentCount)
+					if ((Count.Valid == Count.Total) || PartialOnError)
 					{
-						m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = CacheRecord.GetBuffer().AsIoBuffer()});
+						ZenCacheValue CacheValue;
+						CacheValue.Value = CacheRecord.GetBuffer().AsIoBuffer();
+						CacheValue.Value.SetContentType(ZenContentType::kCbObject);
+
+						m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, CacheValue);
 
 						if (SkipAttachments)
 						{
-							CbPackage PackageWithoutAttachments;
-							PackageWithoutAttachments.SetObject(CacheRecord);
+							Package.Reset();
+							Package.SetObject(CacheRecord);
+						}
 
-							BinaryWriter MemStream;
-							PackageWithoutAttachments.Save(MemStream);
+						BinaryWriter MemStream;
+						Package.Save(MemStream);
 
-							UpstreamCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
-						}
+						UpstreamCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
+						UpstreamCacheValue.Value.SetContentType(ZenContentType::kCbPackage);
 					}
 					else
 					{