tyx_AI_xhs/supervisor/Dada.Supervisor/ManagedChildProcess.cs

using System.Diagnostics;
using System.IO.Pipes;
using System.Text;

namespace Dada.Supervisor;

internal sealed class ManagedChildProcess : IAsyncDisposable
{
    private readonly NamedPipeServerStream controlPipe;
    private readonly StreamWriter controlWriter;
    private bool stopping;

    private ManagedChildProcess(Process process, NamedPipeServerStream controlPipe, StreamWriter controlWriter)
    {
        Process = process;
        this.controlPipe = controlPipe;
        this.controlWriter = controlWriter;
    }

    internal Process Process { get; }
    internal bool IsStopping => stopping;
    internal static TimeSpan ShutdownDeadline { get; } = TimeSpan.FromSeconds(15);

    internal static async Task<ManagedChildProcess> StartAsync(
        ProcessStartInfo startInfo,
        ChildRole role,
        ICredentialStore credentialStore,
        CancellationToken cancellationToken = default)
    {
        var pipeName = $"Dada.P0A.Control.{role}.{Guid.NewGuid():N}";
        var pipe = new NamedPipeServerStream(
            pipeName,
            PipeDirection.InOut,
            1,
            PipeTransmissionMode.Byte,
            PipeOptions.Asynchronous);
        startInfo.ArgumentList.Add("--dada-control-pipe");
        startInfo.ArgumentList.Add(pipeName);
        startInfo.ArgumentList.Add("--dada-credential-stdin");
        Process? process = null;
        try
        {
            process = await CredentialProcessLauncher.StartAsync(startInfo, role, credentialStore, cancellationToken);
            using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            timeout.CancelAfter(TimeSpan.FromSeconds(10));
            await pipe.WaitForConnectionAsync(timeout.Token);
            var reader = new StreamReader(pipe, Encoding.UTF8, false, leaveOpen: true);
            var ready = await reader.ReadLineAsync(timeout.Token);
            if (!string.Equals(ready, "ready", StringComparison.Ordinal)) throw new InvalidOperationException("Managed child did not report ready.");
            var writer = new StreamWriter(pipe, new UTF8Encoding(false), leaveOpen: true) { AutoFlush = true };
            return new ManagedChildProcess(process, pipe, writer);
        }
        catch
        {
            pipe.Dispose();
            if (process is { HasExited: false }) process.Kill(entireProcessTree: true);
            process?.Dispose();
            throw;
        }
    }

    internal async Task StopAsync(TimeSpan? timeout = null, CancellationToken cancellationToken = default)
    {
        if (stopping || Process.HasExited) return;
        stopping = true;
        await controlWriter.WriteLineAsync("shutdown".AsMemory(), cancellationToken);
        var exited = Process.WaitForExitAsync(cancellationToken);
        var deadline = Task.Delay(timeout ?? ShutdownDeadline, cancellationToken);
        if (await Task.WhenAny(exited, deadline) != exited && !Process.HasExited)
        {
            Process.Kill(entireProcessTree: true);
            await Process.WaitForExitAsync(cancellationToken);
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (!Process.HasExited) await StopAsync();
        controlWriter.Dispose();
        controlPipe.Dispose();
        Process.Dispose();
    }
}

internal sealed class ManagedComponentSupervisor : IAsyncDisposable
{
    private readonly Func<CancellationToken, Task<ManagedChildProcess>> start;
    private readonly SupervisorState degradedState;
    private ManagedChildProcess? child;
    private int completedRestarts;
    private DateTimeOffset firstFailure;
    private bool stopping;

    internal ManagedComponentSupervisor(Func<CancellationToken, Task<ManagedChildProcess>> start, SupervisorState degradedState)
    {
        this.start = start;
        this.degradedState = degradedState;
    }

    internal event Action<SupervisorState>? StateChanged;

    internal async Task StartAsync(CancellationToken cancellationToken = default)
    {
        child = await start(cancellationToken);
        child.Process.EnableRaisingEvents = true;
        child.Process.Exited += OnChildExited;
    }

    private async void OnChildExited(object? sender, EventArgs eventArgs)
    {
        if (stopping || child?.IsStopping == true) return;
        if (completedRestarts == 0) firstFailure = DateTimeOffset.UtcNow;
        if (!RestartPolicy.CanRestart(completedRestarts, DateTimeOffset.UtcNow - firstFailure))
        {
            StateChanged?.Invoke(degradedState);
            return;
        }
        var delay = RestartPolicy.Delays[completedRestarts++];
        await Task.Delay(delay);
        try
        {
            child = await start(CancellationToken.None);
            child.Process.EnableRaisingEvents = true;
            child.Process.Exited += OnChildExited;
        }
        catch
        {
            OnChildExited(null, EventArgs.Empty);
        }
    }

    public async ValueTask DisposeAsync()
    {
        stopping = true;
        if (child is not null) await child.DisposeAsync();
    }
}