DodoSSH/src/DodoSSH.Client.Terminal/TerminalSessionPump.cs

using System.Buffers;
using System.Text;
using System.Threading.Channels;
using DodoSSH.Client.Ssh;

namespace DodoSSH.Client.Terminal;

/// <summary>Where terminal frames are sent.</summary>
/// <remarks>
/// Implementations must serialise sends: a WebSocket does not permit concurrent writes, and the pump
/// deliberately does not know whether its transport is a socket, a test double or something else.
/// </remarks>
public interface ITerminalTransport
{
    /// <summary>Sends one binary frame.</summary>
    ValueTask SendAsync(ReadOnlyMemory<byte> frame, CancellationToken cancellationToken);
}

/// <summary>Tuning for one session's output path.</summary>
public sealed class TerminalPumpOptions
{
    /// <summary>How many bytes to read from the channel at once.</summary>
    public int ReadBufferBytes { get; init; } = 32 * 1024;

    /// <summary>
    /// How long to accumulate output before sending it.
    /// </summary>
    /// <remarks>
    /// xterm cannot render faster than the display refreshes, so flushing more often than once a frame
    /// is work whose result is overwritten before anyone sees it. 16 ms is one frame at 60 Hz, and the
    /// added latency on an echoed keystroke is below the threshold of perception.
    /// </remarks>
    public TimeSpan FlushInterval { get; init; } = TimeSpan.FromMilliseconds(16);

    /// <summary>How far behind the renderer may fall, in bytes.</summary>
    public int WindowBytes { get; init; } = CreditWindow.DefaultWindowBytes;
}

/// <summary>
/// Moves bytes between one SSH shell channel and the renderer, under flow control.
/// </summary>
/// <remarks>
/// <para>
/// Credit is reserved before reading, never after. That ordering is the whole design: because the pump
/// cannot read more than the renderer has room for, the coalescing buffer is bounded by the credit
/// window rather than by how fast the remote can talk. Reserving after reading would leave an
/// unbounded queue between the socket and the screen, which is the failure mode this exists to
/// prevent.
/// </para>
/// <para>
/// When credit runs out the pump stops reading. SSH's own receive window then closes, the remote
/// <c>sshd</c> blocks on write, and the process producing output blocks in turn — backpressure all the
/// way to the source, with no custom protocol.
/// </para>
/// </remarks>
public sealed class TerminalSessionPump : IAsyncDisposable
{
    private readonly uint sessionId;
    private readonly ISshShellSession session;
    private readonly ITerminalTransport transport;
    private readonly TimeProvider clock;
    private readonly TerminalPumpOptions options;
    private readonly Channel<byte[]> pending = Channel.CreateUnbounded<byte[]>(
        new UnboundedChannelOptions { SingleReader = true, SingleWriter = true });

    private readonly CancellationTokenSource lifetime = new();
    private int disposed;

    /// <param name="sessionId">Identifies this session in every frame.</param>
    /// <param name="session">The shell channel.</param>
    /// <param name="transport">Where frames go.</param>
    /// <param name="clock">Time source, so the flush interval is testable.</param>
    /// <param name="options">Tuning, or null for the defaults.</param>
    public TerminalSessionPump(
        uint sessionId,
        ISshShellSession session,
        ITerminalTransport transport,
        TimeProvider clock,
        TerminalPumpOptions? options = null)
    {
        ArgumentNullException.ThrowIfNull(session);
        ArgumentNullException.ThrowIfNull(transport);
        ArgumentNullException.ThrowIfNull(clock);

        this.sessionId = sessionId;
        this.session = session;
        this.transport = transport;
        this.clock = clock;
        this.options = options ?? new TerminalPumpOptions();

        Credits = new CreditWindow(this.options.WindowBytes);
    }

    /// <summary>This session's flow-control window.</summary>
    public CreditWindow Credits { get; }

    /// <summary>Total bytes read from the remote, for the throughput harness.</summary>
    public long BytesRead { get; private set; }

    /// <summary>Total frames sent to the renderer, for the throughput harness.</summary>
    public long FramesSent { get; private set; }

    /// <summary>
    /// Runs until the remote closes the channel or the token is cancelled.
    /// </summary>
    /// <remarks>
    /// The read and flush loops are separate tasks because a read blocks until bytes arrive: combining
    /// them would mean output sitting unflushed until the next byte happened to show up, so a prompt
    /// would appear only after the user pressed a key.
    /// </remarks>
    public async Task RunAsync(CancellationToken cancellationToken)
    {
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(
            cancellationToken,
            lifetime.Token);

        await SendAsync(TerminalServerOpcode.SessionOpened, default, linked.Token).ConfigureAwait(false);

        var reader = ReadLoopAsync(linked.Token);
        var flusher = FlushLoopAsync(linked.Token);

        string reason;
        try
        {
            await reader.ConfigureAwait(false);
            reason = "The remote closed the session.";
        }
        catch (OperationCanceledException)
        {
            reason = "The session was closed.";
        }
        catch (Exception exception)
        {
            reason = exception.Message;
        }

        // Stop the flusher, but only after draining what the reader already produced — the last thing
        // a remote writes is often the most important, and dropping it makes a clean exit look like a
        // crash.
        pending.Writer.TryComplete();

        try
        {
            await flusher.ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Cancelled during shutdown; the drain below is best-effort anyway.
        }

        await SendAsync(
                TerminalServerOpcode.SessionClosed,
                Encoding.UTF8.GetBytes(reason),
                CancellationToken.None)
            .ConfigureAwait(false);
    }

    /// <summary>Forwards keystrokes to the remote.</summary>
    public ValueTask WriteInputAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken) =>
        session.WriteAsync(data, cancellationToken);

    /// <summary>Tells the remote the terminal was resized.</summary>
    public void Resize(TerminalSize size) => session.Resize(size);

    /// <summary>Returns credit for bytes the renderer reported rendering.</summary>
    public void Acknowledge(uint rendered) =>
        Credits.Return(rendered > int.MaxValue ? int.MaxValue : (int)rendered);

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref disposed, 1) == 1)
        {
            return;
        }

        await lifetime.CancelAsync().ConfigureAwait(false);

        // Unblocks anything waiting on credit that will now never be acknowledged.
        Credits.Reset();

        pending.Writer.TryComplete();
        lifetime.Dispose();

        await session.DisposeAsync().ConfigureAwait(false);
    }

    private async Task ReadLoopAsync(CancellationToken cancellationToken)
    {
        var buffer = ArrayPool<byte>.Shared.Rent(options.ReadBufferBytes);

        try
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                await Credits.WaitForCreditAsync(cancellationToken).ConfigureAwait(false);

                var granted = Credits.TryReserve(options.ReadBufferBytes);
                if (granted == 0)
                {
                    continue;
                }

                int read;
                try
                {
                    read = await session
                        .ReadAsync(buffer.AsMemory(0, granted), cancellationToken)
                        .ConfigureAwait(false);
                }
                catch
                {
                    Credits.Return(granted);
                    throw;
                }

                // Hand back what was reserved but not used, so a short read does not permanently
                // shrink the window.
                Credits.Return(granted - read);

                if (read == 0)
                {
                    return;
                }

                BytesRead += read;
                await pending.Writer.WriteAsync(buffer[..read], cancellationToken).ConfigureAwait(false);
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    private async Task FlushLoopAsync(CancellationToken cancellationToken)
    {
        var segments = new List<byte[]>();

        while (await pending.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
        {
            if (!pending.Reader.TryRead(out var first))
            {
                continue;
            }

            segments.Clear();
            segments.Add(first);

            // Coalesce for one frame's worth of time, then send everything at once. Whatever the
            // remote produced in that window becomes a single write to the terminal.
            await Task.Delay(options.FlushInterval, clock, cancellationToken).ConfigureAwait(false);

            while (pending.Reader.TryRead(out var more))
            {
                segments.Add(more);
            }

            await SendOutputAsync(segments, cancellationToken).ConfigureAwait(false);
        }

        // The channel completed. Anything the reader wrote before finishing still has to go out.
        segments.Clear();
        while (pending.Reader.TryRead(out var trailing))
        {
            segments.Add(trailing);
        }

        if (segments.Count > 0)
        {
            await SendOutputAsync(segments, CancellationToken.None).ConfigureAwait(false);
        }
    }

    private async ValueTask SendOutputAsync(List<byte[]> segments, CancellationToken cancellationToken)
    {
        var total = 0;
        foreach (var segment in segments)
        {
            total += segment.Length;
        }

        var frame = new byte[TerminalFrame.HeaderLength + total];
        TerminalFrame.Write(frame, (byte)TerminalServerOpcode.Output, sessionId, default);

        var offset = TerminalFrame.HeaderLength;
        foreach (var segment in segments)
        {
            segment.CopyTo(frame, offset);
            offset += segment.Length;
        }

        FramesSent++;
        await transport.SendAsync(frame, cancellationToken).ConfigureAwait(false);
    }

    private async ValueTask SendAsync(
        TerminalServerOpcode opcode,
        ReadOnlyMemory<byte> payload,
        CancellationToken cancellationToken)
    {
        FramesSent++;

        await transport
            .SendAsync(TerminalFrame.Create((byte)opcode, sessionId, payload.Span), cancellationToken)
            .ConfigureAwait(false);
    }
}