Added batch API to upstream endpoints.

1 files changed, 276 insertions, 70 deletions

diff --git a/zenserver/upstream/upstreamcache.cpp b/zenserver/upstream/upstreamcache.cpp
index 7ef0caf62..0a0706656 100644
--- a/zenserver/upstream/upstreamcache.cpp
+++ b/zenserver/upstream/upstreamcache.cpp
@@ -70,7 +70,7 @@ namespace detail {
 
 		virtual std::string_view DisplayName() const override { return m_DisplayName; }
 
-		virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
+		virtual GetUpstreamCacheResult GetCacheRecord(CacheKey CacheKey, ZenContentType Type) override
 		{
 			try
 			{
@@ -144,12 +144,54 @@ namespace detail {
 			}
 		}
 
-		virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
+		virtual GetUpstreamCacheBatchResult GetCacheRecords(std::span<CacheKey>		   CacheKeys,
+															std::span<size_t>		   KeyIndex,
+															const CacheRecordPolicy&   Policy,
+															OnCacheRecordGetComplete&& OnComplete) override
+		{
+			ZEN_UNUSED(Policy);
+
+			CloudCacheSession Session(m_Client);
+
+			for (size_t Index : KeyIndex)
+			{
+				const CacheKey&	 CacheKey = CacheKeys[Index];
+				CloudCacheResult Result	  = Session.GetRef(CacheKey.Bucket, CacheKey.Hash, ZenContentType::kCbObject);
+
+				CbPackage	 Package;
+				CbObjectView Record;
+
+				if (Result.Success)
+				{
+					const CbValidateError ValidationResult = ValidateCompactBinary(Result.Response, CbValidateMode::All);
+					if (ValidationResult == CbValidateError::None)
+					{
+						Record = CbObjectView(Result.Response.GetData());
+						Record.IterateAttachments([&Session, &Result, &Package](CbFieldView AttachmentHash) {
+							CloudCacheResult AttachmentResult = Session.GetCompressedBlob(AttachmentHash.AsHash());
+							if (AttachmentResult.Success)
+							{
+								if (CompressedBuffer Chunk = CompressedBuffer::FromCompressed(SharedBuffer(AttachmentResult.Response)))
+								{
+									Package.AddAttachment(CbAttachment(Chunk));
+								}
+							}
+						});
+					}
+				}
+
+				OnComplete({.CacheKey = CacheKey, .KeyIndex = Index, .Record = Record, .Package = Package});
+			}
+
+			return {};
+		}
+
+		virtual GetUpstreamCacheResult GetCachePayload(const CacheKey&, const IoHash& PayloadId) override
 		{
 			try
 			{
 				CloudCacheSession	   Session(m_Client);
-				const CloudCacheResult Result = Session.GetCompressedBlob(PayloadKey.PayloadId);
+				const CloudCacheResult Result = Session.GetCompressedBlob(PayloadId);
 
 				if (Result.ErrorCode == 0)
 				{
@@ -171,6 +213,29 @@ namespace detail {
 			}
 		}
 
+		virtual GetUpstreamCacheBatchResult GetCachePayloads(std::span<CacheChunkRequest> CacheChunkRequests,
+															 std::span<size_t>			  RequestIndex,
+															 OnCachePayloadGetComplete&&  OnComplete) override final
+		{
+			CloudCacheSession Session(m_Client);
+
+			for (size_t Index : RequestIndex)
+			{
+				const CacheChunkRequest& Request = CacheChunkRequests[Index];
+				const CloudCacheResult	 Result	 = Session.GetCompressedBlob(Request.ChunkId);
+
+				OnComplete({.Request = Request, .RequestIndex = Index, .Payload = Result.Response});
+
+				if (Result.ErrorCode)
+				{
+					m_HealthOk = false;
+					break;
+				}
+			}
+
+			return {};
+		}
+
 		virtual PutUpstreamCacheResult PutCacheRecord(const UpstreamCacheRecord& CacheRecord,
 													  IoBuffer					 RecordValue,
 													  std::span<IoBuffer const>	 Payloads) override
@@ -419,7 +484,7 @@ namespace detail {
 
 		virtual std::string_view DisplayName() const override { return m_DisplayName; }
 
-		virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
+		virtual GetUpstreamCacheResult GetCacheRecord(CacheKey CacheKey, ZenContentType Type) override
 		{
 			try
 			{
@@ -446,13 +511,81 @@ namespace detail {
 			}
 		}
 
-		virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
+		virtual GetUpstreamCacheBatchResult GetCacheRecords(std::span<CacheKey>		   CacheKeys,
+															std::span<size_t>		   KeyIndex,
+															const CacheRecordPolicy&   Policy,
+															OnCacheRecordGetComplete&& OnComplete) override
+		{
+			std::vector<size_t> IndexMap;
+			IndexMap.reserve(KeyIndex.size());
+
+			CbObjectWriter BatchRequest;
+			BatchRequest << "Method"sv
+						 << "GetCacheRecords";
+
+			BatchRequest.BeginObject("Params"sv);
+			{
+				BatchRequest.BeginArray("CacheKeys"sv);
+				for (size_t Index : KeyIndex)
+				{
+					const CacheKey& Key = CacheKeys[Index];
+					IndexMap.push_back(Index);
+
+					BatchRequest.BeginObject();
+					BatchRequest << "Bucket"sv << Key.Bucket;
+					BatchRequest << "Hash"sv << Key.Hash;
+					BatchRequest.EndObject();
+				}
+				BatchRequest.EndArray();
+
+				BatchRequest.BeginObject("Policy"sv);
+				CacheRecordPolicy::Save(Policy, BatchRequest);
+				BatchRequest.EndObject();
+			}
+			BatchRequest.EndObject();
+
+			CbPackage BatchResponse;
+			bool	  Success = false;
+
+			{
+				ZenStructuredCacheSession Session(*m_Client);
+				ZenCacheResult			  Result = Session.InvokeRpc(BatchRequest.Save());
+				if (Result.Success)
+				{
+					Success = BatchResponse.TryLoad(Result.Response);
+				}
+				else if (Result.ErrorCode)
+				{
+					Success = m_HealthOk = false;
+				}
+			}
+
+			if (!Success)
+			{
+				for (size_t Index : KeyIndex)
+				{
+					OnComplete({.CacheKey = CacheKeys[Index], .KeyIndex = Index, .Record = CbObjectView(), .Package = CbPackage()});
+				}
+
+				return {};
+			}
+
+			for (size_t LocalIndex = 0; CbFieldView Record : BatchResponse.GetObject()["Result"sv])
+			{
+				const size_t Index = IndexMap[LocalIndex++];
+				OnComplete({.CacheKey = CacheKeys[Index], .KeyIndex = Index, .Record = Record.AsObjectView(), .Package = BatchResponse});
+			}
+
+			return {};
+		}
+
+		virtual GetUpstreamCacheResult GetCachePayload(const CacheKey& CacheKey, const IoHash& PayloadId) override
 		{
 			try
 			{
 				ZenStructuredCacheSession Session(*m_Client);
 				const ZenCacheResult	  Result =
-					Session.GetCachePayload(PayloadKey.CacheKey.Bucket, PayloadKey.CacheKey.Hash, PayloadKey.PayloadId);
+					Session.GetCachePayload(CacheKey.Bucket, CacheKey.Hash, PayloadId);
 
 				if (Result.ErrorCode == 0)
 				{
@@ -474,6 +607,91 @@ namespace detail {
 			}
 		}
 
+		virtual GetUpstreamCacheBatchResult GetCachePayloads(std::span<CacheChunkRequest> CacheChunkRequests,
+															 std::span<size_t>			  RequestIndex,
+															 OnCachePayloadGetComplete&&  OnComplete) override final
+		{
+			std::vector<size_t> IndexMap;
+			IndexMap.reserve(RequestIndex.size());
+
+			CbObjectWriter BatchRequest;
+			BatchRequest << "Method"sv
+						 << "GetCachePayloads";
+
+			BatchRequest.BeginObject("Params"sv);
+			{
+				BatchRequest.BeginArray("ChunkRequests"sv);
+				{
+					for (size_t Index : RequestIndex)
+					{
+						const CacheChunkRequest& Request = CacheChunkRequests[Index];
+						IndexMap.push_back(Index);
+
+						BatchRequest.BeginObject();
+						{
+							BatchRequest.BeginObject("Key"sv);
+							BatchRequest << "Bucket"sv << Request.Key.Bucket;
+							BatchRequest << "Hash"sv << Request.Key.Hash;
+							BatchRequest.EndObject();
+
+							BatchRequest.AddObjectId("PayloadId"sv, Request.PayloadId);
+							BatchRequest << "ChunkId"sv << Request.ChunkId;
+							BatchRequest << "RawOffset"sv << Request.RawOffset;
+							BatchRequest << "RawSize"sv << Request.RawSize;
+							BatchRequest << "Policy"sv << static_cast<uint32_t>(Request.Policy);
+						}
+						BatchRequest.EndObject();
+					}
+				}
+				BatchRequest.EndArray();
+			}
+			BatchRequest.EndObject();
+
+			CbPackage BatchResponse;
+			bool	  Success = false;
+
+			{
+				ZenStructuredCacheSession Session(*m_Client);
+				ZenCacheResult			  Result = Session.InvokeRpc(BatchRequest.Save());
+				if (Result.Success)
+				{
+					Success = BatchResponse.TryLoad(Result.Response);
+				}
+				else if (Result.ErrorCode)
+				{
+					m_HealthOk = false;
+				}
+			}
+
+			if (!Success)
+			{
+				for (size_t Index : RequestIndex)
+				{
+					OnComplete({.Request = CacheChunkRequests[Index], .RequestIndex = Index, .Payload = IoBuffer()});
+				}
+
+				return {};
+			}
+
+			for (int32_t LocalIndex = 0; CbFieldView AttachmentHash : BatchResponse.GetObject()["Result"sv])
+			{
+				const size_t Index = IndexMap[LocalIndex++];
+				IoBuffer	 Payload;
+
+				if (const CbAttachment* Attachment = BatchResponse.FindAttachment(AttachmentHash.AsHash()))
+				{
+					if (const CompressedBuffer& Compressed = Attachment->AsCompressedBinary())
+					{
+						Payload = Compressed.GetCompressed().Flatten().AsIoBuffer();
+					}
+				}
+
+				OnComplete({.Request = CacheChunkRequests[Index], .RequestIndex = Index, .Payload = std::move(Payload)});
+			}
+
+			return {};
+		}
+
 		virtual PutUpstreamCacheResult PutCacheRecord(const UpstreamCacheRecord& CacheRecord,
 													  IoBuffer					 RecordValue,
 													  std::span<IoBuffer const>	 Payloads) override
@@ -758,7 +976,7 @@ public:
 
 	virtual void RegisterEndpoint(std::unique_ptr<UpstreamEndpoint> Endpoint) override { m_Endpoints.emplace_back(std::move(Endpoint)); }
 
-	virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
+	virtual GetUpstreamCacheResult GetCacheRecord(CacheKey CacheKey, ZenContentType Type) override
 	{
 		if (m_Options.ReadUpstream)
 		{
@@ -780,94 +998,82 @@ public:
 		return {};
 	}
 
-	virtual GetUpstreamCacheBatchResult GetCacheRecords(std::span<CacheKey> CacheKeys,
-														std::span<size_t>	KeyIndex,
-														OnCacheGetComplete	OnComplete) override final
+	virtual void GetCacheRecords(std::span<CacheKey>		CacheKeys,
+								 std::span<size_t>			KeyIndex,
+								 const CacheRecordPolicy&	Policy,
+								 OnCacheRecordGetComplete&& OnComplete) override final
 	{
-		if (!m_Options.ReadUpstream)
-		{
-			return {.Missing = std::vector<size_t>(KeyIndex.begin(), KeyIndex.end())};
-		}
+		std::vector<size_t> MissingKeys(KeyIndex.begin(), KeyIndex.end());
 
-		GetUpstreamCacheBatchResult Result;
-
-		for (size_t Idx : KeyIndex)
+		if (m_Options.ReadUpstream)
 		{
-			const UpstreamCacheKey CacheKey = {CacheKeys[Idx].Bucket, CacheKeys[Idx].Hash};
-
-			GetUpstreamCacheResult CacheResult;
 			for (auto& Endpoint : m_Endpoints)
 			{
-				if (Endpoint->IsHealthy())
+				if (Endpoint->IsHealthy() && !MissingKeys.empty())
 				{
-					CacheResult = Endpoint->GetCacheRecord(CacheKey, ZenContentType::kCbPackage);
-					m_Stats.Add(m_Log, *Endpoint, CacheResult, m_Endpoints);
+					std::vector<size_t> Missing;
 
-					if (CacheResult.Success)
-					{
-						break;
-					}
-				}
-			}
+					auto EndpointResult =
+						Endpoint->GetCacheRecords(CacheKeys, MissingKeys, Policy, [&](CacheRecordGetCompleteParams&& Params) {
+							if (Params.Record)
+							{
+								OnComplete(std::forward<CacheRecordGetCompleteParams>(Params));
+							}
+							else
+							{
+								Missing.push_back(Params.KeyIndex);
+							}
+						});
 
-			if (CacheResult.Success)
-			{
-				OnComplete(Idx, CacheResult.Value);
-			}
-			else
-			{
-				Result.Missing.push_back(Idx);
+					MissingKeys = std::move(Missing);
+				}
 			}
 		}
 
-		return Result;
-	}
-
-	virtual GetUpstreamCacheBatchResult GetCachePayloads(std::span<CacheChunkRequest> CacheChunkRequests,
-														 std::span<size_t>			  ChunkIndex,
-														 OnCacheGetComplete			  OnComplete) override final
-	{
-		if (!m_Options.ReadUpstream)
+		for (size_t Index : MissingKeys)
 		{
-			return {.Missing = std::vector<size_t>(ChunkIndex.begin(), ChunkIndex.end())};
+			OnComplete({.CacheKey = CacheKeys[Index], .KeyIndex = Index, .Record = CbObjectView(), .Package = CbPackage()});
 		}
+	}
 
-		GetUpstreamCacheBatchResult Result;
+	virtual void GetCachePayloads(std::span<CacheChunkRequest> CacheChunkRequests,
+								  std::span<size_t>			   RequestIndex,
+								  OnCachePayloadGetComplete&&  OnComplete) override final
+	{
+		std::vector<size_t> MissingPayloads(RequestIndex.begin(), RequestIndex.end());
 
-		for (size_t Idx : ChunkIndex)
+		if (m_Options.ReadUpstream)
 		{
-			const CacheChunkRequest& Chunk = CacheChunkRequests[Idx];
-			UpstreamPayloadKey		 PayloadKey{{Chunk.Key.Bucket, Chunk.Key.Hash}, Chunk.ChunkId};
-
-			GetUpstreamCacheResult CacheResult;
 			for (auto& Endpoint : m_Endpoints)
 			{
-				if (Endpoint->IsHealthy())
+				if (Endpoint->IsHealthy() && !MissingPayloads.empty())
 				{
-					CacheResult = Endpoint->GetCachePayload(PayloadKey);
-					m_Stats.Add(m_Log, *Endpoint, CacheResult, m_Endpoints);
+					std::vector<size_t> Missing;
 
-					if (CacheResult.Success)
-					{
-						break;
-					}
-				}
-			}
+					auto EndpointResult =
+						Endpoint->GetCachePayloads(CacheChunkRequests, MissingPayloads, [&](CachePayloadGetCompleteParams&& Params) {
+							if (Params.Payload)
+							{
+								OnComplete(std::forward<CachePayloadGetCompleteParams>(Params));
+							}
+							else
+							{
+								Missing.push_back(Params.RequestIndex);
+							}
+						});
 
-			if (CacheResult.Success)
-			{
-				OnComplete(Idx, CacheResult.Value);
-			}
-			else
-			{
-				Result.Missing.push_back(Idx);
+					MissingPayloads = std::move(Missing);
+				}
 			}
 		}
 
-		return Result;
+		for (size_t Index : MissingPayloads)
+		{
+			OnComplete({.Request = CacheChunkRequests[Index], .RequestIndex = Index, .Payload = IoBuffer()});
+		}
 	}
 
-	virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
+	virtual GetUpstreamCacheResult GetCachePayload(const CacheKey& CacheKey, const IoHash& PayloadId) override
 	{
 		if (m_Options.ReadUpstream)
 		{
@@ -875,7 +1081,7 @@ public:
 			{
 				if (Endpoint->IsHealthy())
 				{
-					const GetUpstreamCacheResult Result = Endpoint->GetCachePayload(PayloadKey);
+					const GetUpstreamCacheResult Result = Endpoint->GetCachePayload(CacheKey, PayloadId);
 					m_Stats.Add(m_Log, *Endpoint, Result, m_Endpoints);
 
 					if (Result.Success)