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>
/// How one session ended, and the sentence that says so.
/// </summary>
/// <param name="End">Which of the three ways it was.</param>
/// <param name="Description">
/// What to show a person. Written into the pane by the renderer, and available to the shell beside the
/// <paramref name="End"/> it belongs to.
/// </param>
/// <remarks>
/// The pair travels together because either alone is not enough. The enum is what code branches on — whether
/// to offer a reconnect, whether a tab is worth keeping — and the sentence is the detail underneath it, which
/// for a lost connection is the only place the actual cause appears.
/// </remarks>
public sealed record TerminalSessionOutcome(SshSessionEnd End, string Description);

/// <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>
    /// How this session ended and what to say about it, once <see cref="RunAsync"/> has returned.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Meaningless before the run completes, and deliberately not nullable for it: a session that has not
    /// ended is not a session that ended in an unknown way, and every caller here already knows whether the
    /// run is over — the workspace awaits it before reading this.
    /// </para>
    /// <para>
    /// The same description that goes to the renderer in the <see cref="TerminalServerOpcode.SessionClosed"/>
    /// frame, so the line written into the pane and whatever the shell draws around the tab are the one
    /// sentence rather than two that can disagree.
    /// </para>
    /// </remarks>
    public TerminalSessionOutcome Outcome { get; private set; } =
        new(SshSessionEnd.ClosedLocally, "The session was closed.");

    /// <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);

        TerminalSessionOutcome outcome;
        try
        {
            await reader.ConfigureAwait(false);

            // The read loop finished of its own accord, which means the shell is over. Which of the two ways
            // it is over is not a question asked of the transport now — by now the server has very likely hung
            // up behind a clean exit, and asking would call `exit` a dropped connection — but of what the shell
            // latched at the moment it happened. See ISshShellSession.TransportFailure.
            outcome = session.TransportFailure is { Length: > 0 } failure
                ? new TerminalSessionOutcome(SshSessionEnd.ConnectionLost, failure)
                : new TerminalSessionOutcome(SshSessionEnd.RemoteClosed, "The remote closed the session.");
        }
        catch (OperationCanceledException)
        {
            outcome = new TerminalSessionOutcome(SshSessionEnd.ClosedLocally, "The session was closed.");
        }
        catch (Exception exception)
        {
            // A read that threw rather than returning zero. Whatever it was, it was not this client asking —
            // that arrives as the cancellation above — so it is the connection going out from under the read.
            outcome = new TerminalSessionOutcome(
                SshSessionEnd.ConnectionLost,
                session.TransportFailure is { Length: > 0 } failure ? failure : exception.Message);
        }

        // Before the drain and the last frame, so that anything reading this off a completed run — the
        // workspace's own announcement — cannot see the default.
        Outcome = outcome;

        // 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(outcome.Description),
                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);
    }
}