Initial support for cache policies.

3 files changed, 220 insertions, 19 deletions

diff --git a/zenserver-test/zenserver-test.cpp b/zenserver-test/zenserver-test.cpp
index 48bbacd80..8db056bc4 100644
--- a/zenserver-test/zenserver-test.cpp
+++ b/zenserver-test/zenserver-test.cpp
@@ -670,7 +670,7 @@ main(int argc, char** argv)
 	zen::logging::InitializeLogging();
 
 	spdlog::set_level(spdlog::level::debug);
-	spdlog::set_formatter(std::make_unique<::logging::full_formatter>("test", std::chrono::system_clock::now()));
+	spdlog::set_formatter(std::make_unique< ::logging::full_formatter>("test", std::chrono::system_clock::now()));
 
 	std::filesystem::path ProgramBaseDir = std::filesystem::path(argv[0]).parent_path();
 	std::filesystem::path TestBaseDir	 = ProgramBaseDir.parent_path().parent_path() / ".test";
@@ -1402,6 +1402,85 @@ TEST_CASE("zcache.cbpackage")
 	}
 }
 
+TEST_CASE("zcache.policy")
+{
+	using namespace std::literals;
+
+	struct ZenConfig
+	{
+		std::filesystem::path DataDir;
+		uint16_t			  Port;
+		std::string			  BaseUri;
+		std::string			  Args;
+
+		static ZenConfig New(uint16_t Port = 13337, std::string Args = "")
+		{
+			return ZenConfig{.DataDir = TestEnv.CreateNewTestDir(),
+							 .Port	  = Port,
+							 .BaseUri = "http://localhost:{}/z$"_format(Port),
+							 .Args	  = std::move(Args)};
+		}
+
+		static ZenConfig NewWithUpstream(uint16_t UpstreamPort)
+		{
+			return New(13337, "--upstream-thread-count=0 --upstream-zen-url=http://localhost:{}"_format(UpstreamPort));
+		}
+
+		void Spawn(ZenServerInstance& Inst)
+		{
+			Inst.SetTestDir(DataDir);
+			Inst.SpawnServer(Port, Args);
+			Inst.WaitUntilReady();
+		}
+	};
+
+	auto GenerateData = [](uint64_t Size, zen::IoHash& OutHash) -> zen::UniqueBuffer {
+		auto	 Buf  = zen::UniqueBuffer::Alloc(Size);
+		uint8_t* Data = reinterpret_cast<uint8_t*>(Buf.GetData());
+		for (uint64_t Idx = 0; Idx < Size; Idx++)
+		{
+			Data[Idx] = Idx % 256;
+		}
+		OutHash = zen::IoHash::HashBuffer(Data, Size);
+		return Buf;
+	};
+
+	SUBCASE("query - 'local' does not query upstream")
+	{
+		ZenConfig		  UpstreamCfg = ZenConfig::New(13338);
+		ZenServerInstance UpstreamInst(TestEnv);
+		ZenConfig		  LocalCfg = ZenConfig::NewWithUpstream(13338);
+		ZenServerInstance LocalInst(TestEnv);
+		const auto		  Bucket = "legacy"sv;
+
+		UpstreamCfg.Spawn(UpstreamInst);
+		LocalCfg.Spawn(LocalInst);
+
+		zen::IoHash Key;
+		auto		BinaryValue = GenerateData(1024, Key);
+
+		// Store binary cache value upstream
+		{
+			cpr::Response Result = cpr::Put(cpr::Url{"{}/{}/{}"_format(UpstreamCfg.BaseUri, Bucket, Key)},
+											cpr::Body{(const char*)BinaryValue.GetData(), BinaryValue.GetSize()},
+											cpr::Header{{"Content-Type", "application/octet-stream"}});
+			CHECK(Result.status_code == 201);
+		}
+
+		{
+			cpr::Response Result = cpr::Get(cpr::Url{"{}/{}/{}?query=local"_format(LocalCfg.BaseUri, Bucket, Key)},
+											cpr::Header{{"Accept", "application/octet-stream"}});
+			CHECK(Result.status_code == 404);
+		}
+
+		{
+			cpr::Response Result = cpr::Get(cpr::Url{"{}/{}/{}?query=local,remote"_format(LocalCfg.BaseUri, Bucket, Key)},
+											cpr::Header{{"Accept", "application/octet-stream"}});
+			CHECK(Result.status_code == 200);
+		}
+	}
+}
+
 struct RemoteExecutionRequest
 {
 	RemoteExecutionRequest(std::string_view Host, int Port, std::filesystem::path& TreePath)
@@ -1773,5 +1852,5 @@ TEST_CASE("lifetime.owner.2")
 }
 #	endif
 
-}
+}  // namespace zen::tests
 #endif
diff --git a/zenserver/cache/structuredcache.cpp b/zenserver/cache/structuredcache.cpp
index da2e8850e..c66b1f98d 100644
--- a/zenserver/cache/structuredcache.cpp
+++ b/zenserver/cache/structuredcache.cpp
@@ -25,12 +25,125 @@
 #include <queue>
 #include <thread>
 
+#include <gsl/gsl-lite.hpp>
+
 namespace zen {
 
 using namespace std::literals;
 
 //////////////////////////////////////////////////////////////////////////
 
+namespace detail { namespace cacheopt {
+	constexpr std::string_view Local	   = "local"sv;
+	constexpr std::string_view Remote	   = "remote"sv;
+	constexpr std::string_view Data		   = "data"sv;
+	constexpr std::string_view Meta		   = "meta"sv;
+	constexpr std::string_view Value	   = "value"sv;
+	constexpr std::string_view Attachments = "attachments"sv;
+}}	// namespace detail::cacheopt
+
+//////////////////////////////////////////////////////////////////////////
+
+enum class CachePolicy : uint8_t
+{
+	None			= 0,
+	QueryLocal		= 1 << 0,
+	QueryRemote		= 1 << 1,
+	Query			= QueryLocal | QueryRemote,
+	StoreLocal		= 1 << 2,
+	StoreRemote		= 1 << 3,
+	Store			= StoreLocal | StoreRemote,
+	SkipMeta		= 1 << 4,
+	SkipValue		= 1 << 5,
+	SkipAttachments = 1 << 6,
+	SkipData		= SkipMeta | SkipValue | SkipAttachments,
+	SkipLocalCopy	= 1 << 7,
+	Local			= QueryLocal | StoreLocal,
+	Remote			= QueryRemote | StoreRemote,
+	Default			= Query | Store,
+	Disable			= None,
+};
+
+gsl_DEFINE_ENUM_BITMASK_OPERATORS(CachePolicy);
+
+CachePolicy
+ParseCachePolicy(const zen::HttpServerRequest::QueryParams& QueryParams)
+{
+	CachePolicy QueryPolicy = CachePolicy::Query;
+
+	{
+		std::string_view Opts = QueryParams.GetValue("query"sv);
+		if (!Opts.empty())
+		{
+			QueryPolicy = CachePolicy::None;
+			zen::ForEachStrTok(Opts, ',', [&QueryPolicy](const std::string_view& Opt) {
+				if (Opt == detail::cacheopt::Local)
+				{
+					QueryPolicy |= CachePolicy::QueryLocal;
+				}
+				if (Opt == detail::cacheopt::Remote)
+				{
+					QueryPolicy |= CachePolicy::QueryRemote;
+				}
+				return true;
+			});
+		}
+	}
+
+	CachePolicy StorePolicy = CachePolicy::Store;
+
+	{
+		std::string_view Opts = QueryParams.GetValue("store"sv);
+		if (!Opts.empty())
+		{
+			StorePolicy = CachePolicy::None;
+			zen::ForEachStrTok(Opts, ',', [&StorePolicy](const std::string_view& Opt) {
+				if (Opt == detail::cacheopt::Local)
+				{
+					StorePolicy |= CachePolicy::StoreLocal;
+				}
+				if (Opt == detail::cacheopt::Remote)
+				{
+					StorePolicy |= CachePolicy::StoreRemote;
+				}
+				return true;
+			});
+		}
+	}
+
+	CachePolicy SkipPolicy = CachePolicy::None;
+
+	{
+		std::string_view Opts = QueryParams.GetValue("skip"sv);
+		if (!Opts.empty())
+		{
+			zen::ForEachStrTok(Opts, ',', [&SkipPolicy](const std::string_view& Opt) {
+				if (Opt == detail::cacheopt::Meta)
+				{
+					SkipPolicy |= CachePolicy::SkipMeta;
+				}
+				if (Opt == detail::cacheopt::Value)
+				{
+					SkipPolicy |= CachePolicy::SkipValue;
+				}
+				if (Opt == detail::cacheopt::Attachments)
+				{
+					SkipPolicy |= CachePolicy::SkipAttachments;
+				}
+				if (Opt == detail::cacheopt::Data)
+				{
+					SkipPolicy |= CachePolicy::SkipData;
+				}
+				return true;
+			});
+		}
+	}
+
+	return QueryPolicy | StorePolicy | SkipPolicy;
+}
+
+//////////////////////////////////////////////////////////////////////////
+
 HttpStructuredCacheService::HttpStructuredCacheService(::ZenCacheStore&				  InCacheStore,
 													   zen::CasStore&				  InStore,
 													   zen::CidStore&				  InCidStore,
@@ -78,13 +191,16 @@ HttpStructuredCacheService::HandleRequest(zen::HttpServerRequest& Request)
 		return Request.WriteResponse(zen::HttpResponseCode::BadRequest);  // invalid URL
 	}
 
+	const auto	QueryParams = Request.GetQueryParams();
+	CachePolicy Policy		= ParseCachePolicy(QueryParams);
+
 	if (Ref.PayloadId == IoHash::Zero)
 	{
-		return HandleCacheRecordRequest(Request, Ref);
+		return HandleCacheRecordRequest(Request, Ref, Policy);
 	}
 	else
 	{
-		return HandleCachePayloadRequest(Request, Ref);
+		return HandleCachePayloadRequest(Request, Ref, Policy);
 	}
 
 	return;
@@ -121,7 +237,7 @@ HttpStructuredCacheService::HandleCacheBucketRequest(zen::HttpServerRequest& Req
 }
 
 void
-HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Request, CacheRef& Ref)
+HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Request, CacheRef& Ref, CachePolicy Policy)
 {
 	switch (auto Verb = Request.RequestVerb())
 	{
@@ -136,7 +252,10 @@ HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Req
 				bool		  Success		  = m_CacheStore.Get(Ref.BucketSegment, Ref.HashKey, /* out */ Value);
 				bool		  InUpstreamCache = false;
 
-				if (!Success && m_UpstreamCache)
+				const bool QueryUpstream =
+					!Success && m_UpstreamCache && (zen::CachePolicy::QueryRemote == (Policy & zen::CachePolicy::QueryRemote));
+
+				if (QueryUpstream)
 				{
 					const ZenContentType CacheRecordType = Ref.BucketSegment == "legacy"sv			  ? ZenContentType::kBinary
 														   : AcceptType == ZenContentType::kCbPackage ? ZenContentType::kCbPackage
@@ -326,13 +445,15 @@ HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Req
 
 				const HttpContentType ContentType = Request.RequestContentType();
 
+				const bool StoreUpstream = m_UpstreamCache && (zen::CachePolicy::StoreRemote == (Policy & zen::CachePolicy::StoreRemote));
+
 				if (ContentType == HttpContentType::kBinary || ContentType == HttpContentType::kUnknownContentType)
 				{
 					// TODO: create a cache record and put value in CAS?
 					m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = Body});
 					ZEN_DEBUG("PUT - binary '{}/{}' {}", Ref.BucketSegment, Ref.HashKey, NiceBytes(Body.Size()));
 
-					if (m_UpstreamCache)
+					if (StoreUpstream)
 					{
 						auto Result = m_UpstreamCache->EnqueueUpstream(
 							{.Type = ZenContentType::kBinary, .CacheKey = {Ref.BucketSegment, Ref.HashKey}});
@@ -397,15 +518,12 @@ HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Req
 							  MissingRefs.size(),
 							  References.size());
 
-					if (MissingRefs.empty())
+					if (MissingRefs.empty() && StoreUpstream)
 					{
-						// Only enqueue valid cache records, i.e. all referenced payloads exists
-						if (m_UpstreamCache)
-						{
-							auto Result = m_UpstreamCache->EnqueueUpstream({.Type		= ZenContentType::kCbObject,
-																			.CacheKey	= {Ref.BucketSegment, Ref.HashKey},
-																			.PayloadIds = std::move(References)});
-						}
+						ZEN_ASSERT(m_UpstreamCache);
+						auto Result = m_UpstreamCache->EnqueueUpstream({.Type		= ZenContentType::kCbObject,
+																		.CacheKey	= {Ref.BucketSegment, Ref.HashKey},
+																		.PayloadIds = std::move(References)});
 
 						return Request.WriteResponse(zen::HttpResponseCode::Created);
 					}
@@ -502,8 +620,9 @@ HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Req
 					ZenCacheValue CacheValue{.Value = CacheRecordChunk};
 					m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, CacheValue);
 
-					if (m_UpstreamCache)
+					if (StoreUpstream)
 					{
+						ZEN_ASSERT(m_UpstreamCache);
 						auto Result = m_UpstreamCache->EnqueueUpstream({.Type		= ZenContentType::kCbPackage,
 																		.CacheKey	= {Ref.BucketSegment, Ref.HashKey},
 																		.PayloadIds = std::move(PayloadIds)});
@@ -536,7 +655,7 @@ HttpStructuredCacheService::HandleCacheRecordRequest(zen::HttpServerRequest& Req
 }
 
 void
-HttpStructuredCacheService::HandleCachePayloadRequest(zen::HttpServerRequest& Request, CacheRef& Ref)
+HttpStructuredCacheService::HandleCachePayloadRequest(zen::HttpServerRequest& Request, CacheRef& Ref, CachePolicy Policy)
 {
 	// Note: the URL references the uncompressed payload hash - so this maintains the mapping
 	// from uncompressed CAS identity (aka CID/Content ID) to the stored payload hash
@@ -544,6 +663,8 @@ HttpStructuredCacheService::HandleCachePayloadRequest(zen::HttpServerRequest& Re
 	// this is a PITA but a consequence of the fact that the client side code is not able to
 	// address data by compressed hash
 
+	ZEN_UNUSED(Policy);
+
 	switch (auto Verb = Request.RequestVerb())
 	{
 		using enum zen::HttpVerb;
diff --git a/zenserver/cache/structuredcache.h b/zenserver/cache/structuredcache.h
index 8289fd700..796b21d1f 100644
--- a/zenserver/cache/structuredcache.h
+++ b/zenserver/cache/structuredcache.h
@@ -17,6 +17,7 @@ namespace zen {
 class CasStore;
 class CidStore;
 class UpstreamCache;
+enum class CachePolicy : uint8_t;
 
 /**
  * Structured cache service. Imposes constraints on keys, supports blobs and
@@ -70,8 +71,8 @@ private:
 	};
 
 	[[nodiscard]] bool ValidateKeyUri(zen::HttpServerRequest& Request, CacheRef& OutRef);
-	void			   HandleCacheRecordRequest(zen::HttpServerRequest& Request, CacheRef& Ref);
-	void			   HandleCachePayloadRequest(zen::HttpServerRequest& Request, CacheRef& Ref);
+	void			   HandleCacheRecordRequest(zen::HttpServerRequest& Request, CacheRef& Ref, CachePolicy Policy);
+	void			   HandleCachePayloadRequest(zen::HttpServerRequest& Request, CacheRef& Ref, CachePolicy Policy);
 	void			   HandleCacheBucketRequest(zen::HttpServerRequest& Request, std::string_view Bucket);
 
 	spdlog::logger&				   Log() { return m_Log; }