1 files changed, 229 insertions, 83 deletions

diff --git a/src/zenserver/hub/storageserverinstance.cpp b/src/zenserver/hub/storageserverinstance.cpp
index 8e71e7aca..0c9354990 100644
--- a/src/zenserver/hub/storageserverinstance.cpp
+++ b/src/zenserver/hub/storageserverinstance.cpp
@@ -28,6 +28,7 @@ void
 StorageServerInstance::SpawnServerProcess()
 {
 	ZEN_ASSERT_FORMAT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
+	m_ServerInstance.ResetDeadProcess();
 
 	m_ServerInstance.SetServerExecutablePath(GetRunningExecutablePath());
 	m_ServerInstance.SetDataDir(m_BaseDir);
@@ -56,160 +57,305 @@ StorageServerInstance::SpawnServerProcess()
 	m_ServerInstance.EnableShutdownOnDestroy();
 }
 
-void
-StorageServerInstance::Provision()
+ProcessMetrics
+StorageServerInstance::GetProcessMetrics() const
 {
-	RwLock::ExclusiveLockScope _(m_Lock);
+	ProcessMetrics Metrics;
+	if (m_ServerInstance.IsRunning())
+	{
+		zen::GetProcessMetrics(m_ServerInstance.GetProcessHandle(), Metrics);
+	}
+	return Metrics;
+}
 
-	if (m_IsProvisioned)
+void
+StorageServerInstance::ProvisionLocked()
+{
+	if (m_ServerInstance.IsRunning())
 	{
 		ZEN_WARN("Storage server instance for module '{}' is already provisioned", m_ModuleId);
-
 		return;
 	}
 
-	if (m_IsHibernated)
+	ZEN_INFO("Provisioning storage server instance for module '{}', at '{}'", m_ModuleId, m_BaseDir);
+	try
 	{
-		WakeLocked();
+		Hydrate();
+		SpawnServerProcess();
 	}
-	else
+	catch (const std::exception& Ex)
 	{
-		ZEN_INFO("Provisioning storage server instance for module '{}', at '{}'", m_ModuleId, m_BaseDir);
-
-		Hydrate();
+		ZEN_WARN("Failed spawning server instance for module '{}', at '{}' during provisioning. Reason: {}",
+				 m_ModuleId,
+				 m_BaseDir,
+				 Ex.what());
+		throw;
+	}
+}
 
-		SpawnServerProcess();
+void
+StorageServerInstance::DeprovisionLocked()
+{
+	if (m_ServerInstance.IsRunning())
+	{
+		// m_ServerInstance.Shutdown() never throws.
+		m_ServerInstance.Shutdown();
 	}
 
-	m_IsProvisioned = true;
+	// Crashed or Hibernated: process already dead; skip Shutdown.
+	// Dehydrate preserves instance state for future re-provisioning. Failure means saved state
+	// may be stale or absent, but the process is already dead so the slot can still be released.
+	// Swallow the exception and proceed with cleanup rather than leaving the module stuck.
+	try
+	{
+		Dehydrate();
+	}
+	catch (const std::exception& Ex)
+	{
+		ZEN_WARN("Dehydration of module {} failed during deprovisioning, current state not saved. Reason: {}", m_ModuleId, Ex.what());
+	}
 }
 
 void
-StorageServerInstance::Deprovision()
+StorageServerInstance::HibernateLocked()
 {
-	RwLock::ExclusiveLockScope _(m_Lock);
+	// Signal server to shut down, but keep data around for later wake
 
-	if (!m_IsProvisioned)
+	if (!m_ServerInstance.IsRunning())
 	{
-		ZEN_WARN("Attempted to deprovision storage server instance for module '{}' which is not provisioned", m_ModuleId);
-
 		return;
 	}
 
-	ZEN_INFO("Deprovisioning storage server instance for module '{}'", m_ModuleId);
-
+	// m_ServerInstance.Shutdown() never throws.
 	m_ServerInstance.Shutdown();
-
-	Dehydrate();
-
-	m_IsProvisioned = false;
 }
 
 void
-StorageServerInstance::Hibernate()
+StorageServerInstance::WakeLocked()
 {
-	// Signal server to shut down, but keep data around for later wake
-
-	RwLock::ExclusiveLockScope _(m_Lock);
+	// Start server in-place using existing data
 
-	if (!m_IsProvisioned)
+	if (m_ServerInstance.IsRunning())
 	{
-		ZEN_WARN("Attempted to hibernate storage server instance for module '{}' which is not provisioned", m_ModuleId);
-
 		return;
 	}
 
-	if (m_IsHibernated)
+	try
 	{
-		ZEN_WARN("Storage server instance for module '{}' is already hibernated", m_ModuleId);
-
-		return;
+		SpawnServerProcess();
 	}
-
-	if (!m_ServerInstance.IsRunning())
+	catch (const std::exception& Ex)
 	{
-		ZEN_WARN("Attempted to hibernate storage server instance for module '{}' which is not running", m_ModuleId);
+		ZEN_WARN("Failed spawning server instance for module '{}', at '{}' during waking. Reason: {}", m_ModuleId, m_BaseDir, Ex.what());
+		throw;
+	}
+}
 
-		// This is an unexpected state. Should consider the instance invalid?
+void
+StorageServerInstance::Hydrate()
+{
+	HydrationConfig Config{.ServerStateDir		= m_BaseDir,
+						   .TempDir				= m_TempDir,
+						   .ModuleId			= m_ModuleId,
+						   .TargetSpecification = m_Config.HydrationTargetSpecification,
+						   .Options				= m_Config.HydrationOptions};
 
-		return;
-	}
+	std::unique_ptr<HydrationStrategyBase> Hydrator = CreateHydrator(Config);
 
-	try
+	Hydrator->Hydrate();
+}
+
+void
+StorageServerInstance::Dehydrate()
+{
+	HydrationConfig Config{.ServerStateDir		= m_BaseDir,
+						   .TempDir				= m_TempDir,
+						   .ModuleId			= m_ModuleId,
+						   .TargetSpecification = m_Config.HydrationTargetSpecification,
+						   .Options				= m_Config.HydrationOptions};
+
+	std::unique_ptr<HydrationStrategyBase> Hydrator = CreateHydrator(Config);
+
+	Hydrator->Dehydrate();
+}
+
+StorageServerInstance::SharedLockedPtr::SharedLockedPtr() : m_Lock(nullptr), m_Instance(nullptr)
+{
+}
+
+StorageServerInstance::SharedLockedPtr::SharedLockedPtr(RwLock& Lock, StorageServerInstance* Instance, bool Wait)
+: m_Lock(nullptr)
+, m_Instance(nullptr)
+{
+	ZEN_ASSERT(Instance != nullptr);
+	if (Wait)
 	{
-		m_ServerInstance.Shutdown();
+		Lock.AcquireShared();
+		m_Lock	   = &Lock;
+		m_Instance = Instance;
+	}
+	else
+	{
+		if (Lock.TryAcquireShared())
+		{
+			m_Lock	   = &Lock;
+			m_Instance = Instance;
+		}
+	}
+}
 
-		m_IsHibernated	= true;
-		m_IsProvisioned = false;
+StorageServerInstance::SharedLockedPtr::SharedLockedPtr(SharedLockedPtr&& Rhs) : m_Lock(Rhs.m_Lock), m_Instance(Rhs.m_Instance)
+{
+	Rhs.m_Lock	   = nullptr;
+	Rhs.m_Instance = nullptr;
+}
 
-		return;
+StorageServerInstance::SharedLockedPtr::~SharedLockedPtr()
+{
+	if (m_Lock != nullptr)
+	{
+		m_Lock->ReleaseShared();
+		m_Lock = nullptr;
 	}
-	catch (const std::exception& Ex)
+	m_Instance = nullptr;
+}
+
+StorageServerInstance::SharedLockedPtr&
+StorageServerInstance::SharedLockedPtr::operator=(SharedLockedPtr&& Rhs)
+{
+	if (m_Lock)
 	{
-		ZEN_ERROR("Failed to hibernate storage server instance for module '{}': {}", m_ModuleId, Ex.what());
+		m_Lock->ReleaseShared();
+		m_Lock	   = nullptr;
+		m_Instance = nullptr;
 	}
+	m_Lock		   = Rhs.m_Lock;
+	m_Instance	   = Rhs.m_Instance;
+	Rhs.m_Lock	   = nullptr;
+	Rhs.m_Instance = nullptr;
+	return *this;
 }
 
-void
-StorageServerInstance::Wake()
+std::string_view
+StorageServerInstance::SharedLockedPtr::GetModuleId() const
+{
+	ZEN_ASSERT(m_Instance != nullptr);
+	return m_Instance->m_ModuleId;
+}
+
+bool
+StorageServerInstance::SharedLockedPtr::IsRunning() const
 {
-	RwLock::ExclusiveLockScope _(m_Lock);
-	WakeLocked();
+	ZEN_ASSERT(m_Instance != nullptr);
+	return m_Instance->m_ServerInstance.IsRunning();
 }
 
+#if ZEN_WITH_TESTS
 void
-StorageServerInstance::WakeLocked()
+StorageServerInstance::SharedLockedPtr::TerminateForTesting() const
 {
-	// Start server in-place using existing data
+	ZEN_ASSERT(m_Instance != nullptr);
+	m_Instance->m_ServerInstance.Terminate();
+}
+#endif
 
-	if (!m_IsHibernated)
-	{
-		ZEN_WARN("Attempted to wake storage server instance for module '{}' which is not hibernated", m_ModuleId);
+StorageServerInstance::ExclusiveLockedPtr::ExclusiveLockedPtr() : m_Lock(nullptr), m_Instance(nullptr)
+{
+}
 
-		return;
+StorageServerInstance::ExclusiveLockedPtr::ExclusiveLockedPtr(RwLock& Lock, StorageServerInstance* Instance, bool Wait)
+: m_Lock(nullptr)
+, m_Instance(nullptr)
+{
+	ZEN_ASSERT(Instance != nullptr);
+	if (Wait)
+	{
+		Lock.AcquireExclusive();
+		m_Lock	   = &Lock;
+		m_Instance = Instance;
+	}
+	else
+	{
+		if (Lock.TryAcquireExclusive())
+		{
+			m_Lock	   = &Lock;
+			m_Instance = Instance;
+		}
 	}
+}
 
-	ZEN_ASSERT_FORMAT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
+StorageServerInstance::ExclusiveLockedPtr::ExclusiveLockedPtr(ExclusiveLockedPtr&& Rhs) : m_Lock(Rhs.m_Lock), m_Instance(Rhs.m_Instance)
+{
+	Rhs.m_Lock	   = nullptr;
+	Rhs.m_Instance = nullptr;
+}
 
-	try
+StorageServerInstance::ExclusiveLockedPtr::~ExclusiveLockedPtr()
+{
+	if (m_Lock != nullptr)
 	{
-		SpawnServerProcess();
-		m_IsHibernated = false;
+		m_Lock->ReleaseExclusive();
+		m_Lock = nullptr;
 	}
-	catch (const std::exception& Ex)
-	{
-		ZEN_ERROR("Failed to wake storage server instance for module '{}': {}", m_ModuleId, Ex.what());
+	m_Instance = nullptr;
+}
 
-		// TODO: this instance should be marked as invalid
+StorageServerInstance::ExclusiveLockedPtr&
+StorageServerInstance::ExclusiveLockedPtr::operator=(ExclusiveLockedPtr&& Rhs)
+{
+	if (m_Lock)
+	{
+		m_Lock->ReleaseExclusive();
+		m_Lock	   = nullptr;
+		m_Instance = nullptr;
 	}
+	m_Lock		   = Rhs.m_Lock;
+	m_Instance	   = Rhs.m_Instance;
+	Rhs.m_Lock	   = nullptr;
+	Rhs.m_Instance = nullptr;
+	return *this;
 }
 
-void
-StorageServerInstance::Hydrate()
+std::string_view
+StorageServerInstance::ExclusiveLockedPtr::GetModuleId() const
 {
-	HydrationConfig Config{.ServerStateDir		= m_BaseDir,
-						   .TempDir				= m_TempDir,
-						   .ModuleId			= m_ModuleId,
-						   .TargetSpecification = WideToUtf8(m_Config.FileHydrationPath.native())};
+	ZEN_ASSERT(m_Instance != nullptr);
+	return m_Instance->m_ModuleId;
+}
 
-	std::unique_ptr<HydrationStrategyBase> Hydrator = CreateFileHydrator();
+bool
+StorageServerInstance::ExclusiveLockedPtr::IsRunning() const
+{
+	ZEN_ASSERT(m_Instance != nullptr);
+	return m_Instance->m_ServerInstance.IsRunning();
+}
 
-	Hydrator->Configure(Config);
-	Hydrator->Hydrate();
+void
+StorageServerInstance::ExclusiveLockedPtr::Provision()
+{
+	ZEN_ASSERT(m_Instance != nullptr);
+	m_Instance->ProvisionLocked();
 }
 
 void
-StorageServerInstance::Dehydrate()
+StorageServerInstance::ExclusiveLockedPtr::Deprovision()
 {
-	HydrationConfig Config{.ServerStateDir		= m_BaseDir,
-						   .TempDir				= m_TempDir,
-						   .ModuleId			= m_ModuleId,
-						   .TargetSpecification = WideToUtf8(m_Config.FileHydrationPath.native())};
+	ZEN_ASSERT(m_Instance != nullptr);
+	m_Instance->DeprovisionLocked();
+}
 
-	std::unique_ptr<HydrationStrategyBase> Hydrator = CreateFileHydrator();
+void
+StorageServerInstance::ExclusiveLockedPtr::Hibernate()
+{
+	ZEN_ASSERT(m_Instance != nullptr);
+	m_Instance->HibernateLocked();
+}
 
-	Hydrator->Configure(Config);
-	Hydrator->Dehydrate();
+void
+StorageServerInstance::ExclusiveLockedPtr::Wake()
+{
+	ZEN_ASSERT(m_Instance != nullptr);
+	m_Instance->WakeLocked();
 }
 
 }  // namespace zen