1 files changed, 46 insertions, 23 deletions

diff --git a/src/zenutil/process/subprocessmanager.cpp b/src/zenutil/process/subprocessmanager.cpp
index 3a91b0a61..b053ac6bd 100644
--- a/src/zenutil/process/subprocessmanager.cpp
+++ b/src/zenutil/process/subprocessmanager.cpp
@@ -196,18 +196,6 @@ ManagedProcess::GetCpuUsagePercent() const
 	return m_Impl->m_CpuUsagePercent.load();
 }
 
-void
-ManagedProcess::SetStdoutCallback(ProcessDataCallback Callback)
-{
-	m_Impl->m_StdoutCallback = std::move(Callback);
-}
-
-void
-ManagedProcess::SetStderrCallback(ProcessDataCallback Callback)
-{
-	m_Impl->m_StderrCallback = std::move(Callback);
-}
-
 std::string
 ManagedProcess::GetCapturedStdout() const
 {
@@ -288,7 +276,9 @@ struct SubprocessManager::Impl
 	ManagedProcess* Spawn(const std::filesystem::path& Executable,
 						  std::string_view			   CommandLine,
 						  CreateProcOptions&		   Options,
-						  ProcessExitCallback		   OnExit);
+						  ProcessExitCallback		   OnExit,
+						  ProcessDataCallback		   OnStdout,
+						  ProcessDataCallback		   OnStderr);
 	ManagedProcess* Adopt(ProcessHandle&& Handle, ProcessExitCallback OnExit);
 	void			Remove(int Pid);
 	void			RemoveAll();
@@ -462,7 +452,9 @@ ManagedProcess*
 SubprocessManager::Impl::Spawn(const std::filesystem::path& Executable,
 							   std::string_view				CommandLine,
 							   CreateProcOptions&			Options,
-							   ProcessExitCallback			OnExit)
+							   ProcessExitCallback			OnExit,
+							   ProcessDataCallback			OnStdout,
+							   ProcessDataCallback			OnStderr)
 {
 	bool HasStdout = Options.StdoutPipe != nullptr;
 	bool HasStderr = Options.StderrPipe != nullptr;
@@ -476,6 +468,16 @@ SubprocessManager::Impl::Spawn(const std::filesystem::path& Executable,
 	ImplPtr->m_Handle.Initialize(static_cast<int>(Result));
 #endif
 
+	// Install callbacks before starting async readers so no data is missed.
+	if (OnStdout)
+	{
+		ImplPtr->m_StdoutCallback = std::move(OnStdout);
+	}
+	if (OnStderr)
+	{
+		ImplPtr->m_StderrCallback = std::move(OnStderr);
+	}
+
 	auto Proc = std::unique_ptr<ManagedProcess>(new ManagedProcess(std::move(ImplPtr)));
 
 	ManagedProcess* Ptr = AddProcess(std::move(Proc));
@@ -719,10 +721,12 @@ ManagedProcess*
 SubprocessManager::Spawn(const std::filesystem::path& Executable,
 						 std::string_view			  CommandLine,
 						 CreateProcOptions&			  Options,
-						 ProcessExitCallback		  OnExit)
+						 ProcessExitCallback		  OnExit,
+						 ProcessDataCallback		  OnStdout,
+						 ProcessDataCallback		  OnStderr)
 {
 	ZEN_TRACE_CPU("SubprocessManager::Spawn");
-	return m_Impl->Spawn(Executable, CommandLine, Options, std::move(OnExit));
+	return m_Impl->Spawn(Executable, CommandLine, Options, std::move(OnExit), std::move(OnStdout), std::move(OnStderr));
 }
 
 ManagedProcess*
@@ -835,7 +839,9 @@ struct ProcessGroup::Impl
 	ManagedProcess* Spawn(const std::filesystem::path& Executable,
 						  std::string_view			   CommandLine,
 						  CreateProcOptions&		   Options,
-						  ProcessExitCallback		   OnExit);
+						  ProcessExitCallback		   OnExit,
+						  ProcessDataCallback		   OnStdout,
+						  ProcessDataCallback		   OnStderr);
 	ManagedProcess* Adopt(ProcessHandle&& Handle, ProcessExitCallback OnExit);
 	void			Remove(int Pid);
 	void			KillAll();
@@ -884,7 +890,9 @@ ManagedProcess*
 ProcessGroup::Impl::Spawn(const std::filesystem::path& Executable,
 						  std::string_view			   CommandLine,
 						  CreateProcOptions&		   Options,
-						  ProcessExitCallback		   OnExit)
+						  ProcessExitCallback		   OnExit,
+						  ProcessDataCallback		   OnStdout,
+						  ProcessDataCallback		   OnStderr)
 {
 	bool HasStdout = Options.StdoutPipe != nullptr;
 	bool HasStderr = Options.StderrPipe != nullptr;
@@ -917,6 +925,16 @@ ProcessGroup::Impl::Spawn(const std::filesystem::path& Executable,
 	}
 #endif
 
+	// Install callbacks before starting async readers so no data is missed.
+	if (OnStdout)
+	{
+		ImplPtr->m_StdoutCallback = std::move(OnStdout);
+	}
+	if (OnStderr)
+	{
+		ImplPtr->m_StderrCallback = std::move(OnStderr);
+	}
+
 	auto Proc = std::unique_ptr<ManagedProcess>(new ManagedProcess(std::move(ImplPtr)));
 
 	ManagedProcess* Ptr = AddProcess(std::move(Proc));
@@ -1077,10 +1095,12 @@ ManagedProcess*
 ProcessGroup::Spawn(const std::filesystem::path& Executable,
 					std::string_view			 CommandLine,
 					CreateProcOptions&			 Options,
-					ProcessExitCallback			 OnExit)
+					ProcessExitCallback			 OnExit,
+					ProcessDataCallback			 OnStdout,
+					ProcessDataCallback			 OnStderr)
 {
 	ZEN_TRACE_CPU("ProcessGroup::Spawn");
-	return m_Impl->Spawn(Executable, CommandLine, Options, std::move(OnExit));
+	return m_Impl->Spawn(Executable, CommandLine, Options, std::move(OnExit), std::move(OnStdout), std::move(OnStderr));
 }
 
 ManagedProcess*
@@ -1289,9 +1309,12 @@ TEST_CASE("SubprocessManager.StdoutCallback")
 	std::string ReceivedData;
 	bool		Exited = false;
 
-	ManagedProcess* Proc = Manager.Spawn(AppStub, CmdLine, Options, [&](ManagedProcess&, int) { Exited = true; });
-
-	Proc->SetStdoutCallback([&](ManagedProcess&, std::string_view Data) { ReceivedData.append(Data); });
+	ManagedProcess* Proc = Manager.Spawn(
+		AppStub,
+		CmdLine,
+		Options,
+		[&](ManagedProcess&, int) { Exited = true; },
+		[&](ManagedProcess&, std::string_view Data) { ReceivedData.append(Data); });
 
 	IoContext.run_for(5s);