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

using System.Buffers;
using System.Globalization;
using System.Text;
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>
    /// How much unsent output to hold for a renderer that is behind, before dropping the oldest.
    /// </summary>
    /// <remarks>
    /// This is the number that bounds a session's memory, and <see cref="OutputBacklog"/> says at length
    /// why it rather than <see cref="WindowBytes"/> is the one that does. Larger drops less and shows a
    /// flood's tail later; smaller drops sooner and keeps the terminal closer to live.
    /// </remarks>
    public int BacklogBytes { get; init; } = OutputBacklog.DefaultCapacityBytes;
}

/// <summary>
/// Moves bytes between one SSH shell channel and the renderer, under flow control.
/// </summary>
/// <remarks>
/// <para>
/// ◆ <b>The pump reads the shell channel continuously and never stops, which is deliberate rather than
/// careless.</b> It used to stop as soon as the renderer fell behind, on the reasoning that a reader
/// which does not read closes SSH's receive window, blocks the remote <c>sshd</c> on write, and blocks
/// the process producing the output in turn — backpressure all the way to the source, with no custom
/// protocol. That reasoning is wrong for this client's SSH library, and wrong in the direction that costs
/// memory: SSH.NET re-opens the channel window when bytes <em>arrive</em> rather than when they are
/// consumed, and keeps everything it has been handed in a growable array it does not give back. Not
/// reading slowed nothing down; it only moved the flood into a buffer with no ceiling.
/// <see cref="OutputBacklog"/> carries the full account.
/// </para>
/// <para>
/// So there are two bounds now, each enforced where it can be. The backlog bounds what is waiting to be
/// sent, by dropping its oldest bytes and saying so. The credit window bounds what has been sent and not
/// yet drawn — the buffer inside the WebView that nothing on this side can measure — and the flush loop
/// reserves against it before every frame, holding off while there is none. Reading is no longer part of
/// either, because reading was never what the remote was waiting on.
/// </para>
/// </remarks>
public sealed class TerminalSessionPump : IAsyncDisposable
{
    /// <summary>
    /// The most output one frame carries.
    /// </summary>
    /// <remarks>
    /// Under the 85,000-byte threshold at which the runtime allocates on the large object heap, and that is
    /// the whole of why the number is here: a frame is allocated per send, and a session doing megabytes a
    /// second would otherwise put a quarter-megabyte array on the LOH sixty times a second, which is not
    /// compacted and is collected only with a generation 2 pass. Below the threshold the same traffic is
    /// generation 0 garbage that costs nothing to collect. It is a ceiling rather than a target — a quiet
    /// terminal sends whatever it has, which is a few bytes.
    /// </remarks>
    private const int MaximumFrameBytes = 64 * 1024;

    private readonly uint sessionId;
    private readonly ISshShellSession session;
    private readonly ITerminalTransport transport;
    private readonly TimeProvider clock;
    private readonly TerminalPumpOptions options;
    private readonly OutputBacklog backlog;

    private readonly CancellationTokenSource lifetime = new();

    /// <summary>One write to the remote at a time, in the order they were asked for.</summary>
    /// <remarks>
    /// The data plane hands each input frame over without waiting for the one before it — a blocked write
    /// must not hold up the acknowledgements of every other session on the socket — so without this, two
    /// frames' writes run at once. SSH.NET's <c>ShellStream</c> takes no lock around its write buffer, and its
    /// flush sends what is there and then empties it by length: bytes added mid-flush are discarded unsent.
    /// A snippet's RUN is its text and then a carriage return as two frames, which is exactly the shape that lost or
    /// reordered keystrokes. Waiters on a <see cref="SemaphoreSlim"/> are released in the order they arrived,
    /// and the data plane calls this from one receive loop, so arrival order is the order typed.
    /// </remarks>
    private readonly SemaphoreSlim inputGate = new(1, 1);

    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);
        backlog = new OutputBacklog(this.options.BacklogBytes);
    }

    /// <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 bytes this session discarded because the renderer could not keep up with the remote.
    /// </summary>
    /// <remarks>
    /// Zero for every ordinary session, and the number to look at first when somebody reports missing
    /// output. The user is told as it happens — a dropped run puts a line in the pane rather than a gap —
    /// so this is for a test or a report rather than for the interface.
    /// </remarks>
    public long BytesDropped => backlog.DroppedTotal;

    /// <summary>Bytes read from the remote and not yet sent to the renderer.</summary>
    public int BytesPending => backlog.Buffered;

    /// <summary>The last of what the remote wrote, for the phone's command history.</summary>
    /// <remarks>See <see cref="OutputTail"/> for why a client keeps any of it.</remarks>
    public OutputTail Recent { get; } = new();

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

        var outcome = await ReadUntilItEndsAsync(reader).ConfigureAwait(false);

        // 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.
        backlog.Complete();

        // ◆ AND RELEASE THE FLUSHER IF IT IS PARKED ON CREDIT, WHICH IS NOT THE SAME SIGNAL. Completing
        // the backlog wakes anything waiting for output; it says nothing to a flush loop that has output
        // and is waiting for the renderer to make room for it. Against a renderer that has stopped
        // acknowledging — the page's process killed, the window gone — that acknowledgement never comes,
        // and without this the await below would never return: the shell would have ended, the tab would
        // never be told, and the session would never be torn down. Resetting rather than waiting is
        // right because the window describes a renderer's remaining room and there is nothing left to
        // protect it from: what follows is one backlog's worth of tail and then the session is over.
        Credits.Reset();

        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>
    /// Waits for the read loop to finish, and says which of the three ways it did.
    /// </summary>
    /// <remarks>
    /// Split out of <see cref="RunAsync"/> for length rather than for reuse, and it is the half of that
    /// method with no side effects: three branches, one answer, nothing touched.
    /// </remarks>
    /// <param name="reader">The read loop's task.</param>
    private async Task<TerminalSessionOutcome> ReadUntilItEndsAsync(Task reader)
    {
        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.
            return session.TransportFailure is { Length: > 0 } failure
                ? new TerminalSessionOutcome(SshSessionEnd.ConnectionLost, failure)
                : new TerminalSessionOutcome(SshSessionEnd.RemoteClosed, "The remote closed the session.");
        }
        catch (OperationCanceledException)
        {
            return 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.
            return new TerminalSessionOutcome(
                SshSessionEnd.ConnectionLost,
                session.TransportFailure is { Length: > 0 } failure ? failure : exception.Message);
        }
    }

    /// <summary>Forwards keystrokes to the remote, after any still being written.</summary>
    public async ValueTask WriteInputAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
    {
        await inputGate.WaitAsync(cancellationToken).ConfigureAwait(false);

        try
        {
            await session.WriteAsync(data, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            inputGate.Release();
        }
    }

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

        backlog.Complete();
        lifetime.Dispose();

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

    /// <remarks>
    /// No credit is consulted here, and the class remark says at length why: the remote is not waiting on
    /// this loop, so pausing it buys nothing and costs a buffer inside SSH.NET that grows without bound.
    /// What this loop owes is to keep that buffer empty, which means reading as fast as the channel
    /// delivers and putting the result somewhere with a ceiling.
    /// </remarks>
    private async Task ReadLoopAsync(CancellationToken cancellationToken)
    {
        var buffer = ArrayPool<byte>.Shared.Rent(options.ReadBufferBytes);

        try
        {
            while (true)
            {
                // Thrown rather than tested in the loop condition, and the difference is the whole of how
                // this session's ending is described. RunAsync reads a read loop that RETURNS as
                // the remote having closed the channel, and one that throws OperationCanceledException as
                // this client having ended it. A condition on the token would leave by the first door for
                // a reason that belongs to the second, and a tab the user closed would report itself as a
                // shell that exited on its own.
                cancellationToken.ThrowIfCancellationRequested();

                var read = await session
                    .ReadAsync(buffer.AsMemory(0, options.ReadBufferBytes), cancellationToken)
                    .ConfigureAwait(false);

                if (read == 0)
                {
                    return;
                }

                BytesRead += read;
                Recent.Append(buffer.AsSpan(0, read));

                // Copied out of the rented buffer, because the backlog holds it until the renderer has
                // room and the next read is already about to overwrite this one.
                backlog.Append(buffer[..read]);
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    private async Task FlushLoopAsync(CancellationToken cancellationToken)
    {
        while (await backlog.WaitForOutputAsync(cancellationToken).ConfigureAwait(false))
        {
            if (backlog.IsCompleted)
            {
                // The reader has finished and what is left is the tail, which goes out regardless of
                // credit. Leaving by this door rather than going round again is load-bearing: with the
                // backlog complete, SendWhatCreditAllowsAsync declines to wait for an acknowledgement
                // that may never come, so a renderer that has stopped would put this loop into a spin —
                // output to send, no credit to send it with, and a wait for more output that returns at
                // once because there is some.
                break;
            }

            // Coalesce for one frame's worth of time before the first send, so a burst arriving at an idle
            // terminal becomes a single write. Under sustained output the loop below never comes back
            // here, and does not need to: credit is what paces it then, and a frame is as large as the
            // renderer has room for.
            await Task.Delay(options.FlushInterval, clock, cancellationToken).ConfigureAwait(false);

            await SendWhatCreditAllowsAsync(cancellationToken).ConfigureAwait(false);
        }

        await DrainAsync().ConfigureAwait(false);
    }

    /// <summary>Sends until the backlog is empty or the renderer has no room left.</summary>
    private async Task SendWhatCreditAllowsAsync(CancellationToken cancellationToken)
    {
        while (true)
        {
            var wanted = Math.Min(backlog.Buffered, MaximumFrameBytes);
            if (wanted == 0)
            {
                return;
            }

            var granted = Credits.TryReserve(wanted);

            if (granted == 0)
            {
                // The renderer is behind. Nothing more goes out until it acknowledges, and meanwhile the
                // backlog holds what it can and drops its oldest for what it cannot — which is the whole
                // of how this session's memory stays bounded.
                if (backlog.IsCompleted)
                {
                    // Except at the end, where waiting on an acknowledgement that may never come would
                    // hang the session's teardown behind a renderer that has already gone.
                    return;
                }

                await Credits.WaitForCreditAsync(cancellationToken).ConfigureAwait(false);
                continue;
            }

            var sent = await SendFromBacklogAsync(granted, cancellationToken).ConfigureAwait(false);

            if (granted > sent)
            {
                // Whatever was reserved and not used, so a backlog that shrank between the reservation
                // and the take does not permanently narrow the window. Guarded rather than called with a
                // zero, because returning credit wakes every waiter and replaces the signal they were
                // waiting on, and this is the hot path.
                Credits.Return(granted - sent);
            }
        }
    }

    /// <summary>
    /// Puts the last of the output out, regardless of credit.
    /// </summary>
    /// <remarks>
    /// The session is over by the time this runs, so there is no renderer left to protect from a flood —
    /// and what remains is bounded by the backlog's own capacity rather than by anything the remote can
    /// still do. Waiting on credit here instead would mean a shell that exited against a renderer which
    /// had gone never finished tearing down.
    /// </remarks>
    private async Task DrainAsync()
    {
        while (await SendFromBacklogAsync(MaximumFrameBytes, CancellationToken.None)
            .ConfigureAwait(false) > 0)
        {
            // Until it is empty.
        }
    }

    /// <summary>
    /// Moves up to <paramref name="maxBytes"/> out of the backlog and onto the wire.
    /// </summary>
    /// <returns>How many bytes of output were sent, which is zero once the backlog is empty.</returns>
    private async ValueTask<int> SendFromBacklogAsync(int maxBytes, CancellationToken cancellationToken)
    {
        // Taken straight into the frame, header space and all. A second buffer here would be an
        // allocation and a copy of every byte the terminal ever shows.
        var frame = new byte[TerminalFrame.HeaderLength + maxBytes];
        var taken = backlog.TakeInto(frame.AsSpan(TerminalFrame.HeaderLength), out var dropped);

        if (dropped > 0)
        {
            // Sent immediately before the bytes it explains, and under its own opcode — see
            // TerminalServerOpcode.OutputDropped for why this is not simply spliced into the output.
            await SendAsync(TerminalServerOpcode.OutputDropped, DroppedNotice(dropped), cancellationToken)
                .ConfigureAwait(false);
        }

        if (taken == 0)
        {
            return 0;
        }

        TerminalFrame.Write(frame, (byte)TerminalServerOpcode.Output, sessionId, default);

        FramesSent++;

        await transport
            .SendAsync(frame.AsMemory(0, TerminalFrame.HeaderLength + taken), cancellationToken)
            .ConfigureAwait(false);

        return taken;
    }

    /// <summary>
    /// The line a user sees where output was dropped.
    /// </summary>
    /// <remarks>
    /// <para>
    /// It opens with a reset rather than the colour, and that is the load-bearing part: the bytes were cut
    /// at an arbitrary offset, so whatever escape sequence the remote was in the middle of has been left
    /// half-said. Resetting first means the notice reads as a notice and the output after it starts from a
    /// known state, instead of inheriting a colour or a mode from a sequence nobody finished.
    /// </para>
    /// <para>
    /// Dim grey, matching the renderer's own banners for a session that closed and a view that
    /// reconnected — this is the same kind of thing, a sentence from the application rather than from the
    /// remote, and it should not read like part of the transcript.
    /// </para>
    /// </remarks>
    private static byte[] DroppedNotice(long dropped) =>
        Encoding.UTF8.GetBytes(string.Create(
            CultureInfo.InvariantCulture,
            $"\u001b[0m\r\n\u001b[38;5;244m── {Sized(dropped)} of output was dropped; "
            + $"the terminal could not keep up ──\u001b[0m\r\n"));

    /// <remarks>
    /// One decimal place from a megabyte up, because the number is an order of magnitude rather than an
    /// accounting: what a person needs from it is whether they lost a screenful or a session's worth.
    /// </remarks>
    private static string Sized(long bytes) => bytes switch
    {
        < 1024 => string.Create(CultureInfo.CurrentCulture, $"{bytes} bytes"),
        < 1024 * 1024 => string.Create(CultureInfo.CurrentCulture, $"{bytes / 1024.0:0.#} KB"),
        < 1024L * 1024 * 1024 => string.Create(CultureInfo.CurrentCulture, $"{bytes / (1024.0 * 1024):0.#} MB"),
        _ => string.Create(CultureInfo.CurrentCulture, $"{bytes / (1024.0 * 1024 * 1024):0.#} GB"),
    };

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