4 files changed, 33 insertions, 12 deletions

diff --git a/src/zenserver/config.cpp b/src/zenserver/config.cpp
index bedab7049..0108e8b9f 100644
--- a/src/zenserver/config.cpp
+++ b/src/zenserver/config.cpp
@@ -593,6 +593,9 @@ ParseCliOptions(int argc, char* argv[], ZenServerOptions& ServerOptions)
 	options.add_options()("detach",
 						  "Indicate whether zenserver should detach from parent process group",
 						  cxxopts::value<bool>(ServerOptions.Detach)->default_value("true"));
+	options.add_options()("malloc",
+						  "Configure memory allocator subsystem",
+						  cxxopts::value(ServerOptions.MemoryOptions)->default_value("mimalloc"));
 
 	// clang-format off
 	options.add_options("logging")
@@ -713,6 +716,13 @@ ParseCliOptions(int argc, char* argv[], ZenServerOptions& ServerOptions)
 #if ZEN_WITH_TRACE
 	options.add_option("ue-trace",
 					   "",
+					   "trace",
+					   "Specify which trace channels should be enabled",
+					   cxxopts::value<std::string>(ServerOptions.TraceChannels)->default_value(""),
+					   "");
+
+	options.add_option("ue-trace",
+					   "",
 					   "tracehost",
 					   "Hostname to send the trace to",
 					   cxxopts::value<std::string>(ServerOptions.TraceHost)->default_value(""),
diff --git a/src/zenserver/config.h b/src/zenserver/config.h
index 5c56695f3..c7781aada 100644
--- a/src/zenserver/config.h
+++ b/src/zenserver/config.h
@@ -176,9 +176,11 @@ struct ZenServerOptions
 	std::string				 Loggers[zen::logging::level::LogLevelCount];
 	std::string				 ScrubOptions;
 #if ZEN_WITH_TRACE
-	std::string TraceHost;	// Host name or IP address to send trace data to
-	std::string TraceFile;	// Path of a file to write a trace
+	std::string TraceChannels;	// Trace channels to enable
+	std::string TraceHost;		// Host name or IP address to send trace data to
+	std::string TraceFile;		// Path of a file to write a trace
 #endif
+	std::string MemoryOptions;	// Memory allocation options
 	std::string CommandLine;
 };
 
diff --git a/src/zenserver/diag/logging.cpp b/src/zenserver/diag/logging.cpp
index 595be70cb..0d96cd8d6 100644
--- a/src/zenserver/diag/logging.cpp
+++ b/src/zenserver/diag/logging.cpp
@@ -6,6 +6,7 @@
 
 #include <zencore/filesystem.h>
 #include <zencore/fmtutils.h>
+#include <zencore/memory/llm.h>
 #include <zencore/session.h>
 #include <zencore/string.h>
 #include <zenutil/logging.h>
@@ -20,6 +21,8 @@ namespace zen {
 void
 InitializeServerLogging(const ZenServerOptions& InOptions)
 {
+	UE_MEMSCOPE(ELLMTag::Logging);
+
 	const LoggingOptions LogOptions = {.IsDebug			= InOptions.IsDebug,
 									   .IsVerbose		= false,
 									   .IsTest			= InOptions.IsTest,
@@ -79,6 +82,8 @@ InitializeServerLogging(const ZenServerOptions& InOptions)
 void
 ShutdownServerLogging()
 {
+	UE_MEMSCOPE(ELLMTag::Logging);
+
 	zen::ShutdownLogging();
 }
 
diff --git a/src/zenserver/main.cpp b/src/zenserver/main.cpp
index 2fb01ebf1..4444241cc 100644
--- a/src/zenserver/main.cpp
+++ b/src/zenserver/main.cpp
@@ -17,16 +17,15 @@
 #include <zencore/trace.h>
 #include <zenhttp/httpserver.h>
 
+#include <zencore/memory/fmalloc.h>
+#include <zencore/memory/memory.h>
+#include <zencore/memory/memorytrace.h>
+#include <zencore/memory/newdelete.h>
+
 #include "config.h"
 #include "diag/logging.h"
 #include "sentryintegration.h"
 
-#if ZEN_USE_MIMALLOC
-ZEN_THIRD_PARTY_INCLUDES_START
-#	include <mimalloc-new-delete.h>
-ZEN_THIRD_PARTY_INCLUDES_END
-#endif
-
 #if ZEN_PLATFORM_WINDOWS
 #	include <zencore/windows.h>
 #	include "windows/service.h"
@@ -354,9 +353,6 @@ test_main(int argc, char** argv)
 int
 main(int argc, char* argv[])
 {
-#if ZEN_USE_MIMALLOC
-	mi_version();
-#endif
 	using namespace zen;
 
 	if (argc >= 2)
@@ -433,9 +429,17 @@ main(int argc, char* argv[])
 		{
 			TraceInit("zenserver");
 		}
-		atexit(TraceShutdown);
 #endif	// ZEN_WITH_TRACE
 
+#if ZEN_WITH_MEMTRACK
+		FMalloc* TraceMalloc = MemoryTrace_Create(GMalloc);
+		if (TraceMalloc != GMalloc)
+		{
+			GMalloc = TraceMalloc;
+			MemoryTrace_Initialize();
+		}
+#endif
+
 #if ZEN_PLATFORM_WINDOWS
 		if (ServerOptions.InstallService)
 		{