perf(desktop): hash payload receipts during transport IO

crates/rds-desktop/src/client.rs

@@ -1028,19 +1028,24 @@ async fn read_one_with_receipt(
         // discards its partial buffer on cancellation and hid whether even a
         // header or any media bytes arrived during observed freezes.
         let mut body = Vec::new();
+        let mut hasher = receipts.as_ref().map(|_| blake3::Hasher::new());
         let mut chunk = [0u8; 16 * 1024];
         while let Some(count) = stream.read(&mut chunk).await.ok()? {
             if count > MAX_FRAME_BYTES.saturating_sub(body.len()) {
                 return None;
             }
+            if let Some(hasher) = hasher.as_mut() {
+                hasher.update(&chunk[..count]);
+            }
             body.extend_from_slice(&chunk[..count]);
             body_bytes = body.len();
         }
         if let Some(receipts) = receipts {
             if body.is_empty() {
                 return None;
             }
-            let digest = *blake3::hash(&body).as_bytes();
+            let hasher = hasher?;
+            let digest = *hasher.finalize().as_bytes();
             let obsolete = header.seq < next_seq.load(Ordering::Relaxed);
             // The proof belongs to the bounded wire reader, before ordered
             // decode or local IPC can backpressure completed payloads.

crates/rds-desktop/src/receipts.rs

@@ -30,7 +30,7 @@ pub(crate) async fn wait(
     }
 }
 
-type Pending = BTreeMap<u64, ([u8; 32], Option<oneshot::Sender<Evidence>>)>;
+type Pending = BTreeMap<u64, (Option<[u8; 32]>, Option<oneshot::Sender<Evidence>>)>;
 
 #[derive(Clone)]
 pub(crate) struct Receipts(Arc<Mutex<Pending>>);
@@ -39,27 +39,43 @@ impl Receipts {
     pub(crate) fn new() -> Self {
         Self(Arc::new(Mutex::new(BTreeMap::new())))
     }
-    pub(crate) fn register(&self, seq: u64, digest: [u8; 32]) -> Option<Ticket> {
+    /// Reserve one receipt slot before the payload is sent. The digest is
+    /// filled incrementally by the sender as bytes actually leave the
+    /// stream, avoiding a second full-payload hash pass.
+    pub(crate) fn register(&self, seq: u64) -> Option<Ticket> {
         let mut pending = self.0.lock().ok()?;
         if pending.len() >= crate::session::MAX_PENDING_FRAME_ACKS || pending.contains_key(&seq) {
             return None;
         }
         let (send, receive) = oneshot::channel();
-        pending.insert(seq, (digest, Some(send)));
+        pending.insert(seq, (None, Some(send)));
         Some(Ticket {
             owner: self.clone(),
             seq,
             receive,
         })
     }
+    pub(crate) fn set_digest(&self, seq: u64, digest: [u8; 32]) -> bool {
+        let Ok(mut pending) = self.0.lock() else {
+            return false;
+        };
+        let Some((expected, _)) = pending.get_mut(&seq) else {
+            return false;
+        };
+        if expected.is_some() {
+            return false;
+        }
+        *expected = Some(digest);
+        true
+    }
     pub(crate) fn confirm(&self, seq: u64, digest: &[u8; 32], obsolete: bool) -> bool {
         let Ok(mut pending) = self.0.lock() else {
             return false;
         };
         let Some((expected, send)) = pending.get_mut(&seq) else {
             return false;
         };
-        if expected != digest {
+        if expected.as_ref() != Some(digest) {
             return false;
         }
         send.take().is_some_and(|send| {
@@ -94,11 +110,18 @@ impl Drop for Ticket {
 #[cfg(test)]
 mod tests {
     use super::*;
+
+    fn registered(receipts: &Receipts, seq: u64, digest: [u8; 32]) -> Ticket {
+        let ticket = receipts.register(seq).unwrap();
+        assert!(receipts.set_digest(seq, digest));
+        ticket
+    }
+
     #[tokio::test]
     async fn only_an_exact_live_payload_proof_completes_once() {
         let receipts = Receipts::new();
         let digest = *blake3::hash(b"controlled payload").as_bytes();
-        let mut ticket = receipts.register(7, digest).unwrap();
+        let mut ticket = registered(&receipts, 7, digest);
         assert!(!receipts.confirm(6, &digest, false));
         assert!(!receipts.confirm(7, &[0; 32], false));
         assert!(
@@ -111,29 +134,40 @@ mod tests {
         ticket.received().await.unwrap();
         drop(ticket);
         assert!(!receipts.confirm(7, &digest, false));
-        assert!(receipts.register(7, digest).is_some());
+        assert!(receipts.register(7).is_some());
+    }
+
+    #[test]
+    fn payload_proof_stays_unconfirmable_until_sender_sets_digest() {
+        let receipts = Receipts::new();
+        let digest = *blake3::hash(b"payload").as_bytes();
+        let _ticket = receipts.register(7).unwrap();
+        assert!(!receipts.confirm(7, &digest, false));
+        assert!(receipts.set_digest(7, digest));
+        assert!(receipts.confirm(7, &digest, false));
+        assert!(!receipts.set_digest(7, digest));
     }
     #[test]
     fn pending_proofs_are_bounded_and_cancellation_frees_the_exact_slot() {
         let receipts = Receipts::new();
         let mut tickets: Vec<_> = (0..crate::session::MAX_PENDING_FRAME_ACKS)
-            .map(|seq| receipts.register(seq as u64, [seq as u8; 32]).unwrap())
+            .map(|seq| receipts.register(seq as u64).unwrap())
             .collect();
-        assert!(receipts.register(0, [9; 32]).is_none());
-        assert!(receipts.register(100, [9; 32]).is_none());
+        assert!(receipts.register(0).is_none());
+        assert!(receipts.register(100).is_none());
         let retired = tickets.pop().unwrap();
         let seq = retired.seq;
         drop(retired);
         assert!(!receipts.confirm(seq, &[seq as u8; 32], false));
-        assert!(receipts.register(100, [9; 32]).is_some());
+        assert!(receipts.register(100).is_some());
         drop(tickets);
         assert!(receipts.0.lock().unwrap().is_empty());
     }
     #[tokio::test]
     async fn a_delayed_transport_fin_cannot_hold_a_validated_payload_ticket() {
         let receipts = Receipts::new();
         let digest = *blake3::hash(b"already received frame").as_bytes();
-        let mut ticket = Some(receipts.register(1796, digest).unwrap());
+        let mut ticket = Some(registered(&receipts, 1796, digest));
         let transport = std::future::pending::<Result<Evidence, String>>();
         assert!(receipts.confirm(1796, &digest, false));
         assert_eq!(
@@ -152,7 +186,7 @@ mod tests {
     #[tokio::test]
     async fn transport_fin_alone_is_not_a_negotiated_payload_proof() {
         let receipts = Receipts::new();
-        let mut ticket = Some(receipts.register(1, [1; 32]).unwrap());
+        let mut ticket = Some(registered(&receipts, 1, [1; 32]));
         assert!(
             tokio::time::timeout(
                 std::time::Duration::ZERO,
@@ -180,15 +214,15 @@ mod tests {
     async fn obsolete_payload_proofs_and_session_retirement_do_not_confirm_fresh_work() {
         let old = Receipts::new();
         let digest = [9; 32];
-        let mut ticket = Some(old.register(0, digest).unwrap());
+        let mut ticket = Some(registered(&old, 0, digest));
         assert!(old.confirm(0, &digest, true));
         assert_eq!(
             wait(&mut ticket, std::future::pending()).await.unwrap(),
             Evidence::Obsolete
         );
         drop(ticket);
         let fresh = Receipts::new();
-        let mut ticket = fresh.register(0, digest).unwrap();
+        let mut ticket = registered(&fresh, 0, digest);
         assert!(!old.confirm(0, &digest, false));
         assert!(
             tokio::time::timeout(std::time::Duration::ZERO, ticket.received())

crates/rds-desktop/src/session.rs

@@ -8,8 +8,10 @@
 //! on the bi-directional control stream, which outranks every frame
 //! stream.
 
+use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::task::{Context, Poll};
 use std::time::{Duration, Instant};
 
 use bytes::Bytes;
@@ -1959,11 +1961,9 @@ async fn send_frame_inner(
         generation,
         mut key,
     } = produced;
-    let mut payload_ticket = match &delivery.payload_receipts {
-        Some(receipts) => match receipts.register(
-            produced.header.seq,
-            *blake3::hash(&produced.payload).as_bytes(),
-        ) {
+    let payload_receipts = delivery.payload_receipts.clone();
+    let mut payload_ticket = match &payload_receipts {
+        Some(receipts) => match receipts.register(produced.header.seq) {
             Some(ticket) => Some(ticket),
             None => {
                 tracing::warn!(
@@ -2008,7 +2008,8 @@ async fn send_frame_inner(
         tracing::debug!("frame header write failed: {e}");
         return SendOutcome::Failed;
     }
-    let outcome = match send_payload(stream, &produced, rx).await {
+    let mut payload_hasher = payload_receipts.as_ref().map(|_| blake3::Hasher::new());
+    let outcome = match send_payload(stream, &produced, rx, payload_hasher.as_mut()).await {
         Ok(PayloadOutcome::Abandoned(next)) => {
             return SendOutcome::Superseded(next);
         }
@@ -2023,6 +2024,17 @@ async fn send_frame_inner(
             return SendOutcome::Failed;
         }
     };
+    if matches!(outcome, SendOutcome::Sent | SendOutcome::Done)
+        && let (Some(receipts), Some(hasher)) = (&payload_receipts, payload_hasher)
+    {
+        let digest = *hasher.finalize().as_bytes();
+        if !receipts.set_digest(produced.header.seq, digest) {
+            tracing::debug!(
+                frame_seq = produced.header.seq,
+                "payload receipt digest was retired before send completed"
+            );
+        }
+    }
     if let Err(e) = stream.finish() {
         // finish() reports an erased ClosedStream. Only a ready, exact peer
         // disposition can classify it as obsolete; never wait on other errors.
@@ -2138,11 +2150,51 @@ enum PayloadOutcome<T> {
     Abandoned(T),
 }
 
+/// Hash payload bytes as the async writer actually accepts them. Receipt
+/// hashing used to scan the complete payload before the write and therefore
+/// paid a second Full HD memory pass on every frame.
+struct DigestingWriter<'a, W> {
+    inner: &'a mut W,
+    hasher: Option<&'a mut blake3::Hasher>,
+}
+
+impl<W: AsyncWrite + Unpin> AsyncWrite for DigestingWriter<'_, W> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        let this = self.get_mut();
+        match Pin::new(&mut *this.inner).poll_write(cx, buf) {
+            Poll::Ready(Ok(written)) => {
+                if let Some(hasher) = this.hasher.as_deref_mut() {
+                    hasher.update(&buf[..written]);
+                }
+                Poll::Ready(Ok(written))
+            }
+            other => other,
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.get_mut().inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.get_mut().inner).poll_shutdown(cx)
+    }
+}
+
 async fn send_payload<W: AsyncWrite + Unpin, T: Borrow<Produced>>(
     stream: &mut W,
     produced: &Produced,
     rx: &mut mpsc::Receiver<T>,
+    hasher: Option<&mut blake3::Hasher>,
 ) -> std::io::Result<PayloadOutcome<T>> {
+    let mut stream = DigestingWriter {
+        inner: stream,
+        hasher,
+    };
     // write_all is not cancellation-safe: keep its progress alive across a
     // producer event. Starting another write_all would duplicate the prefix.
     let writing = stream.write_all(&produced.payload);
@@ -2449,7 +2501,7 @@ mod tests {
             let frame = produced(0, keyframe);
             let mut wire = vec![0; 17];
             {
-                let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx));
+                let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None));
                 // Poll until the tiny stream buffer is full, then consume only
                 // a prefix. The producer event must interrupt a partial write.
                 poll_fn(|cx| {
@@ -2513,7 +2565,7 @@ mod tests {
             let (mut writer, reader) = tokio::io::duplex(64);
             let (tx, mut rx) = mpsc::channel(1);
             let frame = produced(0, true);
-            let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx));
+            let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None));
             poll_fn(|cx| {
                 assert!(sending.as_mut().poll(cx).is_pending());
                 Poll::Ready(())
@@ -2548,7 +2600,7 @@ mod tests {
             let (mut writer, _reader) = tokio::io::duplex(64);
             let (tx, mut rx) = mpsc::channel(1);
             let frame = produced(0, false);
-            let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx));
+            let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None));
             poll_fn(|cx| {
                 assert!(sending.as_mut().poll(cx).is_pending());
                 Poll::Ready(())