mxaccessgw/src/MxGateway.Server/Workers/WorkerProcessLauncher.cs

using System.Diagnostics;
using Microsoft.Extensions.Options;
using MxGateway.Server.Configuration;
using MxGateway.Server.Metrics;

namespace MxGateway.Server.Workers;

public sealed class WorkerProcessLauncher : IWorkerProcessLauncher
{
    public const string WorkerNonceEnvironmentVariableName = "MXGATEWAY_WORKER_NONCE";

    private readonly IWorkerProcessFactory _processFactory;
    private readonly IWorkerStartupProbe _startupProbe;
    private readonly GatewayMetrics _metrics;
    private readonly TimeProvider _timeProvider;
    private readonly WorkerOptions _workerOptions;

    public WorkerProcessLauncher(
        IOptions<GatewayOptions> gatewayOptions,
        IWorkerProcessFactory processFactory,
        IWorkerStartupProbe startupProbe,
        GatewayMetrics metrics,
        TimeProvider? timeProvider = null)
    {
        ArgumentNullException.ThrowIfNull(gatewayOptions);
        ArgumentNullException.ThrowIfNull(processFactory);
        ArgumentNullException.ThrowIfNull(startupProbe);
        ArgumentNullException.ThrowIfNull(metrics);

        _workerOptions = gatewayOptions.Value.Worker;
        _processFactory = processFactory;
        _startupProbe = startupProbe;
        _metrics = metrics;
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    public async Task<WorkerProcessHandle> LaunchAsync(
        WorkerProcessLaunchRequest request,
        CancellationToken cancellationToken = default)
    {
        try
        {
            return await LaunchCoreAsync(request, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            request.PipeReservation?.Dispose();
            throw;
        }
    }

    private async Task<WorkerProcessHandle> LaunchCoreAsync(
        WorkerProcessLaunchRequest request,
        CancellationToken cancellationToken)
    {
        ValidateRequest(request);

        DateTimeOffset startedAt = _timeProvider.GetUtcNow();
        ProcessStartInfo startInfo = CreateStartInfo(request, out WorkerProcessCommandLine commandLine);

        IWorkerProcess process;
        try
        {
            process = _processFactory.Start(startInfo);
        }
        catch (Exception exception) when (exception is not WorkerProcessLaunchException)
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.StartFailed,
                "Worker process failed to start.",
                exception);
        }

        try
        {
            using CancellationTokenSource startupTimeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            startupTimeout.CancelAfter(TimeSpan.FromSeconds(_workerOptions.StartupTimeoutSeconds));

            await _startupProbe
                .WaitUntilReadyAsync(process, request, startupTimeout.Token)
                .ConfigureAwait(false);

            _metrics.WorkerStarted(_timeProvider.GetUtcNow() - startedAt);

            return new WorkerProcessHandle(process, commandLine, startedAt);
        }
        catch (OperationCanceledException exception) when (!cancellationToken.IsCancellationRequested)
        {
            KillAndDispose(process, "StartupTimeout");
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.StartupTimeout,
                "Worker process did not complete startup before the configured timeout.",
                exception);
        }
        catch (OperationCanceledException)
        {
            KillAndDispose(process, "LaunchCanceled");
            throw;
        }
        catch (Exception exception) when (exception is not WorkerProcessLaunchException)
        {
            KillAndDispose(process, "StartupFailed");
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.StartupFailed,
                "Worker process failed during startup.",
                exception);
        }
        catch (WorkerProcessLaunchException)
        {
            KillAndDispose(process, "StartupFailed");
            throw;
        }
    }

    private ProcessStartInfo CreateStartInfo(
        WorkerProcessLaunchRequest request,
        out WorkerProcessCommandLine commandLine)
    {
        string executablePath = ResolveExecutablePath();
        string workingDirectory = ResolveWorkingDirectory(executablePath);
        string[] arguments =
        [
            "--session-id",
            request.SessionId,
            "--pipe-name",
            request.PipeName,
            "--protocol-version",
            request.ProtocolVersion.ToString(System.Globalization.CultureInfo.InvariantCulture),
        ];

        ProcessStartInfo startInfo = new()
        {
            FileName = executablePath,
            WorkingDirectory = workingDirectory,
            UseShellExecute = false,
            CreateNoWindow = true,
            ErrorDialog = false,
        };

        foreach (string argument in arguments)
        {
            startInfo.ArgumentList.Add(argument);
        }

        startInfo.Environment[WorkerNonceEnvironmentVariableName] = request.Nonce;

        commandLine = new WorkerProcessCommandLine(executablePath, arguments);

        return startInfo;
    }

    private string ResolveExecutablePath()
    {
        string executablePath;
        try
        {
            executablePath = Path.GetFullPath(_workerOptions.ExecutablePath);
        }
        catch (Exception exception) when (exception is ArgumentException or NotSupportedException or PathTooLongException)
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidExecutable,
                "Worker executable path is not a valid filesystem path.",
                exception);
        }

        if (!string.Equals(Path.GetExtension(executablePath), ".exe", StringComparison.OrdinalIgnoreCase))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidExecutable,
                "Worker executable path must point to a .exe file.");
        }

        if (!File.Exists(executablePath))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.ExecutableNotFound,
                "Worker executable does not exist.");
        }

        WorkerExecutableValidator.Validate(executablePath, _workerOptions.RequiredArchitecture);

        return executablePath;
    }

    private string ResolveWorkingDirectory(string executablePath)
    {
        if (string.IsNullOrWhiteSpace(_workerOptions.WorkingDirectory))
        {
            return Path.GetDirectoryName(executablePath) ?? Environment.CurrentDirectory;
        }

        string workingDirectory;
        try
        {
            workingDirectory = Path.GetFullPath(_workerOptions.WorkingDirectory);
        }
        catch (Exception exception) when (exception is ArgumentException or NotSupportedException or PathTooLongException)
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidWorkingDirectory,
                "Worker working directory is not a valid filesystem path.",
                exception);
        }

        if (!Directory.Exists(workingDirectory))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidWorkingDirectory,
                "Worker working directory does not exist.");
        }

        return workingDirectory;
    }

    private void KillAndDispose(IWorkerProcess process, string reason)
    {
        try
        {
            if (!process.HasExited)
            {
                process.Kill(entireProcessTree: true);
                _metrics.WorkerKilled(reason);
            }
        }
        finally
        {
            process.Dispose();
        }
    }

    private static void ValidateRequest(WorkerProcessLaunchRequest request)
    {
        if (string.IsNullOrWhiteSpace(request.SessionId))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidRequest,
                "Worker launch requires a session id.");
        }

        if (string.IsNullOrWhiteSpace(request.PipeName))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidRequest,
                "Worker launch requires a pipe name.");
        }

        if (request.ProtocolVersion == 0)
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidRequest,
                "Worker launch requires a non-zero protocol version.");
        }

        if (string.IsNullOrWhiteSpace(request.Nonce))
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidRequest,
                "Worker launch requires a nonce.");
        }
    }
}