diff --git a/LocalMultiplayerAgent/MultiplayerServerManager.cs b/LocalMultiplayerAgent/MultiplayerServerManager.cs index b0130d2..d374180 100644 --- a/LocalMultiplayerAgent/MultiplayerServerManager.cs +++ b/LocalMultiplayerAgent/MultiplayerServerManager.cs @@ -84,7 +84,9 @@ public async Task CreateAndStartContainerWaitForExit(SessionHostsStartInfo start _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); diff --git a/VmAgent.Core.UnitTests/DockerContainerEngineTests.cs b/VmAgent.Core.UnitTests/DockerContainerEngineTests.cs index ac39665..e4c385e 100644 --- a/VmAgent.Core.UnitTests/DockerContainerEngineTests.cs +++ b/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(); + var mockContainerOperations = new Mock(); + + long expectedStatusCode = 137; + mockContainerOperations + .Setup(x => x.WaitContainerAsync(It.IsAny(), 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); + } } } \ No newline at end of file diff --git a/VmAgent.Core.UnitTests/ProcessWrapperTests.cs b/VmAgent.Core.UnitTests/ProcessWrapperTests.cs new file mode 100644 index 0000000..d23cb2a --- /dev/null +++ b/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 + { + /// + /// 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. + /// + [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(); + + wrapper.Dispose(); + } + + /// + /// Verifies that WaitForProcessExit returns the correct exit code + /// and removes the tracked process. + /// + [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(); + + wrapper.Dispose(); + } + + /// + /// Verifies that exit code is captured correctly even for non-zero codes. + /// + [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(); + } + + /// + /// Verifies that Kill handles an already-exited process gracefully + /// (no exceptions thrown) and still cleans up the dictionary entry. + /// + [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(); + } + + /// + /// Verifies that Dispose cleans up any remaining tracked processes. + /// + [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 */ } + } + + /// + /// Verifies that StartWithEventHandler tracks the process and Kill + /// cleans it up properly. + /// + [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(); + + 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 + }; + } +} diff --git a/VmAgent.Core/Interfaces/BaseSessionHostRunner.cs b/VmAgent.Core/Interfaces/BaseSessionHostRunner.cs index f5be125..335edbd 100644 --- a/VmAgent.Core/Interfaces/BaseSessionHostRunner.cs +++ b/VmAgent.Core/Interfaces/BaseSessionHostRunner.cs @@ -50,6 +50,6 @@ protected BaseSessionHostRunner( abstract public Task TryDelete(string id); - abstract public Task WaitOnServerExit(string containerId); + abstract public Task WaitOnServerExit(string containerId); } } diff --git a/VmAgent.Core/Interfaces/DockerContainerEngine.cs b/VmAgent.Core/Interfaces/DockerContainerEngine.cs index 255ef07..a86021b 100644 --- a/VmAgent.Core/Interfaces/DockerContainerEngine.cs +++ b/VmAgent.Core/Interfaces/DockerContainerEngine.cs @@ -129,10 +129,12 @@ public override string GetVmAgentIpAddress() } } - public override async Task WaitOnServerExit(string containerId) + public override async Task 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 CreateContainer( diff --git a/VmAgent.Core/Interfaces/IProcessWrapper.cs b/VmAgent.Core/Interfaces/IProcessWrapper.cs index eaa16f6..b7907e6 100644 --- a/VmAgent.Core/Interfaces/IProcessWrapper.cs +++ b/VmAgent.Core/Interfaces/IProcessWrapper.cs @@ -17,6 +17,6 @@ public interface IProcessWrapper IEnumerable List(); - void WaitForProcessExit(int id); + int WaitForProcessExit(int id); } } diff --git a/VmAgent.Core/Interfaces/ISessionHostRunner.cs b/VmAgent.Core/Interfaces/ISessionHostRunner.cs index f087685..edafd60 100644 --- a/VmAgent.Core/Interfaces/ISessionHostRunner.cs +++ b/VmAgent.Core/Interfaces/ISessionHostRunner.cs @@ -22,7 +22,7 @@ public interface ISessionHostRunner string GetVmAgentIpAddress(); - Task WaitOnServerExit(string containerId); + Task WaitOnServerExit(string containerId); Task CollectLogs(string id, string logsFolder, ISessionHostManager sessionHostManager); diff --git a/VmAgent.Core/Interfaces/ProcessRunner.cs b/VmAgent.Core/Interfaces/ProcessRunner.cs index e1dcd01..4183e89 100644 --- a/VmAgent.Core/Interfaces/ProcessRunner.cs +++ b/VmAgent.Core/Interfaces/ProcessRunner.cs @@ -166,10 +166,11 @@ public override Task> List() return Task.FromResult(_processWrapper.List().Select(x => x.ToString())); } - public override Task WaitOnServerExit(string containerId) + public override Task 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() diff --git a/VmAgent.Core/Interfaces/ProcessWrapper.cs b/VmAgent.Core/Interfaces/ProcessWrapper.cs index 5b1fe98..2064d89 100644 --- a/VmAgent.Core/Interfaces/ProcessWrapper.cs +++ b/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 _trackedProcesses = new ConcurrentDictionary(); + 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 StdOutputHandler, Action ErrorOutputHandler, Action ProcessExitedHandler) @@ -49,6 +64,7 @@ public int StartWithEventHandler(ProcessStartInfo startInfo, Action 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(); + } + } } } }