mxaccessgw/src/MxGateway.Worker/Ipc/WorkerPipeClient.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Pipes;
using System.Threading;
using System.Threading.Tasks;
using MxGateway.Worker.Bootstrap;
using Polly;
using Polly.Retry;

namespace MxGateway.Worker.Ipc;

public sealed class WorkerPipeClient : IWorkerPipeClient
{
    public const int DefaultConnectTimeoutMilliseconds = 30000;
    public const int DefaultConnectAttemptTimeoutMilliseconds = 2000;
    public const string ConnectAttemptTimeoutEnvironmentVariableName =
        "MXGATEWAY_WORKER_PIPE_CONNECT_ATTEMPT_TIMEOUT_MS";

    private readonly int _connectTimeoutMilliseconds;
    private readonly int _connectAttemptTimeoutMilliseconds;
    private readonly Func<Stream, WorkerFrameProtocolOptions, IWorkerLogger?, WorkerPipeSession> _sessionFactory;
    private readonly IWorkerLogger? _logger;

    public WorkerPipeClient()
        : this(null, DefaultConnectTimeoutMilliseconds)
    {
    }

    public WorkerPipeClient(IWorkerLogger? logger)
        : this(logger, DefaultConnectTimeoutMilliseconds)
    {
    }

    public WorkerPipeClient(int connectTimeoutMilliseconds)
        : this(null, connectTimeoutMilliseconds)
    {
    }

    public WorkerPipeClient(
        int connectTimeoutMilliseconds,
        Func<Stream, WorkerFrameProtocolOptions, WorkerPipeSession> sessionFactory)
        : this(
            null,
            connectTimeoutMilliseconds,
            ResolveDefaultConnectAttemptTimeoutMilliseconds(),
            (stream, frameOptions, _) => sessionFactory(stream, frameOptions))
    {
    }

    public WorkerPipeClient(
        IWorkerLogger? logger,
        int connectTimeoutMilliseconds)
        : this(
            logger,
            connectTimeoutMilliseconds,
            ResolveDefaultConnectAttemptTimeoutMilliseconds(),
            (stream, frameOptions, workerLogger) => new WorkerPipeSession(stream, frameOptions, workerLogger))
    {
    }

    public WorkerPipeClient(
        IWorkerLogger? logger,
        int connectTimeoutMilliseconds,
        Func<Stream, WorkerFrameProtocolOptions, IWorkerLogger?, WorkerPipeSession> sessionFactory)
        : this(
            logger,
            connectTimeoutMilliseconds,
            ResolveDefaultConnectAttemptTimeoutMilliseconds(),
            sessionFactory)
    {
    }

    public WorkerPipeClient(
        IWorkerLogger? logger,
        int connectTimeoutMilliseconds,
        int connectAttemptTimeoutMilliseconds,
        Func<Stream, WorkerFrameProtocolOptions, IWorkerLogger?, WorkerPipeSession> sessionFactory)
    {
        if (connectTimeoutMilliseconds <= 0)
        {
            throw new ArgumentOutOfRangeException(
                nameof(connectTimeoutMilliseconds),
                "Worker pipe connect timeout must be greater than zero.");
        }

        if (connectAttemptTimeoutMilliseconds <= 0)
        {
            throw new ArgumentOutOfRangeException(
                nameof(connectAttemptTimeoutMilliseconds),
                "Worker pipe connect attempt timeout must be greater than zero.");
        }

        _logger = logger;
        _sessionFactory = sessionFactory ?? throw new ArgumentNullException(nameof(sessionFactory));
        _connectTimeoutMilliseconds = connectTimeoutMilliseconds;
        _connectAttemptTimeoutMilliseconds = connectAttemptTimeoutMilliseconds;
    }

    public async Task RunAsync(
        WorkerOptions options,
        CancellationToken cancellationToken = default)
    {
        if (options is null)
        {
            throw new ArgumentNullException(nameof(options));
        }

        WorkerFrameProtocolOptions frameOptions = new(options);

        using NamedPipeClientStream pipe = await ConnectWithRetryAsync(options.PipeName, cancellationToken)
            .ConfigureAwait(false);

        WorkerPipeSession session = _sessionFactory(pipe, frameOptions, _logger);
        await session.RunAsync(cancellationToken).ConfigureAwait(false);
    }

    private async Task<NamedPipeClientStream> ConnectWithRetryAsync(
        string pipeName,
        CancellationToken cancellationToken)
    {
        int retryAttempts = Math.Max(
            0,
            (_connectTimeoutMilliseconds / Math.Min(_connectTimeoutMilliseconds, _connectAttemptTimeoutMilliseconds)) - 1);

        ResiliencePipeline<NamedPipeClientStream> pipeline = new ResiliencePipelineBuilder<NamedPipeClientStream>()
            .AddRetry(new RetryStrategyOptions<NamedPipeClientStream>
            {
                MaxRetryAttempts = retryAttempts,
                BackoffType = DelayBackoffType.Exponential,
                UseJitter = true,
                Delay = TimeSpan.FromMilliseconds(250),
                MaxDelay = TimeSpan.FromSeconds(2),
                ShouldHandle = new PredicateBuilder<NamedPipeClientStream>()
                    .Handle<Exception>(exception => exception is TimeoutException or IOException),
                OnRetry = args =>
                {
                    args.Outcome.Result?.Dispose();
                    _logger?.Information(
                        "WorkerPipeConnectRetry",
                        new Dictionary<string, object?>
                        {
                            ["attempt"] = args.AttemptNumber + 1,
                            ["pipe_name"] = pipeName,
                        });
                    return default;
                },
            })
            .Build();

        using CancellationTokenSource connectDeadline =
            CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        connectDeadline.CancelAfter(_connectTimeoutMilliseconds);

        try
        {
            return await pipeline.ExecuteAsync(
                    async token => await ConnectSingleAttemptAsync(pipeName, token).ConfigureAwait(false),
                    connectDeadline.Token)
                .ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            throw new TimeoutException(
                $"Worker pipe {pipeName} did not connect within {_connectTimeoutMilliseconds}ms.");
        }
    }

    private async Task<NamedPipeClientStream> ConnectSingleAttemptAsync(
        string pipeName,
        CancellationToken cancellationToken)
    {
        NamedPipeClientStream pipe = new(
            ".",
            pipeName,
            PipeDirection.InOut,
            PipeOptions.Asynchronous);

        try
        {
            using CancellationTokenSource attemptTimeout =
                CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            attemptTimeout.CancelAfter(_connectAttemptTimeoutMilliseconds);

            await Task.Run(
                    () =>
                    {
                        attemptTimeout.Token.ThrowIfCancellationRequested();
                        pipe.Connect(_connectAttemptTimeoutMilliseconds);
                    },
                    attemptTimeout.Token)
                .ConfigureAwait(false);

            return pipe;
        }
        catch
        {
            pipe.Dispose();
            throw;
        }
    }

    private static int ResolveDefaultConnectAttemptTimeoutMilliseconds()
    {
        string? configuredValue = Environment.GetEnvironmentVariable(ConnectAttemptTimeoutEnvironmentVariableName);
        return int.TryParse(configuredValue, out int milliseconds) && milliseconds > 0
            ? milliseconds
            : DefaultConnectAttemptTimeoutMilliseconds;
    }
}