From 97fff220777aaccae985e4780aab91d780bcba85 Mon Sep 17 00:00:00 2001 From: "@permadeath.com" Date: Fri, 11 Sep 2026 21:31:02 -0400 Subject: [PATCH] perf(serve): read a blob off the async worker, not on it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `getBlob` hashed the whole file against its CID and then read it chunk by chunk, both with `std::io::Read` on an async worker — so on the two-worker host `infra/pds` deploys onto, two blob reads were every route waiting. The fetch and the reading both run on the blocking pool now, with a four-chunk channel between the reader and the socket. Co-Authored-By: Claude Opus 5 (1M context) Change-Id: Ic72837020995f8d29ef393c18e09c2ff2b229c62 --- crates/didbot-serve/src/blobs.rs | 156 ++++++++++++++++++++++++------ crates/didbot-serve/src/routes.rs | 2 +- 2 files changed, 129 insertions(+), 29 deletions(-) diff --git a/crates/didbot-serve/src/blobs.rs b/crates/didbot-serve/src/blobs.rs index 1d030a28..20a49d8a 100644 --- a/crates/didbot-serve/src/blobs.rs +++ b/crates/didbot-serve/src/blobs.rs @@ -151,7 +151,20 @@ pub async fn get_blob( Ok(query) => query, Err(rejection) => return ApiError::bad_request(rejection.body_text()).into_response(), }; - match state.registry.fetch_blob(&query.did, &query.cid) { + // Off the async worker: `FileBlobStore::fetch` hashes the whole file + // against the CID it is filed under before it hands back a single byte, + // so this call is as long as the blob is. + let registry = state.registry.clone(); + let (did, cid) = (query.did.clone(), query.cid.clone()); + let fetched = tokio::task::spawn_blocking(move || registry.fetch_blob(&did, &cid)).await; + let fetched = match fetched { + Ok(fetched) => fetched, + Err(join_err) => { + return crate::routes::blocking_task_failed("com.atproto.sync.getBlob", join_err) + .into_response() + } + }; + match fetched { Ok(Fetch::Stream { mime_type, size, @@ -288,44 +301,131 @@ pub fn blob_error(err: &ProvisionError) -> ApiError { /// the store's interface one trait rather than one per runtime. The sentence /// to come back to when this server is asked to hold a realistic number of /// readers, which `plan/pds-writes.md` already has as its own item. +/// How many chunks the reading task may run ahead of the socket. +/// +/// The back-pressure: the reader fills this and then waits, so a slow client +/// costs one task and [`BLOB_READAHEAD`] × [`CHUNK`] of buffer rather than +/// the whole blob in memory. +const BLOB_READAHEAD: usize = 4; + struct BlobStream { - reader: Box, - buffer: Vec, - done: bool, + chunks: tokio::sync::mpsc::Receiver>, } impl BlobStream { - fn new(reader: Box) -> Self { - Self { - reader, - buffer: vec![0u8; CHUNK], - done: false, - } + /// Reads `reader` on the blocking pool, handing chunks over a bounded + /// channel. + /// + /// `std::io::Read` on a file blocks the thread that calls it, and before + /// this that thread was an async worker: a deployment with two of them + /// (the `t3.small` `infra/pds` deploys onto) had every route, the + /// firehose included, waiting behind two blob reads. The read happens on + /// a blocking thread now, and the channel is what keeps it from running + /// ahead of the socket. + fn new(mut reader: Box) -> Self { + let (chunks, rx) = tokio::sync::mpsc::channel(BLOB_READAHEAD); + tokio::task::spawn_blocking(move || { + let mut buffer = vec![0u8; CHUNK]; + loop { + let chunk = match reader.read(&mut buffer) { + Ok(0) => return, + Ok(read) => Ok(Bytes::copy_from_slice(&buffer[..read])), + // The stream ends here rather than repeating the failure: + // a body that has already begun cannot become an error + // response, so the honest thing left is a short one that + // the declared `Content-Length` makes visible as + // truncated. + Err(err) => { + let _ = chunks.blocking_send(Err(err)); + return; + } + }; + // A closed channel is a client that went away mid-body, and + // the read stops with it rather than pulling the rest of the + // blob off the disk for nobody. + if chunks.blocking_send(chunk).is_err() { + return; + } + } + }); + Self { chunks: rx } } } impl Stream for BlobStream { type Item = Result; - fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - if this.done { - return Poll::Ready(None); - } - match this.reader.read(&mut this.buffer) { - Ok(0) => { - this.done = true; - Poll::Ready(None) - } - Ok(read) => Poll::Ready(Some(Ok(Bytes::copy_from_slice(&this.buffer[..read])))), - Err(err) => { - // The stream ends here rather than repeating the failure: a - // body that has already begun cannot become an error - // response, so the honest thing left is a short one that the - // declared `Content-Length` makes visible as truncated. - this.done = true; - Poll::Ready(Some(Err(err))) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.get_mut().chunks.poll_recv(cx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + /// A reader whose first `read` parks until it is released, standing in + /// for a disk that is slow — or, for `FileBlobStore`, for the whole-file + /// re-hash it does before it hands back a byte. + struct ParkedReader { + go: Option>, + } + + impl std::io::Read for ParkedReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + match self.go.take() { + Some(go) => { + let _ = go.recv(); + buf[..3].copy_from_slice(b"hi!"); + Ok(3) + } + None => Ok(0), } } } + + /// **What a blob read used to cost.** `std::io::Read` blocks the thread + /// that calls it, and `poll_next` was calling it on an async worker — so + /// on the two-worker host `infra/pds` deploys onto, two slow blob reads + /// were every route waiting, the firehose included. + /// + /// One worker here, so the property is exact rather than probabilistic: + /// while the read is parked, another task still runs. + #[tokio::test(flavor = "current_thread")] + async fn a_blob_body_is_read_off_the_async_worker() { + let (release, go) = std::sync::mpsc::channel(); + let mut stream = Box::pin(BlobStream::new(Box::new(ParkedReader { go: Some(go) }))); + + // The read is parked inside `ParkedReader::read`. The worker is free + // only if that call is not on it. + let ran = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let flag = ran.clone(); + tokio::spawn(async move { + flag.store(true, std::sync::atomic::Ordering::SeqCst); + }); + let polled = std::future::poll_fn(|cx| { + let _ = stream.as_mut().poll_next(cx); + Poll::Ready(()) + }); + tokio::time::timeout(std::time::Duration::from_secs(5), polled) + .await + .expect("polling a parked blob body must not block the worker"); + tokio::task::yield_now().await; + assert!( + ran.load(std::sync::atomic::Ordering::SeqCst), + "the parked read held the only async worker" + ); + + let _ = release.send(()); + let chunk = tokio::time::timeout( + std::time::Duration::from_secs(5), + std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)), + ) + .await + .expect("the released read answers") + .expect("a chunk") + .expect("read without error"); + assert_eq!(chunk.as_ref(), b"hi!"); + } } diff --git a/crates/didbot-serve/src/routes.rs b/crates/didbot-serve/src/routes.rs index 4dba8c07..63298a0a 100644 --- a/crates/didbot-serve/src/routes.rs +++ b/crates/didbot-serve/src/routes.rs @@ -819,7 +819,7 @@ fn parse_body(body: &Bytes) -> Result { /// becomes a response rather than a connection that hangs up with nothing /// said. `route` names the call that panicked, for the log line a caller /// never sees. -fn blocking_task_failed(route: &str, error: tokio::task::JoinError) -> ApiError { +pub(crate) fn blocking_task_failed(route: &str, error: tokio::task::JoinError) -> ApiError { ApiError::new( StatusCode::INTERNAL_SERVER_ERROR, "InternalServerError", -- 2.51.2