natsdotnet/tests/NATS.Server.Benchmark.Tests/Infrastructure/DotNetServerProcess.cs

using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using System.Text;

namespace NATS.Server.Benchmark.Tests.Infrastructure;

/// <summary>
/// Manages a NATS.Server.Host child process for benchmark testing.
/// </summary>
public sealed class DotNetServerProcess : IAsyncDisposable
{
    private Process? _process;
    private readonly StringBuilder _output = new();
    private readonly object _outputLock = new();
    private readonly string? _configContent;
    private string? _configFilePath;

    public int Port { get; }

    public string Output
    {
        get
        {
            lock (_outputLock)
                return _output.ToString();
        }
    }

    public DotNetServerProcess(string? configContent = null)
    {
        Port = PortAllocator.AllocateFreePort();
        _configContent = configContent;
    }

    public async Task StartAsync()
    {
        var hostDll = ResolveHostDll();

        if (_configContent is not null)
        {
            _configFilePath = Path.Combine(Path.GetTempPath(), $"nats-bench-dotnet-{Guid.NewGuid():N}.conf");
            await File.WriteAllTextAsync(_configFilePath, _configContent);
        }

        var args = new StringBuilder($"exec \"{hostDll}\" -p {Port}");
        if (_configFilePath is not null)
            args.Append($" -c \"{_configFilePath}\"");

        var psi = new ProcessStartInfo
        {
            FileName = "dotnet",
            Arguments = args.ToString(),
            UseShellExecute = false,
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            CreateNoWindow = true,
        };

        _process = new Process { StartInfo = psi, EnableRaisingEvents = true };

        _process.OutputDataReceived += (_, e) =>
        {
            if (e.Data is not null)
                lock (_outputLock) _output.AppendLine(e.Data);
        };
        _process.ErrorDataReceived += (_, e) =>
        {
            if (e.Data is not null)
                lock (_outputLock) _output.AppendLine(e.Data);
        };

        _process.Start();
        _process.BeginOutputReadLine();
        _process.BeginErrorReadLine();

        await WaitForTcpReadyAsync();
    }

    public async ValueTask DisposeAsync()
    {
        if (_process is not null)
        {
            if (!_process.HasExited)
            {
                _process.Kill(entireProcessTree: true);

                using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
                try
                {
                    await _process.WaitForExitAsync(cts.Token);
                }
                catch (OperationCanceledException ex) when (!_process.HasExited)
                {
                    throw new InvalidOperationException(
                        "NATS .NET server process did not exit within 5s after kill.", ex);
                }
            }

            _process.Dispose();
            _process = null;
        }

        if (_configFilePath is not null && File.Exists(_configFilePath))
        {
            File.Delete(_configFilePath);
            _configFilePath = null;
        }
    }

    private async Task WaitForTcpReadyAsync()
    {
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(100));
        SocketException? lastError = null;

        while (await timer.WaitForNextTickAsync(timeout.Token).ConfigureAwait(false))
        {
            try
            {
                using var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                await socket.ConnectAsync(new IPEndPoint(IPAddress.Loopback, Port), timeout.Token);
                return;
            }
            catch (SocketException ex)
            {
                lastError = ex;
            }
        }

        throw new TimeoutException(
            $"NATS .NET server did not become ready on port {Port} within 10s. Last error: {lastError?.Message}\n\nServer output:\n{Output}");
    }

    private static string ResolveHostDll()
    {
        var dir = new DirectoryInfo(AppContext.BaseDirectory);
        while (dir is not null)
        {
            if (File.Exists(Path.Combine(dir.FullName, "NatsDotNet.slnx")))
            {
                var dll = Path.Combine(dir.FullName, "src", "NATS.Server.Host", "bin", "Release", "net10.0", "NATS.Server.Host.dll");
                if (File.Exists(dll))
                    return dll;

                var build = Process.Start(new ProcessStartInfo
                {
                    FileName = "dotnet",
                    Arguments = "build src/NATS.Server.Host/NATS.Server.Host.csproj -c Release",
                    WorkingDirectory = dir.FullName,
                    UseShellExecute = false,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                });
                build!.WaitForExit();

                if (build.ExitCode != 0)
                    throw new InvalidOperationException(
                        $"Failed to build NATS.Server.Host:\n{build.StandardError.ReadToEnd()}");

                if (File.Exists(dll))
                    return dll;

                throw new FileNotFoundException($"Built NATS.Server.Host but DLL not found at: {dll}");
            }

            dir = dir.Parent;
        }

        throw new FileNotFoundException(
            "Could not find solution root (NatsDotNet.slnx) walking up from " + AppContext.BaseDirectory);
    }
}