From 22d25f59c7ead3de0b5d335684242b7364bce8f1 Mon Sep 17 00:00:00 2001
From: Per Larsson <p3rl@protonmail.com>
Date: Sat, 2 Oct 2021 16:31:52 +0200
Subject: Added support for choosing best ZEN upstream endpoint based on
 latency.

---
 zenserver/upstream/upstreamcache.cpp | 100 ++++++++++++++++++++++++++++++++---
 1 file changed, 93 insertions(+), 7 deletions(-)

(limited to 'zenserver/upstream/upstreamcache.cpp')

diff --git a/zenserver/upstream/upstreamcache.cpp b/zenserver/upstream/upstreamcache.cpp
index b1966e299..a183e7ebf 100644
--- a/zenserver/upstream/upstreamcache.cpp
+++ b/zenserver/upstream/upstreamcache.cpp
@@ -111,6 +111,8 @@ namespace detail {
 
 		virtual ~JupiterUpstreamEndpoint() = default;
 
+		virtual UpstreamEndpointHealth Initialize() override { return CheckHealth(); }
+
 		virtual bool IsHealthy() const override { return m_HealthOk.load(); }
 
 		virtual UpstreamEndpointHealth CheckHealth() override
@@ -400,22 +402,70 @@ namespace detail {
 
 	class ZenUpstreamEndpoint final : public UpstreamEndpoint
 	{
+		struct ZenEndpoint
+		{
+			std::string Url;
+			std::string Reason;
+			double		Latency{};
+			bool		Ok = false;
+
+			bool operator<(const ZenEndpoint& RHS) const { return Ok && RHS.Ok ? Latency < RHS.Latency : Ok; }
+		};
+
 	public:
-		ZenUpstreamEndpoint(std::string_view ServiceUrl)
+		ZenUpstreamEndpoint(std::span<std::string const> Urls) : m_Log(zen::logging::Get("upstream")), m_DisplayName("ZEN")
 		{
-			using namespace fmt::literals;
-			m_DisplayName = "Zen - {}"_format(ServiceUrl);
-			m_Client	  = new ZenStructuredCacheClient(ServiceUrl);
+			for (const auto& Url : Urls)
+			{
+				m_Endpoints.push_back({.Url = Url});
+			}
 		}
 
 		~ZenUpstreamEndpoint() = default;
 
+		virtual UpstreamEndpointHealth Initialize() override
+		{
+			using namespace fmt::literals;
+
+			const ZenEndpoint& Ep = GetEndpoint();
+			if (Ep.Ok)
+			{
+				m_ServiceUrl  = Ep.Url;
+				m_DisplayName = "ZEN - {}"_format(m_ServiceUrl);
+				m_Client	  = new ZenStructuredCacheClient(m_ServiceUrl);
+
+				m_HealthOk = true;
+				return {.Ok = true};
+			}
+
+			m_HealthOk = false;
+			return {.Reason = Ep.Reason};
+		}
+
 		virtual bool IsHealthy() const override { return m_HealthOk; }
 
 		virtual UpstreamEndpointHealth CheckHealth() override
 		{
+			using namespace fmt::literals;
+
 			try
 			{
+				if (m_Client.IsNull())
+				{
+					const ZenEndpoint& Ep = GetEndpoint();
+					if (Ep.Ok)
+					{
+						m_ServiceUrl  = Ep.Url;
+						m_DisplayName = "ZEN - {}"_format(m_ServiceUrl);
+						m_Client	  = new ZenStructuredCacheClient(m_ServiceUrl);
+
+						m_HealthOk = true;
+						return {.Ok = true};
+					}
+
+					return {.Reason = Ep.Reason};
+				}
+
 				ZenStructuredCacheSession Session(*m_Client);
 				ZenCacheResult			  Result;
 
@@ -591,6 +641,42 @@ namespace detail {
 		virtual UpstreamEndpointStats& Stats() override { return m_Stats; }
 
 	private:
+		const ZenEndpoint& GetEndpoint()
+		{
+			for (ZenEndpoint& Ep : m_Endpoints)
+			{
+				ZenStructuredCacheClient  Client(Ep.Url);
+				ZenStructuredCacheSession Session(Client);
+				const int32_t			  SampleCount = 2;
+
+				Ep.Ok	   = false;
+				Ep.Latency = {};
+
+				for (int32_t Sample = 0; Sample < SampleCount; ++Sample)
+				{
+					ZenCacheResult Result = Session.CheckHealth();
+					Ep.Ok				  = Result.Success;
+					Ep.Reason			  = std::move(Result.Reason);
+					Ep.Latency += Result.ElapsedSeconds;
+				}
+				Ep.Latency /= double(SampleCount);
+			}
+
+			std::sort(std::begin(m_Endpoints), std::end(m_Endpoints));
+
+			for (const auto& Ep : m_Endpoints)
+			{
+				ZEN_INFO("ping ZEN endpoint '{}' latency '{:.3}s' {}", Ep.Url, Ep.Latency, Ep.Ok ? "OK" : Ep.Reason);
+			}
+
+			return m_Endpoints.front();
+		}
+
+		spdlog::logger& Log() { return m_Log; }
+
+		spdlog::logger&					 m_Log;
+		std::string						 m_ServiceUrl;
+		std::vector<ZenEndpoint>		 m_Endpoints;
 		std::string						 m_DisplayName;
 		RefPtr<ZenStructuredCacheClient> m_Client;
 		UpstreamEndpointStats			 m_Stats;
@@ -711,7 +797,7 @@ public:
 	{
 		for (auto& Endpoint : m_Endpoints)
 		{
-			const UpstreamEndpointHealth Health = Endpoint->CheckHealth();
+			const UpstreamEndpointHealth Health = Endpoint->Initialize();
 			if (Health.Ok)
 			{
 				ZEN_INFO("initialize endpoint '{}' OK", Endpoint->DisplayName());
@@ -993,9 +1079,9 @@ MakeJupiterUpstreamEndpoint(const CloudCacheClientOptions& Options)
 }
 
 std::unique_ptr<UpstreamEndpoint>
-MakeZenUpstreamEndpoint(std::string_view Url)
+MakeZenUpstreamEndpoint(std::span<std::string const> Urls)
 {
-	return std::make_unique<detail::ZenUpstreamEndpoint>(Url);
+	return std::make_unique<detail::ZenUpstreamEndpoint>(Urls);
 }
 
 }  // namespace zen
-- 
cgit v1.2.3