Stop reading SSH output the webview has not caught up with

core/src/output.rs

@@ -4,8 +4,14 @@
 //! to a full 32 KiB window. Delivering every chunk separately makes the fixed
 //! per-message cost of the UI transport dominate bulk output, so chunks are
 //! accumulated into batches instead.
+//!
+//! Delivered batches are also metered: [`OutputCredit`] counts the bytes the UI
+//! layer has not yet acknowledged, so a reader can stop pulling from the
+//! connection while the UI is behind.
 
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::time::{Duration, Instant};
+use tokio::sync::Notify;
 
 /// Longest a partially filled batch waits before delivery. Bounds the delay
 /// added to output that arrives while a batch is already in flight.
@@ -114,6 +120,98 @@ impl OutputBatcher {
     }
 }
 
+/// Unacknowledged output a session may have in flight before its reader stops
+/// pulling from the connection.
+///
+/// Large enough for many full [`MAX_BATCH`] batches to be on their way at once,
+/// so bulk output is not throttled to one round trip to the UI layer and back
+/// per batch; small enough that a UI that stops consuming, or a peer that
+/// produces faster than the UI can, holds a few megabytes per session rather
+/// than whatever the peer chooses to send.
+pub const OUTPUT_WINDOW: u64 = 2 * 1024 * 1024;
+
+/// Count of delivered output bytes the UI layer has not acknowledged yet.
+///
+/// The reader spends what it delivers and stops reading the connection while
+/// the count is at or above [`OUTPUT_WINDOW`]; the UI refunds what it has
+/// received. Pausing the reader is what bounds memory: the SSH library hands
+/// channel data over through a bounded queue and waits when it is full, so an
+/// idle reader stalls the connection and, through the SSH window and TCP, the
+/// peer itself.
+///
+/// The window is a threshold rather than a hard cap: a batch already due is
+/// still delivered, so the count can run past the window by up to one batch
+/// and one protocol packet.
+pub struct OutputCredit {
+    outstanding: AtomicU64,
+    /// Set once the session is closing, after which output is never held back.
+    lifted: AtomicBool,
+    room: Notify,
+}
+
+impl OutputCredit {
+    pub fn new() -> Self {
+        Self {
+            outstanding: AtomicU64::new(0),
+            lifted: AtomicBool::new(false),
+            room: Notify::new(),
+        }
+    }
+
+    /// Records `bytes` delivered to the UI layer.
+    pub fn spend(&self, bytes: usize) {
+        self.outstanding.fetch_add(bytes as u64, Ordering::AcqRel);
+    }
+
+    /// Records `bytes` the UI layer has received.
+    ///
+    /// Saturates at zero: the amount comes from the UI layer, and acknowledging
+    /// more than was delivered must not bank credit for output not yet sent.
+    pub fn refund(&self, bytes: u64) {
+        let _ = self
+            .outstanding
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |outstanding| {
+                Some(outstanding.saturating_sub(bytes))
+            });
+        self.room.notify_one();
+    }
+
+    /// Stops holding output back for good.
+    ///
+    /// A closing session must still reach the end of its stream, which is
+    /// where the reader flushes the trailing output and announces the close.
+    /// Nothing is obliged to keep acknowledging by then, so a reader left
+    /// waiting for credit might never get there.
+    pub fn lift(&self) {
+        self.lifted.store(true, Ordering::Release);
+        self.room.notify_one();
+    }
+
+    /// Whether the reader may pull more output from the connection.
+    pub fn has_room(&self) -> bool {
+        self.lifted.load(Ordering::Acquire)
+            || self.outstanding.load(Ordering::Acquire) < OUTPUT_WINDOW
+    }
+
+    /// Resolves once [`has_room`](Self::has_room) holds.
+    ///
+    /// Meant for a single waiter, the session's reader. `refund` and `lift`
+    /// wake it with `notify_one`, which leaves a permit behind when nobody is
+    /// waiting yet, so a refund that lands between the check and the wait is
+    /// not lost: the wait completes at once and the loop checks again.
+    pub async fn room(&self) {
+        while !self.has_room() {
+            self.room.notified().await;
+        }
+    }
+}
+
+impl Default for OutputCredit {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -315,4 +413,95 @@ mod tests {
         assert!(batcher.is_empty());
         assert_eq!(batcher.deadline(), None);
     }
+
+    #[test]
+    fn credit_runs_out_at_the_window_and_returns_with_a_refund() {
+        let credit = OutputCredit::new();
+        assert!(credit.has_room());
+        credit.spend(OUTPUT_WINDOW as usize - 1);
+        assert!(credit.has_room());
+        credit.spend(1);
+        assert!(!credit.has_room());
+        credit.refund(1);
+        assert!(credit.has_room());
+    }
+
+    #[test]
+    fn an_over_acknowledgement_banks_no_credit() {
+        let credit = OutputCredit::new();
+        credit.spend(MAX_BATCH);
+        credit.refund(u64::MAX);
+        // Had the excess carried over, this would still leave room.
+        credit.spend(OUTPUT_WINDOW as usize);
+        assert!(!credit.has_room());
+    }
+
+    #[test]
+    fn lifting_the_limit_leaves_room_whatever_is_outstanding() {
+        let credit = OutputCredit::new();
+        credit.spend(OUTPUT_WINDOW as usize * 4);
+        assert!(!credit.has_room());
+        credit.lift();
+        assert!(credit.has_room());
+        credit.spend(OUTPUT_WINDOW as usize);
+        assert!(credit.has_room());
+    }
+
+    #[tokio::test]
+    async fn waiting_for_room_returns_at_once_when_there_is_some() {
+        let credit = OutputCredit::new();
+        tokio::time::timeout(Duration::from_secs(1), credit.room())
+            .await
+            .expect("room was available");
+    }
+
+    #[tokio::test]
+    async fn a_refund_wakes_a_waiting_reader() {
+        let credit = std::sync::Arc::new(OutputCredit::new());
+        credit.spend(OUTPUT_WINDOW as usize);
+        let waiter = tokio::spawn({
+            let credit = credit.clone();
+            async move { credit.room().await }
+        });
+        tokio::task::yield_now().await;
+        assert!(!waiter.is_finished());
+
+        // A refund that leaves the count at the window wakes the reader, which
+        // must check again and keep waiting.
+        credit.refund(0);
+        tokio::time::sleep(Duration::from_millis(20)).await;
+        assert!(!waiter.is_finished());
+
+        credit.refund(MAX_BATCH as u64);
+        tokio::time::timeout(Duration::from_secs(1), waiter)
+            .await
+            .expect("the refund did not wake the reader")
+            .expect("waiter");
+    }
+
+    #[tokio::test]
+    async fn a_refund_before_the_wait_starts_is_not_missed() {
+        let credit = OutputCredit::new();
+        credit.spend(OUTPUT_WINDOW as usize);
+        credit.refund(1);
+        tokio::time::timeout(Duration::from_secs(1), credit.room())
+            .await
+            .expect("the earlier refund was lost");
+    }
+
+    #[tokio::test]
+    async fn lifting_the_limit_wakes_a_waiting_reader() {
+        let credit = std::sync::Arc::new(OutputCredit::new());
+        credit.spend(OUTPUT_WINDOW as usize);
+        let waiter = tokio::spawn({
+            let credit = credit.clone();
+            async move { credit.room().await }
+        });
+        tokio::task::yield_now().await;
+        credit.lift();
+        tokio::time::timeout(Duration::from_secs(1), waiter)
+            .await
+            .expect("lifting did not wake the reader")
+            .expect("waiter");
+    }
 }

core/src/session.rs

@@ -1,7 +1,7 @@
 use crate::connection;
 use crate::host::Host;
 use crate::known_hosts::KnownHosts;
-use crate::output::OutputBatcher;
+use crate::output::{OutputBatcher, OutputCredit};
 use crate::vault::Vault;
 use crate::{CoreError, Result};
 use russh::client::Handle;
@@ -129,6 +129,9 @@ struct SshSession {
     /// its own input.
     writes: Mutex<()>,
     resize_tx: mpsc::Sender<(u32, u32)>,
+    /// Output delivered to the UI layer and not yet acknowledged, shared with
+    /// the reader task, which stops reading the channel while it is spent.
+    credit: Arc<OutputCredit>,
     /// Held so the id stays reserved while this entry exists.
     #[allow(dead_code)]
     claim: Arc<ClaimedId>,
@@ -144,8 +147,13 @@ impl SshSession {
     /// Mark the session closed and say goodbye to the peer. The goodbye holds
     /// only the connection handle, so the id is released as soon as the map
     /// entry and the reader are gone, however long the peer takes.
+    ///
+    /// The output limit is lifted as well: a reader stalled on unacknowledged
+    /// output has to keep reading to reach the end of the channel, flush the
+    /// tail and announce the close.
     fn retire(&self) -> impl std::future::Future<Output = ()> + Send + 'static {
         self.closed.send_replace(true);
+        self.credit.lift();
         let handle = self.handle.clone();
         async move { disconnect(&handle).await }
     }
@@ -243,6 +251,7 @@ impl SessionManager {
         }
 
         let (resize_tx, mut resize_rx) = mpsc::channel::<(u32, u32)>(32);
+        let credit = Arc::new(OutputCredit::new());
 
         let held = claim.commit();
         let session = Arc::new(SshSession {
@@ -251,6 +260,7 @@ impl SessionManager {
             channel_id,
             writes: Mutex::new(()),
             resize_tx,
+            credit: credit.clone(),
             claim: held.clone(),
         });
 
@@ -271,16 +281,28 @@ impl SessionManager {
             let mut batcher = OutputBatcher::new(Instant::now());
             let flush = |batcher: &mut OutputBatcher| {
                 if !batcher.is_empty() {
-                    data_cb(&sid, batcher.batch());
+                    let batch = batcher.batch();
+                    // Spent before the batch is handed over, so an
+                    // acknowledgement can never arrive ahead of the spend
+                    // it pays back.
+                    credit.spend(batch.len());
+                    data_cb(&sid, batch);
                     batcher.mark_flushed(Instant::now());
                 }
             };
             loop {
                 // Read the deadline before the select so no branch borrows the
                 // batcher while another branch's handler mutates it.
                 let flush_at = batcher.deadline();
+                // While the UI layer is a window behind, the channel is left
+                // unread, which backs the stall up to the peer. Resizes and a
+                // batch already due are still handled meanwhile; the end of the
+                // stream is seen once credit returns or the session is closed,
+                // which lifts the limit.
+                let reading = credit.has_room();
                 tokio::select! {
-                    msg = channel.wait() => {
+                    () = credit.room(), if !reading => {}
+                    msg = channel.wait(), if reading => {
                         match msg {
                             Some(ChannelMsg::Data { data })
                             | Some(ChannelMsg::ExtendedData { data, .. }) => {
@@ -373,7 +395,9 @@ impl SessionManager {
     /// The entry leaves the map before the disconnect is sent, so the session
     /// stops being addressable at once even if the peer never answers. The
     /// entry's share of the id goes with it; the reader task still holds one
-    /// until it has flushed and announced the close.
+    /// until it has flushed and announced the close. Output credit is lifted
+    /// first, so a reader waiting on the UI layer goes on to the end of the
+    /// stream instead of waiting on acknowledgements that may never come.
     pub async fn close(&self, session_id: &str) -> Result<()> {
         let removed = self.sessions.lock().await.remove(session_id);
         if let Some(session) = removed {
@@ -382,6 +406,17 @@ impl SessionManager {
         Ok(())
     }
 
+    /// Return `bytes` of output credit to a session once the UI layer has
+    /// received them, letting a reader that ran out resume.
+    ///
+    /// An id with no live session is not an error: acknowledgements trail the
+    /// output they cover and can arrive after the session has closed.
+    pub async fn ack(&self, session_id: &str, bytes: u64) {
+        if let Ok(session) = self.get(session_id).await {
+            session.credit.refund(bytes);
+        }
+    }
+
     /// List active session ids.
     pub async fn list(&self) -> Vec<String> {
         self.sessions.lock().await.keys().cloned().collect()