DodoSSH/src/DodoSSH.Client.Terminal/OutputBacklog.cs

namespace DodoSSH.Client.Terminal;

/// <summary>
/// One session's bounded queue of output waiting for the renderer, oldest first.
/// </summary>
/// <remarks>
/// <para>
/// ◆ <b>This is where a session's memory is actually bounded, and it exists because the obvious place
/// cannot do it.</b> The design used to be that <see cref="CreditWindow"/> bounded everything: the pump
/// stopped reading the SSH channel when the renderer fell behind, SSH's own receive window closed, and the
/// remote blocked on write. The middle step is not true of the library this client uses.
/// <c>Renci.SshNet.Channels.Channel.OnData</c> calls its window adjustment <em>before</em> raising
/// <c>DataReceived</c> and re-opens the channel window as soon as it drops below one packet, and
/// <c>ShellStream</c>'s own handler copies the bytes into a growable buffer without blocking and without a
/// ceiling. So a reader that stops reading does not slow the remote down by one byte; it only decides
/// whose buffer the flood accumulates in — and SSH.NET's grows by doubling and does not give the memory
/// back until the stream is disposed. A remote running <c>yes</c> took the process with it.
/// </para>
/// <para>
/// So the pump now reads continuously, which is what keeps SSH.NET's buffer empty, and the backlog it
/// reads into is this one, with a hard capacity. When the renderer cannot keep up the <em>oldest</em>
/// bytes go rather than the newest, because the newest are the ones somebody is waiting to see: the tail
/// of a flood is the prompt that comes back after it. Dropping is announced rather than silent — see
/// <see cref="Take"/> — and it is the honest end of the trade, since output arriving faster than the
/// terminal can draw it was going to scroll off a 5,000-line scrollback regardless.
/// </para>
/// <para>
/// The credit window is still here and still does its own job: it bounds what is in flight <em>to the
/// renderer</em>, which is the buffer inside the WebView that nothing on this side can see. What changed
/// is that it no longer pretends to bound what is in flight from the remote.
/// </para>
/// </remarks>
public sealed class OutputBacklog
{
    /// <summary>
    /// Default capacity: 4 MiB.
    /// </summary>
    /// <remarks>
    /// Sixteen times the credit window, so nothing an ordinary burst produces is ever dropped — a
    /// screenful of dense output is kilobytes, and a build that scrolls for a second is well under this.
    /// It is a per-session cost, and it is the whole of what a runaway remote can take, which is the
    /// number worth choosing against: twenty tabs all flooding at once is 80 MiB and survivable, where the
    /// same twenty unbounded was whatever the machine had.
    /// </remarks>
    public const int DefaultCapacityBytes = 4 * 1024 * 1024;

    private readonly Lock gate = new();
    private readonly Queue<byte[]> segments = new();
    private readonly int capacityBytes;

    /// <summary>Signalled when output arrives at an empty backlog, and when it is completed.</summary>
    private TaskCompletionSource available = CreateSignal();

    /// <summary>How much of the head segment has already been taken.</summary>
    private int head;

    private int buffered;
    private long droppedSinceTake;
    private long droppedTotal;
    private bool completed;

    /// <param name="capacityBytes">How many unsent bytes to hold before dropping the oldest.</param>
    public OutputBacklog(int capacityBytes = DefaultCapacityBytes)
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(capacityBytes, 1);

        this.capacityBytes = capacityBytes;
    }

    /// <summary>How many unsent bytes are held.</summary>
    public int Buffered
    {
        get
        {
            lock (gate)
            {
                return buffered;
            }
        }
    }

    /// <summary>How many bytes this backlog has discarded over its whole life.</summary>
    public long DroppedTotal
    {
        get
        {
            lock (gate)
            {
                return droppedTotal;
            }
        }
    }

    /// <summary>Whether the producer has said there will be no more output.</summary>
    public bool IsCompleted
    {
        get
        {
            lock (gate)
            {
                return completed;
            }
        }
    }

    /// <summary>
    /// Adds output, dropping the oldest already held if that puts the backlog over capacity.
    /// </summary>
    /// <remarks>
    /// Takes ownership of <paramref name="segment"/>; the caller must not write to it afterwards.
    /// Appending to a completed backlog does nothing rather than throwing, because a producer finding out
    /// it has been shut down is a race rather than a mistake.
    /// </remarks>
    /// <param name="segment">Bytes read from the remote, in order.</param>
    public void Append(byte[] segment)
    {
        ArgumentNullException.ThrowIfNull(segment);

        if (segment.Length == 0)
        {
            return;
        }

        lock (gate)
        {
            if (completed)
            {
                return;
            }

            var wasEmpty = buffered == 0;

            segments.Enqueue(segment);
            buffered += segment.Length;

            // Never the segment just added, however large it is. A read bigger than the whole capacity is
            // possible if somebody configures one, and a backlog that answered every take with nothing
            // would stall the session rather than bound it.
            while (buffered > capacityBytes && segments.Count > 1)
            {
                var oldest = segments.Dequeue();
                var discarded = oldest.Length - head;

                head = 0;
                buffered -= discarded;
                droppedSinceTake += discarded;
                droppedTotal += discarded;
            }

            // Only on the transition, because a waiter that found bytes here never captured the signal in
            // the first place — see WaitForOutputAsync, which checks under this same lock. Signalling on
            // every append would allocate a TaskCompletionSource per read of a session doing megabytes a
            // second, for a waiter that is not there.
            if (wasEmpty)
            {
                Signal();
            }
        }
    }

    /// <summary>
    /// Moves the oldest output held into <paramref name="destination"/>, as much of it as will fit.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Into a caller's span rather than into an array of its own, because this is the hot path: the pump
    /// writes straight into the frame it is about to send, so a terminal showing a gigabyte costs one copy
    /// of it here rather than two.
    /// </para>
    /// <para>
    /// <paramref name="dropped"/> is what makes a drop visible instead of silent. It counts the bytes
    /// discarded since the previous take, which is exactly the gap sitting immediately before the bytes
    /// being handed over now. The caller writes a line saying so, and it has to — the alternative is a
    /// terminal quietly showing a different transcript from the one the remote sent. It is reported even
    /// when nothing is taken, so a drop at the very end of a session is still announced.
    /// </para>
    /// </remarks>
    /// <param name="destination">Where the bytes go, oldest first.</param>
    /// <param name="dropped">Bytes discarded since the previous take, and zero on the ordinary path.</param>
    /// <returns>How many bytes were written, which is zero when nothing is held.</returns>
    public int TakeInto(Span<byte> destination, out long dropped)
    {
        lock (gate)
        {
            dropped = droppedSinceTake;
            droppedSinceTake = 0;

            var take = Math.Min(destination.Length, buffered);
            var offset = 0;

            while (offset < take)
            {
                var oldest = segments.Peek();
                var fromThis = Math.Min(oldest.Length - head, take - offset);

                oldest.AsSpan(head, fromThis).CopyTo(destination[offset..]);

                offset += fromThis;
                head += fromThis;
                buffered -= fromThis;

                if (head == oldest.Length)
                {
                    _ = segments.Dequeue();
                    head = 0;
                }
            }

            return take;
        }
    }

    /// <summary>Says there will be no more output, releasing anything waiting for some.</summary>
    public void Complete()
    {
        lock (gate)
        {
            completed = true;
            Signal();
        }
    }

    /// <summary>Waits until there is output to take, or until there will never be any.</summary>
    /// <param name="cancellationToken">Abandons the wait.</param>
    /// <returns><see langword="false"/> once the backlog is completed and empty.</returns>
    public async ValueTask<bool> WaitForOutputAsync(CancellationToken cancellationToken)
    {
        while (true)
        {
            Task signal;

            lock (gate)
            {
                if (buffered > 0)
                {
                    return true;
                }

                if (completed)
                {
                    return false;
                }

                signal = available.Task;
            }

            await signal.WaitAsync(cancellationToken).ConfigureAwait(false);
        }
    }

    /// <remarks>
    /// Completed inside the lock so a waiter cannot miss the transition, and replaced there too so the
    /// next wait is on a signal that has not already fired. Continuations run asynchronously, so nothing
    /// resumes while the lock is held.
    /// </remarks>
    private void Signal()
    {
        _ = available.TrySetResult();
        available = CreateSignal();
    }

    private static TaskCompletionSource CreateSignal() =>
        new(TaskCreationOptions.RunContinuationsAsynchronously);
}