fix(neshub): a pause belongs to its session, not to the encoder

The hub keeps the set of sessions pausing each encoder. An encoder is paused
while the set is not empty; a resume goes only to the encoder that session
paused; a session that drops, resubscribes or sees the primary picture move
releases its pause.

apps/neshub/src/e2e_tests.rs

@@ -577,4 +577,39 @@ mod windows {
         assert_eq!(got.len(), 3);
         assert!(got.iter().all(|g| (g.0, g.1) == (0, 0)));
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn a_session_that_goes_while_paused_releases_its_pause() {
+        let rig = rig().await;
+        tell(&rig, |m| encode_event_message(m, &opened(0, "launch_a"))).await;
+        tell(&rig, |m| encode_window_pid(m, 0, GAME_A)).await;
+        draw(&rig, GAME_A, 0, true).await;
+        let _ = received(&rig.client).await;
+        let me = rig._client_ep.addr().id;
+        let handle = rig.mgr.pause_handle_of(&me).await.expect("a session");
+        handle.pause(GAME_A);
+        assert!(rig.mgr.encoder_paused(GAME_A));
+        let controller = tokio::sync::Mutex::new(Controller::new(Limits::new(4_000)));
+        let (cmd, _rx) = tokio::sync::mpsc::unbounded_channel();
+        rig.mgr.remove_session(&me, &controller, &cmd).await;
+        assert!(!rig.mgr.encoder_paused(GAME_A), "the pause went with the session");
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn a_resubscribe_while_paused_releases_the_old_encoder() {
+        let rig = rig().await;
+        tell(&rig, |m| encode_event_message(m, &opened(0, "launch_a"))).await;
+        tell(&rig, |m| encode_event_message(m, &opened(1, "launch_b"))).await;
+        tell(&rig, |m| encode_window_pid(m, 0, GAME_A)).await;
+        tell(&rig, |m| encode_window_pid(m, 1, GAME_B)).await;
+        draw(&rig, GAME_A, 0, true).await;
+        let _ = received(&rig.client).await;
+        let me = rig._client_ep.addr().id;
+        let handle = rig.mgr.pause_handle_of(&me).await.expect("a session");
+        handle.pause(GAME_A);
+        assert!(rig.mgr.encoder_paused(GAME_A));
+        let (idr, _rx) = tokio::sync::mpsc::unbounded_channel();
+        rig.mgr.subscribe(me, 1, &idr).await;
+        assert!(!rig.mgr.encoder_paused(GAME_A), "no longer shown, no longer held");
+    }
 }

apps/neshub/src/main.rs

@@ -13,6 +13,7 @@ mod session;
 mod ticket;
 mod windows;
 mod capture_cmd;
+mod pauses;
 
 use std::net::{Ipv4Addr, SocketAddrV4};
 use std::path::PathBuf;

apps/neshub/src/pauses.rs

@@ -0,0 +1,193 @@
+//! Who has paused which encoder, ref(d-0087).
+//!
+//! A pause is a client's, not an encoder's: two clients can be shown the same
+//! window, and one pausing it must not freeze the other's picture. So the hub
+//! keeps, for each encoder, the set of sessions pausing it, and the encoder is
+//! stopped only while that set is not empty. A session leaving (its
+//! connection ending, a resubscribe, the primary picture moving) releases its
+//! pause, so a pause can never outlive the client that asked for it.
+//!
+//! Keyed by the session's gamepad number, which is never reused within a hub's
+//! lifetime, so a client that reconnects is a new session and cannot release
+//! the pause of the one it replaced, or the other way round.
+
+use std::collections::{BTreeMap, BTreeSet};
+use std::sync::{Arc, Mutex};
+
+use tokio::sync::mpsc::UnboundedSender;
+
+use crate::capture_cmd::{window_pause, window_resume};
+
+#[derive(Debug, Default)]
+pub struct PauseBook {
+    /// The encoder each pausing session paused. One at a time per session.
+    paused_by: BTreeMap<u32, u32>,
+}
+
+impl PauseBook {
+    fn users(&self, pid: u32) -> BTreeSet<u32> {
+        self.paused_by
+            .iter()
+            .filter(|(_, p)| **p == pid)
+            .map(|(s, _)| *s)
+            .collect()
+    }
+
+    /// `session` pauses `pid`. The commands to send, in order.
+    fn pause(&mut self, session: u32, pid: u32) -> Vec<Vec<u8>> {
+        let mut out = Vec::new();
+        if self.paused_by.get(&session) == Some(&pid) {
+            return out;
+        }
+        // Moving from one encoder to another releases the first.
+        out.extend(self.release(session));
+        let first = self.users(pid).is_empty();
+        self.paused_by.insert(session, pid);
+        if first {
+            out.push(window_pause(Some(pid)));
+        }
+        out
+    }
+
+    /// `session` stops pausing. Resumes the encoder it paused, and only that
+    /// one, and only when no other session still pauses it.
+    fn release(&mut self, session: u32) -> Vec<Vec<u8>> {
+        let Some(pid) = self.paused_by.remove(&session) else {
+            return Vec::new();
+        };
+        if self.users(pid).is_empty() {
+            vec![window_resume(Some(pid))]
+        } else {
+            Vec::new()
+        }
+    }
+
+    pub fn paused(&self, pid: u32) -> bool {
+        !self.users(pid).is_empty()
+    }
+}
+
+/// One session's way of pausing and resuming, shared with its control reader.
+#[derive(Clone)]
+pub struct PauseHandle {
+    book: Arc<Mutex<PauseBook>>,
+    session: u32,
+    cmd_tx: UnboundedSender<Vec<u8>>,
+}
+
+impl PauseHandle {
+    pub fn new(book: Arc<Mutex<PauseBook>>, session: u32, cmd_tx: UnboundedSender<Vec<u8>>) -> Self {
+        Self { book, session, cmd_tx }
+    }
+
+    fn send(&self, cmds: Vec<Vec<u8>>) {
+        for c in cmds {
+            let _ = self.cmd_tx.send(c);
+        }
+    }
+
+    /// Pause `pid` for this session. `pid` is the encoder it is shown, which
+    /// is never 0 here: the caller does not pause what it cannot name.
+    pub fn pause(&self, pid: u32) {
+        let cmds = self.book.lock().unwrap_or_else(|e| e.into_inner()).pause(self.session, pid);
+        self.send(cmds);
+    }
+
+    /// Resume what this session paused. Never anything else, and nothing at
+    /// all when it paused nothing.
+    pub fn resume(&self) {
+        self.release();
+    }
+
+    /// Let go of whatever this session paused: it is leaving, or looking at
+    /// another picture.
+    pub fn release(&self) {
+        let cmds = self.book.lock().unwrap_or_else(|e| e.into_inner()).release(self.session);
+        self.send(cmds);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn rig() -> (Arc<Mutex<PauseBook>>, UnboundedSender<Vec<u8>>, tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>) {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        (Arc::new(Mutex::new(PauseBook::default())), tx, rx)
+    }
+
+    fn sent(rx: &mut tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>) -> Vec<Vec<u8>> {
+        std::iter::from_fn(|| rx.try_recv().ok()).collect()
+    }
+
+    #[test]
+    fn one_of_two_sessions_pausing_a_shared_encoder_pauses_it_until_both_let_go() {
+        let (book, tx, mut rx) = rig();
+        let a = PauseHandle::new(book.clone(), 1, tx.clone());
+        let b = PauseHandle::new(book.clone(), 2, tx);
+        a.pause(10);
+        assert_eq!(sent(&mut rx), vec![window_pause(Some(10))]);
+        // B pauses too: the encoder is already stopped, nothing to send.
+        b.pause(10);
+        assert!(sent(&mut rx).is_empty());
+        // A resumes, B still pauses: the encoder stays stopped.
+        a.resume();
+        assert!(sent(&mut rx).is_empty());
+        assert!(book.lock().unwrap().paused(10));
+        b.resume();
+        assert_eq!(sent(&mut rx), vec![window_resume(Some(10))]);
+        assert!(!book.lock().unwrap().paused(10));
+    }
+
+    #[test]
+    fn a_session_that_never_paused_resumes_nothing() {
+        let (book, tx, mut rx) = rig();
+        let a = PauseHandle::new(book.clone(), 1, tx.clone());
+        let b = PauseHandle::new(book, 2, tx);
+        a.pause(10);
+        sent(&mut rx);
+        // B, shown the same window, did not pause it: its resume must not
+        // start the encoder A stopped.
+        b.resume();
+        b.release();
+        assert!(sent(&mut rx).is_empty());
+    }
+
+    #[test]
+    fn a_session_that_goes_while_paused_resumes_the_encoder() {
+        let (book, tx, mut rx) = rig();
+        let a = PauseHandle::new(book.clone(), 1, tx);
+        a.pause(10);
+        sent(&mut rx);
+        a.release(); // what its drop does
+        assert_eq!(sent(&mut rx), vec![window_resume(Some(10))]);
+        // Releasing twice is not resuming twice.
+        a.release();
+        assert!(sent(&mut rx).is_empty());
+    }
+
+    #[test]
+    fn a_resubscribe_while_paused_releases_the_old_encoder_only() {
+        let (book, tx, mut rx) = rig();
+        let a = PauseHandle::new(book.clone(), 1, tx.clone());
+        let other = PauseHandle::new(book.clone(), 2, tx);
+        a.pause(10);
+        other.pause(20);
+        sent(&mut rx);
+        a.release(); // restart_picture on moving to another window
+        assert_eq!(sent(&mut rx), vec![window_resume(Some(10))]);
+        assert!(book.lock().unwrap().paused(20), "someone else's pause is untouched");
+    }
+
+    #[test]
+    fn pausing_another_encoder_moves_the_pause() {
+        let (book, tx, mut rx) = rig();
+        let a = PauseHandle::new(book, 1, tx);
+        a.pause(10);
+        a.pause(20);
+        assert_eq!(
+            sent(&mut rx),
+            vec![window_pause(Some(10)), window_resume(Some(10)), window_pause(Some(20))]
+        );
+    }
+}

apps/neshub/src/session.rs

@@ -20,7 +20,7 @@ use nesprotocol::{
 use crate::control::{Controller, PathView};
 use crate::negotiation::{Negotiator, effective_caps};
 
-use crate::capture_cmd::{keyframe_request, window_pause, window_resume};
+use crate::capture_cmd::keyframe_request;
 use crate::dgram::{MediaItem, run_datagram_writer};
 use crate::windows::{Source, Subscription, WindowTable, route, wants_cursor};
 use nesprotocol::window::{MSG_WINDOW_SUBSCRIBE, STREAM_WINDOWS, WINDOW_DEFAULT, WindowEvent, WindowId};
@@ -106,6 +106,9 @@ pub struct ClientSession {
     /// The encoder whose frames this client was last sent (0: none yet), so
     /// its keyframe requests go to that encoder alone.
     encoder_pid: Arc<AtomicU32>,
+    /// This session's pause of an encoder, if any (d-0087); released when it
+    /// goes or looks at another picture.
+    pause: crate::pauses::PauseHandle,
     primary_window: Arc<AtomicU32>,
     /// Where a subscribe request goes: the manager, which holds the windows.
     subscribe_tx: tokio::sync::mpsc::UnboundedSender<(iroh::EndpointId, WindowId)>,
@@ -150,6 +153,8 @@ impl Drop for ClientSession {
         for task in &self.tasks {
             task.abort();
         }
+        // A pause is this client's, and does not outlive it.
+        self.pause.release();
     }
 }
 
@@ -169,6 +174,7 @@ impl ClientSession {
         pyrowave_kbps: u32,
         client: iroh::EndpointId,
         negotiator: Arc<Negotiator>,
+        pauses: Arc<std::sync::Mutex<crate::pauses::PauseBook>>,
     ) -> Self {
         let (gamepad_feedback_tx, gamepad_feedback_rx) = tokio::sync::mpsc::unbounded_channel();
         let (video_tx, video_rx) = tokio::sync::mpsc::unbounded_channel::<MediaItem>();
@@ -176,7 +182,9 @@ impl ClientSession {
         let (cursor_tx, cursor_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
         let (stats_tx, stats_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
         let (windows_tx, windows_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
+        let pause = crate::pauses::PauseHandle::new(pauses, gamepad_session, idr_cmd_tx.clone());
         Self {
+            pause,
             subscription: Arc::new(std::sync::Mutex::new(Subscription::Legacy)),
             encoder_pid: Arc::new(AtomicU32::new(encoder_pid)),
             primary_window,
@@ -349,6 +357,7 @@ impl ClientSession {
                     client: self.client,
                     negotiator: self.negotiator.clone(),
                 };
+                let pause = self.pause.clone();
                 let subscribe = (
                     self.client,
                     self.subscribe_tx.clone(),
@@ -356,7 +365,7 @@ impl ClientSession {
                 );
                 self.tasks.push(tokio::spawn(async move {
                     run_control_reader(
-                        conn, idr, reports, controller, awaiting, pyrowave, caps, subscribe,
+                        conn, idr, reports, controller, awaiting, pyrowave, caps, subscribe, pause,
                     )
                     .await
                 }));
@@ -513,6 +522,9 @@ impl ClientSession {
     /// zero, hold deltas until a keyframe, and forget what it reported about
     /// the old one. The caller asks for the keyframe.
     pub fn restart_picture(&self) {
+        // Another picture: whatever this client paused is no longer the one
+        // it is shown, so it stops holding it.
+        self.pause.release();
         self.awaiting_keyframe.store(true, Ordering::Relaxed);
         if self.pending_video.is_none() {
             let _ = self.send_video.send(MediaItem::Resubscribe);
@@ -812,6 +824,7 @@ async fn run_control_reader(
         tokio::sync::mpsc::UnboundedSender<(iroh::EndpointId, WindowId)>,
         Arc<AtomicU32>,
     ),
+    pause: crate::pauses::PauseHandle,
 ) {
     run_framed_reader(conn, BIDI_CONTROL, "control", None, |msg_type, payload| {
         let idr_cmd_tx = idr_cmd_tx.clone();
@@ -821,6 +834,7 @@ async fn run_control_reader(
         let pyrowave = pyrowave.clone();
         let caps_sink = caps_sink.clone();
         let subscribe = subscribe.clone();
+        let pause = pause.clone();
         async move {
             match msg_type {
                 MSG_WINDOW_SUBSCRIBE => match nesprotocol::window::decode_subscribe(&payload) {
@@ -844,26 +858,26 @@ async fn run_control_reader(
                     let _ = idr_cmd_tx.send(keyframe_request(pid_or_all(&subscribe.2)));
                 }
                 MSG_WINDOW_PAUSE => {
-                    // To the encoder of the window this client is shown, and
-                    // to nobody else. Until the client has been sent a frame
-                    // the hub does not know that encoder, and "every encoder"
-                    // is the wrong answer for a pause: it would stop windows
-                    // nobody asked about.
+                    // This session's pause of the encoder it is shown, and of
+                    // nothing else. Until it has been sent a frame the hub
+                    // does not know that encoder, and there is nothing to pause.
                     match pid_or_all(&subscribe.2) {
                         Some(pid) => {
                             info!("client pauses encoder {pid}");
-                            let _ = idr_cmd_tx.send(window_pause(Some(pid)));
+                            pause.pause(pid);
                         }
                         None => debug!("pause from a client with no encoder yet; ignored"),
                     }
                 }
                 MSG_WINDOW_RESUME => {
-                    // The same gate as a keyframe request: what arrives before
-                    // the keyframe cannot be decoded and starves it.
+                    // Only what this session paused. A resume from one that
+                    // paused nothing starts nothing; the keyframe is asked for
+                    // either way, since the client restarted its picture.
                     awaiting_keyframe.store(true, Ordering::Relaxed);
-                    let pid = pid_or_all(&subscribe.2);
-                    info!("client resumes encoder {pid:?}");
-                    let _ = idr_cmd_tx.send(window_resume(pid));
+                    pause.resume();
+                    if let Some(pid) = pid_or_all(&subscribe.2) {
+                        let _ = idr_cmd_tx.send(keyframe_request(Some(pid)));
+                    }
                 }
                 MSG_CLIENT_CAPS => {
                     // Passed through untouched, and in particular *without*
@@ -1230,6 +1244,8 @@ pub struct SessionManager {
         std::sync::Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<(iroh::EndpointId, WindowId)>>>,
     /// The encoder of the primary picture (0: untagged or none yet).
     primary_pid: Arc<AtomicU32>,
+    /// Which sessions have paused which encoders, ref(d-0087).
+    pauses: Arc<std::sync::Mutex<crate::pauses::PauseBook>>,
     /// Where commands for the encoders go; set by `serve_subscriptions`.
     cmd_tx: std::sync::OnceLock<tokio::sync::mpsc::UnboundedSender<Vec<u8>>>,
 }
@@ -1261,9 +1277,24 @@ impl SessionManager {
             relay_ms: Arc::new(AtomicU32::new(0)),
             pyrowave_kbps,
             negotiator: Arc::new(Negotiator::new()),
+            pauses: Default::default(),
         }
     }
 
+    /// A session's pause handle, for tests.
+    #[cfg(test)]
+    pub(crate) async fn pause_handle_of(
+        &self,
+        id: &iroh::EndpointId,
+    ) -> Option<crate::pauses::PauseHandle> {
+        self.sessions.lock().await.get(id).map(|s| s.pause.clone())
+    }
+
+    #[cfg(test)]
+    pub(crate) fn encoder_paused(&self, pid: u32) -> bool {
+        self.pauses.lock().unwrap().paused(pid)
+    }
+
     /// Take on one of a client's connections, creating the session if this is
     /// the first to arrive.
     ///
@@ -1295,6 +1326,7 @@ impl SessionManager {
                 self.pyrowave_kbps,
                 id,
                 self.negotiator.clone(),
+                self.pauses.clone(),
             )
         });
         session.attach(carrier, conn);