add host discovery and zen cache support for oplog import (#601)

* add host discovery and zen cache support for oplog import

2 files changed, 162 insertions, 8 deletions

diff --git a/src/zenremotestore/builds/buildstoragecache.cpp b/src/zenremotestore/builds/buildstoragecache.cpp
index 2e9cf9954..d7442d7c5 100644
--- a/src/zenremotestore/builds/buildstoragecache.cpp
+++ b/src/zenremotestore/builds/buildstoragecache.cpp
@@ -359,19 +359,16 @@ public:
 		while (true)
 		{
 			intptr_t Remaining = m_PendingBackgroundWorkCount.Remaining();
-			if (UpdateCallback(Remaining))
+			if (!UpdateCallback(Remaining))
 			{
-				if (m_PendingBackgroundWorkCount.Wait(UpdateIntervalMS))
-				{
-					UpdateCallback(0);
-					return;
-				}
+				m_CancelBackgroundWork.store(true);
 			}
-			else
+			if (m_PendingBackgroundWorkCount.Wait(UpdateIntervalMS))
 			{
-				m_CancelBackgroundWork.store(true);
+				break;
 			}
 		}
+		UpdateCallback(0);
 	}
 
 private:
diff --git a/src/zenremotestore/builds/buildstorageutil.cpp b/src/zenremotestore/builds/buildstorageutil.cpp
new file mode 100644
index 000000000..998529714
--- /dev/null
+++ b/src/zenremotestore/builds/buildstorageutil.cpp
@@ -0,0 +1,157 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include <zenremotestore/builds/buildstorageutil.h>
+
+#include <zenremotestore/builds/buildstoragecache.h>
+#include <zenremotestore/builds/jupiterbuildstorage.h>
+#include <zenremotestore/jupiter/jupiterhost.h>
+#include <zenutil/zenserverprocess.h>
+
+namespace zen {
+
+BuildStorageResolveResult
+ResolveBuildStorage(const HttpClientSettings& ClientSettings,
+					std::string_view		  Host,
+					std::string_view		  OverrideHost,
+					std::string_view		  ZenCacheHost,
+					ZenCacheResolveMode		  ZenResolveMode)
+{
+	bool AllowZenCacheDiscovery = ZenResolveMode == ZenCacheResolveMode::Discovery || ZenResolveMode == ZenCacheResolveMode::All;
+	bool AllowLocalZenCache		= ZenResolveMode == ZenCacheResolveMode::LocalHost || ZenResolveMode == ZenCacheResolveMode::All;
+
+	auto GetHostNameFromUrl = [](std::string_view Url) -> std::string_view {
+		std::string::size_type HostnameStart  = 0;
+		std::string::size_type HostnameLength = std::string::npos;
+		if (auto StartPos = Url.find("//"); StartPos != std::string::npos)
+		{
+			HostnameStart = StartPos + 2;
+		}
+		if (auto EndPos = Url.find("/", HostnameStart); EndPos != std::string::npos)
+		{
+			HostnameLength = EndPos - HostnameStart;
+		}
+		if (auto EndPos = Url.find(":", HostnameStart); EndPos != std::string::npos)
+		{
+			HostnameLength = EndPos - HostnameStart;
+		}
+		return Url.substr(HostnameStart, HostnameLength);
+	};
+
+	std::string HostUrl;
+	std::string HostName;
+
+	std::string CacheUrl;
+	std::string CacheName;
+	bool		HostAssumeHttp2	 = ClientSettings.AssumeHttp2;
+	bool		CacheAssumeHttp2 = ClientSettings.AssumeHttp2;
+
+	JupiterServerDiscovery DiscoveryResponse;
+	const std::string_view DiscoveryHost = Host.empty() ? OverrideHost : Host;
+
+	if (OverrideHost.empty() || (ZenCacheHost.empty() && AllowZenCacheDiscovery))
+	{
+		DiscoveryResponse = DiscoverJupiterEndpoints(DiscoveryHost, ClientSettings);
+	}
+
+	if (!OverrideHost.empty())
+	{
+		if (JupiterEndpointTestResult TestResult = TestJupiterEndpoint(OverrideHost, HostAssumeHttp2); TestResult.Success)
+		{
+			HostUrl	 = OverrideHost;
+			HostName = GetHostNameFromUrl(OverrideHost);
+		}
+		else
+		{
+			throw std::runtime_error(fmt::format("Host {} could not be reached. Reason: {}", OverrideHost, TestResult.FailureReason));
+		}
+	}
+	else
+	{
+		if (DiscoveryResponse.ServerEndPoints.empty())
+		{
+			throw std::runtime_error(fmt::format("Failed to find any builds hosts at {}", DiscoveryHost));
+		}
+
+		for (const JupiterServerDiscovery::EndPoint& ServerEndpoint : DiscoveryResponse.ServerEndPoints)
+		{
+			if (!ServerEndpoint.BaseUrl.empty())
+			{
+				if (JupiterEndpointTestResult TestResult = TestJupiterEndpoint(ServerEndpoint.BaseUrl, ServerEndpoint.AssumeHttp2);
+					TestResult.Success)
+				{
+					HostUrl			= ServerEndpoint.BaseUrl;
+					HostAssumeHttp2 = ServerEndpoint.AssumeHttp2;
+					HostName		= ServerEndpoint.Name;
+					break;
+				}
+				else
+				{
+					ZEN_DEBUG("Unable to reach host {}. Reason: {}", ServerEndpoint.BaseUrl, TestResult.FailureReason);
+				}
+			}
+		}
+		if (HostUrl.empty())
+		{
+			throw std::runtime_error(fmt::format("Failed to find any usable builds hosts out of {} using {}",
+												 DiscoveryResponse.ServerEndPoints.size(),
+												 DiscoveryHost));
+		}
+	}
+	if (ZenCacheHost.empty())
+	{
+		if (AllowZenCacheDiscovery)
+		{
+			for (const JupiterServerDiscovery::EndPoint& CacheEndpoint : DiscoveryResponse.CacheEndPoints)
+			{
+				if (!CacheEndpoint.BaseUrl.empty())
+				{
+					if (ZenCacheEndpointTestResult TestResult = TestZenCacheEndpoint(CacheEndpoint.BaseUrl, CacheEndpoint.AssumeHttp2);
+						TestResult.Success)
+					{
+						CacheUrl		 = CacheEndpoint.BaseUrl;
+						CacheAssumeHttp2 = CacheEndpoint.AssumeHttp2;
+						CacheName		 = CacheEndpoint.Name;
+						break;
+					}
+				}
+			}
+		}
+		if (CacheUrl.empty() && AllowLocalZenCache)
+		{
+			ZenServerState State;
+			if (State.InitializeReadOnly())
+			{
+				State.Snapshot([&](const ZenServerState::ZenServerEntry& Entry) {
+					if (CacheUrl.empty())
+					{
+						std::string ZenServerLocalHostUrl = fmt::format("http://127.0.0.1:{}", Entry.EffectiveListenPort.load());
+						if (ZenCacheEndpointTestResult TestResult = TestZenCacheEndpoint(ZenServerLocalHostUrl, false); TestResult.Success)
+						{
+							CacheUrl		 = ZenServerLocalHostUrl;
+							CacheAssumeHttp2 = false;
+							CacheName		 = "localhost";
+						}
+					}
+				});
+			}
+		}
+	}
+	else if (ZenCacheEndpointTestResult TestResult = TestZenCacheEndpoint(ZenCacheHost, false); TestResult.Success)
+	{
+		CacheUrl  = ZenCacheHost;
+		CacheName = GetHostNameFromUrl(ZenCacheHost);
+	}
+	else
+	{
+		ZEN_WARN("Unable to reach cache host {}. Reason: {}", ZenCacheHost, TestResult.FailureReason);
+	}
+
+	return BuildStorageResolveResult{.HostUrl		  = HostUrl,
+									 .HostName		  = HostName,
+									 .HostAssumeHttp2 = HostAssumeHttp2,
+
+									 .CacheUrl		   = CacheUrl,
+									 .CacheName		   = CacheName,
+									 .CacheAssumeHttp2 = CacheAssumeHttp2};
+}
+}  // namespace zen