From ec648634377b329d49e7906c62d331db1a58ebb1 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 09:02:16 +0000 Subject: [PATCH] Docker proxies: release a streamed layer's staging space once it is stored Since #679 a proxy serves an uncached layer while it downloads it, and every response held the shared download, which held a reader of the staged file: its unlinked temp file and its whole-size reservation in server.max_staging_bytes. A client that read slowly, or stopped reading with its TCP window at zero, kept both after the layer was verified and stored, so a few such pulls of large layers could use up the staging budget, and every push and first proxy download answered 503 until those clients went away. The staged file is now the download task's. Responses take a handle on it for each read only, and the task drops the download's own handle when it ends, after reporting the layer stored (with its S3 key) or the download failed. A response still on the layer goes on from the stored blob, through the local blob cache like a pull that joins after the commit, or ends short when the download failed (a verified layer that could not be stored included, as in 2.1). So no response keeps staging space past the download, and a stalled one needs no idle timeout. store_fetched returns the key it stored the blob under. A unit test stalls a reader after verification and checks the reservation is back once the layer is stored; an integration test stalls a real pull of a 32 MiB layer, which fails on the code before with all 32 MiB still held, and then reads the rest of the layer from the store. #679's changelog fragment, docs/docker.md and docs/deployment.md say when the space is given back. Fixes #763 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_012e6DJbjxnhsmaqc5gaMAcx --- crates/server/src/images/download.rs | 373 ++++++++++++++++++++------- crates/server/src/images/proxy.rs | 7 +- crates/server/tests/registry.rs | 62 +++++ docs/changes/2.2.0/679.feature.md | 8 +- docs/deployment.md | 4 +- docs/docker.md | 4 +- 6 files changed, 361 insertions(+), 97 deletions(-) diff --git a/crates/server/src/images/download.rs b/crates/server/src/images/download.rs index 65bddd9a..7bbf5e4f 100644 --- a/crates/server/src/images/download.rs +++ b/crates/server/src/images/download.rs @@ -17,6 +17,13 @@ //! * The download is its task's, not a request's: it is finished and cached //! when its clients go away, within the upstream client's download deadline //! (`proxy.download_timeout_seconds`). +//! * So is the staged file: it is deleted, and its staging space +//! (`server.max_staging_bytes`) given back, when the task ends +//! ([`Download::end`]), however far its responses are. A response still on +//! the layer then goes on from the stored layer, through the local blob +//! cache like a pull that comes after it, or ends short when the download +//! failed. A client that reads slowly, or stops reading, holds no staging +//! space past the download. //! * A download is registered from its start until its layer is linked, or //! it failed. The proxy starts and joins downloads under the layer's //! download gate, so a pull finds the layer either linked or downloading; @@ -29,7 +36,7 @@ use std::sync::{Arc, LazyLock, Mutex, PoisonError}; use bytes::Bytes; use forge_core::net::upstream::UpstreamHttp; use forge_db::model::Repository; -use forge_storage::{CasError, SizedStage, StageReader}; +use forge_storage::{BlobRef, BlobStore, BlobStream, CasError, SizedStage, StageReader}; use futures_util::{Stream, StreamExt}; use tokio::sync::watch; use tracing::Instrument; @@ -48,18 +55,45 @@ static DOWNLOADS: LazyLock>>> = LazyLock::ne pub(crate) struct Download { /// The layer's size, as the upstream declared it. pub size: u64, - reader: StageReader, + /// The staged file, until the download's task ends ([`Self::end`]). + /// Responses take a handle for each read only, so the file and its + /// staging space go with the task, not with the slowest response. + staged: Mutex>, + /// Where the responses still on the layer read it from once it is + /// stored. + store: BlobStore, progress: watch::Receiver, } -/// How far a download is. -#[derive(Debug, Clone, Copy, Default)] -struct Progress { - /// The bytes in the staged file. - written: u64, - /// Whether all of them are there and match the layer's digest. A - /// download whose task ended without it failed. - verified: bool, +/// How far a download is. A download whose task ended before +/// [`Progress::Stored`] failed. +#[derive(Debug, Clone)] +enum Progress { + /// Arriving: the bytes in the staged file so far. + Staging { written: u64 }, + /// All of them are there, and they match the layer's digest. + Verified, + /// Stored and linked. The staged file is released right after. + Stored(Arc), +} + +/// Where a downloaded layer is stored. +#[derive(Debug)] +struct StoredLayer { + sha256: String, + s3_key: String, +} + +/// Where a response is in the layer. +enum Reading { + /// Following the download: `at` is the next offset it sends. + Download { + download: Arc, + progress: watch::Receiver, + at: u64, + }, + /// Reading the stored layer, once the download stored it. + Stored(BlobStream), } impl std::fmt::Debug for Download { @@ -84,6 +118,23 @@ impl Drop for Registered { } impl Download { + /// A download of `size` bytes into `stage`, stored in `store`, and the + /// progress its task reports. + fn new( + size: u64, + stage: &SizedStage, + store: BlobStore, + ) -> (Arc, watch::Sender) { + let (progress, watched) = watch::channel(Progress::Staging { written: 0 }); + let download = Arc::new(Download { + size, + staged: Mutex::new(Some(stage.reader())), + store, + progress: watched, + }); + (download, progress) + } + /// What the download of `digest` of `image` in `repo` is registered /// under, and the layer's download gate. By upstream too: a pull after /// the proxy was pointed elsewhere never joins a download from its @@ -106,18 +157,34 @@ impl Download { .cloned() } - /// How much of the layer may be served: what is written, its last byte - /// only once the layer matches its digest. - fn readable(&self, progress: Progress) -> u64 { - if progress.verified { - self.size - } else { - progress.written.min(self.size.saturating_sub(1)) + /// A handle on the staged file, while the download's task runs. + fn staged(&self) -> Option { + self.staged + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + /// End the download, as its task does once it is unregistered: with its + /// layer `stored`, or failed (`None`). The staged file goes, and its + /// staging space with it, whoever still reads the layer: those + /// responses go on from the store, or end short. + fn end(&self, progress: watch::Sender, stored: Option) { + if let Some(layer) = stored { + progress.send_replace(Progress::Stored(Arc::new(layer))); } + let staged = self + .staged + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + drop(staged); + drop(progress); } - /// The bytes `range` of the layer, as they arrive. The stream fails, - /// cutting its response short, when the download does. + /// The bytes `range` of the layer, as they arrive, and from the store + /// once the layer is stored. The stream fails, cutting its response + /// short, when the download does. pub(crate) fn stream( self: Arc, range: Range, @@ -125,33 +192,108 @@ impl Download { let end = range.end; let progress = self.progress.clone(); futures_util::stream::try_unfold( - (self, progress, range.start), - move |(download, mut progress, at)| async move { - if at >= end { - return Ok(None); - } - loop { - let ready = download.readable(*progress.borrow_and_update()).min(end); - if at < ready { - let len = usize::try_from(ready - at) - .unwrap_or(usize::MAX) - .min(READ_CHUNK); - let bytes = download.reader.read_at(at, len).await?; - if bytes.is_empty() { - return Err(std::io::Error::other("the staged layer ended early")); - } - let next = at + bytes.len() as u64; - return Ok(Some((bytes, (download, progress, next)))); - } - if progress.changed().await.is_err() && !progress.borrow().verified { - return Err(std::io::Error::other( - "the upstream download of the layer failed", - )); - } + Reading::Download { + download: self, + progress, + at: range.start, + }, + move |reading| async move { + match reading { + Reading::Download { + download, + progress, + at, + } => download.next(progress, at, end).await, + Reading::Stored(stored) => next_stored(stored).await, } }, ) } + + /// The next bytes of a response at `at` that ends at `end`: from the + /// staged file as they arrive (its last byte only once the layer matches + /// its digest), from the store once the layer is stored. + async fn next( + self: Arc, + mut progress: watch::Receiver, + at: u64, + end: u64, + ) -> std::io::Result> { + if at >= end { + return Ok(None); + } + let failed = || std::io::Error::other("the upstream download of the layer failed"); + loop { + let readable = match &*progress.borrow_and_update() { + Progress::Staging { written } => Ok((*written).min(self.size.saturating_sub(1))), + Progress::Verified => Ok(self.size), + Progress::Stored(layer) => Err(Arc::clone(layer)), + }; + let ready = match readable { + Ok(readable) => readable.min(end), + Err(layer) => return self.read_stored(&layer, at, end).await, + }; + if at < ready { + let Some(staged) = self.staged() else { + // The task ended: it stored the layer, which its + // progress says by now, or it failed. + if matches!(*progress.borrow(), Progress::Stored(_)) { + continue; + } + return Err(failed()); + }; + let len = usize::try_from(ready - at) + .unwrap_or(usize::MAX) + .min(READ_CHUNK); + let bytes = staged.read_at(at, len).await?; + if bytes.is_empty() { + return Err(std::io::Error::other("the staged layer ended early")); + } + let next = at + bytes.len() as u64; + return Ok(Some(( + bytes, + Reading::Download { + download: self, + progress, + at: next, + }, + ))); + } + if progress.changed().await.is_err() + && !matches!(*progress.borrow(), Progress::Stored(_)) + { + return Err(failed()); + } + } + } + + /// The first bytes of `at..end` of the stored layer, and the stream of + /// the rest. + async fn read_stored( + &self, + layer: &StoredLayer, + at: u64, + end: u64, + ) -> std::io::Result> { + let blob = BlobRef { + sha256: &layer.sha256, + s3_key: &layer.s3_key, + size: self.size, + }; + let stored = self.store.open(blob, Some(at..end)).await.map_err(|e| { + tracing::warn!(error = %e, "registry proxy: reading a layer it just stored failed"); + std::io::Error::other(e) + })?; + next_stored(stored).await + } +} + +/// The next bytes of a response from the stored layer. +async fn next_stored(mut stored: BlobStream) -> std::io::Result> { + match stored.next().await { + Some(bytes) => Ok(Some((bytes?, Reading::Stored(stored)))), + None => Ok(None), + } } impl Registry { @@ -170,12 +312,7 @@ impl Registry { http: UpstreamHttp, ) -> ImgResult> { let stage = self.inner.cas.stage_sized(size)?; - let (progress, watched) = watch::channel(Progress::default()); - let download = Arc::new(Download { - size, - reader: stage.reader(), - progress: watched, - }); + let (download, progress) = Download::new(size, &stage, self.inner.store.clone()); let key = Download::key(repo, image, digest); DOWNLOADS .lock() @@ -190,21 +327,26 @@ impl Registry { image = %image, digest = %digest, ); + let ending = Arc::clone(&download); tokio::spawn( async move { let downloaded = reg .download(&repo, &image, &digest, stage, response, &http, &progress) .await; - // Unregistered before its readers learn that it failed: a - // pull after that starts a new download, never joins this one. + // Unregistered before its readers learn how it ended: a pull + // after a failure starts a new download, never joins this one. drop(registered); - drop(progress); - if let Err(failure) = downloaded { - tracing::warn!( - error = %failure.message(), - "registry proxy: downloading a layer failed; nothing cached" - ); - } + let stored = match downloaded { + Ok(stored) => Some(stored), + Err(failure) => { + tracing::warn!( + error = %failure.message(), + "registry proxy: downloading a layer failed; nothing cached" + ); + None + } + }; + ending.end(progress, stored); } .instrument(span), ); @@ -223,7 +365,7 @@ impl Registry { response: reqwest::Response, http: &UpstreamHttp, progress: &watch::Sender, - ) -> ImgResult<()> { + ) -> ImgResult { let staging = |e: CasError| match e { CasError::SizeMismatch { declared } => Failure::unavailable(format!( "the upstream served a blob of another size than the {declared} bytes it declared" @@ -234,9 +376,8 @@ impl Registry { while let Some(chunk) = body.next().await { let chunk = chunk.map_err(|e| upstream_unreachable(&e))?; stage.write(&chunk).await.map_err(staging)?; - progress.send_replace(Progress { + progress.send_replace(Progress::Staging { written: stage.written(), - verified: false, }); } let staged = stage.finish().await.map_err(staging)?; @@ -247,11 +388,10 @@ impl Registry { "the upstream served a blob that does not match its digest", )); } - progress.send_replace(Progress { - written: staged.size, - verified: true, - }); - self.store_fetched(repo, image, digest, staged).await + progress.send_replace(Progress::Verified); + let sha256 = staged.sha256.clone(); + let s3_key = self.store_fetched(repo, image, digest, staged).await?; + Ok(StoredLayer { sha256, s3_key }) } } @@ -259,12 +399,16 @@ impl Registry { mod tests { use std::time::Duration; - use forge_storage::{BlobStore, Cas, ConditionalWrites, S3Config, S3Credentials}; + use forge_storage::{Cas, ConditionalWrites, S3Config, S3Credentials, Staging}; use super::*; - /// A CAS that stages only: its store and pool never connect. - fn cas() -> Cas { + /// The staging budget of a test's CAS. + const BUDGET: u64 = 1024 * 1024; + + /// A store and a CAS that stages only, with a budget of [`BUDGET`]: + /// neither the store nor the pool ever connects. + fn cas() -> (BlobStore, Cas) { let store = BlobStore::new(S3Config { endpoint: "http://127.0.0.1:9".into(), region: "us-east-1".into(), @@ -277,27 +421,28 @@ mod tests { let pool = sqlx::postgres::PgPoolOptions::new() .connect_lazy("postgres://unused@127.0.0.1:9/unused") .expect("lazy pool"); - Cas::new(store, pool) + let staging = Staging::new(Staging::default_dir(), Some(BUDGET)).expect("staging"); + (store.clone(), Cas::new(store, pool).with_staging(staging)) } /// A download of `size` bytes, its stage and its progress, as its task - /// holds them. - fn download(size: u64) -> (Arc, SizedStage, watch::Sender) { - let stage = cas().stage_sized(size).unwrap(); - let (progress, watched) = watch::channel(Progress::default()); - let download = Arc::new(Download { - size, - reader: stage.reader(), - progress: watched, - }); - (download, stage, progress) + /// holds them, and the staging space it counts against. + fn download(size: u64) -> (Arc, SizedStage, watch::Sender, Staging) { + let (store, cas) = cas(); + let stage = cas.stage_sized(size).unwrap(); + let (download, progress) = Download::new(size, &stage, store); + (download, stage, progress, cas.staging().clone()) + } + + /// The bytes `staging` holds. + fn in_use(staging: &Staging) -> u64 { + staging.usage().expect("a budget").0 } async fn write(stage: &mut SizedStage, progress: &watch::Sender, bytes: &[u8]) { stage.write(bytes).await.unwrap(); - progress.send_replace(Progress { + progress.send_replace(Progress::Staging { written: stage.written(), - verified: false, }); } @@ -320,7 +465,7 @@ mod tests { #[tokio::test] async fn readers_follow_the_download_and_get_its_last_byte_once_verified() { - let (download, mut stage, progress) = download(10); + let (download, mut stage, progress, _staging) = download(10); let mut whole = Box::pin(Arc::clone(&download).stream(0..10)); let mut range = Box::pin(Arc::clone(&download).stream(4..6)); assert!(ready(&mut whole).await.is_none(), "nothing arrived yet"); @@ -338,12 +483,8 @@ mod tests { ); assert!(ready(&mut whole).await.is_none(), "the last byte waits"); - // Verified, and the task ended. - progress.send_replace(Progress { - written: 10, - verified: true, - }); - drop(progress); + // Verified; the task stores it. + progress.send_replace(Progress::Verified); assert_eq!(chunk(&mut whole).await, "9"); assert!(matches!(ready(&mut whole).await, Some(None))); // A reader that comes later reads it all. @@ -354,15 +495,67 @@ mod tests { #[tokio::test] async fn a_failed_download_ends_its_readers_short() { - let (download, mut stage, progress) = download(10); + let (download, mut stage, progress, staging) = download(10); let mut whole = Box::pin(Arc::clone(&download).stream(0..10)); write(&mut stage, &progress, b"0123456789").await; assert_eq!(chunk(&mut whole).await, "012345678"); // The task ended without verifying the layer (it did not match its // digest, the upstream broke off). - drop(progress); + drop(stage); + download.end(progress, None); + assert_eq!(in_use(&staging), 0, "the staged file is gone"); assert!(matches!(ready(&mut whole).await, Some(Some(Err(_))))); let mut late = Box::pin(Arc::clone(&download).stream(9..10)); assert!(matches!(ready(&mut late).await, Some(Some(Err(_))))); } + + #[tokio::test] + async fn a_verified_download_that_is_not_stored_ends_its_readers_short() { + let (download, mut stage, progress, staging) = download(10); + let mut whole = Box::pin(Arc::clone(&download).stream(0..10)); + write(&mut stage, &progress, b"0123456789").await; + assert_eq!(chunk(&mut whole).await, "012345678"); + progress.send_replace(Progress::Verified); + // Storing it failed (the object store is down, the proxy was + // repointed meanwhile). + drop(stage); + download.end(progress, None); + assert_eq!(in_use(&staging), 0, "the staged file is gone"); + assert!(matches!(ready(&mut whole).await, Some(Some(Err(_))))); + } + + /// A client that stops reading once the layer is verified, with its + /// response not done, holds none of the staging space once the layer is + /// stored: the response goes on from the store when it reads again + /// (#763). + #[tokio::test] + async fn a_reader_that_stalls_after_verification_holds_no_staging_space_once_stored() { + let (download, mut stage, progress, staging) = download(10); + let mut stalled = Box::pin(Arc::clone(&download).stream(0..10)); + write(&mut stage, &progress, b"0123").await; + assert_eq!(chunk(&mut stalled).await, "0123"); + write(&mut stage, &progress, b"456789").await; + let staged = stage.finish().await.unwrap(); + progress.send_replace(Progress::Verified); + assert_eq!(in_use(&staging), 10, "the whole layer is reserved"); + + // Stored: the task lets go of the staged blob, and ends. + drop(staged); + assert_eq!(in_use(&staging), 10, "the download still holds it"); + download.end( + progress, + Some(StoredLayer { + sha256: "unused".into(), + s3_key: "unused".into(), + }), + ); + assert_eq!( + in_use(&staging), + 0, + "released, though a response is still on the layer" + ); + assert!(download.staged().is_none()); + assert!(matches!(&*download.progress.borrow(), Progress::Stored(_))); + drop(stalled); + } } diff --git a/crates/server/src/images/proxy.rs b/crates/server/src/images/proxy.rs index 470894f7..5a5c6d47 100644 --- a/crates/server/src/images/proxy.rs +++ b/crates/server/src/images/proxy.rs @@ -1022,14 +1022,15 @@ impl Registry { } /// Store a blob `repo`'s upstream served for `image`, verified against - /// its digest, and link it into the image. + /// its digest, and link it into the image. Returns the S3 key it is + /// stored under. pub(crate) async fn store_fetched( &self, repo: &Repository, image: &str, digest: &ContentDigest, staged: StagedBlob, - ) -> ImgResult<()> { + ) -> ImgResult { let stored = self.inner.cas.commit(staged).await?; let mut tx = forge_db::begin(&self.inner.pool).await?; still_upstream(&mut tx, repo).await?; @@ -1046,7 +1047,7 @@ impl Registry { .map_err(Failure::from) .and_then(require_linked)?; tx.commit().await?; - Ok(()) + Ok(stored.s3_key) } /// A proxy caches a manifest before its config: once the config blob diff --git a/crates/server/tests/registry.rs b/crates/server/tests/registry.rs index 08e273e8..89a2b4ef 100644 --- a/crates/server/tests/registry.rs +++ b/crates/server/tests/registry.rs @@ -2944,6 +2944,68 @@ async fn a_downloading_layer_that_does_not_match_its_digest_is_cut_short() { assert!(linked_soon(&srv, &proxy).await, "cached"); } +/// A client that stops reading a layer a proxy serves while it downloads it +/// holds none of the staging space once the layer is stored, and reads the +/// rest from the store when it goes on (#763). +#[tokio::test] +async fn a_stalled_pull_of_a_downloading_layer_releases_its_staging_space_once_stored() { + let srv = common::start().await; + let mock = MockRegistry::start().await; + // Far more than the sockets between the server and the client buffer: + // the server cannot end its response to a client that stops reading. + let layer = big_layer(32); + mock.put_blob(&layer); + let hold = Arc::new(tokio::sync::Semaphore::new(0)); + *mock.blob_hold.lock().unwrap() = Some(hold.clone()); + let proxy = repo_name("proxy"); + create_repo( + &srv, + json!({ + "name": proxy, "kind": "proxy", "format": "docker", "remote_url": mock.base(), + "upstream_username": "upstream", "upstream_password": "secret", + }), + ) + .await; + let oci = admin_oci(&srv).await; + let name = format!("{proxy}/library/stalled"); + let path = format!("blobs/{}", sha(&layer)); + let staged = || srv.state.cas.staging().usage().expect("a staging budget").0; + + let mut stalled = oci.get(&name, &path).await; + assert_eq!(stalled.status(), 200); + let mut got = read_some(&mut stalled, 1024 * 1024).await; + assert!( + staged() >= layer.len() as u64, + "the layer's size is reserved while it downloads" + ); + + // The client stops reading; the upstream sends the rest, and the layer + // is verified and stored. + hold.add_permits(1); + assert!(linked_soon(&srv, &proxy).await, "cached"); + let mut released = false; + for _ in 0..500 { + if staged() == 0 { + released = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + assert!( + released, + "the staging space is released while the client has not read the layer: {} bytes held", + staged() + ); + + // The client reads on, from the stored layer. + got.extend_from_slice(&stalled.bytes().await.expect("the whole layer")); + assert!( + got == layer, + "the layer the client receives is the upstream's" + ); + assert_eq!(mock.blob_gets.load(Ordering::SeqCst), 1, "one download"); +} + #[tokio::test] async fn proxies_cache_verify_and_serve_stale_tags() { let srv = common::start().await; diff --git a/docs/changes/2.2.0/679.feature.md b/docs/changes/2.2.0/679.feature.md index 3031232f..cba81f52 100644 --- a/docs/changes/2.2.0/679.feature.md +++ b/docs/changes/2.2.0/679.feature.md @@ -5,5 +5,9 @@ on the same instance meanwhile are served from that one download. A layer is still cached only once it matches its digest, and its last byte is sent only then: a layer that does not match ends the response short, an error the client retries. The download finishes and is cached also when its -client goes away. Image configs (at most 1 MiB) and layers whose upstream -sends no `Content-Length` are downloaded before they are served, as before. +client goes away. The temp space it takes (`server.max_staging_bytes`, +reserved whole when it starts) is given back once the layer is stored: a +client still reading the layer goes on from the store, so one that reads +slowly, or stops reading, does not hold it. Image configs (at most 1 MiB) +and layers whose upstream sends no `Content-Length` are downloaded before +they are served, as before. diff --git a/docs/deployment.md b/docs/deployment.md index 0b93978e..c8a11c53 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -1625,7 +1625,9 @@ notices within about 15 s, drops its in-memory cache and resubscribes. `0` removes the limit); past it, pushes and downloads are answered `503` with `Retry-After` until space frees up. A layer a docker proxy serves while it downloads it reserves its whole size when the download starts, - until the download is stored and the last response reading it has ended. + until the layer is stored or the download fails; the responses still + reading it then go on from the stored layer (or end short), so a client + that reads slowly, or stops reading, holds none of it. A repository's upload limit may not exceed the budget (less `nuget.max_nuspec_bytes`), so raise the budget before giving a repository a larger limit. An explicit budget below diff --git a/docs/docker.md b/docs/docker.md index 57583789..8f6810eb 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -236,7 +236,9 @@ docker pull forge.example.com/docker-hub/nginx:1.27 on the same instance meanwhile are served from that one download. Its last byte is sent once the whole layer matches its digest; a layer that does not ends the response short, an error the client retries, and is not - cached. The download is cached even when its client goes away. Blobs of at + cached. The download is cached even when its client goes away. Once the + layer is stored, its temp space is given back and the clients still + reading it go on from the store. Blobs of at most 1 MiB (image configs), and layers whose upstream sends no `Content-Length`, are downloaded and verified before they are served. - Multi-arch indexes are cached as they are; each platform's manifest and