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