Merge branch 'main' into lm/restrict-content-type

1 files changed, 59 insertions, 9 deletions

diff --git a/src/zenserver/hub/hubservice.cpp b/src/zenserver/hub/hubservice.cpp
index 4d9da3a57..7b999ae20 100644
--- a/src/zenserver/hub/hubservice.cpp
+++ b/src/zenserver/hub/hubservice.cpp
@@ -4,10 +4,12 @@
 
 #include "hydration.h"
 
+#include <zencore/assertfmt.h>
 #include <zencore/compactbinarybuilder.h>
 #include <zencore/filesystem.h>
 #include <zencore/fmtutils.h>
 #include <zencore/logging.h>
+#include <zencore/process.h>
 #include <zencore/scopeguard.h>
 #include <zencore/system.h>
 #include <zenutil/zenserverprocess.h>
@@ -150,7 +152,12 @@ struct StorageServerInstance
 
 	inline uint16_t GetBasePort() const { return m_ServerInstance.GetBasePort(); }
 
+#if ZEN_PLATFORM_WINDOWS
+	void SetJobObject(JobObject* InJobObject) { m_JobObject = InJobObject; }
+#endif
+
 private:
+	void				  WakeLocked();
 	RwLock				  m_Lock;
 	std::string			  m_ModuleId;
 	std::atomic<bool>	  m_IsProvisioned{false};
@@ -160,6 +167,9 @@ private:
 	std::filesystem::path m_TempDir;
 	std::filesystem::path m_HydrationPath;
 	ResourceMetrics		  m_ResourceMetrics;
+#if ZEN_PLATFORM_WINDOWS
+	JobObject* m_JobObject = nullptr;
+#endif
 
 	void SpawnServerProcess();
 
@@ -186,10 +196,13 @@ StorageServerInstance::~StorageServerInstance()
 void
 StorageServerInstance::SpawnServerProcess()
 {
-	ZEN_ASSERT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
+	ZEN_ASSERT_FORMAT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
 
 	m_ServerInstance.SetServerExecutablePath(GetRunningExecutablePath());
 	m_ServerInstance.SetDataDir(m_BaseDir);
+#if ZEN_PLATFORM_WINDOWS
+	m_ServerInstance.SetJobObject(m_JobObject);
+#endif
 	const uint16_t BasePort = m_ServerInstance.SpawnServerAndWaitUntilReady();
 
 	ZEN_DEBUG("Storage server instance for module '{}' started, listening on port {}", m_ModuleId, BasePort);
@@ -211,7 +224,7 @@ StorageServerInstance::Provision()
 
 	if (m_IsHibernated)
 	{
-		Wake();
+		WakeLocked();
 	}
 	else
 	{
@@ -294,9 +307,14 @@ StorageServerInstance::Hibernate()
 void
 StorageServerInstance::Wake()
 {
-	// Start server in-place using existing data
-
 	RwLock::ExclusiveLockScope _(m_Lock);
+	WakeLocked();
+}
+
+void
+StorageServerInstance::WakeLocked()
+{
+	// Start server in-place using existing data
 
 	if (!m_IsHibernated)
 	{
@@ -305,7 +323,7 @@ StorageServerInstance::Wake()
 		return;
 	}
 
-	ZEN_ASSERT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
+	ZEN_ASSERT_FORMAT(!m_ServerInstance.IsRunning(), "Storage server instance for module '{}' is already running", m_ModuleId);
 
 	try
 	{
@@ -374,6 +392,21 @@ struct HttpHubService::Impl
 		// flexibility, and to allow running multiple hubs on the same host if
 		// necessary.
 		m_RunEnvironment.SetNextPortNumber(21000);
+
+#if ZEN_PLATFORM_WINDOWS
+		if (m_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
 	}
 
 	void Cleanup()
@@ -416,6 +449,12 @@ struct HttpHubService::Impl
 				IsNewInstance = true;
 				auto NewInstance =
 					std::make_unique<StorageServerInstance>(m_RunEnvironment, ModuleId, m_FileHydrationPath, m_HydrationTempPath);
+#if ZEN_PLATFORM_WINDOWS
+				if (m_JobObject.IsValid())
+				{
+					NewInstance->SetJobObject(&m_JobObject);
+				}
+#endif
 				Instance = NewInstance.get();
 				m_Instances.emplace(std::string(ModuleId), std::move(NewInstance));
 
@@ -573,10 +612,15 @@ struct HttpHubService::Impl
 	inline int GetInstanceLimit() { return m_InstanceLimit; }
 	inline int GetMaxInstanceCount() { return m_MaxInstanceCount; }
 
+	bool m_UseJobObject = true;
+
 private:
-	ZenServerEnvironment													m_RunEnvironment;
-	std::filesystem::path													m_FileHydrationPath;
-	std::filesystem::path													m_HydrationTempPath;
+	ZenServerEnvironment  m_RunEnvironment;
+	std::filesystem::path m_FileHydrationPath;
+	std::filesystem::path m_HydrationTempPath;
+#if ZEN_PLATFORM_WINDOWS
+	JobObject m_JobObject;
+#endif
 	RwLock																	m_Lock;
 	std::unordered_map<std::string, std::unique_ptr<StorageServerInstance>> m_Instances;
 	std::unordered_set<std::string>											m_DeprovisioningModules;
@@ -802,7 +846,7 @@ HttpHubService::HttpHubService(std::filesystem::path HubBaseDir, std::filesystem
 			Obj << "currentInstanceCount" << m_Impl->GetInstanceCount();
 			Obj << "maxInstanceCount" << m_Impl->GetMaxInstanceCount();
 			Obj << "instanceLimit" << m_Impl->GetInstanceLimit();
-			Req.ServerRequest().WriteResponse(HttpResponseCode::OK);
+			Req.ServerRequest().WriteResponse(HttpResponseCode::OK, Obj.Save());
 		},
 		HttpVerb::kGet);
 }
@@ -811,6 +855,12 @@ HttpHubService::~HttpHubService()
 {
 }
 
+void
+HttpHubService::SetUseJobObject(bool Enable)
+{
+	m_Impl->m_UseJobObject = Enable;
+}
+
 const char*
 HttpHubService::BaseUri() const
 {