diff --git a/crates/didbot-serve/src/blobs.rs b/crates/didbot-serve/src/blobs.rs index 20a49d8a..f8d0f522 100644 --- a/crates/didbot-serve/src/blobs.rs +++ b/crates/didbot-serve/src/blobs.rs @@ -45,6 +45,22 @@ use crate::wire::{GetBlobQuery, ListBlobsQuery}; /// in memory at once. const CHUNK: usize = 64 * 1024; +/// How many `com.atproto.sync.getBlob` calls may hold a blocking-pool thread +/// at once, across every caller. +/// +/// One read takes a thread twice: once to hash the file against the CID it is +/// filed under, and once for the task that reads the body out, which +/// [`BLOB_READAHEAD`]'s back-pressure parks for as long as the client is slow +/// to take it. That pool is the same one session creation, provisioning, +/// deletion and firehose publishing draw from, and it holds fewer threads +/// than [`crate::MAX_CONCURRENT_REQUESTS`] can put in front of it. One permit +/// covers both threads of one read, from the hash until the body ends. +/// +/// A caller that finds every permit taken waits for one, and +/// [`crate::routes::REQUEST_TIMEOUT`]'s layer answers `503` if that wait +/// outlasts the request. +static BLOB_READ_CONCURRENCY: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(32); + /// `POST /xrpc/com.atproto.repo.uploadBlob` /// /// The body is streamed into the store frame by frame and never assembled. @@ -151,6 +167,16 @@ pub async fn get_blob( Ok(query) => query, Err(rejection) => return ApiError::bad_request(rejection.body_text()).into_response(), }; + // The permit is acquired before `spawn_blocking`, not inside it: what it + // bounds is blocking-pool threads (see `BLOB_READ_CONCURRENCY`), and + // claiming one of those first would just hold it idle while this request + // waits its turn. It is released at the end of this function on every + // branch that does not stream, and moved into `BlobStream` on the one + // that does. + let permit = BLOB_READ_CONCURRENCY + .acquire() + .await + .expect("this semaphore is never closed"); // 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. @@ -181,7 +207,7 @@ pub async fn get_blob( "public, max-age=31536000, immutable".to_owned(), ), ], - Body::from_stream(BlobStream::new(body)), + Body::from_stream(BlobStream::new(body, permit)), ) .into_response(), Ok(Fetch::Redirect { url, expires_in }) => ( @@ -292,15 +318,6 @@ pub fn blob_error(err: &ProvisionError) -> ApiError { } } -/// A blob on its way out, a chunk at a time. -/// -/// The store hands back a plain [`std::io::Read`], and this turns it into the -/// stream `axum::body::Body` wants. The read is blocking, which is the same -/// posture the rest of this server takes towards its disk — the write-ahead -/// log is written from inside async handlers too — and it is what keeps -/// 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 @@ -308,6 +325,15 @@ pub fn blob_error(err: &ProvisionError) -> ApiError { /// the whole blob in memory. const BLOB_READAHEAD: usize = 4; +/// A blob on its way out, a chunk at a time. +/// +/// The store hands back a plain [`std::io::Read`], and this turns it into the +/// stream `axum::body::Body` wants. The read is blocking, which is the same +/// posture the rest of this server takes towards its disk — the write-ahead +/// log is written from inside async handlers too — and it is what keeps +/// the store's interface one trait rather than one per runtime. +/// [`BLOB_READ_CONCURRENCY`] is what bounds how many of those reads run at +/// once. struct BlobStream { chunks: tokio::sync::mpsc::Receiver>, } @@ -322,9 +348,18 @@ impl BlobStream { /// 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 { + /// + /// `permit` is the caller's [`BLOB_READ_CONCURRENCY`] permit, held for as + /// long as the reading task holds its thread and given back the moment + /// that task ends — whether the blob ran out, the read failed, or the + /// client went away. + fn new( + mut reader: Box, + permit: tokio::sync::SemaphorePermit<'static>, + ) -> Self { let (chunks, rx) = tokio::sync::mpsc::channel(BLOB_READAHEAD); tokio::task::spawn_blocking(move || { + let _permit = permit; let mut buffer = vec![0u8; CHUNK]; loop { let chunk = match reader.read(&mut buffer) { @@ -372,6 +407,13 @@ mod tests { go: Option>, } + /// A one-permit semaphore of this test's own, standing in for + /// [`BLOB_READ_CONCURRENCY`] so that what one test counts cannot be + /// another test's blob read. + fn own_semaphore() -> &'static tokio::sync::Semaphore { + Box::leak(Box::new(tokio::sync::Semaphore::new(1))) + } + impl std::io::Read for ParkedReader { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { match self.go.take() { @@ -395,7 +437,12 @@ mod tests { #[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) }))); + let semaphore = own_semaphore(); + let permit = semaphore.try_acquire().expect("a free permit"); + let mut stream = Box::pin(BlobStream::new( + Box::new(ParkedReader { go: Some(go) }), + permit, + )); // The read is parked inside `ParkedReader::read`. The worker is free // only if that call is not on it. @@ -428,4 +475,36 @@ mod tests { .expect("read without error"); assert_eq!(chunk.as_ref(), b"hi!"); } + + /// **A download that ends early still gives its thread back.** The permit + /// bounds blocking-pool threads and is held for the whole body, so the + /// case to get right is the body that never finishes: a client that + /// disconnects mid-download leaves the reading task parked in a read it + /// only abandons on its next send. + #[tokio::test] + async fn a_blob_body_gives_its_permit_back_when_the_client_goes_away() { + let (release, go) = std::sync::mpsc::channel(); + let semaphore = own_semaphore(); + let permit = semaphore.try_acquire().expect("a free permit"); + let stream = BlobStream::new(Box::new(ParkedReader { go: Some(go) }), permit); + assert_eq!( + semaphore.available_permits(), + 0, + "a blob read in flight holds its permit" + ); + + drop(stream); + let _ = release.send(()); + + let returned = tokio::time::timeout(std::time::Duration::from_secs(5), async { + while semaphore.available_permits() == 0 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await; + assert!( + returned.is_ok(), + "the permit outlived the blob read it bounds" + ); + } }