From 0fe6cc364195f2bef274a388b1d1b6f0c31fd6cb Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Thu, 3 Sep 2026 02:45:24 +0800 Subject: [PATCH] feat(ecstore): add on-demand migration write-back pipeline (#7079) * feat(ecstore): add on-demand migration pull queue and write-back pipeline Background pull queue per bucket (bounded by pull_queue_capacity, concurrency via the state's pull slot), OdmWriteBack/PullSource traits, single-part and multipart write-back with a pumped body that enforces idle timeout, cancel and content length, retry policy for retryable source errors, inline commit helper, and stats accounting (rustfs/backlog#2153). * feat(object): implement on-demand migration write-back over internal put OdmWriteBack impl mapping source heads onto InternalPutContext (content-header allowlist, x-amz-meta copy, tags, dual-prefix odm-* provenance, ETag policy), injected into OnDemandMigrationSys at startup; removes the dead-code gates left by ODM-06a (rustfs/backlog#2153). --- crates/ecstore/src/api/mod.rs | 5 + .../src/bucket/on_demand_migration/mod.rs | 6 + .../src/bucket/on_demand_migration/pull.rs | 1697 +++++++++++++++++ .../src/bucket/on_demand_migration/sys.rs | 36 +- rustfs/src/app/object/internal_put.rs | 2 +- rustfs/src/app/object/mod.rs | 5 +- .../src/app/object/on_demand_migration_put.rs | 863 +++++++++ rustfs/src/app/object/put.rs | 2 - rustfs/src/app/storage_api.rs | 14 + rustfs/src/startup_bucket_metadata.rs | 12 +- 10 files changed, 2630 insertions(+), 12 deletions(-) create mode 100644 crates/ecstore/src/bucket/on_demand_migration/pull.rs create mode 100644 rustfs/src/app/object/on_demand_migration_put.rs diff --git a/crates/ecstore/src/api/mod.rs b/crates/ecstore/src/api/mod.rs index 6fed68316..e5c10d288 100644 --- a/crates/ecstore/src/api/mod.rs +++ b/crates/ecstore/src/api/mod.rs @@ -159,6 +159,11 @@ pub mod bucket { OnDemandMigrationConfig, OnDemandMigrationConfigError, PathStyle, PolicyConfig, Provider, RangeGetPolicy, SourceConfig, SourceCredentials, SourceErrorPolicy, SourceTimeout, TlsConfig, ValidationContext, }; + pub use crate::bucket::on_demand_migration::{ + EnqueueOutcome, LocalObject, MAX_MULTIPART_PARTS, OdmWriteBack, PULL_MAX_RETRIES, PULL_RETRY_BASE_DELAYS, + PullCompletion, PullQueue, PullReason, PullSource, SourceBody, WriteBackBody, WriteBackError, WriteBackOutcome, + WriteBackPart, WriteBackRequest, commit_inline, commit_inline_with, + }; pub mod source_client { pub use crate::bucket::on_demand_migration::source_client::{ SourceClient, SourceClientSpec, SourceError, SourceGet, SourceHead, SourceObject, SourcePage, SourceProbe, diff --git a/crates/ecstore/src/bucket/on_demand_migration/mod.rs b/crates/ecstore/src/bucket/on_demand_migration/mod.rs index b2b4c1a08..78d92c038 100644 --- a/crates/ecstore/src/bucket/on_demand_migration/mod.rs +++ b/crates/ecstore/src/bucket/on_demand_migration/mod.rs @@ -23,6 +23,7 @@ pub mod breaker; pub mod config; pub mod negative_cache; +pub mod pull; pub mod source_client; pub mod stats; pub mod sys; @@ -37,6 +38,11 @@ pub use config::{ SourceCredentials, SourceErrorPolicy, SourceTimeout, TlsConfig, ValidationContext, }; pub use negative_cache::{NEGATIVE_CACHE_MAX_ENTRIES, NegativeCache}; +pub use pull::{ + EnqueueOutcome, LocalObject, MAX_MULTIPART_PARTS, OdmWriteBack, PULL_MAX_RETRIES, PULL_RETRY_BASE_DELAYS, PullCompletion, + PullQueue, PullReason, PullSource, SourceBody, WriteBackBody, WriteBackError, WriteBackOutcome, WriteBackPart, + WriteBackRequest, commit_inline, commit_inline_with, +}; pub use stats::{ GaugeGuard, LastSourceError, LatencyBucketSnapshot, OdmOp, OdmOutcome, OdmStats, OdmStatsSnapshot, PullFailureReason, PullPath, SOURCE_LATENCY_BUCKET_BOUNDS_MS, SourceLatencySnapshot, diff --git a/crates/ecstore/src/bucket/on_demand_migration/pull.rs b/crates/ecstore/src/bucket/on_demand_migration/pull.rs new file mode 100644 index 000000000..93caa572b --- /dev/null +++ b/crates/ecstore/src/bucket/on_demand_migration/pull.rs @@ -0,0 +1,1697 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Write-back pipeline of on-demand migration (rustfs/backlog#2153). +//! +//! Two paths store a source object locally: +//! +//! - **inline**: the GET handler tees the source body to the client and to +//! [`commit_inline`], which writes the copy through the injected +//! [`OdmWriteBack`] in one pass; +//! - **background**: [`BucketOdmState::enqueue_pull`] puts a key on the +//! bucket's bounded [`PullQueue`]; a dispatcher takes one job at a time, +//! waits for a pull slot (singleflight plus `max_concurrent_pulls`, shared +//! with the inline path) and runs the pull in its own task: local +//! existence check, HEAD, GET, single-part or multipart write-back. +//! +//! The local write must look like an ordinary client PUT (bucket default +//! SSE, quota, versioning, Object Lock, replication scheduling, events). +//! Those policies live in the `rustfs` app layer, which this crate cannot +//! call, so the write is delegated to an [`OdmWriteBack`] trait object that +//! the binary injects at startup through +//! [`OnDemandMigrationSys::set_write_back`](super::sys::OnDemandMigrationSys::set_write_back). +//! Source reads go through [`PullSource`], implemented by [`SourceClient`], +//! so the pipeline is testable without HTTP. +//! +//! The source body never reaches the write-back directly: a pump task +//! copies it into a bounded channel while enforcing the idle timeout, the +//! cancellation token and the advertised content length. A truncated or +//! oversized body therefore fails the local write before it can commit, +//! independently of the digest check the write-back performs. + +use super::source_client::{SourceClient, SourceError, SourceHead}; +use super::stats::{PullFailureReason, PullPath}; +use super::sys::{BucketOdmState, OnDemandMigrationSys, PullError, PullOutcome, PullSlot}; +use async_trait::async_trait; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use parking_lot::Mutex; +use rand::RngExt; +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::io; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant}; +use time::OffsetDateTime; +use tokio::sync::mpsc::{self, error::TrySendError}; +use tokio::sync::watch; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::{debug, trace}; + +const EVENT_ODM_PULL_STORED: &str = "odm_pull_stored"; +const EVENT_ODM_PULL_SKIPPED: &str = "odm_pull_skipped"; +const EVENT_ODM_PULL_FAILED: &str = "odm_pull_failed"; +const EVENT_ODM_PULL_QUEUE_STOPPED: &str = "odm_pull_queue_stopped"; +const LOG_COMPONENT_ECSTORE: &str = "ecstore"; +const LOG_SUBSYSTEM_ON_DEMAND_MIGRATION: &str = "on_demand_migration"; + +/// Retries after the first attempt for a retryable source failure. +pub const PULL_MAX_RETRIES: usize = 3; +/// Base delay before retry `n` (a random jitter of up to 25% is added). +pub const PULL_RETRY_BASE_DELAYS: [Duration; PULL_MAX_RETRIES] = + [Duration::from_secs(1), Duration::from_secs(4), Duration::from_secs(16)]; +/// S3 multipart upload limit; larger objects need a larger part size. +pub const MAX_MULTIPART_PARTS: u64 = 10_000; +/// Chunk size the source body is read with. +const SOURCE_READ_CHUNK_BYTES: usize = 256 * 1024; +/// Chunks the pump may run ahead of the write-back. +const PUMP_CHANNEL_CHUNKS: usize = 8; + +/// Why a background pull was requested; selects the `pulled_objects_total` +/// label. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PullReason { + /// A Range GET was served from the source; fetch the whole object. + RangeGet, + /// The object exceeded `inline_max_bytes`. + LargeObject, + /// The backfill job listed the key on the source. + Backfill, +} + +impl PullReason { + pub fn as_str(self) -> &'static str { + match self { + PullReason::RangeGet => "range_get", + PullReason::LargeObject => "large_object", + PullReason::Backfill => "backfill", + } + } + + pub fn path(self) -> PullPath { + match self { + PullReason::RangeGet | PullReason::LargeObject => PullPath::Background, + PullReason::Backfill => PullPath::Backfill, + } + } +} + +/// Result of [`PullQueue::enqueue`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum EnqueueOutcome { + /// A new job was queued. + Enqueued, + /// The key is already queued or being pulled; nothing was added. + Coalesced, + /// `pull_queue_capacity` jobs are waiting; the caller only counts this. + QueueFull, + /// The bucket has no usable source client or write-back, or its state + /// was torn down. + Unavailable, +} + +/// Body a source read produces; consumed inside the pump task only. +pub type SourceBody = Pin> + Send + 'static>>; + +/// Body handed to the write-back; `Sync` because the app-layer put path +/// wraps it into an S3 streaming blob. +pub type WriteBackBody = Pin> + Send + Sync + 'static>>; + +/// Read-only view of the source a pull needs. +#[async_trait] +pub trait PullSource: Send + Sync { + async fn head_object(&self, key: &str) -> Result; + + /// Unranged GET: the returned head describes the whole object. + async fn get_object(&self, key: &str) -> Result<(SourceHead, SourceBody), SourceError>; + + async fn get_object_tagging(&self, key: &str) -> Result, SourceError>; +} + +#[async_trait] +impl PullSource for SourceClient { + async fn head_object(&self, key: &str) -> Result { + SourceClient::head_object(self, key).await + } + + async fn get_object(&self, key: &str) -> Result<(SourceHead, SourceBody), SourceError> { + let get = SourceClient::get_object(self, key, None).await?; + let body = tokio_util::io::ReaderStream::with_capacity(get.body.into_async_read(), SOURCE_READ_CHUNK_BYTES); + Ok((get.head, Box::pin(body))) + } + + async fn get_object_tagging(&self, key: &str) -> Result, SourceError> { + SourceClient::get_object_tagging(self, key).await + } +} + +/// What the write-back must store. Everything the app layer needs to build +/// the provenance keys, the metadata allowlist and the ETag policy. +#[derive(Clone, Debug)] +pub struct WriteBackRequest { + pub bucket: String, + pub key: String, + /// Source HEAD/GET of the whole object. + pub head: SourceHead, + /// `:`, stored under `odm-source`. + pub source_label: String, + pub pulled_at: OffsetDateTime, + /// `policy.preserve_etag`. + pub preserve_etag: bool, + /// `policy.emit_events`. + pub emit_events: bool, + /// Source tags to copy (`policy.copy_tags`), `None` to skip. + pub tags: Option>, +} + +impl WriteBackRequest { + pub fn new(state: &BucketOdmState, key: &str, head: SourceHead, tags: Option>) -> Self { + let config = state.config(); + Self { + bucket: state.bucket().to_string(), + key: key.to_string(), + head, + source_label: format!("{}:{}", config.source.provider.as_str(), config.source.bucket), + pulled_at: OffsetDateTime::now_utc(), + preserve_etag: config.policy.preserve_etag, + emit_events: config.policy.emit_events, + tags, + } + } +} + +/// The committed local object. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WriteBackOutcome { + pub etag: Option, + pub size: u64, + pub version_id: Option, +} + +/// One staged part of an internal multipart upload. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WriteBackPart { + pub part_number: usize, + pub etag: String, +} + +/// The current local version of a key, when one exists. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LocalObject { + pub etag: Option, + pub size: u64, + pub delete_marker: bool, +} + +/// Why the local write did not commit. `reason()` maps it onto the +/// `pull_failures_total` label set fixed by ODM-05. +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum WriteBackError { + /// The body did not hash to the ETag the source advertised. + #[error("source bytes did not match the source ETag")] + Integrity, + /// The bucket quota rejected the write. + #[error("bucket quota exceeded: {0}")] + Quota(String), + /// The object cannot be represented locally (for example too many parts). + #[error("unsupported source object: {0}")] + Unsupported(String), + /// Any other local failure. + #[error("local write failed: {0}")] + Local(String), +} + +impl WriteBackError { + pub fn reason(&self) -> PullFailureReason { + match self { + WriteBackError::Integrity => PullFailureReason::EtagMismatch, + WriteBackError::Unsupported(_) => PullFailureReason::SourceUnsupported, + WriteBackError::Quota(_) | WriteBackError::Local(_) => PullFailureReason::LocalWrite, + } + } +} + +/// Local store facet of the pull pipeline: an ordinary app-layer write plus +/// the existence lookup that guards it. Implemented in the `rustfs` binary +/// on top of its internal put entry points and injected at startup. +#[async_trait] +pub trait OdmWriteBack: Send + Sync { + /// Current version of `key`, `None` when the key has no readable + /// current version (absent or delete marker latest). + async fn local_object(&self, bucket: &str, key: &str) -> Result, WriteBackError>; + + /// Single-object write of exactly `request.head.size` bytes. + async fn put_object(&self, request: &WriteBackRequest, body: WriteBackBody) -> Result; + + async fn create_multipart_upload(&self, request: &WriteBackRequest) -> Result; + + async fn upload_part( + &self, + request: &WriteBackRequest, + upload_id: &str, + part_number: usize, + size: u64, + body: WriteBackBody, + ) -> Result; + + async fn complete_multipart_upload( + &self, + request: &WriteBackRequest, + upload_id: &str, + parts: Vec, + ) -> Result; + + async fn abort_multipart_upload(&self, bucket: &str, key: &str, upload_id: &str) -> Result<(), WriteBackError>; +} + +/// Why the pump stopped feeding the write-back before EOF. +#[derive(Debug)] +enum PumpFailure { + Source(SourceError), + Canceled, +} + +#[derive(Debug, Default)] +struct PumpState { + failure: Mutex>, +} + +impl PumpState { + fn fail(&self, failure: PumpFailure) -> io::Error { + let message = match &failure { + PumpFailure::Source(err) => err.to_string(), + PumpFailure::Canceled => "pull canceled".to_string(), + }; + let kind = match &failure { + PumpFailure::Source(SourceError::Timeout) => io::ErrorKind::TimedOut, + PumpFailure::Source(SourceError::Connect(_)) => io::ErrorKind::UnexpectedEof, + PumpFailure::Source(_) => io::ErrorKind::Other, + PumpFailure::Canceled => io::ErrorKind::Interrupted, + }; + let mut slot = self.failure.lock(); + if slot.is_none() { + *slot = Some(failure); + } + io::Error::new(kind, message) + } + + fn take(&self) -> Option { + self.failure.lock().take() + } +} + +/// Copies the source body into a bounded channel, enforcing `idle_timeout` +/// per chunk, `cancel`, and the advertised `expected_size`. +fn spawn_pump( + mut body: SourceBody, + expected_size: u64, + idle_timeout: Duration, + cancel: CancellationToken, +) -> (mpsc::Receiver>, Arc) { + let (tx, rx) = mpsc::channel(PUMP_CHANNEL_CHUNKS); + let state = Arc::new(PumpState::default()); + let pump_state = Arc::clone(&state); + tokio::spawn(async move { + let mut delivered: u64 = 0; + loop { + let next = tokio::select! { + _ = cancel.cancelled() => Err(pump_state.fail(PumpFailure::Canceled)), + next = tokio::time::timeout(idle_timeout, body.next()) => match next { + Err(_elapsed) => Err(pump_state.fail(PumpFailure::Source(SourceError::Timeout))), + Ok(None) if delivered < expected_size => Err(pump_state.fail(PumpFailure::Source(SourceError::Connect( + format!("source body ended after {delivered} of {expected_size} bytes"), + )))), + Ok(None) => return, + Ok(Some(Err(err))) => { + let failure = if err.kind() == io::ErrorKind::TimedOut { + SourceError::Timeout + } else { + SourceError::Connect(format!("source body read failed: {err}")) + }; + Err(pump_state.fail(PumpFailure::Source(failure))) + } + Ok(Some(Ok(chunk))) => { + let len = u64::try_from(chunk.len()).unwrap_or(u64::MAX); + delivered = delivered.saturating_add(len); + if delivered > expected_size { + Err(pump_state.fail(PumpFailure::Source(SourceError::Other(format!( + "source body exceeded the advertised {expected_size} bytes" + ))))) + } else { + Ok(chunk) + } + } + }, + }; + let stop = next.is_err(); + if tx.send(next).await.is_err() || stop { + return; + } + } + }); + (rx, state) +} + +/// The pumped body plus the tail of the last chunk that crossed a part +/// boundary. Parts are written sequentially, so the mutex is uncontended; +/// it is only ever held inside a poll. +struct SharedSource { + rx: mpsc::Receiver>, + leftover: Option, +} + +type SharedSourceHandle = Arc>; + +/// Exactly `remaining` bytes of the shared source, then EOF. A source EOF +/// before that is an `UnexpectedEof` error, never a short part. +struct PartBody { + source: SharedSourceHandle, + remaining: u64, +} + +impl PartBody { + fn boxed(source: &SharedSourceHandle, len: u64) -> WriteBackBody { + Box::pin(Self { + source: Arc::clone(source), + remaining: len, + }) + } +} + +impl Stream for PartBody { + type Item = io::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.remaining == 0 { + return Poll::Ready(None); + } + let mut chunk = { + let mut source = self.source.lock(); + match source.leftover.take() { + Some(leftover) => leftover, + None => match source.rx.poll_recv(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + return Poll::Ready(Some(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("source body ended with {} bytes outstanding", self.remaining), + )))); + } + Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err))), + Poll::Ready(Some(Ok(chunk))) => chunk, + }, + } + }; + let take = usize::try_from(self.remaining).map_or(chunk.len(), |remaining| remaining.min(chunk.len())); + if take < chunk.len() { + self.source.lock().leftover = Some(chunk.split_off(take)); + } + self.remaining = self.remaining.saturating_sub(u64::try_from(take).unwrap_or(u64::MAX)); + Poll::Ready(Some(Ok(chunk))) + } +} + +/// How one pull ended without error. +#[derive(Debug)] +pub enum PullCompletion { + /// The write-back committed a new local object. + Stored(WriteBackOutcome), + /// A current local version already existed; nothing was pulled. + AlreadyPresent(LocalObject), +} + +impl PullCompletion { + fn outcome(&self) -> PullOutcome { + match self { + PullCompletion::Stored(outcome) => PullOutcome { + etag: outcome.etag.clone(), + size: outcome.size, + }, + PullCompletion::AlreadyPresent(local) => PullOutcome { + etag: local.etag.clone(), + size: local.size, + }, + } + } +} + +struct AttemptError { + error: PullError, + retryable: bool, +} + +impl AttemptError { + fn source(err: &SourceError) -> Self { + Self { + error: PullError::from(err), + retryable: err.is_retryable(), + } + } + + fn write_back(err: &WriteBackError) -> Self { + Self { + error: PullError::new(err.reason(), err.to_string()), + retryable: false, + } + } + + fn canceled() -> Self { + Self { + error: PullError::canceled("bucket on-demand migration state was removed"), + retryable: false, + } + } +} + +struct PullContext<'a> { + state: &'a Arc, + source: &'a Arc, + write_back: &'a Arc, + key: &'a str, + cancel: &'a CancellationToken, +} + +impl PullContext<'_> { + async fn observe(&self, call: impl Future>) -> Result { + let started = Instant::now(); + let result = call.await; + self.state.observe_source(started.elapsed(), self.key, result.as_ref().err()); + result.map_err(|err| AttemptError::source(&err)) + } + + async fn local_object(&self) -> Result, AttemptError> { + self.write_back + .local_object(self.state.bucket(), self.key) + .await + .map_err(|err| AttemptError::write_back(&err)) + } +} + +fn retry_delay(retry: usize) -> Duration { + let base = PULL_RETRY_BASE_DELAYS[retry.min(PULL_RETRY_BASE_DELAYS.len() - 1)]; + let jitter_cap = base.as_millis() / 4; + let jitter = rand::rng().random_range(0..=jitter_cap); + base + Duration::from_millis(u64::try_from(jitter).unwrap_or(0)) +} + +/// One attempt: local check, HEAD, optional tags, GET, second local check, +/// then the write-back. Errors carry whether a retry may help. +async fn pull_once(ctx: &PullContext<'_>) -> Result { + if ctx.cancel.is_cancelled() { + return Err(AttemptError::canceled()); + } + if let Some(local) = ctx.local_object().await? + && !local.delete_marker + { + return Ok(PullCompletion::AlreadyPresent(local)); + } + + let policy = &ctx.state.config().policy; + ctx.observe(ctx.source.head_object(ctx.key)).await?; + let tags = if policy.copy_tags { + Some(ctx.observe(ctx.source.get_object_tagging(ctx.key)).await?) + } else { + None + }; + let (head, body) = ctx.observe(ctx.source.get_object(ctx.key)).await?; + + // Re-checked after the source round trips: a client PUT that landed in + // between must not be overwritten by the older source copy. + if let Some(local) = ctx.local_object().await? + && !local.delete_marker + { + return Ok(PullCompletion::AlreadyPresent(local)); + } + + let idle_timeout = Duration::from_millis(policy.source_timeout.idle_ms); + let (rx, pump) = spawn_pump(body, head.size, idle_timeout, ctx.cancel.clone()); + let shared: SharedSourceHandle = Arc::new(Mutex::new(SharedSource { rx, leftover: None })); + let request = WriteBackRequest::new(ctx.state, ctx.key, head, tags); + let part_size = policy.multipart_part_size_bytes.max(1); + let written = if request.head.size > part_size { + write_multipart(ctx.write_back, &request, &shared, part_size).await + } else { + ctx.write_back + .put_object(&request, PartBody::boxed(&shared, request.head.size)) + .await + }; + match written { + Ok(outcome) => Ok(PullCompletion::Stored(outcome)), + Err(err) => Err(match pump.take() { + Some(PumpFailure::Canceled) => AttemptError::canceled(), + Some(PumpFailure::Source(source_err)) => AttemptError::source(&source_err), + None => AttemptError::write_back(&err), + }), + } +} + +async fn write_multipart( + write_back: &Arc, + request: &WriteBackRequest, + shared: &SharedSourceHandle, + part_size: u64, +) -> Result { + let size = request.head.size; + let part_count = size.div_ceil(part_size); + if part_count > MAX_MULTIPART_PARTS { + return Err(WriteBackError::Unsupported(format!( + "object of {size} bytes needs {part_count} parts of {part_size} bytes; the limit is {MAX_MULTIPART_PARTS}" + ))); + } + let upload_id = write_back.create_multipart_upload(request).await?; + let staged = stage_parts(write_back, request, &upload_id, shared, part_size, part_count).await; + let completed = match staged { + Ok(parts) => write_back.complete_multipart_upload(request, &upload_id, parts).await, + Err(err) => Err(err), + }; + if completed.is_err() + && let Err(abort_err) = write_back + .abort_multipart_upload(&request.bucket, &request.key, &upload_id) + .await + { + debug!( + event = EVENT_ODM_PULL_FAILED, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_ON_DEMAND_MIGRATION, + bucket = %request.bucket, + key = %request.key, + error = %abort_err, + "On-demand migration multipart abort failed after a write-back error" + ); + } + completed +} + +async fn stage_parts( + write_back: &Arc, + request: &WriteBackRequest, + upload_id: &str, + shared: &SharedSourceHandle, + part_size: u64, + part_count: u64, +) -> Result, WriteBackError> { + let size = request.head.size; + let mut parts = Vec::with_capacity(usize::try_from(part_count).unwrap_or(0)); + let mut offset = 0; + for part_number in 1..=part_count { + let len = part_size.min(size - offset); + let part_number = usize::try_from(part_number) + .map_err(|_| WriteBackError::Unsupported(format!("part number {part_number} does not fit this platform")))?; + let part = write_back + .upload_part(request, upload_id, part_number, len, PartBody::boxed(shared, len)) + .await?; + parts.push(part); + offset += len; + } + Ok(parts) +} + +/// Full pull with the retry policy: transient source failures are retried +/// up to [`PULL_MAX_RETRIES`] times with [`PULL_RETRY_BASE_DELAYS`] plus +/// jitter; everything else fails immediately. +async fn pull_object(ctx: &PullContext<'_>) -> Result { + let mut retries = 0; + loop { + match pull_once(ctx).await { + Ok(completion) => return Ok(completion), + Err(AttemptError { error, retryable }) => { + if !retryable || retries >= PULL_MAX_RETRIES { + return Err(error); + } + let delay = retry_delay(retries); + retries += 1; + tokio::select! { + _ = tokio::time::sleep(delay) => {} + _ = ctx.cancel.cancelled() => return Err(AttemptError::canceled().error), + } + } + } + } +} + +/// Counts and logs the end of one pull. +fn record_completion(state: &BucketOdmState, key: &str, path: PullPath, result: &Result) { + let stats = state.stats(); + match result { + Ok(PullCompletion::Stored(outcome)) => { + stats.record_pulled_object(path); + stats.record_pulled_bytes(outcome.size); + trace!( + event = EVENT_ODM_PULL_STORED, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_ON_DEMAND_MIGRATION, + path = path.as_str(), + bucket = %state.bucket(), + key = %key, + size = outcome.size, + "On-demand migration stored a source object locally" + ); + } + Ok(PullCompletion::AlreadyPresent(_)) => { + debug!( + event = EVENT_ODM_PULL_SKIPPED, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_ON_DEMAND_MIGRATION, + path = path.as_str(), + bucket = %state.bucket(), + key = %key, + "On-demand migration pull skipped; a local object already exists" + ); + } + Err(err) => { + stats.record_pull_failure(err.reason); + debug!( + event = EVENT_ODM_PULL_FAILED, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_ON_DEMAND_MIGRATION, + path = path.as_str(), + reason = err.reason.as_str(), + bucket = %state.bucket(), + key = %key, + "On-demand migration pull failed" + ); + } + } +} + +/// Inline write-back of a body the GET handler is already streaming (the +/// tee secondary). No retry: the bytes cannot be re-read. The caller owns +/// the singleflight slot; this only writes and accounts. +pub async fn commit_inline( + state: &Arc, + key: &str, + head: SourceHead, + tags: Option>, + body: WriteBackBody, +) -> Result { + let Some(write_back) = state.write_back() else { + let error = PullError::new(PullFailureReason::LocalWrite, "on-demand migration write-back is not installed"); + record_completion(state, key, PullPath::Inline, &Err(error.clone())); + return Err(error); + }; + commit_inline_with(state, write_back, key, head, tags, body).await +} + +/// [`commit_inline`] with an explicit write-back (tests and embedders). +pub async fn commit_inline_with( + state: &Arc, + write_back: &Arc, + key: &str, + head: SourceHead, + tags: Option>, + body: WriteBackBody, +) -> Result { + let request = WriteBackRequest::new(state, key, head, tags); + let result = write_back + .put_object(&request, body) + .await + .map_err(|err| PullError::new(err.reason(), err.to_string())); + let completion = result + .as_ref() + .map(|outcome| PullCompletion::Stored(outcome.clone())) + .map_err(Clone::clone); + record_completion(state, key, PullPath::Inline, &completion); + result +} + +struct PullJob { + key: String, + reason: PullReason, +} + +/// Bounded per-bucket queue of background pulls. Keys are unique while +/// queued or in flight; capacity is `pull_queue_capacity`. +pub struct PullQueue { + bucket: String, + tx: mpsc::Sender, + /// Keys queued or running; the job removes its key when it ends. + pending: Mutex>, + capacity: usize, + cancel: CancellationToken, + stats: Arc, + stopped: watch::Receiver, +} + +impl fmt::Debug for PullQueue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PullQueue") + .field("bucket", &self.bucket) + .field("capacity", &self.capacity) + .field("pending", &self.pending.lock().len()) + .field("stopped", &self.is_stopped()) + .finish() + } +} + +/// Removes the job's key from the pending set however the job ends. +struct PendingKeyGuard { + queue: Arc, + key: String, +} + +impl Drop for PendingKeyGuard { + fn drop(&mut self) { + self.queue.pending.lock().remove(&self.key); + } +} + +impl PullQueue { + /// Starts the dispatcher for `state`. Requires a Tokio runtime. + pub fn start(state: Arc, source: Arc, write_back: Arc) -> Arc { + let capacity = usize::try_from(state.config().policy.pull_queue_capacity.max(1)).unwrap_or(usize::MAX); + let (tx, rx) = mpsc::channel(capacity); + let (stopped_tx, stopped_rx) = watch::channel(false); + let queue = Arc::new(Self { + bucket: state.bucket().to_string(), + tx, + pending: Mutex::new(HashSet::new()), + capacity, + cancel: state.cancel_token(), + stats: Arc::clone(state.stats()), + stopped: stopped_rx, + }); + tokio::spawn(dispatch(Arc::clone(&queue), state, source, write_back, rx, stopped_tx)); + queue + } + + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Keys queued or in flight. + pub fn pending_keys(&self) -> usize { + self.pending.lock().len() + } + + pub fn is_stopped(&self) -> bool { + *self.stopped.borrow() + } + + /// Resolves once the dispatcher and every in-flight job have exited. + pub async fn wait_until_stopped(&self) { + let mut stopped = self.stopped.clone(); + // A closed channel means the dispatcher is gone as well. + let _ = stopped.wait_for(|stopped| *stopped).await; + } + + pub fn enqueue(&self, key: &str, reason: PullReason) -> EnqueueOutcome { + if self.cancel.is_cancelled() { + return EnqueueOutcome::Unavailable; + } + let mut pending = self.pending.lock(); + if pending.contains(key) { + return EnqueueOutcome::Coalesced; + } + match self.tx.try_send(PullJob { + key: key.to_string(), + reason, + }) { + Ok(()) => { + pending.insert(key.to_string()); + EnqueueOutcome::Enqueued + } + Err(TrySendError::Full(_)) => { + self.stats.record_pull_failure(PullFailureReason::QueueFull); + EnqueueOutcome::QueueFull + } + Err(TrySendError::Closed(_)) => EnqueueOutcome::Unavailable, + } + } +} + +/// Takes jobs in order, waits for a pull slot for each (this is what bounds +/// concurrency to `max_concurrent_pulls`) and runs the pull in its own task. +/// On cancellation it stops taking jobs, fails the queued ones as +/// `canceled`, and waits for in-flight tasks before reporting stopped. +async fn dispatch( + queue: Arc, + state: Arc, + source: Arc, + write_back: Arc, + mut rx: mpsc::Receiver, + stopped: watch::Sender, +) { + let cancel = state.cancel_token(); + let mut tasks = JoinSet::new(); + loop { + while tasks.try_join_next().is_some() {} + let job = tokio::select! { + _ = cancel.cancelled() => break, + job = rx.recv() => match job { + Some(job) => job, + None => break, + }, + }; + let pending = PendingKeyGuard { + queue: Arc::clone(&queue), + key: job.key.clone(), + }; + let slot = match state.acquire_pull_slot(&job.key).await { + Ok(slot) => slot, + Err(err) => { + let result: Result = Err(err); + record_completion(&state, &job.key, job.reason.path(), &result); + drop(pending); + break; + } + }; + tasks.spawn(run_job( + Arc::clone(&state), + Arc::clone(&source), + Arc::clone(&write_back), + slot, + job, + pending, + )); + } + + rx.close(); + while let Ok(job) = rx.try_recv() { + queue.pending.lock().remove(&job.key); + state.stats().record_pull_failure(PullFailureReason::Canceled); + } + while tasks.join_next().await.is_some() {} + debug!( + event = EVENT_ODM_PULL_QUEUE_STOPPED, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_ON_DEMAND_MIGRATION, + bucket = %state.bucket(), + "On-demand migration pull queue stopped" + ); + stopped.send_replace(true); +} + +async fn run_job( + state: Arc, + source: Arc, + write_back: Arc, + slot: PullSlot, + job: PullJob, + pending: PendingKeyGuard, +) { + let _pending = pending; + let cancel = state.cancel_token(); + match slot { + PullSlot::Follower(follower) => { + // Someone else (inline GET or an earlier job) is pulling the key; + // its result makes this job redundant. + let _ = follower.wait().await; + } + PullSlot::Leader(leader) => { + let ctx = PullContext { + state: &state, + source: &source, + write_back: &write_back, + key: &job.key, + cancel: &cancel, + }; + let result = pull_object(&ctx).await; + record_completion(&state, &job.key, job.reason.path(), &result); + leader.complete(result.map(|completion| completion.outcome())); + } + } +} + +impl BucketOdmState { + /// The bucket's queue, started on first use. `None` when the state is + /// torn down, has no usable client, or no write-back was injected. + pub fn pull_queue(self: &Arc) -> Option> { + if self.is_cancelled() { + return None; + } + if let Some(queue) = self.pull_queue.get() { + return Some(Arc::clone(queue)); + } + let client: Arc = Arc::clone(self.client().ok()?); + let source: Arc = client; + let write_back = Arc::clone(self.write_back()?); + tokio::runtime::Handle::try_current().ok()?; + Some(Arc::clone( + self.pull_queue + .get_or_init(|| PullQueue::start(Arc::clone(self), source, write_back)), + )) + } + + /// Queues a background pull of `key`; see [`EnqueueOutcome`]. + pub fn enqueue_pull(self: &Arc, key: &str, reason: PullReason) -> EnqueueOutcome { + match self.pull_queue() { + Some(queue) => queue.enqueue(key, reason), + None => EnqueueOutcome::Unavailable, + } + } +} + +impl OnDemandMigrationSys { + /// [`BucketOdmState::enqueue_pull`] by bucket name; `Unavailable` when + /// the bucket has no state. + pub fn enqueue_pull(&self, bucket: &str, key: &str, reason: PullReason) -> EnqueueOutcome { + match self.state(bucket) { + Some(state) => state.enqueue_pull(key, reason), + None => EnqueueOutcome::Unavailable, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::bucket::on_demand_migration::config::{ + FilterConfig, OnDemandMigrationConfig, PathStyle as ConfigPathStyle, PolicyConfig, Provider, SourceConfig, + SourceCredentials, TlsConfig, + }; + use std::collections::VecDeque; + use std::sync::atomic::{AtomicUsize, Ordering}; + + const BUCKET: &str = "odm-pull-bucket"; + + fn config() -> OnDemandMigrationConfig { + OnDemandMigrationConfig { + version: 1, + enabled: true, + source: SourceConfig { + provider: Provider::Minio, + endpoint: Some("https://source.example.com:9000".to_string()), + region: "auto".to_string(), + bucket: "legacy".to_string(), + path_style: ConfigPathStyle::Auto, + credentials: Some(SourceCredentials { + access_key: "AK".to_string(), + secret_key: "SK".to_string(), + session_token: None, + }), + tls: TlsConfig::default(), + }, + filter: FilterConfig::default(), + policy: PolicyConfig::default(), + } + } + + async fn enabled_state(sys: &OnDemandMigrationSys, cfg: &OnDemandMigrationConfig) -> Arc { + sys.set_module_enabled(true); + sys.apply(BUCKET, Some(cfg)).await; + sys.state(BUCKET).expect("state must be installed") + } + + fn head(size: u64) -> SourceHead { + SourceHead { + etag: Some("0123456789abcdef0123456789abcdef".to_string()), + size, + content_type: Some("text/plain".to_string()), + ..Default::default() + } + } + + fn body_bytes(len: usize) -> Vec { + (0..len).map(|i| u8::try_from(i % 251).expect("fits")).collect() + } + + #[derive(Clone)] + enum BodyKind { + Bytes(Vec), + Hang, + } + + #[derive(Clone)] + struct MockObject { + head: SourceHead, + body: BodyKind, + } + + #[derive(Default)] + struct MockSource { + objects: Mutex>, + head_failures: Mutex>, + get_failures: Mutex>, + tags: HashMap, + head_calls: AtomicUsize, + get_calls: AtomicUsize, + tag_calls: AtomicUsize, + /// Simulates a concurrent client PUT landing while the source GET + /// is in flight: inserts the key into the write-back's local map. + land_local_on_get: Mutex, String)>>, + } + + impl MockSource { + fn with_object(key: &str, size: u64, body: BodyKind) -> Arc { + let source = Self::default(); + source + .objects + .lock() + .insert(key.to_string(), MockObject { head: head(size), body }); + Arc::new(source) + } + + fn object(&self, key: &str) -> Result { + self.objects.lock().get(key).cloned().ok_or(SourceError::NotFound) + } + } + + #[async_trait] + impl PullSource for MockSource { + async fn head_object(&self, key: &str) -> Result { + self.head_calls.fetch_add(1, Ordering::SeqCst); + if let Some(err) = self.head_failures.lock().pop_front() { + return Err(err); + } + Ok(self.object(key)?.head) + } + + async fn get_object(&self, key: &str) -> Result<(SourceHead, SourceBody), SourceError> { + self.get_calls.fetch_add(1, Ordering::SeqCst); + if let Some(err) = self.get_failures.lock().pop_front() { + return Err(err); + } + let object = self.object(key)?; + if let Some((write_back, local_key)) = self.land_local_on_get.lock().take() { + write_back.local.lock().insert( + local_key, + LocalObject { + etag: Some("client-put".to_string()), + size: 1, + delete_marker: false, + }, + ); + } + let body: SourceBody = match object.body { + BodyKind::Bytes(bytes) => { + let chunks: Vec> = + bytes.chunks(700).map(|chunk| Ok(Bytes::copy_from_slice(chunk))).collect(); + Box::pin(futures::stream::iter(chunks)) + } + BodyKind::Hang => Box::pin(futures::stream::pending()), + }; + Ok((object.head, body)) + } + + async fn get_object_tagging(&self, _key: &str) -> Result, SourceError> { + self.tag_calls.fetch_add(1, Ordering::SeqCst); + Ok(self.tags.clone()) + } + } + + /// `(upload id, part number, size, bytes)` as staged by the mock. + type StagedPart = (String, usize, u64, Vec); + + #[derive(Default)] + struct MockWriteBack { + local: Mutex>, + puts: Mutex)>>, + failed_puts: AtomicUsize, + uploads: Mutex>, + parts: Mutex>, + completed: Mutex)>>, + aborted: Mutex>, + fail_part: Option, + forced_put_error: Mutex>, + upload_seq: AtomicUsize, + } + + async fn drain(mut body: WriteBackBody, expected: u64) -> Result, WriteBackError> { + let mut buf = Vec::new(); + while let Some(chunk) = body.next().await { + let chunk = chunk.map_err(|err| WriteBackError::Local(format!("body read failed: {err}")))?; + buf.extend_from_slice(&chunk); + } + if u64::try_from(buf.len()).expect("fits") != expected { + return Err(WriteBackError::Local(format!("expected {expected} bytes, got {}", buf.len()))); + } + Ok(buf) + } + + #[async_trait] + impl OdmWriteBack for MockWriteBack { + async fn local_object(&self, _bucket: &str, key: &str) -> Result, WriteBackError> { + Ok(self.local.lock().get(key).cloned()) + } + + async fn put_object(&self, request: &WriteBackRequest, body: WriteBackBody) -> Result { + let drained = drain(body, request.head.size).await; + let bytes = match drained { + Ok(bytes) => bytes, + Err(err) => { + self.failed_puts.fetch_add(1, Ordering::SeqCst); + return Err(err); + } + }; + if let Some(err) = self.forced_put_error.lock().take() { + self.failed_puts.fetch_add(1, Ordering::SeqCst); + return Err(err); + } + self.puts.lock().push((request.clone(), bytes)); + self.local.lock().insert( + request.key.clone(), + LocalObject { + etag: request.head.etag.clone(), + size: request.head.size, + delete_marker: false, + }, + ); + Ok(WriteBackOutcome { + etag: request.head.etag.clone(), + size: request.head.size, + version_id: None, + }) + } + + async fn create_multipart_upload(&self, _request: &WriteBackRequest) -> Result { + let id = format!("upload-{}", self.upload_seq.fetch_add(1, Ordering::SeqCst)); + self.uploads.lock().push(id.clone()); + Ok(id) + } + + async fn upload_part( + &self, + _request: &WriteBackRequest, + upload_id: &str, + part_number: usize, + size: u64, + body: WriteBackBody, + ) -> Result { + let bytes = drain(body, size).await?; + if self.fail_part == Some(part_number) { + return Err(WriteBackError::Local("injected part failure".to_string())); + } + self.parts.lock().push((upload_id.to_string(), part_number, size, bytes)); + Ok(WriteBackPart { + part_number, + etag: format!("part-{part_number}"), + }) + } + + async fn complete_multipart_upload( + &self, + request: &WriteBackRequest, + upload_id: &str, + parts: Vec, + ) -> Result { + self.completed.lock().push((upload_id.to_string(), parts)); + self.local.lock().insert( + request.key.clone(), + LocalObject { + etag: request.head.etag.clone(), + size: request.head.size, + delete_marker: false, + }, + ); + Ok(WriteBackOutcome { + etag: request.head.etag.clone(), + size: request.head.size, + version_id: None, + }) + } + + async fn abort_multipart_upload(&self, _bucket: &str, _key: &str, upload_id: &str) -> Result<(), WriteBackError> { + self.aborted.lock().push(upload_id.to_string()); + Ok(()) + } + } + + async fn wait_until(what: &str, condition: impl Fn() -> bool) { + tokio::time::timeout(Duration::from_secs(10), async { + while !condition() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); + } + + fn failures(state: &BucketOdmState) -> std::collections::BTreeMap { + state + .snapshot() + .stats + .pull_failures_total + .into_iter() + .filter(|(_, count)| *count > 0) + .collect() + } + + fn pulled(state: &BucketOdmState, path: PullPath) -> u64 { + state.snapshot().stats.pulled_objects_total[path.as_str()] + } + + async fn pull( + state: &Arc, + source: &Arc, + write_back: &Arc, + key: &str, + ) -> Result { + let cancel = state.cancel_token(); + let ctx = PullContext { + state, + source, + write_back, + key, + cancel: &cancel, + }; + let result = pull_object(&ctx).await; + record_completion(state, key, PullPath::Background, &result); + result + } + + #[tokio::test] + async fn coalesced_enqueues_produce_one_source_get() { + let sys = OnDemandMigrationSys::new(); + let state = enabled_state(&sys, &config()).await; + let body = body_bytes(1000); + let source = MockSource::with_object("a", 1000, BodyKind::Bytes(body.clone())); + let write_back = Arc::new(MockWriteBack::default()); + let queue = PullQueue::start(Arc::clone(&state), source.clone(), write_back.clone()); + assert_eq!(queue.capacity(), 1024); + + let mut outcomes = HashMap::new(); + for _ in 0..100 { + *outcomes.entry(queue.enqueue("a", PullReason::RangeGet)).or_insert(0) += 1; + } + assert_eq!(outcomes.get(&EnqueueOutcome::Enqueued), Some(&1)); + assert_eq!(outcomes.get(&EnqueueOutcome::Coalesced), Some(&99)); + assert_eq!(queue.pending_keys(), 1); + + wait_until("first pull to finish", || queue.pending_keys() == 0).await; + assert_eq!(source.head_calls.load(Ordering::SeqCst), 1); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1); + assert_eq!(source.tag_calls.load(Ordering::SeqCst), 0, "copy_tags is off by default"); + { + let puts = write_back.puts.lock(); + assert_eq!(puts.len(), 1); + let (request, stored) = &puts[0]; + assert_eq!(stored, &body); + assert_eq!(request.bucket, BUCKET); + assert_eq!(request.key, "a"); + assert_eq!(request.source_label, "minio:legacy"); + assert!(request.preserve_etag && request.emit_events, "policy defaults flow into the request"); + assert!(request.tags.is_none()); + } + assert_eq!(pulled(&state, PullPath::Background), 1); + assert_eq!(state.snapshot().stats.pulled_bytes_total, 1000); + assert!(failures(&state).is_empty(), "{:?}", failures(&state)); + + // A second round finds the local copy and does not touch the source. + assert_eq!(queue.enqueue("a", PullReason::Backfill), EnqueueOutcome::Enqueued); + wait_until("second job to finish", || queue.pending_keys() == 0).await; + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1); + assert_eq!(pulled(&state, PullPath::Background), 1); + assert_eq!(pulled(&state, PullPath::Backfill), 0); + assert_eq!(state.inflight_keys(), 0); + + sys.remove(BUCKET); + queue.wait_until_stopped().await; + assert_eq!(queue.enqueue("a", PullReason::RangeGet), EnqueueOutcome::Unavailable); + } + + #[tokio::test] + async fn queue_full_is_reported_and_cancel_drains_without_leaking_tasks() { + let sys = OnDemandMigrationSys::new(); + let mut cfg = config(); + cfg.policy.pull_queue_capacity = 1; + cfg.policy.max_concurrent_pulls = 1; + let state = enabled_state(&sys, &cfg).await; + let source = MockSource::with_object("hang", 10, BodyKind::Hang); + for key in ["b", "c", "d"] { + source.objects.lock().insert( + key.to_string(), + MockObject { + head: head(3), + body: BodyKind::Bytes(vec![1, 2, 3]), + }, + ); + } + let write_back = Arc::new(MockWriteBack::default()); + let queue = PullQueue::start(Arc::clone(&state), source.clone(), write_back.clone()); + + assert_eq!(queue.enqueue("hang", PullReason::LargeObject), EnqueueOutcome::Enqueued); + wait_until("hanging pull to reach the body", || source.get_calls.load(Ordering::SeqCst) == 1).await; + assert_eq!(state.stats().inflight_pulls(), 1); + // The dispatcher takes "b" and blocks on the pull slot; "c" fills the + // single channel slot; "d" has nowhere to go. + assert_eq!(queue.enqueue("b", PullReason::LargeObject), EnqueueOutcome::Enqueued); + wait_until("dispatcher to wait for a slot", || state.stats().queue_depth() == 1).await; + assert_eq!(queue.enqueue("c", PullReason::LargeObject), EnqueueOutcome::Enqueued); + assert_eq!(queue.enqueue("d", PullReason::LargeObject), EnqueueOutcome::QueueFull); + assert_eq!(queue.enqueue("c", PullReason::LargeObject), EnqueueOutcome::Coalesced); + assert_eq!(queue.pending_keys(), 3); + assert_eq!(failures(&state).get("queue_full"), Some(&1)); + assert!(!queue.is_stopped()); + + assert_eq!(sys.remove(BUCKET), crate::bucket::on_demand_migration::ApplyOutcome::Removed); + tokio::time::timeout(Duration::from_secs(5), queue.wait_until_stopped()) + .await + .expect("dispatcher and in-flight job must exit after cancel"); + assert!(queue.is_stopped()); + assert_eq!(queue.pending_keys(), 0); + assert_eq!(state.inflight_keys(), 0); + assert_eq!(state.stats().inflight_pulls(), 0); + assert_eq!(state.stats().queue_depth(), 0); + let failures = failures(&state); + assert_eq!(failures.get("canceled"), Some(&3), "{failures:?}"); + assert!(write_back.puts.lock().is_empty()); + assert_eq!(queue.enqueue("b", PullReason::LargeObject), EnqueueOutcome::Unavailable); + } + + #[tokio::test(start_paused = true)] + async fn retry_policy_retries_transient_source_errors_only() { + let sys = OnDemandMigrationSys::new(); + let state = enabled_state(&sys, &config()).await; + let source = MockSource::with_object("a", 5, BodyKind::Bytes(vec![7; 5])); + source + .head_failures + .lock() + .extend([SourceError::ServerError(500), SourceError::ServerError(503)]); + let write_back = Arc::new(MockWriteBack::default()); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + + let started = tokio::time::Instant::now(); + let completion = pull(&state, &source_dyn, &write_back_dyn, "a") + .await + .expect("two server errors are within the retry budget"); + assert!(matches!(completion, PullCompletion::Stored(_))); + let elapsed = started.elapsed(); + assert!(elapsed >= Duration::from_secs(5), "1s + 4s base backoff, got {elapsed:?}"); + assert!(elapsed <= Duration::from_millis(6_300), "jitter is at most 25%, got {elapsed:?}"); + assert_eq!(source.head_calls.load(Ordering::SeqCst), 3); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1); + assert!(failures(&state).is_empty(), "{:?}", failures(&state)); + assert_eq!(pulled(&state, PullPath::Background), 1); + + write_back.local.lock().clear(); + source.head_failures.lock().push_back(SourceError::AccessDenied); + let err = pull(&state, &source_dyn, &write_back_dyn, "a") + .await + .expect_err("access denied is not retried"); + assert_eq!(err.reason, PullFailureReason::SourceAccessDenied); + assert_eq!(source.head_calls.load(Ordering::SeqCst), 4); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1); + assert_eq!(failures(&state).get("source_access_denied"), Some(&1)); + + // Exhausting the budget: four attempts, then the last error wins. + source.get_failures.lock().extend((0..4).map(|_| SourceError::Timeout)); + let err = pull(&state, &source_dyn, &write_back_dyn, "a") + .await + .expect_err("four timeouts exceed three retries"); + assert_eq!(err.reason, PullFailureReason::SourceTimeout); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 5); + assert_eq!(failures(&state).get("source_timeout"), Some(&1)); + } + + #[tokio::test(start_paused = true)] + async fn truncated_source_body_never_commits_and_is_retried() { + let sys = OnDemandMigrationSys::new(); + let state = enabled_state(&sys, &config()).await; + // Advertises 1000 bytes, delivers 600. + let source = MockSource::with_object("short", 1000, BodyKind::Bytes(body_bytes(600))); + let write_back = Arc::new(MockWriteBack::default()); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + + let err = pull(&state, &source_dyn, &write_back_dyn, "short") + .await + .expect_err("a truncated body must not commit"); + assert_eq!(err.reason, PullFailureReason::SourceConnect, "{err}"); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1 + PULL_MAX_RETRIES); + assert_eq!(write_back.failed_puts.load(Ordering::SeqCst), 1 + PULL_MAX_RETRIES); + assert!(write_back.puts.lock().is_empty()); + assert!(write_back.local.lock().is_empty()); + assert_eq!(failures(&state).get("source_connect"), Some(&1)); + } + + #[tokio::test(start_paused = true)] + async fn stalled_source_body_hits_the_idle_timeout() { + let sys = OnDemandMigrationSys::new(); + let mut cfg = config(); + cfg.policy.source_timeout.idle_ms = 1_000; + let state = enabled_state(&sys, &cfg).await; + let source = MockSource::with_object("stall", 10, BodyKind::Hang); + let write_back = Arc::new(MockWriteBack::default()); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + + let started = tokio::time::Instant::now(); + let err = pull(&state, &source_dyn, &write_back_dyn, "stall") + .await + .expect_err("a stalled body must time out"); + assert_eq!(err.reason, PullFailureReason::SourceTimeout, "{err}"); + assert!(started.elapsed() >= Duration::from_secs(4 + 21), "4 idle timeouts + 3 backoffs"); + assert!(write_back.local.lock().is_empty()); + assert_eq!(failures(&state).get("source_timeout"), Some(&1)); + } + + #[tokio::test] + async fn multipart_path_splits_parts_and_aborts_on_failure() { + let sys = OnDemandMigrationSys::new(); + let mut cfg = config(); + cfg.policy.multipart_part_size_bytes = 1024; + let state = enabled_state(&sys, &cfg).await; + let body = body_bytes(2500); + let source = MockSource::with_object("big", 2500, BodyKind::Bytes(body.clone())); + let write_back = Arc::new(MockWriteBack::default()); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + + let completion = pull(&state, &source_dyn, &write_back_dyn, "big") + .await + .expect("multipart write-back must succeed"); + let PullCompletion::Stored(outcome) = completion else { + panic!("expected a stored object"); + }; + assert_eq!(outcome.size, 2500); + let parts = write_back.parts.lock().clone(); + assert_eq!(parts.iter().map(|(_, n, _, _)| *n).collect::>(), vec![1, 2, 3]); + assert_eq!(parts.iter().map(|(_, _, size, _)| *size).collect::>(), vec![1024, 1024, 452]); + let joined: Vec = parts.iter().flat_map(|(_, _, _, bytes)| bytes.clone()).collect(); + assert_eq!(joined, body); + { + let completed = write_back.completed.lock(); + assert_eq!(completed.len(), 1); + assert_eq!(completed[0].1.len(), 3); + } + assert!(write_back.aborted.lock().is_empty()); + assert!(write_back.puts.lock().is_empty(), "large objects never use the single-part path"); + assert_eq!(state.snapshot().stats.pulled_bytes_total, 2500); + + // A failing part aborts the upload and reports a local write failure. + let failing = Arc::new(MockWriteBack { + fail_part: Some(2), + ..Default::default() + }); + let failing_dyn: Arc = failing.clone(); + let err = pull(&state, &source_dyn, &failing_dyn, "big") + .await + .expect_err("part failure must fail the pull"); + assert_eq!(err.reason, PullFailureReason::LocalWrite); + assert_eq!(failing.aborted.lock().as_slice(), ["upload-0"]); + assert!(failing.completed.lock().is_empty()); + assert!(failing.local.lock().is_empty()); + assert_eq!(failures(&state).get("local_write"), Some(&1)); + + // Too many parts is rejected before any upload is created. + let mut tiny = config(); + tiny.policy.multipart_part_size_bytes = 1; + let sys2 = OnDemandMigrationSys::new(); + let state2 = enabled_state(&sys2, &tiny).await; + let big = Arc::new(MockWriteBack::default()); + let big_dyn: Arc = big.clone(); + let source2 = MockSource::with_object("huge", 20_000, BodyKind::Bytes(body_bytes(20_000))); + let source2_dyn: Arc = source2; + let err = pull(&state2, &source2_dyn, &big_dyn, "huge") + .await + .expect_err("more than 10000 parts is unsupported"); + assert_eq!(err.reason, PullFailureReason::SourceUnsupported); + assert!(big.uploads.lock().is_empty()); + } + + #[tokio::test] + async fn local_object_checks_skip_the_pull_before_and_after_the_source_get() { + let sys = OnDemandMigrationSys::new(); + let state = enabled_state(&sys, &config()).await; + let source = MockSource::with_object("present", 4, BodyKind::Bytes(vec![1; 4])); + let write_back = Arc::new(MockWriteBack::default()); + write_back.local.lock().insert( + "present".to_string(), + LocalObject { + etag: Some("local".to_string()), + size: 4, + delete_marker: false, + }, + ); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + let completion = pull(&state, &source_dyn, &write_back_dyn, "present").await.expect("skip"); + assert!(matches!(completion, PullCompletion::AlreadyPresent(ref local) if local.etag.as_deref() == Some("local"))); + assert_eq!(source.head_calls.load(Ordering::SeqCst), 0); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 0); + assert_eq!(pulled(&state, PullPath::Background), 0); + + // A client PUT that lands during the source GET wins over the copy. + source.objects.lock().insert( + "racy".to_string(), + MockObject { + head: head(4), + body: BodyKind::Bytes(vec![2; 4]), + }, + ); + *source.land_local_on_get.lock() = Some((write_back.clone(), "racy".to_string())); + let completion = pull(&state, &source_dyn, &write_back_dyn, "racy").await.expect("skip"); + assert!(matches!(completion, PullCompletion::AlreadyPresent(ref local) if local.etag.as_deref() == Some("client-put"))); + assert_eq!(source.get_calls.load(Ordering::SeqCst), 1); + assert!(write_back.puts.lock().is_empty(), "the source copy must not overwrite the client write"); + assert!(failures(&state).is_empty()); + } + + #[tokio::test] + async fn copy_tags_policy_fetches_source_tags() { + let sys = OnDemandMigrationSys::new(); + let mut cfg = config(); + cfg.policy.copy_tags = true; + let state = enabled_state(&sys, &cfg).await; + let source = MockSource { + tags: HashMap::from([("env".to_string(), "prod".to_string())]), + ..Default::default() + }; + source.objects.lock().insert( + "tagged".to_string(), + MockObject { + head: head(2), + body: BodyKind::Bytes(vec![9, 9]), + }, + ); + let source = Arc::new(source); + let write_back = Arc::new(MockWriteBack::default()); + let source_dyn: Arc = source.clone(); + let write_back_dyn: Arc = write_back.clone(); + pull(&state, &source_dyn, &write_back_dyn, "tagged").await.expect("stored"); + assert_eq!(source.tag_calls.load(Ordering::SeqCst), 1); + let puts = write_back.puts.lock(); + assert_eq!(puts[0].0.tags.as_ref().and_then(|tags| tags.get("env")).map(String::as_str), Some("prod")); + } + + #[tokio::test] + async fn commit_inline_accounts_the_inline_path_and_maps_write_errors() { + let sys = OnDemandMigrationSys::new(); + let mock = Arc::new(MockWriteBack::default()); + let mock_dyn: Arc = mock.clone(); + sys.set_write_back(mock_dyn.clone()); + let state = enabled_state(&sys, &config()).await; + assert!(state.write_back().is_some(), "states capture the injected write-back"); + + let body = body_bytes(300); + let stream: WriteBackBody = Box::pin(futures::stream::iter(vec![Ok(Bytes::from(body.clone()))])); + let outcome = commit_inline(&state, "inline", head(300), None, stream) + .await + .expect("inline commit succeeds"); + assert_eq!(outcome.size, 300); + assert_eq!(mock.puts.lock()[0].1, body); + assert_eq!(pulled(&state, PullPath::Inline), 1); + assert_eq!(pulled(&state, PullPath::Background), 0); + assert_eq!(state.snapshot().stats.pulled_bytes_total, 300); + + *mock.forced_put_error.lock() = Some(WriteBackError::Integrity); + let stream: WriteBackBody = Box::pin(futures::stream::iter(vec![Ok(Bytes::from(body.clone()))])); + let err = commit_inline_with(&state, &mock_dyn, "inline", head(300), None, stream) + .await + .expect_err("integrity failure surfaces"); + assert_eq!(err.reason, PullFailureReason::EtagMismatch); + assert_eq!(failures(&state).get("etag_mismatch"), Some(&1)); + + // A tee secondary that ends early (primary dropped) never commits. + let stream: WriteBackBody = Box::pin(futures::stream::iter(vec![ + Ok(Bytes::from(body_bytes(100))), + Err(io::Error::new(io::ErrorKind::BrokenPipe, "tee primary dropped")), + ])); + let err = commit_inline(&state, "torn", head(300), None, stream) + .await + .expect_err("a broken secondary must fail"); + assert_eq!(err.reason, PullFailureReason::LocalWrite); + assert!(mock.local.lock().get("torn").is_none()); + + // Without an injected write-back the inline path fails closed. + let bare = OnDemandMigrationSys::new(); + let bare_state = enabled_state(&bare, &config()).await; + assert!(bare_state.write_back().is_none()); + let stream: WriteBackBody = Box::pin(futures::stream::empty()); + let err = commit_inline(&bare_state, "x", head(0), None, stream) + .await + .expect_err("no write-back"); + assert_eq!(err.reason, PullFailureReason::LocalWrite); + assert_eq!(failures(&bare_state).get("local_write"), Some(&1)); + } + + #[tokio::test] + async fn state_and_sys_enqueue_entry_points() { + let sys = OnDemandMigrationSys::new(); + assert_eq!(sys.enqueue_pull("missing", "k", PullReason::Backfill), EnqueueOutcome::Unavailable); + let state = enabled_state(&sys, &config()).await; + // No write-back injected: the lazily started queue is unavailable. + assert!(state.pull_queue().is_none()); + assert_eq!(state.enqueue_pull("k", PullReason::Backfill), EnqueueOutcome::Unavailable); + + let source = MockSource::with_object("k", 3, BodyKind::Bytes(vec![1, 2, 3])); + let write_back = Arc::new(MockWriteBack::default()); + let queue = PullQueue::start(Arc::clone(&state), source.clone(), write_back.clone()); + assert!(state.pull_queue.set(Arc::clone(&queue)).is_ok()); + assert!(Arc::ptr_eq(&state.pull_queue().expect("seeded"), &queue)); + assert_eq!(sys.enqueue_pull(BUCKET, "k", PullReason::Backfill), EnqueueOutcome::Enqueued); + assert_eq!(state.enqueue_pull("k", PullReason::Backfill), EnqueueOutcome::Coalesced); + wait_until("backfill job", || queue.pending_keys() == 0).await; + assert_eq!(pulled(&state, PullPath::Backfill), 1); + assert_eq!( + format!("{queue:?}"), + format!("PullQueue {{ bucket: {BUCKET:?}, capacity: 1024, pending: 0, stopped: false }}") + ); + sys.remove(BUCKET); + queue.wait_until_stopped().await; + } + + #[tokio::test] + async fn part_body_slices_the_shared_source_at_part_boundaries() { + let (tx, rx) = mpsc::channel(4); + tx.send(Ok(Bytes::from(body_bytes(300)))).await.expect("send"); + tx.send(Ok(Bytes::from(body_bytes(500)))).await.expect("send"); + drop(tx); + let shared: SharedSourceHandle = Arc::new(Mutex::new(SharedSource { rx, leftover: None })); + let mut expected = body_bytes(300); + expected.extend(body_bytes(500)); + + let first = drain(PartBody::boxed(&shared, 600), 600).await.expect("first part"); + assert_eq!(first, expected[..600]); + let second = drain(PartBody::boxed(&shared, 200), 200).await.expect("second part"); + assert_eq!(second, expected[600..]); + let err = drain(PartBody::boxed(&shared, 1), 1).await.expect_err("source exhausted"); + assert!( + err.to_string().contains("UnexpectedEof") || err.to_string().contains("outstanding"), + "{err}" + ); + let empty = drain(PartBody::boxed(&shared, 0), 0).await.expect("zero-length part"); + assert!(empty.is_empty()); + } + + #[test] + fn reasons_paths_and_error_classes_are_fixed() { + assert_eq!(PullReason::RangeGet.path(), PullPath::Background); + assert_eq!(PullReason::LargeObject.path(), PullPath::Background); + assert_eq!(PullReason::Backfill.path(), PullPath::Backfill); + assert_eq!(PullReason::RangeGet.as_str(), "range_get"); + assert_eq!(WriteBackError::Integrity.reason(), PullFailureReason::EtagMismatch); + assert_eq!(WriteBackError::Quota("full".into()).reason(), PullFailureReason::LocalWrite); + assert_eq!(WriteBackError::Local("x".into()).reason(), PullFailureReason::LocalWrite); + assert_eq!(WriteBackError::Unsupported("x".into()).reason(), PullFailureReason::SourceUnsupported); + for (retry, base) in PULL_RETRY_BASE_DELAYS.iter().enumerate() { + let delay = retry_delay(retry); + assert!(delay >= *base && delay <= *base + *base / 4, "{delay:?}"); + } + } +} diff --git a/crates/ecstore/src/bucket/on_demand_migration/sys.rs b/crates/ecstore/src/bucket/on_demand_migration/sys.rs index 366b1d5a9..3e5ae9cc5 100644 --- a/crates/ecstore/src/bucket/on_demand_migration/sys.rs +++ b/crates/ecstore/src/bucket/on_demand_migration/sys.rs @@ -35,12 +35,16 @@ //! The module switch (`RUSTFS_ON_DEMAND_MIGRATION_ENABLED`, default off) is //! injected by the `rustfs` binary through [`OnDemandMigrationSys::set_module_enabled`] //! before bucket metadata loads; this crate never reads the environment. +//! The same startup step injects the [`OdmWriteBack`] the pull pipeline +//! (`pull.rs`) stores objects with; each bucket state captures it at build +//! time together with its lazily started [`PullQueue`]. use super::breaker::{Breaker, BreakerState, BreakerTransition, BreakerVerdict}; use super::config::{ ON_DEMAND_MIGRATION_CONFIG_HOOK, OnDemandMigrationConfig, PathStyle as ConfigPathStyle, Provider, SourceConfig, }; use super::negative_cache::NegativeCache; +use super::pull::{OdmWriteBack, PullQueue}; use super::source_client::{SourceClient, SourceClientSpec, SourceError, SourceProvider, SourceTimeouts}; use super::stats::{GaugeGuard, OdmStats, OdmStatsSnapshot, PullFailureReason}; use crate::bucket::remote_s3_client::{PathStyle as ClientPathStyle, RemoteCredentials, RemoteS3ClientError}; @@ -269,6 +273,9 @@ pub struct BucketOdmState { stats: Arc, cancel: CancellationToken, last_source_error_logged_at: Mutex>, + write_back: Option>, + /// Started by `pull::BucketOdmState::pull_queue` on first enqueue. + pub(super) pull_queue: OnceLock>, } impl fmt::Debug for BucketOdmState { @@ -286,7 +293,12 @@ impl fmt::Debug for BucketOdmState { } impl BucketOdmState { - async fn build(bucket: &str, config: &OnDemandMigrationConfig, stats: Arc) -> Arc { + async fn build( + bucket: &str, + config: &OnDemandMigrationConfig, + stats: Arc, + write_back: Option>, + ) -> Arc { let spec = source_client_spec(config); let client = if config.source.credentials.is_none() { Err(OdmStateError::AnonymousUnsupported) @@ -310,6 +322,8 @@ impl BucketOdmState { stats, cancel: CancellationToken::new(), last_source_error_logged_at: Mutex::new(None), + write_back, + pull_queue: OnceLock::new(), }) } @@ -346,6 +360,11 @@ impl BucketOdmState { &self.stats } + /// The write-back injected by the binary when this state was built. + pub fn write_back(&self) -> Option<&Arc> { + self.write_back.as_ref() + } + /// Fires when this state is replaced or removed; background pulls /// started for it must exit. pub fn cancel_token(&self) -> CancellationToken { @@ -601,6 +620,7 @@ pub struct OnDemandMigrationSys { module_enabled: AtomicBool, buckets: RwLock>, generation: AtomicU64, + write_back: RwLock>>, } impl fmt::Debug for OnDemandMigrationSys { @@ -625,6 +645,7 @@ impl OnDemandMigrationSys { module_enabled: AtomicBool::new(false), buckets: RwLock::new(HashMap::new()), generation: AtomicU64::new(0), + write_back: RwLock::new(None), } } @@ -641,6 +662,17 @@ impl OnDemandMigrationSys { self.module_enabled.load(Ordering::Relaxed) } + /// Installs the local write path used by every bucket state built from + /// now on (states built earlier keep what they captured). The binary + /// calls this before bucket metadata loads. + pub fn set_write_back(&self, write_back: Arc) { + *self.write_back.write() = Some(write_back); + } + + pub fn write_back(&self) -> Option> { + self.write_back.read().clone() + } + /// Registers `publish` as the bucket-metadata publish hook. Returns /// `false` when a hook was already registered. pub fn register_config_hook(&'static self) -> bool { @@ -701,7 +733,7 @@ impl OnDemandMigrationSys { return ApplyOutcome::Unchanged; } let stats = self.state(bucket).map(|state| Arc::clone(&state.stats)).unwrap_or_default(); - let state = BucketOdmState::build(bucket, config, stats).await; + let state = BucketOdmState::build(bucket, config, stats, self.write_back()).await; let (outcome, previous) = { let mut buckets = self.buckets.write(); diff --git a/rustfs/src/app/object/internal_put.rs b/rustfs/src/app/object/internal_put.rs index dcc406c1f..0455ca192 100644 --- a/rustfs/src/app/object/internal_put.rs +++ b/rustfs/src/app/object/internal_put.rs @@ -163,7 +163,7 @@ impl InternalPutObjectEvent { }) } - fn builder(event_name: EventName, bucket: &str, key: &str, principal_id: &'static str) -> EventArgsBuilder { + pub(super) fn builder(event_name: EventName, bucket: &str, key: &str, principal_id: &'static str) -> EventArgsBuilder { // The object is a placeholder until `object()` supplies the committed // ObjectInfo, matching the S3 helper. let placeholder = ObjectInfo { diff --git a/rustfs/src/app/object/mod.rs b/rustfs/src/app/object/mod.rs index a5ca06d7e..367c2aff0 100644 --- a/rustfs/src/app/object/mod.rs +++ b/rustfs/src/app/object/mod.rs @@ -191,10 +191,8 @@ mod delete; mod extract; mod get; mod head; -// Consumed by the on-demand migration write-back (rustfs/backlog#2153); until -// that lands only tests construct the internal entry points. -#[cfg_attr(not(test), expect(dead_code, reason = "wired by the on-demand migration write-back"))] mod internal_put; +mod on_demand_migration_put; mod put; mod restore; mod shared; @@ -207,6 +205,7 @@ pub(crate) use self::delete::*; pub(crate) use self::extract::*; pub(crate) use self::get::*; pub(crate) use self::internal_put::*; +pub(crate) use self::on_demand_migration_put::*; use self::put::*; pub(crate) use self::put::{guard_put_object_body_read_timeout, put_object_body_read_timeout}; pub(crate) use self::shared::*; diff --git a/rustfs/src/app/object/on_demand_migration_put.rs b/rustfs/src/app/object/on_demand_migration_put.rs new file mode 100644 index 000000000..89ed3addb --- /dev/null +++ b/rustfs/src/app/object/on_demand_migration_put.rs @@ -0,0 +1,863 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! On-demand migration write-back (rustfs/backlog#2153): the app-layer +//! [`OdmWriteBack`] the ecstore pull pipeline stores source objects with. +//! +//! Every write goes through the internal put entry points, so a pulled +//! object is indistinguishable from a client PUT: bucket default SSE, quota, +//! versioning, Object Lock defaults, replication scheduling and creation +//! events all apply. This module only maps a [`WriteBackRequest`] onto an +//! [`InternalPutContext`]: the content-header allowlist, the `x-amz-meta-*` +//! copy, optional tags, the five `odm-*` provenance keys (dual prefix), and +//! the ETag policy. +//! +//! ETag policy: a single-part source ETag (32 hex digits) of an unencrypted +//! source object is the plaintext MD5 and doubles as the integrity check; +//! `policy.preserve_etag` keeps the source ETag (multipart ETags included, +//! display only) unless the bucket encrypts by default, where the override +//! is dropped and the SSE write path decides the ETag, exactly like +//! replication receive. The source ETag is always recorded under +//! `odm-source-etag`. + +use super::*; + +use crate::app::storage_api::multipart_usecase::contract::multipart::CompletePart; +use crate::app::storage_api::object_usecase::on_demand_migration::{ + LocalObject, OdmWriteBack, SourceHead, WriteBackBody, WriteBackError, WriteBackOutcome, WriteBackPart, WriteBackRequest, + is_multipart_etag, +}; +use rustfs_utils::http::{ + SUFFIX_ODM_PULLED_AT, SUFFIX_ODM_SOURCE, SUFFIX_ODM_SOURCE_ETAG, SUFFIX_ODM_SOURCE_LAST_MODIFIED, + SUFFIX_ODM_SOURCE_VERSION_ID, +}; + +/// `userIdentity.principalId` of every creation event a write-back emits. +pub(crate) const ON_DEMAND_MIGRATION_PRINCIPAL_ID: &str = "rustfs-on-demand-migration"; + +/// [`OdmWriteBack`] over [`DefaultObjectUsecase`]'s internal put entry +/// points. The ambient app context is resolved per call: the write-back is +/// installed at startup, before the context is published. +#[derive(Debug, Default)] +pub(crate) struct OnDemandMigrationWriteBack; + +impl OnDemandMigrationWriteBack { + pub(crate) fn new() -> Self { + Self + } + + fn usecase(&self) -> DefaultObjectUsecase { + DefaultObjectUsecase::from_global() + } + + fn store(&self) -> Result, WriteBackError> { + self.usecase() + .object_store() + .ok_or_else(|| WriteBackError::Local("object store is not initialized".to_string())) + } +} + +fn rfc3339(time: OffsetDateTime) -> String { + time.format(&Rfc3339).unwrap_or_default() +} + +/// Standard object headers copied from the source. `storage_class` and the +/// source `Last-Modified` are deliberately absent: the local class follows +/// the bucket and the local mtime is the write time. +pub(super) fn content_headers(head: &SourceHead) -> HashMap { + let mut headers = HashMap::with_capacity(6); + for (name, value) in [ + ("Content-Type", &head.content_type), + ("Content-Encoding", &head.content_encoding), + ("Content-Disposition", &head.content_disposition), + ("Content-Language", &head.content_language), + ("Cache-Control", &head.cache_control), + ("Expires", &head.expires), + ] { + if let Some(value) = value.as_deref().map(str::trim).filter(|value| !value.is_empty()) { + headers.insert(name.to_string(), value.to_string()); + } + } + headers +} + +/// The five `odm-*` provenance keys under both internal prefixes. Absent +/// source values are stored as empty strings so the key set is constant. +pub(super) fn provenance_metadata(request: &WriteBackRequest) -> HashMap { + let head = &request.head; + let mut metadata = HashMap::with_capacity(10); + insert_str(&mut metadata, SUFFIX_ODM_SOURCE, request.source_label.clone()); + insert_str(&mut metadata, SUFFIX_ODM_SOURCE_ETAG, head.etag.clone().unwrap_or_default()); + insert_str( + &mut metadata, + SUFFIX_ODM_SOURCE_LAST_MODIFIED, + head.last_modified.map(OffsetDateTime::from).map(rfc3339).unwrap_or_default(), + ); + insert_str(&mut metadata, SUFFIX_ODM_SOURCE_VERSION_ID, head.version_id.clone().unwrap_or_default()); + insert_str(&mut metadata, SUFFIX_ODM_PULLED_AT, rfc3339(request.pulled_at)); + metadata +} + +/// The source ETag as the expected plaintext MD5: only a bare 32-digit hex +/// ETag of an unencrypted source object is one. +pub(super) fn expected_md5_hex(head: &SourceHead) -> Option { + if head.sse.is_some() { + return None; + } + let etag = head.etag.as_deref()?; + if etag.len() != 32 || is_multipart_etag(etag) || !etag.bytes().all(|byte| byte.is_ascii_hexdigit()) { + return None; + } + Some(etag.to_ascii_lowercase()) +} + +/// `x-amz-tagging` form of the source tags, sorted by key for a stable +/// stored value. +pub(super) fn encode_tags(tags: &HashMap) -> Option { + if tags.is_empty() { + return None; + } + let mut pairs: Vec<(&String, &String)> = tags.iter().collect(); + pairs.sort(); + let mut encoded = url::form_urlencoded::Serializer::new(String::new()); + for (key, value) in pairs { + encoded.append_pair(key, value); + } + Some(encoded.finish()) +} + +async fn bucket_encrypts_by_default(bucket: &str) -> bool { + metadata_sys::get_sse_config(bucket).await.is_ok_and(|(config, _)| { + config + .rules + .iter() + .any(|rule| rule.apply_server_side_encryption_by_default.is_some()) + }) +} + +/// Builds the internal put context of a write-back. `single_part` enables +/// the MD5 integrity check, which only the single-object path can honor. +pub(super) async fn write_back_context(request: &WriteBackRequest, single_part: bool) -> InternalPutContext { + let head = &request.head; + let preserve_etag = if request.preserve_etag && head.etag.is_some() && !bucket_encrypts_by_default(&request.bucket).await { + head.etag.clone() + } else { + None + }; + InternalPutContext { + bucket: request.bucket.clone(), + key: request.key.clone(), + size: Some(head.size), + expected_md5_hex: single_part.then(|| expected_md5_hex(head)).flatten(), + preserve_etag, + content_headers: content_headers(head), + user_metadata: head.user_metadata.clone(), + tags: request.tags.as_ref().and_then(encode_tags), + internal_metadata: provenance_metadata(request), + emit_events: request.emit_events, + principal_id: ON_DEMAND_MIGRATION_PRINCIPAL_ID, + } +} + +/// Maps an internal put failure onto the write-back error classes. A +/// digest mismatch is the only integrity signal; both quota producers +/// (admission and the durable reservation) say "quota exceeded". +pub(super) fn write_back_error(err: ApiError) -> WriteBackError { + if err.code == S3ErrorCode::BadDigest { + return WriteBackError::Integrity; + } + if err.message.to_ascii_lowercase().contains("quota exceeded") { + return WriteBackError::Quota(err.message); + } + WriteBackError::Local(format!("{}: {}", err.code.as_str(), err.message)) +} + +fn outcome(info: ObjectInfo) -> WriteBackOutcome { + WriteBackOutcome { + etag: info.etag, + size: u64::try_from(info.size).unwrap_or(0), + version_id: info.version_id.map(|version_id| version_id.to_string()), + } +} + +#[async_trait::async_trait] +impl OdmWriteBack for OnDemandMigrationWriteBack { + async fn local_object(&self, bucket: &str, key: &str) -> Result, WriteBackError> { + let store = self.store()?; + match store.get_object_info(bucket, key, &ObjectOptions::default()).await { + Ok(info) => Ok(Some(LocalObject { + etag: info.etag.clone(), + size: u64::try_from(info.size).unwrap_or(0), + delete_marker: info.delete_marker, + })), + Err(err) if is_err_object_not_found(&err) || is_err_version_not_found(&err) => Ok(None), + Err(err) => Err(WriteBackError::Local(err.to_string())), + } + } + + async fn put_object(&self, request: &WriteBackRequest, body: WriteBackBody) -> Result { + let ctx = write_back_context(request, true).await; + self.usecase() + .internal_put_object(ctx, body) + .await + .map(outcome) + .map_err(write_back_error) + } + + async fn create_multipart_upload(&self, request: &WriteBackRequest) -> Result { + let ctx = write_back_context(request, false).await; + self.usecase() + .internal_create_multipart_upload(&ctx) + .await + .map_err(write_back_error) + } + + async fn upload_part( + &self, + request: &WriteBackRequest, + upload_id: &str, + part_number: usize, + size: u64, + body: WriteBackBody, + ) -> Result { + let ctx = write_back_context(request, false).await; + let part = self + .usecase() + .internal_upload_part(&ctx, upload_id, part_number, size, None, body) + .await + .map_err(write_back_error)?; + Ok(WriteBackPart { + part_number: part.part_num, + etag: part.etag.unwrap_or_default(), + }) + } + + async fn complete_multipart_upload( + &self, + request: &WriteBackRequest, + upload_id: &str, + parts: Vec, + ) -> Result { + let ctx = write_back_context(request, false).await; + let parts = parts + .into_iter() + .map(|part| CompletePart { + part_num: part.part_number, + etag: Some(part.etag), + ..Default::default() + }) + .collect(); + self.usecase() + .internal_complete_multipart_upload(&ctx, upload_id, parts) + .await + .map(outcome) + .map_err(write_back_error) + } + + async fn abort_multipart_upload(&self, bucket: &str, key: &str, upload_id: &str) -> Result<(), WriteBackError> { + self.usecase() + .internal_abort_multipart_upload(bucket, key, upload_id) + .await + .map_err(write_back_error) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::app::storage_api::multipart_usecase::contract::multipart::MultipartOperations as _; + use crate::app::storage_api::object_usecase::on_demand_migration::{PullFailureReason, SourceSse}; + use crate::app::storage_api::test::bucket::utils::serialize; + use crate::app::storage_api::test::contract::bucket::{BucketOperations as _, MakeBucketOptions}; + use crate::app::storage_api::test::{get_global_bucket_metadata_sys, set_bucket_metadata}; + use http::Method; + use rustfs_utils::http::{MINIO_INTERNAL_PREFIX, RUSTFS_INTERNAL_PREFIX, get_str}; + use s3s::dto::{ + BucketVersioningStatus, DeleteMarkerReplication, DeleteMarkerReplicationStatus, Destination, ReplicationConfiguration, + ReplicationRule, ReplicationRuleStatus, ServerSideEncryptionByDefault, ServerSideEncryptionConfiguration, + ServerSideEncryptionRule, VersioningConfiguration, + }; + use sha2::{Digest as Sha256Digest, Sha256}; + use std::time::SystemTime; + use tokio::io::AsyncReadExt; + + const SOURCE_LABEL: &str = "s3:legacy-bucket"; + + fn md5_hex(body: &[u8]) -> String { + hex_simd::encode_to_string(Md5::digest(body), hex_simd::AsciiCase::Lower) + } + + fn sha256_hex(body: &[u8]) -> String { + hex_simd::encode_to_string(Sha256::digest(body), hex_simd::AsciiCase::Lower) + } + + fn stream(chunks: Vec>) -> WriteBackBody { + Box::pin(futures::stream::iter(chunks)) + } + + fn body_stream(body: &[u8]) -> WriteBackBody { + stream(body.chunks(1 << 20).map(|chunk| Ok(Bytes::copy_from_slice(chunk))).collect()) + } + + fn source_head(body: &[u8]) -> SourceHead { + SourceHead { + etag: Some(md5_hex(body)), + size: body.len() as u64, + last_modified: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000)), + content_type: Some("text/plain; charset=utf-8".to_string()), + cache_control: Some("max-age=60".to_string()), + content_language: Some("en".to_string()), + user_metadata: HashMap::from([("origin".to_string(), "legacy".to_string())]), + version_id: Some("src-v1".to_string()), + storage_class: Some("STANDARD_IA".to_string()), + ..Default::default() + } + } + + fn request(bucket: &str, key: &str, head: SourceHead) -> WriteBackRequest { + WriteBackRequest { + bucket: bucket.to_string(), + key: key.to_string(), + head, + source_label: SOURCE_LABEL.to_string(), + pulled_at: OffsetDateTime::from_unix_timestamp(1_756_800_000).expect("valid timestamp"), + preserve_etag: true, + emit_events: true, + tags: Some(HashMap::from([ + ("team".to_string(), "storage".to_string()), + ("env".to_string(), "prod".to_string()), + ])), + } + } + + async fn write_back_test_bucket(prefix: &str, versioned: bool) -> (Arc, String) { + let store = crate::app::gating_test_env::shared_gating_ecstore().await; + crate::app::runtime_sources::install_test_app_context(Arc::clone(&store)).await; + let bucket = format!("{prefix}-{}", Uuid::new_v4().simple()); + store + .make_bucket( + &bucket, + &MakeBucketOptions { + versioning_enabled: versioned, + ..Default::default() + }, + ) + .await + .expect("create write-back test bucket"); + (store, bucket) + } + + async fn stored_object(store: &Arc, bucket: &str, key: &str) -> ObjectInfo { + store + .get_object_info(bucket, key, &ObjectOptions::default()) + .await + .expect("write-back must leave a readable object") + } + + async fn assert_nothing_left(store: &Arc, bucket: &str, key: &str) { + let lookup = store.get_object_info(bucket, key, &ObjectOptions::default()).await; + assert!( + lookup.as_ref().is_err_and(is_err_object_not_found), + "a failed write-back must not leave an object: {lookup:?}" + ); + let uploads = store + .list_multipart_uploads(bucket, key, None, None, None, 100) + .await + .expect("list multipart uploads"); + assert!( + uploads.uploads.is_empty(), + "a failed write-back must not leave uploads: {:?}", + uploads.uploads + ); + } + + async fn raw_object_bytes(store: &Arc, bucket: &str, key: &str) -> Vec { + let mut reader = (**store) + .get_object_reader(bucket, key, None, HeaderMap::new(), &ObjectOptions::default()) + .await + .expect("read object"); + let mut buf = Vec::new(); + reader.stream.read_to_end(&mut buf).await.expect("drain object reader"); + buf + } + + async fn get_via_app(bucket: &str, key: &str) -> Vec { + let input = GetObjectInput::builder() + .bucket(bucket.to_string()) + .key(key.to_string()) + .build() + .expect("GET input must build"); + let req = build_request(input, Method::GET); + let mut response = DefaultObjectUsecase::from_global() + .execute_get_object(req) + .await + .expect("app-layer GET must succeed"); + let mut body = response.output.body.take().expect("GET response must include a body"); + let mut actual = Vec::new(); + while let Some(chunk) = body.next().await { + actual.extend_from_slice(&chunk.expect("GET body chunk")); + } + actual + } + + fn assert_provenance(metadata: &HashMap, head: &SourceHead) { + for suffix in [ + SUFFIX_ODM_SOURCE, + SUFFIX_ODM_SOURCE_ETAG, + SUFFIX_ODM_SOURCE_LAST_MODIFIED, + SUFFIX_ODM_SOURCE_VERSION_ID, + SUFFIX_ODM_PULLED_AT, + ] { + assert!( + metadata.contains_key(&format!("{RUSTFS_INTERNAL_PREFIX}{suffix}")), + "missing rustfs {suffix}" + ); + assert!( + metadata.contains_key(&format!("{MINIO_INTERNAL_PREFIX}{suffix}")), + "missing minio {suffix}" + ); + } + assert_eq!(get_str(metadata, SUFFIX_ODM_SOURCE).as_deref(), Some(SOURCE_LABEL)); + assert_eq!(get_str(metadata, SUFFIX_ODM_SOURCE_ETAG), head.etag); + assert_eq!( + get_str(metadata, SUFFIX_ODM_SOURCE_LAST_MODIFIED).as_deref(), + Some("2023-11-14T22:13:20Z") + ); + assert_eq!(get_str(metadata, SUFFIX_ODM_SOURCE_VERSION_ID).as_deref(), Some("src-v1")); + assert_eq!(get_str(metadata, SUFFIX_ODM_PULLED_AT).as_deref(), Some("2025-09-02T08:00:00Z")); + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_commits_the_source_object_with_provenance_and_source_etag() { + let (store, bucket) = write_back_test_bucket("odm-wb", false).await; + let write_back = OnDemandMigrationWriteBack::new(); + assert!( + write_back + .local_object(&bucket, "dir/obj.txt") + .await + .expect("lookup") + .is_none() + ); + + let body = b"pulled from the legacy bucket".to_vec(); + let head = source_head(&body); + let outcome = write_back + .put_object(&request(&bucket, "dir/obj.txt", head.clone()), body_stream(&body)) + .await + .expect("write-back must commit"); + assert_eq!(outcome.etag, head.etag, "single-part source ETag is preserved"); + assert_eq!(outcome.size, body.len() as u64); + + let stored = stored_object(&store, &bucket, "dir/obj.txt").await; + assert_eq!(stored.etag, head.etag); + assert_eq!(stored.size, body.len() as i64); + let metadata = &stored.user_defined; + assert_provenance(metadata, &head); + assert_eq!(metadata.get("content-type").map(String::as_str), Some("text/plain; charset=utf-8")); + assert_eq!(metadata.get("cache-control").map(String::as_str), Some("max-age=60")); + assert_eq!(metadata.get("content-language").map(String::as_str), Some("en")); + assert_eq!(metadata.get("origin").map(String::as_str), Some("legacy")); + assert!( + !metadata.keys().any(|key| key.eq_ignore_ascii_case("x-amz-storage-class")), + "source storage class is not copied: {metadata:?}" + ); + assert_eq!(stored.user_tags.as_str(), "env=prod&team=storage"); + assert_eq!(raw_object_bytes(&store, &bucket, "dir/obj.txt").await, body); + + let local = write_back + .local_object(&bucket, "dir/obj.txt") + .await + .expect("lookup") + .expect("object now exists"); + assert_eq!(local.etag, head.etag); + assert_eq!(local.size, body.len() as u64); + assert!(!local.delete_marker); + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_integrity_failure_leaves_nothing_behind() { + let (store, bucket) = write_back_test_bucket("odm-wb-etag", false).await; + let body = b"the source lied about this body".to_vec(); + let mut head = source_head(&body); + head.etag = Some(md5_hex(b"a different body")); + + let err = OnDemandMigrationWriteBack::new() + .put_object(&request(&bucket, "wrong.bin", head), body_stream(&body)) + .await + .expect_err("an ETag mismatch must fail the write-back"); + assert_eq!(err, WriteBackError::Integrity); + assert_eq!(err.reason(), PullFailureReason::EtagMismatch); + assert_nothing_left(&store, &bucket, "wrong.bin").await; + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_truncated_stream_leaves_nothing_behind() { + let (store, bucket) = write_back_test_bucket("odm-wb-trunc", false).await; + let body = vec![0x5a; 200 * 1024]; + let head = source_head(&body); + + // The tee secondary reports the source failure mid-stream. + let torn = stream(vec![ + Ok(Bytes::copy_from_slice(&body[..64 * 1024])), + Err(io::Error::new(io::ErrorKind::BrokenPipe, "tee primary dropped before EOF")), + ]); + let err = OnDemandMigrationWriteBack::new() + .put_object(&request(&bucket, "torn.bin", head.clone()), torn) + .await + .expect_err("a broken stream must fail the write-back"); + assert_ne!(err, WriteBackError::Integrity, "{err}"); + assert_nothing_left(&store, &bucket, "torn.bin").await; + + // A clean EOF short of the advertised size is just as fatal. + let short = stream(vec![Ok(Bytes::copy_from_slice(&body[..64 * 1024]))]); + let err = OnDemandMigrationWriteBack::new() + .put_object(&request(&bucket, "short.bin", head), short) + .await + .expect_err("a short body must fail the write-back"); + assert!(matches!(err, WriteBackError::Local(_) | WriteBackError::Integrity), "{err}"); + assert_nothing_left(&store, &bucket, "short.bin").await; + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_multipart_path_matches_the_source_digest() { + const PART_SIZE: usize = 5 * 1024 * 1024; + let (store, bucket) = write_back_test_bucket("odm-wb-mpu", false).await; + let body: Vec = (0..PART_SIZE + 4096).map(|i| (i % 253) as u8).collect(); + let mut head = source_head(&body); + head.etag = Some(format!("{}-2", md5_hex(&body))); + head.is_multipart_etag = true; + let request = request(&bucket, "big/object.bin", head.clone()); + let write_back = OnDemandMigrationWriteBack::new(); + + let upload_id = write_back.create_multipart_upload(&request).await.expect("create"); + let mut parts = Vec::new(); + for (index, chunk) in body.chunks(PART_SIZE).enumerate() { + let part = write_back + .upload_part(&request, &upload_id, index + 1, chunk.len() as u64, body_stream(chunk)) + .await + .expect("stage part"); + assert_eq!(part.part_number, index + 1); + assert!(!part.etag.is_empty()); + parts.push(part); + } + let outcome = write_back + .complete_multipart_upload(&request, &upload_id, parts) + .await + .expect("complete"); + assert_eq!(outcome.size, body.len() as u64); + assert_eq!(outcome.etag, head.etag, "multipart source ETag is preserved for display"); + + let stored = stored_object(&store, &bucket, "big/object.bin").await; + assert_eq!(stored.parts.len(), 2, "HEAD parts_count"); + assert_eq!(stored.size, body.len() as i64); + assert_eq!(stored.etag, head.etag); + assert_provenance(&stored.user_defined, &head); + assert_eq!( + stored.user_defined.get("content-type").map(String::as_str), + Some("text/plain; charset=utf-8") + ); + assert_eq!(sha256_hex(&raw_object_bytes(&store, &bucket, "big/object.bin").await), sha256_hex(&body)); + + // A failed upload is aborted and leaves no residue. + let aborted = write_back.create_multipart_upload(&request).await.expect("create"); + write_back + .upload_part(&request, &aborted, 1, 4096, body_stream(&body[..4096])) + .await + .expect("stage part"); + write_back + .abort_multipart_upload(&bucket, "big/object.bin", &aborted) + .await + .expect("abort"); + let uploads = store + .list_multipart_uploads(&bucket, "big/object.bin", None, None, None, 100) + .await + .expect("list uploads"); + assert!(uploads.uploads.iter().all(|upload| upload.upload_id != aborted)); + } + + async fn install_bucket_default_sse(bucket: &str) { + let sys = get_global_bucket_metadata_sys().expect("bucket metadata system"); + let metadata = { + let sys = sys.read().await; + sys.get(bucket).await.expect("bucket metadata cached") + }; + let mut metadata = (*metadata).clone(); + let config = ServerSideEncryptionConfiguration { + rules: vec![ServerSideEncryptionRule { + apply_server_side_encryption_by_default: Some(ServerSideEncryptionByDefault { + sse_algorithm: ServerSideEncryption::from_static(ServerSideEncryption::AES256), + kms_master_key_id: None, + }), + blocked_encryption_types: None, + bucket_key_enabled: None, + }], + }; + metadata.encryption_config_xml = serialize(&config).expect("sse config serializes"); + metadata.sse_config = Some(config); + set_bucket_metadata(bucket.to_string(), metadata) + .await + .expect("install bucket default SSE"); + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_under_bucket_default_sse_stores_ciphertext_and_records_source_etag() { + let local_sse_master_key = base64_simd::STANDARD.encode_to_string([0x42u8; 32]); + temp_env::async_with_vars([("RUSTFS_SSE_S3_MASTER_KEY", Some(local_sse_master_key))], async { + let (store, bucket) = write_back_test_bucket("odm-wb-sse", false).await; + // The gating store adopts the bootstrap context; server startup is + // what normally installs the read-side decryption resolver on it. + let _ = crate::app::storage_api::test::bootstrap_instance_ctx(); + install_bucket_default_sse(&bucket).await; + assert!(bucket_encrypts_by_default(&bucket).await); + + let body = b"plaintext that must be encrypted at rest".to_vec(); + let head = source_head(&body); + let request = request(&bucket, "secret.txt", head.clone()); + // The source ETag is not forced onto an encrypted object; the + // local ETag is whatever the SSE write path computes. + assert_eq!(write_back_context(&request, true).await.preserve_etag, None); + let outcome = OnDemandMigrationWriteBack::new() + .put_object(&request, body_stream(&body)) + .await + .expect("write-back under SSE must commit"); + assert!(outcome.etag.is_some()); + + let stored = stored_object(&store, &bucket, "secret.txt").await; + assert_eq!(stored.etag, outcome.etag); + assert!( + stored.user_defined.contains_key("x-amz-server-side-encryption"), + "{:?}", + stored.user_defined + ); + assert_eq!(get_str(&stored.user_defined, SUFFIX_ODM_SOURCE_ETAG), head.etag); + assert_provenance(&stored.user_defined, &head); + assert!( + stored + .user_defined + .keys() + .any(|key| key.starts_with("x-rustfs-encryption-") || key.starts_with("x-minio-encryption-")), + "disk holds ciphertext under a managed key: {:?}", + stored.user_defined + ); + assert_eq!(get_via_app(&bucket, "secret.txt").await, body, "GET returns the plaintext"); + }) + .await; + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_reports_a_full_bucket_quota() { + let (store, bucket) = crate::app::gating_test_env::durable_quota_test_bucket("odm-wb-quota", 64).await; + let body = vec![0x71; 4096]; + let err = OnDemandMigrationWriteBack::new() + .put_object(&request(&bucket, "over.bin", source_head(&body)), body_stream(&body)) + .await + .expect_err("a full quota must reject the write-back"); + assert!(matches!(err, WriteBackError::Quota(_)), "{err}"); + // ODM-05 fixed the failure label set without a quota label; quota + // failures are accounted as local writes until it grows one. + assert_eq!(err.reason(), PullFailureReason::LocalWrite); + assert_nothing_left(&store, &bucket, "over.bin").await; + } + + async fn install_replication_rule(bucket: &str) { + let sys = get_global_bucket_metadata_sys().expect("bucket metadata system"); + let metadata = { + let sys = sys.read().await; + sys.get(bucket).await.expect("bucket metadata cached") + }; + let mut metadata = (*metadata).clone(); + metadata.versioning_config_xml = b"Enabled".to_vec(); + metadata.versioning_config = Some(VersioningConfiguration { + status: Some(BucketVersioningStatus::from_static(BucketVersioningStatus::ENABLED)), + ..Default::default() + }); + let config = ReplicationConfiguration { + role: String::new(), + rules: vec![ReplicationRule { + delete_marker_replication: Some(DeleteMarkerReplication { + status: Some(DeleteMarkerReplicationStatus::from_static(DeleteMarkerReplicationStatus::DISABLED)), + }), + delete_replication: None, + destination: Destination { + bucket: "arn:aws:s3:::target-bucket".to_string(), + ..Default::default() + }, + existing_object_replication: None, + filter: None, + id: Some("odm".to_string()), + prefix: Some(String::new()), + priority: Some(1), + source_selection_criteria: None, + status: ReplicationRuleStatus::from_static(ReplicationRuleStatus::ENABLED), + }], + }; + metadata.replication_config_xml = serialize(&config).expect("replication config serializes"); + metadata.replication_config = Some(config); + set_bucket_metadata(bucket.to_string(), metadata) + .await + .expect("install replication rule"); + } + + #[tokio::test] + #[serial_test::serial] + async fn write_back_schedules_replication_and_names_the_migration_principal() { + let (store, bucket) = write_back_test_bucket("odm-wb-repl", true).await; + install_replication_rule(&bucket).await; + + let body = b"replicate me".to_vec(); + let head = source_head(&body); + let request = request(&bucket, "replicated.txt", head); + let ctx = write_back_context(&request, true).await; + assert!(ctx.emit_events, "policy.emit_events reaches the creation event"); + assert_eq!(ctx.principal_id, ON_DEMAND_MIGRATION_PRINCIPAL_ID); + assert_eq!(ctx.principal_id, "rustfs-on-demand-migration"); + let event = + InternalPutObjectEvent::builder(EventName::ObjectCreatedPut, &bucket, "replicated.txt", ctx.principal_id).build(); + assert_eq!(event.event_name, EventName::ObjectCreatedPut); + assert_eq!( + event.req_params.get("principalId").map(String::as_str), + Some("rustfs-on-demand-migration") + ); + + let outcome = OnDemandMigrationWriteBack::new() + .put_object(&request, body_stream(&body)) + .await + .expect("write-back must commit"); + assert!(outcome.version_id.is_some(), "versioned bucket yields a version id"); + let stored = stored_object(&store, &bucket, "replicated.txt").await; + // Stored per target as `=;`; the S3 header is derived from it. + let status = get_str(&stored.user_defined, SUFFIX_REPLICATION_STATUS).unwrap_or_default(); + assert!( + status.contains("arn:aws:s3:::target-bucket=PENDING;") || status.contains("arn:aws:s3:::target-bucket=COMPLETED;"), + "write-back must enter the replication queue: {status:?}" + ); + } + + #[test] + fn content_headers_follow_the_allowlist() { + let head = SourceHead { + content_type: Some("image/png".to_string()), + content_encoding: Some("gzip".to_string()), + content_disposition: Some("attachment".to_string()), + content_language: Some(" ".to_string()), + cache_control: None, + expires: Some("Thu, 01 Jan 2026 00:00:00 GMT".to_string()), + storage_class: Some("GLACIER".to_string()), + ..Default::default() + }; + let headers = content_headers(&head); + assert_eq!( + headers, + HashMap::from([ + ("Content-Type".to_string(), "image/png".to_string()), + ("Content-Encoding".to_string(), "gzip".to_string()), + ("Content-Disposition".to_string(), "attachment".to_string()), + ("Expires".to_string(), "Thu, 01 Jan 2026 00:00:00 GMT".to_string()), + ]) + ); + } + + #[test] + fn expected_md5_only_for_bare_single_part_unencrypted_etags() { + let mut head = source_head(b"abc"); + assert_eq!(expected_md5_hex(&head), Some(md5_hex(b"abc"))); + head.etag = Some(md5_hex(b"abc").to_ascii_uppercase()); + assert_eq!(expected_md5_hex(&head), Some(md5_hex(b"abc")), "normalized to lowercase"); + head.etag = Some(format!("{}-2", md5_hex(b"abc"))); + assert_eq!(expected_md5_hex(&head), None, "multipart ETag"); + head.etag = Some("not-hex-not-hex-not-hex-not-hex-".to_string()); + assert_eq!(expected_md5_hex(&head), None, "non-hex"); + head.etag = Some(md5_hex(b"abc")); + head.sse = Some(SourceSse::S3); + assert_eq!(expected_md5_hex(&head), None, "encrypted source"); + head.sse = None; + head.etag = None; + assert_eq!(expected_md5_hex(&head), None); + } + + #[test] + fn provenance_and_tags_are_stable() { + let mut request = request("b", "k", source_head(b"x")); + let metadata = provenance_metadata(&request); + assert_eq!(metadata.len(), 10, "five keys under two prefixes"); + assert_provenance(&metadata, &request.head); + request.head.etag = None; + request.head.version_id = None; + request.head.last_modified = None; + let metadata = provenance_metadata(&request); + assert_eq!(metadata.len(), 10, "absent values keep the key set constant"); + assert_eq!(get_str(&metadata, SUFFIX_ODM_SOURCE_ETAG).as_deref(), Some("")); + assert_eq!(get_str(&metadata, SUFFIX_ODM_SOURCE_VERSION_ID).as_deref(), Some("")); + assert_eq!(get_str(&metadata, SUFFIX_ODM_SOURCE_LAST_MODIFIED).as_deref(), Some("")); + + assert_eq!(encode_tags(&HashMap::new()), None); + let tags = HashMap::from([("b key".to_string(), "v&2".to_string()), ("a".to_string(), "1".to_string())]); + assert_eq!(encode_tags(&tags).as_deref(), Some("a=1&b+key=v%262")); + } + + #[tokio::test] + async fn write_back_context_applies_the_etag_and_event_policy() { + let body = b"context".to_vec(); + let mut request = request("no-such-bucket", "k", source_head(&body)); + let ctx = write_back_context(&request, true).await; + assert_eq!(ctx.expected_md5_hex, Some(md5_hex(&body))); + assert_eq!(ctx.preserve_etag, Some(md5_hex(&body))); + assert_eq!(ctx.size, Some(body.len() as u64)); + assert_eq!(ctx.tags.as_deref(), Some("env=prod&team=storage")); + assert_eq!(ctx.user_metadata.get("origin").map(String::as_str), Some("legacy")); + assert!(ctx.emit_events); + assert_eq!(ctx.principal_id, ON_DEMAND_MIGRATION_PRINCIPAL_ID); + + let multipart = write_back_context(&request, false).await; + assert_eq!(multipart.expected_md5_hex, None, "parts cannot be checked against the object ETag"); + assert_eq!(multipart.preserve_etag, Some(md5_hex(&body))); + + request.preserve_etag = false; + request.emit_events = false; + request.tags = None; + let ctx = write_back_context(&request, true).await; + assert_eq!(ctx.preserve_etag, None); + assert_eq!( + ctx.expected_md5_hex, + Some(md5_hex(&body)), + "integrity check is independent of preservation" + ); + assert!(!ctx.emit_events); + assert_eq!(ctx.tags, None); + } + + #[test] + fn write_back_error_classes_follow_the_api_error() { + let bad_digest = ApiError { + code: S3ErrorCode::BadDigest, + message: "digest".to_string(), + source: None, + }; + assert_eq!(write_back_error(bad_digest), WriteBackError::Integrity); + let quota = ApiError::invalid_request("Bucket quota exceeded. Current usage: 1 bytes, limit: 1 bytes"); + assert!(matches!(write_back_error(quota), WriteBackError::Quota(_))); + let other = ApiError { + code: S3ErrorCode::InternalError, + message: "disk".to_string(), + source: None, + }; + assert_eq!(write_back_error(other), WriteBackError::Local("InternalError: disk".to_string())); + } +} diff --git a/rustfs/src/app/object/put.rs b/rustfs/src/app/object/put.rs index 8f506306a..6fe1c06d7 100644 --- a/rustfs/src/app/object/put.rs +++ b/rustfs/src/app/object/put.rs @@ -933,7 +933,6 @@ pub(super) enum PutObjectContentMd5 { /// `Content-MD5` request header value. Base64(String), /// Lowercase hex digest, as an internal caller already holds it. - #[cfg_attr(not(test), expect(dead_code, reason = "constructed by the internal put entry point"))] Hex(String), } @@ -950,7 +949,6 @@ pub(super) enum PutObjectOrigin<'a> { /// request and no credential: managed-SSE authorization treats the write /// as internal, and the creation event, when requested, names /// `principal_id` instead of an access key. - #[cfg_attr(not(test), expect(dead_code, reason = "constructed by the internal put entry point"))] Internal { principal_id: &'static str, emit_events: bool }, } diff --git a/rustfs/src/app/storage_api.rs b/rustfs/src/app/storage_api.rs index c9b1b1c7a..4d9e6e577 100644 --- a/rustfs/src/app/storage_api.rs +++ b/rustfs/src/app/storage_api.rs @@ -1157,6 +1157,19 @@ pub(crate) mod bucket_usecase { pub(crate) mod object_usecase { pub(crate) use super::storage_contracts::BUCKET_LIFECYCLE_LOCK_OBJECT; + pub(crate) mod on_demand_migration { + #[cfg(test)] + pub(crate) use crate::storage::storage_api::ecstore_bucket::on_demand_migration::PullFailureReason; + #[cfg(test)] + pub(crate) use crate::storage::storage_api::ecstore_bucket::on_demand_migration::source_client::SourceSse; + pub(crate) use crate::storage::storage_api::ecstore_bucket::on_demand_migration::source_client::{ + SourceHead, is_multipart_etag, + }; + pub(crate) use crate::storage::storage_api::ecstore_bucket::on_demand_migration::{ + LocalObject, OdmWriteBack, WriteBackBody, WriteBackError, WriteBackOutcome, WriteBackPart, WriteBackRequest, + }; + } + pub(crate) mod object_cache { #[cfg(test)] pub(crate) use crate::storage::storage_api::ecstore_object::GetObjectBodySource; @@ -1285,6 +1298,7 @@ pub(crate) mod test { pub(crate) mod data_usage { pub(crate) use super::super::data_usage::*; } + pub(crate) use crate::storage::storage_api::bootstrap_instance_ctx; pub(crate) use crate::storage::storage_api::ecstore_bucket::install_all_v6_fleet_capability_proof; pub(crate) use crate::storage::storage_api::test_consumer::{get_global_bucket_metadata_sys, set_bucket_metadata}; pub(crate) use crate::storage::storage_api::{ diff --git a/rustfs/src/startup_bucket_metadata.rs b/rustfs/src/startup_bucket_metadata.rs index fc336e990..1c3cdb270 100644 --- a/rustfs/src/startup_bucket_metadata.rs +++ b/rustfs/src/startup_bucket_metadata.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::app::object::OnDemandMigrationWriteBack; use crate::module_switches::{on_demand_migration_enabled_from_env, set_on_demand_migration_module_enabled}; use crate::storage_api::startup::bucket_metadata::contract::bucket::{BucketOperations, BucketOptions}; use crate::storage_api::startup::bucket_metadata::{ @@ -90,15 +91,18 @@ pub(crate) async fn init_bucket_metadata_runtime(store: Arc, ctx: Cance Ok(buckets) } -/// Publishes the on-demand migration module switch and registers the -/// runtime's config hook before bucket metadata is loaded, so every cache -/// install path (initial load included) reaches `OnDemandMigrationSys` -/// (rustfs/backlog#2152). Idempotent across embedded and server startups. +/// Publishes the on-demand migration module switch, installs the app-layer +/// write-back the pull pipeline stores objects with (rustfs/backlog#2153), +/// and registers the runtime's config hook before bucket metadata is +/// loaded, so every cache install path (initial load included) reaches +/// `OnDemandMigrationSys` with a usable write-back (rustfs/backlog#2152). +/// Idempotent across embedded and server startups. fn init_on_demand_migration_runtime() { let enabled = on_demand_migration_enabled_from_env(); set_on_demand_migration_module_enabled(enabled); let sys = OnDemandMigrationSys::get(); sys.set_module_enabled(enabled); + sys.set_write_back(Arc::new(OnDemandMigrationWriteBack::new())); let hook_registered = sys.register_config_hook(); tracing::info!( event = EVENT_ON_DEMAND_MIGRATION_RUNTIME_INITIALIZED,