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

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

namespace NATS.Server.Benchmark.Tests.Infrastructure;

/// <summary>
/// Manages a Go nats-server child process for benchmark testing.
/// Builds the Go binary if not present and skips gracefully if Go is not installed.
/// </summary>
public sealed class GoServerProcess : 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 GoServerProcess(string? configContent = null)
    {
        Port = PortAllocator.AllocateFreePort();
        _configContent = configContent;
    }

    /// <summary>
    /// Returns true if Go is installed and the Go NATS server source exists.
    /// </summary>
    public static bool IsAvailable()
    {
        try
        {
            var goVersion = Process.Start(new ProcessStartInfo
            {
                FileName = "go",
                Arguments = "version",
                UseShellExecute = false,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                CreateNoWindow = true,
            });
            goVersion!.WaitForExit();
            if (goVersion.ExitCode != 0)
                return false;

            var solutionRoot = FindSolutionRoot();
            return solutionRoot is not null &&
                   Directory.Exists(Path.Combine(solutionRoot, "golang", "nats-server"));
        }
        catch
        {
            return false;
        }
    }

    public async Task StartAsync()
    {
        var binary = ResolveGoBinary();

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

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

        var psi = new ProcessStartInfo
        {
            FileName = binary,
            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(
                        "Go nats-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(15));
        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(
            $"Go nats-server did not become ready on port {Port} within 15s. Last error: {lastError?.Message}\n\nServer output:\n{Output}");
    }

    private static string ResolveGoBinary()
    {
        var solutionRoot = FindSolutionRoot()
            ?? throw new FileNotFoundException(
                "Could not find solution root (NatsDotNet.slnx) walking up from " + AppContext.BaseDirectory);

        var goServerDir = Path.Combine(solutionRoot, "golang", "nats-server");
        if (!Directory.Exists(goServerDir))
            throw new DirectoryNotFoundException($"Go nats-server source not found at: {goServerDir}");

        var binaryName = OperatingSystem.IsWindows() ? "nats-server.exe" : "nats-server";
        var binaryPath = Path.Combine(goServerDir, binaryName);

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

        // Build the Go binary
        var build = Process.Start(new ProcessStartInfo
        {
            FileName = "go",
            Arguments = "build -o " + binaryName,
            WorkingDirectory = goServerDir,
            UseShellExecute = false,
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            CreateNoWindow = true,
        });
        build!.WaitForExit();

        if (build.ExitCode != 0)
            throw new InvalidOperationException(
                $"Failed to build Go nats-server:\n{build.StandardError.ReadToEnd()}");

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

        throw new FileNotFoundException($"Built Go nats-server but binary not found at: {binaryPath}");
    }

    internal static string? FindSolutionRoot()
    {
        var dir = new DirectoryInfo(AppContext.BaseDirectory);
        while (dir is not null)
        {
            if (File.Exists(Path.Combine(dir.FullName, "NatsDotNet.slnx")))
                return dir.FullName;
            dir = dir.Parent;
        }

        return null;
    }
}