Handle cache record policy.

4 files changed, 87 insertions, 21 deletions

diff --git a/zenserver/cache/cachekey.cpp b/zenserver/cache/cachekey.cpp
index 57993d424..85d57745d 100644
--- a/zenserver/cache/cachekey.cpp
+++ b/zenserver/cache/cachekey.cpp
@@ -2,6 +2,7 @@
 
 #include "cachekey.h"
 
+#include <zencore/compactbinary.h>
 #include <zencore/string.h>
 
 namespace zen {
@@ -100,6 +101,44 @@ ParseSkipCachePolicy(std::string_view SkipPolicy, CachePolicy Default)
 	return Result;
 }
 
+CachePolicy
+CacheRecordPolicy::GetPayloadPolicy(const Oid& PayloadId) const
+{
+	if (const auto It = m_PayloadPolicies.find(PayloadId); It != m_PayloadPolicies.end())
+	{
+		return It->second;
+	}
+
+	return m_DefaultPayloadPolicy;
+}
+
+bool
+CacheRecordPolicy::FromCompactBinary(CbObjectView RecordPolicyObject, CacheRecordPolicy& OutRecordPolicy)
+{
+	using namespace std::literals;
+
+	const uint32_t RecordPolicy = RecordPolicyObject["recordpolicy"sv].AsUInt32(static_cast<uint32_t>(CachePolicy::Default));
+	const uint32_t DefaultPayloadPolicy =
+		RecordPolicyObject["defaultpayloadpolicy"sv].AsUInt32(static_cast<uint32_t>(CachePolicy::Default));
+
+	OutRecordPolicy.m_RecordPolicy		   = static_cast<CachePolicy>(RecordPolicy);
+	OutRecordPolicy.m_DefaultPayloadPolicy = static_cast<CachePolicy>(DefaultPayloadPolicy);
+
+	for (CbFieldView PayloadPolicyView : RecordPolicyObject["payloadpolicies"sv])
+	{
+		CbObjectView   PayloadPolicyObject = PayloadPolicyView.AsObjectView();
+		const Oid	   PayloadId		   = PayloadPolicyObject["payloadid"sv].AsObjectId();
+		const uint32_t PayloadPolicy	   = PayloadPolicyObject["policy"sv].AsUInt32();
+
+		if (PayloadId != Oid::Zero && PayloadPolicy != 0)
+		{
+			OutRecordPolicy.m_PayloadPolicies.emplace(PayloadId, static_cast<CachePolicy>(PayloadPolicy));
+		}
+	}
+
+	return true;
+}
+
 const CacheKey CacheKey::Empty = CacheKey{.Bucket = std::string(), .Hash = IoHash()};
 
 }  // namespace zen
diff --git a/zenserver/cache/cachekey.h b/zenserver/cache/cachekey.h
index 5b67b0261..6ce5d3aab 100644
--- a/zenserver/cache/cachekey.h
+++ b/zenserver/cache/cachekey.h
@@ -7,8 +7,12 @@
 #include <zencore/uid.h>
 #include <gsl/gsl-lite.hpp>
 
+#include <unordered_map>
+
 namespace zen {
 
+class CbObjectView;
+
 enum class CachePolicy : uint8_t
 {
 	None			= 0,
@@ -37,6 +41,25 @@ CachePolicy ParseStoreCachePolicy(std::string_view StorePolicy, CachePolicy Defa
 
 CachePolicy ParseSkipCachePolicy(std::string_view SkipPolicy, CachePolicy Default = CachePolicy::None);
 
+class CacheRecordPolicy
+{
+public:
+	CacheRecordPolicy() = default;
+
+	CachePolicy GetRecordPolicy() const { return m_RecordPolicy; }
+	CachePolicy GetPayloadPolicy(const Oid& PayloadId) const;
+	CachePolicy GetDefaultPayloadPolicy() const { return m_DefaultPayloadPolicy; }
+
+	static bool FromCompactBinary(CbObjectView RecordPolicyObject, CacheRecordPolicy& OutRecordPolicy);
+
+private:
+	using PayloadPolicyMap = std::unordered_map<Oid, CachePolicy, Oid::Hasher>;
+
+	CachePolicy		 m_RecordPolicy			= CachePolicy::Default;
+	CachePolicy		 m_DefaultPayloadPolicy = CachePolicy::Default;
+	PayloadPolicyMap m_PayloadPolicies;
+};
+
 struct CacheKey
 {
 	std::string Bucket;
diff --git a/zenserver/cache/structuredcache.cpp b/zenserver/cache/structuredcache.cpp
index decad2f04..073192c12 100644
--- a/zenserver/cache/structuredcache.cpp
+++ b/zenserver/cache/structuredcache.cpp
@@ -114,9 +114,7 @@ HttpStructuredCacheService::HandleRequest(HttpServerRequest& Request)
 
 		if (Key == "$batch")
 		{
-			const auto	QueryParams = Request.GetQueryParams();
-			CachePolicy Policy		= ParseCachePolicy(QueryParams);
-			return HandleBatchRequest(Request, Policy);
+			return HandleBatchRequest(Request);
 		}
 
 		if (std::all_of(begin(Key), end(Key), [](const char c) { return std::isalnum(c); }))
@@ -782,10 +780,8 @@ HttpStructuredCacheService::ValidateKeyUri(HttpServerRequest& Request, CacheRef&
 }
 
 void
-HttpStructuredCacheService::HandleBatchRequest(zen::HttpServerRequest& Request, CachePolicy Policy)
+HttpStructuredCacheService::HandleBatchRequest(zen::HttpServerRequest& Request)
 {
-	using namespace fmt::literals;
-
 	switch (auto Verb = Request.RequestVerb())
 	{
 		using enum HttpVerb;
@@ -805,11 +801,11 @@ HttpStructuredCacheService::HandleBatchRequest(zen::HttpServerRequest& Request,
 
 				if (Method == "get-cache-records"sv)
 				{
-					HandleBatchGetCacheRecords(Request, BatchRequest, Policy);
+					HandleBatchGetCacheRecords(Request, BatchRequest);
 				}
 				else if (Method == "get-cache-chunks"sv)
 				{
-					HandleBatchGetCachePayloads(Request, BatchRequest, Policy);
+					HandleBatchGetCachePayloads(Request, BatchRequest);
 				}
 				else
 				{
@@ -824,19 +820,21 @@ HttpStructuredCacheService::HandleBatchRequest(zen::HttpServerRequest& Request,
 }
 
 void
-HttpStructuredCacheService::HandleBatchGetCacheRecords(zen::HttpServerRequest& Request, CbObjectView BatchRequest, CachePolicy Policy)
+HttpStructuredCacheService::HandleBatchGetCacheRecords(zen::HttpServerRequest& Request, CbObjectView BatchRequest)
 {
 	using namespace fmt::literals;
 
-	const bool SkipData		   = (Policy & CachePolicy::SkipData) == CachePolicy::SkipData;
-	const bool SkipAttachments = (Policy & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
-	const bool QueryUpstream   = m_UpstreamCache && ((Policy & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);
-
 	const std::string_view Method = BatchRequest["method"sv].AsString();
 	ZEN_ASSERT(Method == "get-cache-records"sv);
 
 	CbObjectView Params = BatchRequest["params"sv].AsObjectView();
 
+	CacheRecordPolicy Policy;
+	CacheRecordPolicy::FromCompactBinary(Params["policy"sv].AsObjectView(), Policy);
+
+	const bool SkipAttachments = (Policy.GetRecordPolicy() & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
+	const bool QueryUpstream   = m_UpstreamCache && ((Policy.GetRecordPolicy() & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);
+
 	std::vector<CacheKey> CacheKeys;
 	std::vector<IoBuffer> CacheValues;
 	std::vector<IoBuffer> Payloads;
@@ -855,6 +853,11 @@ HttpStructuredCacheService::HandleBatchGetCacheRecords(zen::HttpServerRequest& R
 
 	CacheValues.resize(CacheKeys.size());
 
+	if (!SkipAttachments)
+	{
+		Payloads.reserve(CacheKeys.size());
+	}
+
 	for (size_t Idx = 0; const CacheKey& Key : CacheKeys)
 	{
 		ZenCacheValue CacheValue;
@@ -971,7 +974,7 @@ HttpStructuredCacheService::HandleBatchGetCacheRecords(zen::HttpServerRequest& R
 }
 
 void
-HttpStructuredCacheService::HandleBatchGetCachePayloads(zen::HttpServerRequest& Request, CbObjectView BatchRequest, CachePolicy Policy)
+HttpStructuredCacheService::HandleBatchGetCachePayloads(zen::HttpServerRequest& Request, CbObjectView BatchRequest)
 {
 	using namespace fmt::literals;
 
@@ -999,15 +1002,16 @@ HttpStructuredCacheService::HandleBatchGetCachePayloads(zen::HttpServerRequest&
 		return IoHash::Zero;
 	};
 
-	const bool SkipData		   = (Policy & CachePolicy::SkipData) == CachePolicy::SkipData;
-	const bool SkipAttachments = (Policy & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
-	const bool QueryUpstream   = m_UpstreamCache && ((Policy & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);
-
 	const std::string_view Method = BatchRequest["method"sv].AsString();
 	ZEN_ASSERT(Method == "get-cache-chunks"sv);
 
 	CbObjectView Params = BatchRequest["params"sv].AsObjectView();
 
+	CacheRecordPolicy Policy;
+	CacheRecordPolicy::FromCompactBinary(Params["policy"sv].AsObjectView(), Policy);
+
+	const bool QueryUpstream = m_UpstreamCache && ((Policy.GetRecordPolicy() & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);
+
 	std::vector<CacheChunkRequest> ChunkRequests;
 	std::vector<size_t>			   Missing;
 
diff --git a/zenserver/cache/structuredcache.h b/zenserver/cache/structuredcache.h
index 1ff4f28c9..9efcc05fa 100644
--- a/zenserver/cache/structuredcache.h
+++ b/zenserver/cache/structuredcache.h
@@ -89,9 +89,9 @@ private:
 	void			   HandleCachePayloadRequest(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy);
 	void			   HandleGetCachePayload(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy);
 	void			   HandlePutCachePayload(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy);
-	void			   HandleBatchRequest(zen::HttpServerRequest& Request, CachePolicy Policy);
-	void			   HandleBatchGetCacheRecords(zen::HttpServerRequest& Request, CbObjectView BatchRequest, CachePolicy Policy);
-	void			   HandleBatchGetCachePayloads(zen::HttpServerRequest& Request, CbObjectView BatchRequest, CachePolicy Policy);
+	void			   HandleBatchRequest(zen::HttpServerRequest& Request);
+	void			   HandleBatchGetCacheRecords(zen::HttpServerRequest& Request, CbObjectView BatchRequest);
+	void			   HandleBatchGetCachePayloads(zen::HttpServerRequest& Request, CbObjectView BatchRequest);
 	void			   HandleCacheBucketRequest(zen::HttpServerRequest& Request, std::string_view Bucket);
 	virtual void	   HandleStatsRequest(zen::HttpServerRequest& Request) override;
 	virtual void	   HandleStatusRequest(zen::HttpServerRequest& Request) override;