diff --git a/src/ipc/drm.rs b/src/ipc/drm.rs index 831543481..1521cc570 100644 --- a/src/ipc/drm.rs +++ b/src/ipc/drm.rs @@ -1204,8 +1204,10 @@ pub struct DrmConn { /// The raw stream. Obtained from `connect_drm` (client) or the accepted `_drm` listener stream /// (service). All framing is done by hand on this fd; there is no `Framed` codec. stream: tokio::net::UnixStream, - /// Grow-once accumulation buffer for `recv_msg`/`next_raw` length-prefixed reads (preallocated - /// model: it grows to the largest frame seen and is then reused, never per-frame reallocated). + /// Grow-once accumulation buffer for `recv_msg` length-prefixed reads (preallocated model: it + /// grows to the largest message seen and is then reused, never per-message reallocated). Raw + /// bodies do not use it: `next_raw_into` reads into a buffer the caller owns, so a whole frame + /// can be recycled between the receive path and the slot it is published to. read_buf: Vec, /// Set by `drm_read_full` once the current read has consumed at least one byte off the socket. /// `recv_msg` clears it before reading, and `recv_msg_timeout2` reads it to tell a spurious @@ -1562,9 +1564,11 @@ impl DrmConn { drm_send_frame(&self.stream, &data, None).await } - /// Receive a raw length-prefixed body. Parity with `ConnectionTmpl::next_raw`. A raw body never - /// carries an fd; a stray fd (protocol desync) is collected by `drm_read_full` and dropped/closed. - pub async fn next_raw(&mut self) -> ResultType { + /// Receive a raw length-prefixed body INTO `out`, replacing its contents. Parity with + /// `ConnectionTmpl::next_raw`, except that the caller owns the buffer so it can be recycled + /// across frames. A raw body never carries an fd; a stray fd (protocol desync) is collected by + /// `drm_read_full` and dropped/closed. + pub async fn next_raw_into(&mut self, out: &mut Vec) -> ResultType<()> { // next_raw is not called through recv_msg_timeout2, so its progress flag is unused; pass the // field for signature parity (recv_msg clears it before its own reads). let mut prefix = [0u8; 4]; @@ -1578,10 +1582,14 @@ impl DrmConn { if len > MAX_DRM_RAW_BYTES { bail!("drm: raw body length {len} exceeds cap {MAX_DRM_RAW_BYTES}"); } - let mut out = bytes::BytesMut::new(); + // Read straight into the caller's buffer, reusing its allocation. A CPU-fallback frame is a + // whole packed-BGRA scanout, so allocating and zeroing a fresh one here and handing back a + // copy of it was two full-frame passes per frame, ~250 MB/s of pure overhead at 4K30 on top + // of the copy that actually moves the pixels. `resize` does nothing at all once the caller + // has been through one frame of the same size, which is the steady state. out.resize(len, 0); drm_read_full(&self.stream, &mut out[..], false, &mut self.consumed).await?; - Ok(out) + Ok(()) } } @@ -1711,8 +1719,15 @@ mod drm_conn_tests { let mut rx = DrmConn::new(b); let body = Bytes::from(vec![7u8; 5000]); tx.send_raw(body.clone()).await.unwrap(); - let got = rx.next_raw().await.unwrap(); + let mut got = Vec::new(); + rx.next_raw_into(&mut got).await.unwrap(); assert_eq!(&got[..], &body[..]); + // The buffer is reused across bodies, including a SHORTER one: a stale tail from the + // previous frame must not survive into it. + let short = Bytes::from(vec![9u8; 10]); + tx.send_raw(short.clone()).await.unwrap(); + rx.next_raw_into(&mut got).await.unwrap(); + assert_eq!(&got[..], &short[..]); } // A forged length prefix past the JSON cap is rejected at the prefix, before any body allocation. diff --git a/src/server/drm_capturer.rs b/src/server/drm_capturer.rs index 4f832b418..bbf9783d5 100644 --- a/src/server/drm_capturer.rs +++ b/src/server/drm_capturer.rs @@ -44,10 +44,27 @@ struct FrameSlot { // assuming BGRA; the CPU-fallback path stores BGRA. The row stride is recoverable from // `pixels.len() / height` (the convert output may carry a padded stride). latest: Option<(usize, usize, Pixfmt, Vec)>, + // A frame buffer no longer in use, handed back for the receive path to fill again: by `frame()` + // when it swaps in a new frame, and by the receive path itself when it supersedes one that was + // never consumed. A scanout is megabytes (33 MB at 4K), so allocating one per frame and freeing + // it a moment later is the kind of churn a 30 fps loop should not be doing. One slot is enough: + // at most one buffer is idle at a time, since the pipeline holds exactly two (the one being + // filled and the one published) plus the one `frame()` is lending to the encoder. + free: Option>, // Set once the stream ends so `frame()` returns a hard error (triggers a capturer rebuild). ended: Option, } +impl FrameSlot { + /// Publish `buf` as the newest frame, recycling whatever it supersedes. + fn publish(&mut self, w: usize, h: usize, fmt: Pixfmt, buf: Vec) { + if let Some((.., old)) = self.latest.take() { + self.free = Some(old); + } + self.latest = Some((w, h, fmt, buf)); + } +} + struct Shared { slot: Mutex, cv: Condvar, @@ -202,6 +219,7 @@ impl IpcDrmCapturer { let shared = Arc::new(Shared { slot: Mutex::new(FrameSlot { latest: None, + free: None, ended: None, }), cv: Condvar::new(), @@ -289,7 +307,11 @@ impl TraitCapturer for IpcDrmCapturer { ), )); } - self.cur = buf; + // Hand the buffer this one replaces back to the receive path instead of freeing it. + // The encoder is done with it: `frame()` takes `&mut self`, so the borrow it lent + // out last time has ended. + let previous = std::mem::replace(&mut self.cur, buf); + self.shared.slot.lock().unwrap().free = Some(previous); self.cur_w = w; self.cur_h = h; self.cur_fmt = fmt; @@ -516,8 +538,15 @@ async fn recv_thread( }; match conv.convert(&mut ddesc, received_fd) { Ok((data, w, h, fmt)) => { + // The convert output is borrowed from the render context and is only valid + // until the next convert, so it must be copied out. Copy into a recycled + // buffer, and outside the slot lock, so a multi-megabyte memcpy never holds + // the encoder off the slot. + let mut buf = shared.slot.lock().unwrap().free.take().unwrap_or_default(); + buf.clear(); + buf.extend_from_slice(data); let mut slot = shared.slot.lock().unwrap(); - slot.latest = Some((w as usize, h as usize, fmt, data.to_vec())); + slot.publish(w as usize, h as usize, fmt, buf); shared.cv.notify_one(); } // Transient convert contention: skip this frame (latest-wins keeps the newest), @@ -554,17 +583,19 @@ async fn recv_thread( let need = (width as usize) .saturating_mul(height as usize) .saturating_mul(4); - match conn.next_raw().await { - Ok(raw) => { - if raw.len() < need { + // Read the body straight into a recycled frame buffer and publish that same buffer: + // the pixels are copied once, by the kernel, on their way out of the socket. + let mut buf = shared.slot.lock().unwrap().free.take().unwrap_or_default(); + match conn.next_raw_into(&mut buf).await { + Ok(()) => { + if buf.len() < need { break format!( "cpu frame: body {} bytes < {need} for {width}x{height}", - raw.len() + buf.len() ); } let mut slot = shared.slot.lock().unwrap(); - slot.latest = - Some((width as usize, height as usize, Pixfmt::BGRA, raw.to_vec())); + slot.publish(width as usize, height as usize, Pixfmt::BGRA, buf); shared.cv.notify_one(); } Err(err) => break format!("frame body: {err}"), @@ -588,8 +619,11 @@ async fn recv_thread( let need = (width as usize) .saturating_mul(height as usize) .saturating_mul(4); - match conn.next_raw().await { - Ok(raw) => { + // A cursor is tiny and changes rarely, so this one keeps its own buffer (the frame + // recycler is for scanout-sized bodies) and hands it straight to the cursor cache. + let mut raw = Vec::new(); + match conn.next_raw_into(&mut raw).await { + Ok(()) => { if raw.len() < need { break format!( "cursor body {} bytes < {need} for {width}x{height}", @@ -605,7 +639,7 @@ async fn recv_thread( height: height as i32, hotx, hoty, - colors: raw.to_vec(), + colors: raw, }, ); } @@ -1416,6 +1450,7 @@ mod drm_capturer_tests { shared: Arc::new(Shared { slot: Mutex::new(FrameSlot { latest: None, + free: None, ended: None, }), cv: Condvar::new(), @@ -1432,9 +1467,14 @@ mod drm_capturer_tests { } } + // Publishes exactly the way the receive path does, so the recycling is exercised too: take a + // free buffer if one is on offer, fill it, publish it. fn put_frame(c: &IpcDrmCapturer, w: usize, h: usize) { + let mut buf = c.shared.slot.lock().unwrap().free.take().unwrap_or_default(); + buf.clear(); + buf.resize(w * h * 4, 0); let mut slot = c.shared.slot.lock().unwrap(); - slot.latest = Some((w, h, Pixfmt::BGRA, vec![0u8; w * h * 4])); + slot.publish(w, h, Pixfmt::BGRA, buf); } #[test] @@ -1513,6 +1553,47 @@ mod drm_capturer_tests { } } + // A scanout is megabytes, so the buffers must circulate rather than be allocated per frame: + // whatever a new frame displaces goes back on offer, both when the encoder consumes one and + // when a frame is superseded before anyone reads it. + #[test] + fn frame_buffers_circulate_instead_of_being_reallocated() { + let mut c = capturer_with(Some((64, 32))); + put_frame(&c, 64, 32); + // Superseded before anyone consumed it: its buffer must come back on offer. + put_frame(&c, 64, 32); + let recycled = c + .shared + .slot + .lock() + .unwrap() + .free + .as_ref() + .map(|b| b.as_ptr()); + assert!( + recycled.is_some(), + "a superseded frame must be handed back, not dropped" + ); + // ...and the next frame must be filled into exactly that allocation. + put_frame(&c, 64, 32); + assert_eq!( + c.shared + .slot + .lock() + .unwrap() + .latest + .as_ref() + .map(|(.., b)| b.as_ptr()), + recycled, + "the receive path must refill the recycled buffer rather than allocate" + ); + assert!(matches!(c.frame(Duration::from_millis(50)), Ok(_))); + assert!( + c.shared.slot.lock().unwrap().free.is_some(), + "the buffer the encoder finished with must be handed back to the receive path" + ); + } + #[test] fn outputs_are_matched_by_name_across_the_drm_naming_difference() { let drm = [drm_display("HDMI-A-1", 1920, 1080), drm_display("DP-1", 2560, 1440)];