path: root/src/zenserver/hub/hub.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "hub.h"

#include "storageserverinstance.h"

#include <zencore/assertfmt.h>
#include <zencore/compactbinary.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/timer.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <EASTL/fixed_vector.h>
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

#if ZEN_WITH_TESTS
#	include <zencore/filesystem.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#	include <zencore/workthreadpool.h>
#	include <zenhttp/httpclient.h>
#endif

#include <numeric>

namespace zen {

///////////////////////////////////////////////////////////////////////////

/**
 * A timeline of events with sequence IDs and timestamps. Used to
 * track significant events for broadcasting to listeners.
 */
class EventTimeline
{
public:
	EventTimeline() { m_Events.reserve(1024); }

	~EventTimeline() {}

	EventTimeline(const EventTimeline&) = delete;
	EventTimeline& operator=(const EventTimeline&) = delete;

	void RecordEvent(std::string_view EventTag, CbObject EventMetadata)
	{
		const uint64_t			   SequenceId = m_NextEventId++;
		const auto				   Now		  = std::chrono::steady_clock::now();
		RwLock::ExclusiveLockScope _(m_Lock);
		m_Events.emplace_back(SequenceId, EventTag, Now, std::move(EventMetadata));
	}

	struct EventRecord
	{
		uint64_t							  SequenceId;
		std::string							  Tag;
		std::chrono::steady_clock::time_point Timestamp;
		CbObject							  EventMetadata;

		EventRecord(uint64_t							  InSequenceId,
					std::string_view					  InTag,
					std::chrono::steady_clock::time_point InTimestamp,
					CbObject							  InEventMetadata = CbObject())
		: SequenceId(InSequenceId)
		, Tag(InTag)
		, Timestamp(InTimestamp)
		, EventMetadata(InEventMetadata)
		{
		}
	};

	/**
	 * Iterate over events that have a SequenceId greater than SinceEventId
	 *
	 * @param Callback A callable that takes a const EventRecord&
	 * @param SinceEventId The SequenceId to compare against
	 */
	void IterateEventsSince(auto&& Callback, uint64_t SinceEventId)
	{
		// Hold the lock for as short a time as possible
		eastl::fixed_vector<EventRecord, 128> EventsToProcess;
		m_Lock.WithSharedLock([&] {
			for (auto& Event : m_Events)
			{
				if (Event.SequenceId > SinceEventId)
				{
					EventsToProcess.push_back(Event);
				}
			}
		});

		// Now invoke the callback outside the lock
		for (auto& Event : EventsToProcess)
		{
			Callback(Event);
		}
	}

	/**
	 * Trim events up to (and including) the given SequenceId. Intended
	 * to be used for cleaning up events which are not longer interesting.
	 *
	 * @param UpToEventId The SequenceId up to which events should be removed
	 */
	void TrimEventsUpTo(uint64_t UpToEventId)
	{
		RwLock::ExclusiveLockScope _(m_Lock);
		auto					   It = std::remove_if(m_Events.begin(), m_Events.end(), [UpToEventId](const EventRecord& Event) {
			  return Event.SequenceId <= UpToEventId;
		  });
		m_Events.erase(It, m_Events.end());
	}

private:
	std::atomic<uint64_t> m_NextEventId{0};

	RwLock					 m_Lock;
	std::vector<EventRecord> m_Events;
};

//////////////////////////////////////////////////////////////////////////

Hub::Hub(const Configuration&		   Config,
		 ZenServerEnvironment&&		   RunEnvironment,
		 ProvisionModuleCallbackFunc&& ProvisionedModuleCallback,
		 ProvisionModuleCallbackFunc&& DeprovisionedModuleCallback)
: m_Config(Config)
, m_RunEnvironment(std::move(RunEnvironment))
, m_ProvisionedModuleCallback(std::move(ProvisionedModuleCallback))
, m_DeprovisionedModuleCallback(std::move(DeprovisionedModuleCallback))
{
	m_HostMetrics					  = GetSystemMetrics();
	m_ResourceLimits.DiskUsageBytes	  = 1000ull * 1024 * 1024 * 1024;
	m_ResourceLimits.MemoryUsageBytes = 16ull * 1024 * 1024 * 1024;

	if (m_Config.HydrationTargetSpecification.empty())
	{
		std::filesystem::path FileHydrationPath = m_RunEnvironment.CreateChildDir("hydration_storage");
		ZEN_INFO("using file hydration path: '{}'", FileHydrationPath);
		m_HydrationTargetSpecification = fmt::format("file://{}", WideToUtf8(FileHydrationPath.native()));
	}
	else
	{
		m_HydrationTargetSpecification = m_Config.HydrationTargetSpecification;
	}

	m_HydrationTempPath = m_RunEnvironment.CreateChildDir("hydration_temp");
	ZEN_INFO("using hydration temp path: '{}'", m_HydrationTempPath);

	ZEN_ASSERT(uint64_t(Config.BasePortNumber) + Config.InstanceLimit <= std::numeric_limits<uint16_t>::max());

	m_InstanceLookup.reserve(Config.InstanceLimit);
	m_ActiveInstances.reserve(Config.InstanceLimit);

	m_FreePorts.resize(Config.InstanceLimit);
	std::iota(m_FreePorts.begin(), m_FreePorts.end(), Config.BasePortNumber);

#if ZEN_PLATFORM_WINDOWS
	if (m_Config.UseJobObject)
	{
		m_JobObject.Initialize();
		if (m_JobObject.IsValid())
		{
			ZEN_INFO("Job object initialized for hub service child process management");
		}
		else
		{
			ZEN_WARN("Failed to initialize job object; child processes will not be auto-terminated on hub crash");
		}
	}
#endif
	m_WatchDog = std::thread([this]() { WatchDog(); });
}

Hub::~Hub()
{
	try
	{
		ZEN_INFO("Hub service shutting down, deprovisioning any current instances");

		m_WatchDogEvent.Set();
		if (m_WatchDog.joinable())
		{
			m_WatchDog.join();
		}

		m_WatchDog = {};

		// WatchDog has been joined; no concurrent access is possible
		m_Lock.WithExclusiveLock([this] {
			for (auto& [ModuleId, ActiveInstanceIndex] : m_InstanceLookup)
			{
				std::unique_ptr<StorageServerInstance>& InstanceRaw = m_ActiveInstances[ActiveInstanceIndex];
				{
					StorageServerInstance::ExclusiveLockedPtr Instance(InstanceRaw->LockExclusive(/*Wait*/ true));

					uint16_t	BasePort = InstanceRaw->GetBasePort();
					std::string BaseUri;  // TODO?

					if (m_DeprovisionedModuleCallback)
					{
						try
						{
							m_DeprovisionedModuleCallback(ModuleId, HubProvisionedInstanceInfo{.BaseUri = BaseUri, .Port = BasePort});
						}
						catch (const std::exception& Ex)
						{
							ZEN_ERROR("Deprovision callback for module {} failed. Reason: '{}'", ModuleId, Ex.what());
						}
					}
					Instance.Deprovision();
				}
				InstanceRaw.reset();
			}
			m_InstanceLookup.clear();
			m_ActiveInstances.clear();
			m_FreeActiveInstanceIndexes.clear();
		});
	}
	catch (const std::exception& e)
	{
		ZEN_WARN("Exception during hub service shutdown: {}", e.what());
	}
}

bool
Hub::Provision(std::string_view ModuleId, HubProvisionedInstanceInfo& OutInfo, std::string& OutReason)
{
	StorageServerInstance::ExclusiveLockedPtr Instance;
	bool									  IsNewInstance = false;
	uint16_t								  AllocatedPort = 0;
	{
		RwLock::ExclusiveLockScope _(m_Lock);
		auto					   RestoreAllocatedPort = MakeGuard([this, ModuleId, &IsNewInstance, &AllocatedPort]() {
			  if (IsNewInstance && AllocatedPort != 0 && !m_InstanceLookup.contains(std::string(ModuleId)))
			  {
				  m_FreePorts.push_back(AllocatedPort);
				  AllocatedPort = 0;
			  }
		  });

		if (auto It = m_InstanceLookup.find(std::string(ModuleId)); It == m_InstanceLookup.end())
		{
			std::string Reason;
			if (!CanProvisionInstance(ModuleId, /* out */ Reason))
			{
				ZEN_WARN("Cannot provision new storage server instance for module '{}': {}", ModuleId, Reason);

				OutReason = Reason;

				return false;
			}

			IsNewInstance = true;

			AllocatedPort = m_FreePorts.front();
			ZEN_ASSERT(AllocatedPort != 0);
			m_FreePorts.pop_front();

			auto NewInstance = std::make_unique<StorageServerInstance>(
				m_RunEnvironment,
				StorageServerInstance::Configuration{.BasePort					   = AllocatedPort,
													 .HydrationTempPath			   = m_HydrationTempPath,
													 .HydrationTargetSpecification = m_HydrationTargetSpecification,
													 .HttpThreadCount			   = m_Config.InstanceHttpThreadCount,
													 .CoreLimit					   = m_Config.InstanceCoreLimit,
													 .ConfigPath				   = m_Config.InstanceConfigPath},
				ModuleId);

#if ZEN_PLATFORM_WINDOWS
			if (m_JobObject.IsValid())
			{
				NewInstance->SetJobObject(&m_JobObject);
			}
#endif

			Instance = NewInstance->LockExclusive(/*Wait*/ true);

			size_t ActiveInstanceIndex = (size_t)-1;
			if (!m_FreeActiveInstanceIndexes.empty())
			{
				ActiveInstanceIndex = m_FreeActiveInstanceIndexes.back();
				m_FreeActiveInstanceIndexes.pop_back();
				ZEN_ASSERT(m_ActiveInstances.size() > ActiveInstanceIndex);
				m_ActiveInstances[ActiveInstanceIndex] = std::move(NewInstance);
			}
			else
			{
				ActiveInstanceIndex = m_ActiveInstances.size();
				m_ActiveInstances.emplace_back(std::move(NewInstance));
			}
			ZEN_ASSERT(ActiveInstanceIndex != (size_t)-1);
			m_InstanceLookup.insert_or_assign(std::string(ModuleId), ActiveInstanceIndex);

			ZEN_INFO("Created new storage server instance for module '{}'", ModuleId);

			const int CurrentInstanceCount = gsl::narrow_cast<int>(m_InstanceLookup.size());
			int		  CurrentMaxCount	   = m_MaxInstanceCount.load();
			const int NewMax			   = Max(CurrentMaxCount, CurrentInstanceCount);

			m_MaxInstanceCount.compare_exchange_weak(CurrentMaxCount, NewMax);
		}
		else
		{
			const size_t ActiveInstanceIndex = It->second;
			ZEN_ASSERT(m_ActiveInstances.size() > ActiveInstanceIndex);

			std::unique_ptr<StorageServerInstance>& InstanceRaw = m_ActiveInstances[ActiveInstanceIndex];
			Instance											= InstanceRaw->LockExclusive(/*Wait*/ true);
			AllocatedPort										= InstanceRaw->GetBasePort();
		}

		m_ProvisioningModules.emplace(std::string(ModuleId));
	}

	ZEN_ASSERT(Instance);

	auto RemoveProvisioningModule = MakeGuard([&] {
		RwLock::ExclusiveLockScope _(m_Lock);
		m_ProvisioningModules.erase(std::string(ModuleId));
		if (IsNewInstance && AllocatedPort != 0 && !m_InstanceLookup.contains(std::string(ModuleId)))
		{
			m_FreePorts.push_back(AllocatedPort);
			AllocatedPort = 0;
		}
	});

	// NOTE: this is done while not holding the hub lock, as provisioning may take time
	// and we don't want to block other operations. We track which modules are being
	// provisioned using m_ProvisioningModules, and reject attempts to provision/deprovision
	// those modules while in this state.

	try
	{
		Instance.Provision();
		Instance = {};
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("Failed to provision storage server instance for module '{}': {}", ModuleId, Ex.what());
		Instance  = {};
		OutReason = Ex.what();

		if (IsNewInstance)
		{
			// Clean up failed instance provisioning
			std::unique_ptr<StorageServerInstance> DestroyInstance;
			{
				RwLock::ExclusiveLockScope _(m_Lock);
				if (auto It = m_InstanceLookup.find(std::string(ModuleId)); It != m_InstanceLookup.end())
				{
					const size_t ActiveInstanceIndex = It->second;
					ZEN_ASSERT(ActiveInstanceIndex < m_ActiveInstances.size());
					DestroyInstance = std::move(m_ActiveInstances[ActiveInstanceIndex]);
					ZEN_ASSERT(DestroyInstance);
					ZEN_ASSERT(!m_ActiveInstances[ActiveInstanceIndex]);
					m_FreeActiveInstanceIndexes.push_back(ActiveInstanceIndex);
					m_InstanceLookup.erase(It);
				}
			}
			try
			{
				DestroyInstance.reset();
			}
			catch (const std::exception& Ex)
			{
				ZEN_ERROR("Failed to destroy instance for failed provision module '{}': {}", ModuleId, Ex.what());
			}
		}
		return false;
	}

	OutReason.clear();
	OutInfo.Port = AllocatedPort;
	// TODO: base URI? Would need to know what host name / IP to use

	if (m_ProvisionedModuleCallback)
	{
		try
		{
			m_ProvisionedModuleCallback(ModuleId, OutInfo);
		}
		catch (const std::exception& Ex)
		{
			ZEN_ERROR("Provision callback for module {} failed. Reason: '{}'", ModuleId, Ex.what());
		}
	}

	return true;
}

bool
Hub::Deprovision(const std::string& ModuleId, std::string& OutReason)
{
	std::unique_ptr<StorageServerInstance>	  RawInstance;
	StorageServerInstance::ExclusiveLockedPtr Instance;

	{
		RwLock::ExclusiveLockScope _(m_Lock);

		if (m_ProvisioningModules.contains(ModuleId))
		{
			OutReason = fmt::format("Module '{}' is currently being provisioned", ModuleId);

			ZEN_WARN("Attempted to deprovision module '{}' which is currently being provisioned", ModuleId);

			return false;
		}

		if (auto It = m_InstanceLookup.find(ModuleId); It == m_InstanceLookup.end())
		{
			ZEN_WARN("Attempted to deprovision non-existent module '{}'", ModuleId);

			OutReason.clear();	// empty = not found (-> 404)
			return false;
		}
		else
		{
			const size_t ActiveInstanceIndex = It->second;
			ZEN_ASSERT(ActiveInstanceIndex < m_ActiveInstances.size());
			RawInstance = std::move(m_ActiveInstances[ActiveInstanceIndex]);
			ZEN_ASSERT(RawInstance != nullptr);
			m_FreeActiveInstanceIndexes.push_back(ActiveInstanceIndex);
			m_InstanceLookup.erase(It);
			m_DeprovisioningModules.emplace(ModuleId);

			Instance = RawInstance->LockExclusive(/*Wait*/ true);
		}
	}

	ZEN_ASSERT(RawInstance);
	ZEN_ASSERT(Instance);

	uint16_t	BasePort = RawInstance->GetBasePort();
	std::string BaseUri;  // TODO?

	if (m_DeprovisionedModuleCallback)
	{
		try
		{
			m_DeprovisionedModuleCallback(ModuleId, HubProvisionedInstanceInfo{.BaseUri = BaseUri, .Port = BasePort});
		}
		catch (const std::exception& Ex)
		{
			ZEN_ERROR("Deprovision callback for module {} failed. Reason: '{}'", ModuleId, Ex.what());
		}
	}

	// The module is deprovisioned outside the hub lock to avoid blocking other operations.
	//
	// To ensure that no new provisioning can occur while we're deprovisioning,
	// we add the module ID to m_DeprovisioningModules and remove it once
	// deprovisioning is complete.

	auto _ = MakeGuard([&] {
		{
			RwLock::ExclusiveLockScope _(m_Lock);
			m_DeprovisioningModules.erase(ModuleId);
			m_FreePorts.push_back(BasePort);
		}
	});

	Instance.Deprovision();

	Instance = {};
	OutReason.clear();

	return true;
}

bool
Hub::Hibernate(const std::string& ModuleId, std::string& OutReason)
{
	StorageServerInstance::ExclusiveLockedPtr Instance;

	{
		RwLock::ExclusiveLockScope _(m_Lock);

		if (m_ProvisioningModules.contains(ModuleId) || m_DeprovisioningModules.contains(ModuleId) ||
			m_HibernatingModules.contains(ModuleId) || m_WakingModules.contains(ModuleId))
		{
			OutReason = fmt::format("Module '{}' is currently changing state", ModuleId);
			return false;
		}

		auto It = m_InstanceLookup.find(ModuleId);
		if (It == m_InstanceLookup.end())
		{
			OutReason.clear();	// empty = not found (-> 404)
			return false;
		}

		const size_t ActiveInstanceIndex = It->second;
		ZEN_ASSERT(ActiveInstanceIndex < m_ActiveInstances.size());
		Instance = m_ActiveInstances[ActiveInstanceIndex]->LockExclusive(/*Wait*/ true);
		m_HibernatingModules.emplace(ModuleId);
	}

	ZEN_ASSERT(Instance);

	auto RemoveHibernatingModule = MakeGuard([&] {
		RwLock::ExclusiveLockScope _(m_Lock);
		m_HibernatingModules.erase(ModuleId);
	});

	// NOTE: done while not holding the hub lock, as hibernation may take time.
	// m_HibernatingModules tracks which modules are being hibernated, blocking
	// concurrent Hibernate/Wake/Provision/Deprovision attempts on the same module.
	try
	{
		if (!Instance.Hibernate())
		{
			OutReason = fmt::format("Module '{}' cannot be hibernated from state '{}'", ModuleId, ToString(Instance.GetState()));
			return false;
		}
		Instance = {};
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("Failed to hibernate storage server instance for module '{}': {}", ModuleId, Ex.what());
		Instance  = {};
		OutReason = Ex.what();
		return false;
	}

	OutReason.clear();
	return true;
}

bool
Hub::Wake(const std::string& ModuleId, std::string& OutReason)
{
	StorageServerInstance::ExclusiveLockedPtr Instance;

	{
		RwLock::ExclusiveLockScope _(m_Lock);

		if (m_ProvisioningModules.contains(ModuleId) || m_DeprovisioningModules.contains(ModuleId) ||
			m_HibernatingModules.contains(ModuleId) || m_WakingModules.contains(ModuleId))
		{
			OutReason = fmt::format("Module '{}' is currently changing state", ModuleId);
			return false;
		}

		auto It = m_InstanceLookup.find(ModuleId);
		if (It == m_InstanceLookup.end())
		{
			OutReason.clear();	// empty = not found (-> 404)
			return false;
		}

		const size_t ActiveInstanceIndex = It->second;
		ZEN_ASSERT(ActiveInstanceIndex < m_ActiveInstances.size());
		Instance = m_ActiveInstances[ActiveInstanceIndex]->LockExclusive(/*Wait*/ true);
		m_WakingModules.emplace(ModuleId);
	}

	ZEN_ASSERT(Instance);

	auto RemoveWakingModule = MakeGuard([&] {
		RwLock::ExclusiveLockScope _(m_Lock);
		m_WakingModules.erase(ModuleId);
	});

	// NOTE: done while not holding the hub lock, as waking may take time.
	// m_WakingModules tracks which modules are being woken, blocking
	// concurrent Hibernate/Wake/Provision/Deprovision attempts on the same module.
	try
	{
		if (!Instance.Wake())
		{
			OutReason = fmt::format("Module '{}' cannot be woken from state '{}'", ModuleId, ToString(Instance.GetState()));
			return false;
		}
		Instance = {};
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("Failed to wake storage server instance for module '{}': {}", ModuleId, Ex.what());
		Instance  = {};
		OutReason = Ex.what();
		return false;
	}

	OutReason.clear();
	return true;
}

bool
Hub::Find(std::string_view ModuleId, InstanceInfo* OutInstanceInfo)
{
	RwLock::SharedLockScope _(m_Lock);
	if (auto It = m_InstanceLookup.find(std::string(ModuleId)); It != m_InstanceLookup.end())
	{
		if (OutInstanceInfo)
		{
			const size_t ActiveInstanceIndex = It->second;
			ZEN_ASSERT(ActiveInstanceIndex < m_ActiveInstances.size());
			const std::unique_ptr<StorageServerInstance>& Instance = m_ActiveInstances[ActiveInstanceIndex];
			ZEN_ASSERT(Instance);
			InstanceInfo Info{
				Instance->GetState(),
				std::chrono::system_clock::now()  // TODO
			};
			Instance->GetProcessMetrics(Info.Metrics);

			*OutInstanceInfo = Info;
		}
		return true;
	}
	return false;
}

void
Hub::EnumerateModules(std::function<void(std::string_view ModuleId, const InstanceInfo&)> Callback)
{
	std::vector<std::pair<std::string, InstanceInfo>> Infos;
	{
		RwLock::SharedLockScope _(m_Lock);
		for (auto& [ModuleId, ActiveInstanceIndex] : m_InstanceLookup)
		{
			const std::unique_ptr<StorageServerInstance>& Instance = m_ActiveInstances[ActiveInstanceIndex];
			ZEN_ASSERT(Instance);
			InstanceInfo Info{
				Instance->GetState(),
				std::chrono::system_clock::now()  // TODO
			};
			Instance->GetProcessMetrics(Info.Metrics);

			Infos.push_back(std::make_pair(std::string(Instance->GetModuleId()), Info));
		}
	}

	for (const std::pair<std::string, InstanceInfo>& Info : Infos)
	{
		Callback(Info.first, Info.second);
	}
}

int
Hub::GetInstanceCount()
{
	return m_Lock.WithSharedLock([this]() { return gsl::narrow_cast<int>(m_InstanceLookup.size()); });
}

void
Hub::UpdateCapacityMetrics()
{
	m_HostMetrics = GetSystemMetrics();

	// TODO: Should probably go into WatchDog and use atomic for update so it can be read without locks...
	// Per-instance stats are already refreshed by WatchDog and are readable via the Find and EnumerateModules
}

void
Hub::UpdateStats()
{
	int CurrentInstanceCount = m_Lock.WithSharedLock([this] { return gsl::narrow_cast<int>(m_InstanceLookup.size()); });
	int CurrentMaxCount		 = m_MaxInstanceCount.load();

	int NewMax = Max(CurrentMaxCount, CurrentInstanceCount);

	m_MaxInstanceCount.compare_exchange_weak(CurrentMaxCount, NewMax);
}

bool
Hub::CanProvisionInstance(std::string_view ModuleId, std::string& OutReason)
{
	if (m_DeprovisioningModules.contains(std::string(ModuleId)))
	{
		OutReason = fmt::format("module '{}' is currently being deprovisioned", ModuleId);

		return false;
	}

	if (m_ProvisioningModules.contains(std::string(ModuleId)))
	{
		OutReason = fmt::format("module '{}' is currently being provisioned", ModuleId);

		return false;
	}

	if (gsl::narrow_cast<int>(m_InstanceLookup.size()) >= m_Config.InstanceLimit)
	{
		OutReason = fmt::format("instance limit ({}) exceeded", m_Config.InstanceLimit);

		return false;
	}

	// Since deprovisioning happens outside the lock and we don't return the port until the instance is fully shut down, we might be below
	// the instance count limit but with no free ports available
	if (m_FreePorts.empty())
	{
		OutReason = fmt::format("no free ports available, deprovisioning of instances might be in flight ({})",
								m_Config.InstanceLimit - m_InstanceLookup.size());

		return false;
	}

	// TODO: handle additional resource metrics

	return true;
}

void
Hub::WatchDog()
{
	constexpr uint64_t WatchDogWakeupTimeMs		= 5000;
	constexpr uint64_t WatchDogProcessingTimeMs = 500;

	size_t CheckInstanceIndex = 0;
	while (!m_WatchDogEvent.Wait(WatchDogWakeupTimeMs))
	{
		try
		{
			size_t MaxCheckCount = m_Lock.WithSharedLock([this]() { return m_InstanceLookup.size(); });

			Stopwatch Timer;
			while (MaxCheckCount-- > 0 && Timer.GetElapsedTimeMs() < WatchDogProcessingTimeMs && !m_WatchDogEvent.Wait(5))
			{
				StorageServerInstance::SharedLockedPtr LockedInstance;
				m_Lock.WithSharedLock([this, &CheckInstanceIndex, &LockedInstance]() {
					if (m_InstanceLookup.empty())
					{
						return;
					}

					size_t				   MaxLoopCount = m_ActiveInstances.size();
					StorageServerInstance* Instance		= nullptr;
					while (MaxLoopCount-- > 0 && !Instance)
					{
						CheckInstanceIndex++;
						if (CheckInstanceIndex >= m_ActiveInstances.size())
						{
							CheckInstanceIndex = 0;
						}
						Instance = (CheckInstanceIndex < m_ActiveInstances.size()) ? m_ActiveInstances[CheckInstanceIndex].get() : nullptr;
					}

					if (Instance)
					{
						LockedInstance = Instance->LockShared(/*Wait*/ false);
					}
				});

				if (LockedInstance)
				{
					if (LockedInstance.IsRunning())
					{
						LockedInstance.UpdateMetrics();
					}
					else if (LockedInstance.GetState() == HubInstanceState::Provisioned)
					{
						// Process is not running but state says it should be - instance died unexpectedly.
						// TODO: Track and attempt recovery.
					}
					// else: transitional state (Provisioning, Deprovisioning, Hibernating, Waking) - expected, skip.
					LockedInstance = {};
				}
			}
		}
		catch (const std::exception& Ex)
		{
			// TODO: Catch specific errors such as asserts, OOM, OOD, system_error etc
			ZEN_ERROR("Hub watchdog threw exception: {}", Ex.what());
		}
	}
}

#if ZEN_WITH_TESTS

TEST_SUITE_BEGIN("server.hub");

static const HttpClientSettings kFastTimeout{.ConnectTimeout = std::chrono::milliseconds(200)};

namespace hub_testutils {

	ZenServerEnvironment MakeHubEnvironment(const std::filesystem::path& BaseDir)
	{
		return ZenServerEnvironment(ZenServerEnvironment::Hub, GetRunningExecutablePath().parent_path(), BaseDir);
	}

	std::unique_ptr<Hub> MakeHub(const std::filesystem::path&	  BaseDir,
								 Hub::Configuration				  Config			  = {},
								 Hub::ProvisionModuleCallbackFunc ProvisionCallback	  = {},
								 Hub::ProvisionModuleCallbackFunc DeprovisionCallback = {})
	{
		return std::make_unique<Hub>(Config, MakeHubEnvironment(BaseDir), std::move(ProvisionCallback), std::move(DeprovisionCallback));
	}

}  // namespace hub_testutils

TEST_CASE("hub.provision_basic")
{
	ScopedTemporaryDirectory TempDir;
	std::unique_ptr<Hub>	 HubInstance = hub_testutils::MakeHub(TempDir.Path());

	CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	CHECK_FALSE(HubInstance->Find("module_a"));

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;
	const bool				   ProvisionResult = HubInstance->Provision("module_a", Info, Reason);
	REQUIRE_MESSAGE(ProvisionResult, Reason);
	CHECK(Reason.empty());
	CHECK_NE(Info.Port, 0);
	CHECK_EQ(HubInstance->GetInstanceCount(), 1);
	Hub::InstanceInfo InstanceInfo;
	REQUIRE(HubInstance->Find("module_a", &InstanceInfo));
	CHECK_EQ(InstanceInfo.State, HubInstanceState::Provisioned);

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(ModClient.Get("/health/"));
	}

	const bool DeprovisionResult = HubInstance->Deprovision("module_a", Reason);
	CHECK(DeprovisionResult);
	CHECK(Reason.empty());
	CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	CHECK_FALSE(HubInstance->Find("module_a"));

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(!ModClient.Get("/health/"));
	}
}

TEST_CASE("hub.provision_config")
{
	ScopedTemporaryDirectory TempDir;
	CreateDirectories(TempDir.Path() / "hub");

	std::string LuaConfig =
		"server = {\n"
		"	buildstore = {\n"
		"		enabled = true,\n"
		"	}\n"
		"}\n";

	WriteFile(TempDir.Path() / "config.lua", IoBuffer(IoBuffer::Wrap, LuaConfig.data(), LuaConfig.length()));

	std::unique_ptr<Hub> HubInstance =
		hub_testutils::MakeHub(TempDir.Path() / "hub", Hub::Configuration{.InstanceConfigPath = TempDir.Path() / "config.lua"});

	CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	CHECK_FALSE(HubInstance->Find("module_a"));

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;
	const bool				   ProvisionResult = HubInstance->Provision("module_a", Info, Reason);
	REQUIRE_MESSAGE(ProvisionResult, Reason);
	CHECK_NE(Info.Port, 0);
	CHECK_EQ(HubInstance->GetInstanceCount(), 1);
	Hub::InstanceInfo InstanceInfo;
	REQUIRE(HubInstance->Find("module_a", &InstanceInfo));
	CHECK_EQ(InstanceInfo.State, HubInstanceState::Provisioned);

	HttpClient			 Client(fmt::format("http://127.0.0.1:{}{}", Info.Port, Info.BaseUri));
	HttpClient::Response TestResponse = Client.Get("/status/builds");
	CHECK(TestResponse.IsSuccess());
	CHECK(TestResponse.AsObject()["ok"].AsBool());

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(ModClient.Get("/health/"));
	}

	const bool DeprovisionResult = HubInstance->Deprovision("module_a", Reason);
	CHECK(DeprovisionResult);
	CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	CHECK_FALSE(HubInstance->Find("module_a"));

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(!ModClient.Get("/health/"));
	}
}

TEST_CASE("hub.provision_callbacks")
{
	ScopedTemporaryDirectory TempDir;

	struct CallbackRecord
	{
		std::string ModuleId;
		uint16_t	Port;
	};
	RwLock						CallbackMutex;
	std::vector<CallbackRecord> ProvisionRecords;
	std::vector<CallbackRecord> DeprovisionRecords;

	auto ProvisionCb = [&](std::string_view ModuleId, const HubProvisionedInstanceInfo& Info) {
		CallbackMutex.WithExclusiveLock([&]() { ProvisionRecords.push_back({std::string(ModuleId), Info.Port}); });
	};
	auto DeprovisionCb = [&](std::string_view ModuleId, const HubProvisionedInstanceInfo& Info) {
		CallbackMutex.WithExclusiveLock([&]() { DeprovisionRecords.push_back({std::string(ModuleId), Info.Port}); });
	};

	std::unique_ptr<Hub> HubInstance = hub_testutils::MakeHub(TempDir.Path(), {}, std::move(ProvisionCb), std::move(DeprovisionCb));

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;

	const bool ProvisionResult = HubInstance->Provision("cb_module", Info, Reason);
	REQUIRE_MESSAGE(ProvisionResult, Reason);

	{
		RwLock::SharedLockScope _(CallbackMutex);
		REQUIRE_EQ(ProvisionRecords.size(), 1u);
		CHECK_EQ(ProvisionRecords[0].ModuleId, "cb_module");
		CHECK_EQ(ProvisionRecords[0].Port, Info.Port);
		CHECK_NE(ProvisionRecords[0].Port, 0);
	}

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(ModClient.Get("/health/"));
	}

	const bool DeprovisionResult = HubInstance->Deprovision("cb_module", Reason);
	CHECK(DeprovisionResult);

	{
		HttpClient ModClient(fmt::format("http://localhost:{}", Info.Port), kFastTimeout);
		CHECK(!ModClient.Get("/health/"));
	}

	{
		RwLock::SharedLockScope _(CallbackMutex);
		REQUIRE_EQ(DeprovisionRecords.size(), 1u);
		CHECK_EQ(DeprovisionRecords[0].ModuleId, "cb_module");
		CHECK_NE(DeprovisionRecords[0].Port, 0);
		CHECK_EQ(ProvisionRecords.size(), 1u);
	}
}

TEST_CASE("hub.instance_limit")
{
	ScopedTemporaryDirectory TempDir;
	Hub::Configuration		 Config;
	Config.InstanceLimit  = 2;
	Config.BasePortNumber = 21500;

	std::unique_ptr<Hub> HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config);

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;

	const bool FirstResult = HubInstance->Provision("limit_a", Info, Reason);
	REQUIRE_MESSAGE(FirstResult, Reason);

	const bool SecondResult = HubInstance->Provision("limit_b", Info, Reason);
	REQUIRE_MESSAGE(SecondResult, Reason);
	CHECK_EQ(HubInstance->GetInstanceCount(), 2);

	Reason.clear();
	const bool ThirdResult = HubInstance->Provision("limit_c", Info, Reason);
	CHECK_FALSE(ThirdResult);
	CHECK_EQ(HubInstance->GetInstanceCount(), 2);
	CHECK_NE(Reason.find("instance limit"), std::string::npos);

	HubInstance->Deprovision("limit_a", Reason);
	CHECK_EQ(HubInstance->GetInstanceCount(), 1);

	Reason.clear();
	const bool FourthResult = HubInstance->Provision("limit_d", Info, Reason);
	CHECK_MESSAGE(FourthResult, Reason);
	CHECK_EQ(HubInstance->GetInstanceCount(), 2);
}

TEST_CASE("hub.enumerate_modules")
{
	ScopedTemporaryDirectory TempDir;
	std::unique_ptr<Hub>	 HubInstance = hub_testutils::MakeHub(TempDir.Path());

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;

	REQUIRE_MESSAGE(HubInstance->Provision("enum_a", Info, Reason), Reason);
	REQUIRE_MESSAGE(HubInstance->Provision("enum_b", Info, Reason), Reason);

	std::vector<std::string> Ids;
	int						 ProvisionedCount = 0;
	HubInstance->EnumerateModules([&](std::string_view ModuleId, const Hub::InstanceInfo& InstanceInfo) {
		Ids.push_back(std::string(ModuleId));
		if (InstanceInfo.State == HubInstanceState::Provisioned)
		{
			ProvisionedCount++;
		}
	});
	CHECK_EQ(Ids.size(), 2u);
	CHECK_EQ(ProvisionedCount, 2);
	const bool FoundA = std::find(Ids.begin(), Ids.end(), "enum_a") != Ids.end();
	const bool FoundB = std::find(Ids.begin(), Ids.end(), "enum_b") != Ids.end();
	CHECK(FoundA);
	CHECK(FoundB);

	HubInstance->Deprovision("enum_a", Reason);
	Ids.clear();
	ProvisionedCount = 0;
	HubInstance->EnumerateModules([&](std::string_view ModuleId, const Hub::InstanceInfo& InstanceInfo) {
		Ids.push_back(std::string(ModuleId));
		if (InstanceInfo.State == HubInstanceState::Provisioned)
		{
			ProvisionedCount++;
		}
	});
	REQUIRE_EQ(Ids.size(), 1u);
	CHECK_EQ(Ids[0], "enum_b");
	CHECK_EQ(ProvisionedCount, 1);
}

TEST_CASE("hub.max_instance_count")
{
	ScopedTemporaryDirectory TempDir;
	std::unique_ptr<Hub>	 HubInstance = hub_testutils::MakeHub(TempDir.Path());

	CHECK_EQ(HubInstance->GetMaxInstanceCount(), 0);

	HubProvisionedInstanceInfo Info;
	std::string				   Reason;

	REQUIRE_MESSAGE(HubInstance->Provision("max_a", Info, Reason), Reason);
	CHECK_GE(HubInstance->GetMaxInstanceCount(), 1);

	REQUIRE_MESSAGE(HubInstance->Provision("max_b", Info, Reason), Reason);
	CHECK_GE(HubInstance->GetMaxInstanceCount(), 2);

	const int MaxAfterTwo = HubInstance->GetMaxInstanceCount();

	HubInstance->Deprovision("max_a", Reason);
	CHECK_EQ(HubInstance->GetInstanceCount(), 1);
	CHECK_EQ(HubInstance->GetMaxInstanceCount(), MaxAfterTwo);
}

TEST_CASE("hub.concurrent_callbacks")
{
	ScopedTemporaryDirectory TempDir;
	Hub::Configuration		 Config;
	Config.BasePortNumber = 22300;
	Config.InstanceLimit  = 10;

	struct CallbackRecord
	{
		std::string ModuleId;
		uint16_t	Port;
	};
	RwLock						CallbackMutex;
	std::vector<CallbackRecord> ProvisionCallbacks;
	std::vector<CallbackRecord> DeprovisionCallbacks;

	auto ProvisionCb = [&](std::string_view ModuleId, const HubProvisionedInstanceInfo& Info) {
		CallbackMutex.WithExclusiveLock([&]() { ProvisionCallbacks.push_back({std::string(ModuleId), Info.Port}); });
	};
	auto DeprovisionCb = [&](std::string_view ModuleId, const HubProvisionedInstanceInfo& Info) {
		CallbackMutex.WithExclusiveLock([&]() { DeprovisionCallbacks.push_back({std::string(ModuleId), Info.Port}); });
	};

	std::unique_ptr<Hub> HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config, std::move(ProvisionCb), std::move(DeprovisionCb));

	constexpr int kHalf = 3;

	// Serially pre-provision kHalf modules and drain the resulting callbacks before the
	// concurrent phase so we start with a clean slate.
	for (int I = 0; I < kHalf; ++I)
	{
		HubProvisionedInstanceInfo Info;
		std::string				   Reason;
		REQUIRE_MESSAGE(HubInstance->Provision(fmt::format("pre_{}", I), Info, Reason), Reason);
	}
	CHECK_EQ(HubInstance->GetInstanceCount(), kHalf);

	{
		RwLock::ExclusiveLockScope _(CallbackMutex);
		REQUIRE_EQ(ProvisionCallbacks.size(), static_cast<size_t>(kHalf));
		ProvisionCallbacks.clear();
	}

	// Concurrently provision kHalf new modules while deprovisioning the pre-provisioned ones.
	std::vector<int>		 ProvisionResults(kHalf, 0);
	std::vector<std::string> ProvisionReasons(kHalf);
	std::vector<int>		 DeprovisionResults(kHalf, 0);

	{
		WorkerThreadPool Provisioners(kHalf, "hub_cbtest_provisioners");
		WorkerThreadPool Deprovisioneers(kHalf, "hub_cbtest_deprovisioneers");

		std::vector<std::future<void>> ProvisionFutures(kHalf);
		std::vector<std::future<void>> DeprovisionFutures(kHalf);

		for (int I = 0; I < kHalf; ++I)
		{
			ProvisionFutures[I] = Provisioners.EnqueueTask(std::packaged_task<void()>([&, I] {
															   HubProvisionedInstanceInfo Info;
															   std::string				  Reason;
															   const bool				  Result =
																   HubInstance->Provision(fmt::format("new_{}", I), Info, Reason);
															   ProvisionResults[I] = Result ? 1 : 0;
															   ProvisionReasons[I] = Reason;
														   }),
														   WorkerThreadPool::EMode::EnableBacklog);

			DeprovisionFutures[I] = Deprovisioneers.EnqueueTask(std::packaged_task<void()>([&, I] {
																	std::string Reason;
																	const bool	Result =
																		HubInstance->Deprovision(fmt::format("pre_{}", I), Reason);
																	DeprovisionResults[I] = Result ? 1 : 0;
																}),
																WorkerThreadPool::EMode::EnableBacklog);
		}

		for (std::future<void>& F : ProvisionFutures)
		{
			F.get();
		}
		for (std::future<void>& F : DeprovisionFutures)
		{
			F.get();
		}
	}

	// All operations must have succeeded
	for (int I = 0; I < kHalf; ++I)
	{
		CHECK_MESSAGE(ProvisionResults[I] != 0, ProvisionReasons[I]);
		CHECK(DeprovisionResults[I] != 0);
	}
	CHECK_EQ(HubInstance->GetInstanceCount(), kHalf);

	// Each new_* module must have triggered exactly one provision callback with a non-zero port.
	// Each pre_* module must have triggered exactly one deprovision callback with a non-zero port.
	{
		RwLock::SharedLockScope _(CallbackMutex);
		REQUIRE_EQ(ProvisionCallbacks.size(), static_cast<size_t>(kHalf));
		REQUIRE_EQ(DeprovisionCallbacks.size(), static_cast<size_t>(kHalf));

		for (const CallbackRecord& Record : ProvisionCallbacks)
		{
			CHECK_NE(Record.Port, 0);
			const bool IsNewModule = Record.ModuleId.rfind("new_", 0) == 0;
			CHECK_MESSAGE(IsNewModule, Record.ModuleId);
		}
		for (const CallbackRecord& Record : DeprovisionCallbacks)
		{
			CHECK_NE(Record.Port, 0);
			const bool IsPreModule = Record.ModuleId.rfind("pre_", 0) == 0;
			CHECK_MESSAGE(IsPreModule, Record.ModuleId);
		}
	}
}

#	if ZEN_PLATFORM_WINDOWS
TEST_CASE("hub.job_object")
{
	SUBCASE("UseJobObject=true")
	{
		ScopedTemporaryDirectory TempDir;
		Hub::Configuration		 Config;
		Config.UseJobObject	  = true;
		Config.BasePortNumber = 22100;

		std::unique_ptr<Hub>	   HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config);
		HubProvisionedInstanceInfo Info;
		std::string				   Reason;

		const bool ProvisionResult = HubInstance->Provision("jobobj_a", Info, Reason);
		REQUIRE_MESSAGE(ProvisionResult, Reason);
		CHECK_NE(Info.Port, 0);

		const bool DeprovisionResult = HubInstance->Deprovision("jobobj_a", Reason);
		CHECK(DeprovisionResult);
		CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	}

	SUBCASE("UseJobObject=false")
	{
		ScopedTemporaryDirectory TempDir;
		Hub::Configuration		 Config;
		Config.UseJobObject	  = false;
		Config.BasePortNumber = 22200;

		std::unique_ptr<Hub>	   HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config);
		HubProvisionedInstanceInfo Info;
		std::string				   Reason;

		const bool ProvisionResult = HubInstance->Provision("nojobobj_a", Info, Reason);
		REQUIRE_MESSAGE(ProvisionResult, Reason);
		CHECK_NE(Info.Port, 0);

		const bool DeprovisionResult = HubInstance->Deprovision("nojobobj_a", Reason);
		CHECK(DeprovisionResult);
		CHECK_EQ(HubInstance->GetInstanceCount(), 0);
	}
}
#	endif	// ZEN_PLATFORM_WINDOWS

TEST_CASE("hub.hibernate_wake")
{
	ScopedTemporaryDirectory TempDir;
	Hub::Configuration		 Config;
	Config.BasePortNumber			 = 22600;
	std::unique_ptr<Hub> HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config);

	HubProvisionedInstanceInfo ProvInfo;
	Hub::InstanceInfo		   Info;
	std::string				   Reason;

	// Provision
	REQUIRE_MESSAGE(HubInstance->Provision("hib_a", ProvInfo, Reason), Reason);
	CHECK(Reason.empty());
	REQUIRE(HubInstance->Find("hib_a", &Info));
	CHECK_EQ(Info.State, HubInstanceState::Provisioned);
	{
		HttpClient ModClient(fmt::format("http://localhost:{}", ProvInfo.Port), kFastTimeout);
		CHECK(ModClient.Get("/health/"));
	}

	// Hibernate
	const bool HibernateResult = HubInstance->Hibernate("hib_a", Reason);
	REQUIRE_MESSAGE(HibernateResult, Reason);
	CHECK(Reason.empty());
	REQUIRE(HubInstance->Find("hib_a", &Info));
	CHECK_EQ(Info.State, HubInstanceState::Hibernated);
	{
		HttpClient ModClient(fmt::format("http://localhost:{}", ProvInfo.Port), kFastTimeout);
		CHECK(!ModClient.Get("/health/"));
	}

	// Wake
	const bool WakeResult = HubInstance->Wake("hib_a", Reason);
	REQUIRE_MESSAGE(WakeResult, Reason);
	CHECK(Reason.empty());
	REQUIRE(HubInstance->Find("hib_a", &Info));
	CHECK_EQ(Info.State, HubInstanceState::Provisioned);
	{
		HttpClient ModClient(fmt::format("http://localhost:{}", ProvInfo.Port), kFastTimeout);
		CHECK(ModClient.Get("/health/"));
	}

	// Deprovision
	const bool DeprovisionResult = HubInstance->Deprovision("hib_a", Reason);
	CHECK(DeprovisionResult);
	CHECK(Reason.empty());
	CHECK_FALSE(HubInstance->Find("hib_a"));
	{
		HttpClient ModClient(fmt::format("http://localhost:{}", ProvInfo.Port), kFastTimeout);
		CHECK(!ModClient.Get("/health/"));
	}
}

TEST_CASE("hub.hibernate_wake_errors")
{
	ScopedTemporaryDirectory TempDir;
	Hub::Configuration		 Config;
	Config.BasePortNumber			 = 22700;
	std::unique_ptr<Hub> HubInstance = hub_testutils::MakeHub(TempDir.Path(), Config);

	HubProvisionedInstanceInfo ProvInfo;
	std::string				   Reason;

	// Hibernate/wake on a non-existent module - should return false with empty reason (-> 404)
	CHECK_FALSE(HubInstance->Hibernate("never_provisioned", Reason));
	CHECK(Reason.empty());

	CHECK_FALSE(HubInstance->Wake("never_provisioned", Reason));
	CHECK(Reason.empty());

	// Double-hibernate: first hibernate succeeds, second returns false with non-empty reason (-> 400)
	REQUIRE_MESSAGE(HubInstance->Provision("err_b", ProvInfo, Reason), Reason);
	CHECK(Reason.empty());
	REQUIRE_MESSAGE(HubInstance->Hibernate("err_b", Reason), Reason);
	CHECK(Reason.empty());

	Reason.clear();
	CHECK_FALSE(HubInstance->Hibernate("err_b", Reason));
	CHECK_FALSE(Reason.empty());

	// Wake on provisioned: succeeds (-> Provisioned), then wake again returns false (-> 400)
	REQUIRE_MESSAGE(HubInstance->Wake("err_b", Reason), Reason);
	CHECK(Reason.empty());

	Reason.clear();
	CHECK_FALSE(HubInstance->Wake("err_b", Reason));
	CHECK_FALSE(Reason.empty());

	// Deprovision not-found - should return false with empty reason (-> 404)
	CHECK_FALSE(HubInstance->Deprovision("never_provisioned", Reason));
	CHECK(Reason.empty());
}

TEST_SUITE_END();

void
hub_forcelink()
{
}

#endif	// ZEN_WITH_TESTS

}  // namespace zen