Capture and surface game server exit codes in LocalMultiplayerAgent

LocalMultiplayerAgent/MultiplayerServerManager.cs

@@ -84,7 +84,9 @@
 
             _logger.LogInformation("Waiting for heartbeats from the game server.....");
 
-            await sessionHostRunner.WaitOnServerExit(typeSpecificId).ConfigureAwait(false);
+            int exitCode = await sessionHostRunner.WaitOnServerExit(typeSpecificId).ConfigureAwait(false);
+            _logger.LogInformation($"Game server exited with exit code {exitCode}.");
+            Console.WriteLine($"Game server exited with exit code {exitCode}.");
             string logFolder = Path.Combine(Globals.VmConfiguration.VmDirectories.GameLogsRootFolderVm, sessionHostInfo.LogFolderId);
             await sessionHostRunner.CollectLogs(typeSpecificId, logFolder, sessionHostManager);
             await sessionHostRunner.TryDelete(typeSpecificId);
@@ -119,7 +121,7 @@
 
                     //we currently only support passwordless certificates
                     //passing a certificate with a password here will throw an exception because it requires a password parameter
                     X509Certificate2 cert = new X509Certificate2(certUserDetails.Path);
                     certs.Add(new CertificateDetail()
                     {
                         Name = certUserDetails.Name,

VmAgent.Core.UnitTests/DockerContainerEngineTests.cs

@@ -9,9 +9,12 @@
 using Microsoft.Azure.Gaming.VmAgent.Model;
 using Microsoft.Azure.Gaming.AgentInterfaces;
 using Microsoft.Extensions.Logging.Abstractions;
+using Docker.DotNet.Models;
 using Moq;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
+using System.Threading.Tasks;
+using ISystemOperations = Microsoft.Azure.Gaming.VmAgent.Core.Interfaces.ISystemOperations;
 
 namespace VmAgent.Core.UnitTests
 {
@@ -172,5 +175,27 @@ public void GetGameWorkingDir_LinuxHost_ShouldReturnNull()
 
             Assert.IsNull(result);
         }
+        [TestMethod]
+        [TestCategory("BVT")]
+        public async Task WaitOnServerExit_ReturnsContainerExitCode()
+        {
+            var mockDockerClient = new Mock<Docker.DotNet.IDockerClient>();
+            var mockContainerOperations = new Mock<Docker.DotNet.IContainerOperations>();
+
+            long expectedStatusCode = 137;
+            mockContainerOperations
+                .Setup(x => x.WaitContainerAsync(It.IsAny<string>(), default))
+                .ReturnsAsync(new ContainerWaitResponse { StatusCode = expectedStatusCode });
+
+            mockDockerClient.Setup(x => x.Containers).Returns(mockContainerOperations.Object);
+
+            var logger = new MultiLogger(NullLogger.Instance);
+            var vmConfiguration = new VmConfiguration(56001, "testVmId", new VmDirectories("root"), true);
+            var engine = new DockerContainerEngine(vmConfiguration, logger, _mockSystemOperations.Object, mockDockerClient.Object);
+
+            int exitCode = await engine.WaitOnServerExit("test-container-id");
+
+            Assert.AreEqual(137, exitCode);
+        }
     }
 }

VmAgent.Core.UnitTests/ProcessWrapperTests.cs

@@ -0,0 +1,181 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace VmAgent.Core.UnitTests
+{
+    using System;
+    using System.Diagnostics;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Gaming.VmAgent.Core.Interfaces;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using FluentAssertions;
+
+    [TestClass]
+    public class ProcessWrapperTests
+    {
+        /// <summary>
+        /// Verifies that Kill() removes the Process from the tracked dictionary
+        /// so that it doesn't leak. After Kill, WaitForProcessExit should throw
+        /// InvalidOperationException because the process is no longer tracked.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void Kill_RemovesProcessFromTrackedDictionary()
+        {
+            var wrapper = new ProcessWrapper();
+            int pid = wrapper.Start(GetSleepProcessStartInfo());
+
+            wrapper.Kill(pid);
+
+            // WaitForProcessExit should throw because the process was removed from tracking by Kill
+            Action act = () => wrapper.WaitForProcessExit(pid);
+            act.Should().Throw<InvalidOperationException>();
+
+            wrapper.Dispose();
+        }
+
+        /// <summary>
+        /// Verifies that WaitForProcessExit returns the correct exit code
+        /// and removes the tracked process.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void WaitForProcessExit_ReturnsExitCodeAndCleansUp()
+        {
+            var wrapper = new ProcessWrapper();
+
+            // Start a process that exits with code 0
+            var startInfo = GetExitProcessStartInfo(exitCode: 0);
+            int pid = wrapper.Start(startInfo);
+
+            int exitCode = wrapper.WaitForProcessExit(pid);
+            exitCode.Should().Be(0);
+
+            // Calling again should throw since it was removed from tracking
+            Action act = () => wrapper.WaitForProcessExit(pid);
+            act.Should().Throw<InvalidOperationException>();
+
+            wrapper.Dispose();
+        }
+
+        /// <summary>
+        /// Verifies that exit code is captured correctly even for non-zero codes.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void WaitForProcessExit_CapturesNonZeroExitCode()
+        {
+            var wrapper = new ProcessWrapper();
+
+            var startInfo = GetExitProcessStartInfo(exitCode: 42);
+            int pid = wrapper.Start(startInfo);
+
+            int exitCode = wrapper.WaitForProcessExit(pid);
+            exitCode.Should().Be(42);
+
+            wrapper.Dispose();
+        }
+
+        /// <summary>
+        /// Verifies that Kill handles an already-exited process gracefully
+        /// (no exceptions thrown) and still cleans up the dictionary entry.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void Kill_AlreadyExitedProcess_DoesNotThrow()
+        {
+            var wrapper = new ProcessWrapper();
+
+            var startInfo = GetExitProcessStartInfo(exitCode: 0);
+            int pid = wrapper.Start(startInfo);
+
+            // Wait for the process to exit on its own
+            Process.GetProcessById(pid).WaitForExit(5000);
+            Thread.Sleep(100); // small buffer
+
+            // Kill should not throw even though the process already exited
+            Action act = () => wrapper.Kill(pid);
+            act.Should().NotThrow();
+
+            wrapper.Dispose();
+        }
+
+        /// <summary>
+        /// Verifies that Dispose cleans up any remaining tracked processes.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void Dispose_CleansUpRemainingTrackedProcesses()
+        {
+            var wrapper = new ProcessWrapper();
+            int pid = wrapper.Start(GetSleepProcessStartInfo());
+
+            // Dispose without Kill or WaitForProcessExit — should not leak
+            wrapper.Dispose();
+
+            // Clean up the actual OS process
+            try { Process.GetProcessById(pid).Kill(true); }
+            catch (ArgumentException) { /* process already exited */ }
+            catch (InvalidOperationException) { /* process already exited */ }
+        }
+
+        /// <summary>
+        /// Verifies that StartWithEventHandler tracks the process and Kill
+        /// cleans it up properly.
+        /// </summary>
+        [TestMethod]
+        [TestCategory("BVT")]
+        public void StartWithEventHandler_KillCleansUpTrackedProcess()
+        {
+            var wrapper = new ProcessWrapper();
+            int pid = wrapper.StartWithEventHandler(
+                GetSleepProcessStartInfo(),
+                (sender, args) => { },
+                (sender, args) => { },
+                (sender, args) => { });
+
+            wrapper.Kill(pid);
+
+            // Process should be removed from tracking
+            Action act = () => wrapper.WaitForProcessExit(pid);
+            act.Should().Throw<InvalidOperationException>();
+
+            wrapper.Dispose();
+        }
+
+        private static ProcessStartInfo GetSleepProcessStartInfo() =>
+            OperatingSystem.IsWindows()
+                ? new ProcessStartInfo
+                {
+                    FileName = "cmd.exe",
+                    Arguments = "/c timeout /t 30 /nobreak >nul",
+                    UseShellExecute = false,
+                    CreateNoWindow = true
+                }
+                : new ProcessStartInfo
+                {
+                    FileName = "sleep",
+                    Arguments = "30",
+                    UseShellExecute = false,
+                    CreateNoWindow = true
+                };
+
+        private static ProcessStartInfo GetExitProcessStartInfo(int exitCode) =>
+            OperatingSystem.IsWindows()
+                ? new ProcessStartInfo
+                {
+                    FileName = "cmd.exe",
+                    Arguments = $"/c exit {exitCode}",
+                    UseShellExecute = false,
+                    CreateNoWindow = true
+                }
+                : new ProcessStartInfo
+                {
+                    FileName = "/bin/sh",
+                    Arguments = $"-c \"exit {exitCode}\"",
+                    UseShellExecute = false,
+                    CreateNoWindow = true
+                };
+    }
+}

VmAgent.Core/Interfaces/BaseSessionHostRunner.cs

@@ -50,6 +50,6 @@ protected BaseSessionHostRunner(
 
         abstract public Task<bool> TryDelete(string id);
 
-        abstract public Task WaitOnServerExit(string containerId);
+        abstract public Task<int> WaitOnServerExit(string containerId);
     }
 }

VmAgent.Core/Interfaces/DockerContainerEngine.cs

@@ -129,10 +129,12 @@
             }
         }
 
-        public override async Task WaitOnServerExit(string containerId)
+        public override async Task<int> WaitOnServerExit(string containerId)
         {
             ContainerWaitResponse containerWaitResponse = await _dockerClient.Containers.WaitContainerAsync(containerId).ConfigureAwait(false);
-            _logger.LogInformation($"Container {containerId} exited with exit code {containerWaitResponse.StatusCode}.");
+            int exitCode = (int)containerWaitResponse.StatusCode;
+            _logger.LogInformation($"Container {containerId} exited with exit code {exitCode}.");
+            return exitCode;
         }
 
         private async Task<string> CreateContainer(
@@ -510,12 +512,12 @@
                     // which the host Windows machine does not have "copy" access to, to get the logs with a FileCopy
                     // this is only supposed to run on LocalMultiplayerAgent running on lcow
                     StringBuilder sb = new StringBuilder();
                     Stream logsStream = await _dockerClient.Containers.GetContainerLogsAsync(containerId,
                         new ContainerLogsParameters() {ShowStdout = true, ShowStderr = true});
                     using (StreamReader sr = new StreamReader(logsStream))
                     {
                         Stopwatch sw = new Stopwatch();
                         while (!sr.EndOfStream)
                         {
                             if (sw.Elapsed.Seconds > 3) // don't flood STDOUT with messages, output one every 3 seconds if logs are too many
                             {

VmAgent.Core/Interfaces/IProcessWrapper.cs

@@ -17,6 +17,6 @@ public interface IProcessWrapper
 
         IEnumerable<int> List();
 
-        void WaitForProcessExit(int id);
+        int WaitForProcessExit(int id);
     }
 }

VmAgent.Core/Interfaces/ISessionHostRunner.cs

@@ -22,7 +22,7 @@ public interface ISessionHostRunner
 
         string GetVmAgentIpAddress();
 
-        Task WaitOnServerExit(string containerId);
+        Task<int> WaitOnServerExit(string containerId);
 
         Task CollectLogs(string id, string logsFolder, ISessionHostManager sessionHostManager);
 

VmAgent.Core/Interfaces/ProcessRunner.cs

@@ -166,10 +166,11 @@ public override Task<IEnumerable<string>> List()
             return Task.FromResult(_processWrapper.List().Select(x => x.ToString()));
         }
 
-        public override Task WaitOnServerExit(string containerId)
+        public override Task<int> WaitOnServerExit(string containerId)
         {
-            _processWrapper.WaitForProcessExit(int.Parse(containerId));
-            return Task.CompletedTask;
+            int exitCode = _processWrapper.WaitForProcessExit(int.Parse(containerId));
+            _logger.LogInformation($"Process {containerId} exited with exit code {exitCode}.");
+            return Task.FromResult(exitCode);
         }
 
         public override string GetVmAgentIpAddress()

VmAgent.Core/Interfaces/ProcessWrapper.cs

@@ -4,18 +4,30 @@
 namespace Microsoft.Azure.Gaming.VmAgent.Core.Interfaces
 {
     using System;
+    using System.Collections.Concurrent;
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.Linq;
 
-    public class ProcessWrapper : IProcessWrapper
+    public class ProcessWrapper : IProcessWrapper, IDisposable
     {
+        private readonly ConcurrentDictionary<int, Process> _trackedProcesses = new ConcurrentDictionary<int, Process>();
+
         public void Kill(int id)
         {
             try
             {
-                Process process = Process.GetProcessById(id);
-                process.Kill(true);
+                // Use the tracked process reference when available to avoid PID reuse
+                // issues and to ensure the dictionary entry is cleaned up.
+                if (!_trackedProcesses.TryRemove(id, out Process process))
+                {
+                    process = Process.GetProcessById(id);
+                }
+
+                using (process)
+                {
+                    process.Kill(true);
+                }
             }
             catch (ArgumentException)
             {
@@ -31,7 +43,10 @@ public void Kill(int id)
 
         public int Start(ProcessStartInfo startInfo)
         {
-            return Process.Start(startInfo).Id;
+            Process process = Process.Start(startInfo)
+                ?? throw new InvalidOperationException("Process.Start returned null for: " + startInfo.FileName);
+            _trackedProcesses[process.Id] = process;
+            return process.Id;
         }
 
         public int StartWithEventHandler(ProcessStartInfo startInfo, Action<object, DataReceivedEventArgs> StdOutputHandler, Action<object, DataReceivedEventArgs> ErrorOutputHandler, Action<object, EventArgs> ProcessExitedHandler)
@@ -49,6 +64,7 @@ public int StartWithEventHandler(ProcessStartInfo startInfo, Action<object, Data
             process.BeginOutputReadLine();
             process.BeginErrorReadLine();
 
+            _trackedProcesses[process.Id] = process;
             return process.Id;
         }
 
@@ -57,11 +73,41 @@ public IEnumerable<int> List()
             return Process.GetProcesses().Select(x => x.Id);
         }
 
-        public void WaitForProcessExit(int id)
+        public int WaitForProcessExit(int id)
+        {
+            if (!_trackedProcesses.TryRemove(id, out Process process))
+            {
+                throw new InvalidOperationException(
+                    $"Process {id} is not tracked. All processes should be started through this wrapper.");
+            }
+
+            using (process)
+            {
+                try
+                {
+                    process.WaitForExit();
+                    return process.ExitCode;
+                }
+                finally
+                {
+                    try { process.CancelOutputRead(); }
+                    catch (InvalidOperationException) { /* expected when output was not redirected */ }
+
+                    try { process.CancelErrorRead(); }
+                    catch (InvalidOperationException) { /* expected when error was not redirected */ }
+                }
+            }
+        }
+
+        public void Dispose()
         {
-            // TODO: this may need a bit more polish, it is currently only used by LocalMultiplayerAgent.
-            Process process = Process.GetProcessById(id);
-            process.WaitForExit();
+            foreach (var kvp in _trackedProcesses)
+            {
+                if (_trackedProcesses.TryRemove(kvp.Key, out Process process))
+                {
+                    process.Dispose();
+                }
+            }
         }
     }
 }