diff --git a/crates/libs/vol/mgr/server/src/executor.rs b/crates/libs/vol/mgr/server/src/executor.rs index 4b417f723..e5bb72b1d 100644 --- a/crates/libs/vol/mgr/server/src/executor.rs +++ b/crates/libs/vol/mgr/server/src/executor.rs @@ -48,6 +48,28 @@ use sha2::{Digest, Sha256}; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; +mod artifact_role; +mod codec_concurrency; +mod decode_sinks; +mod pull; +mod replica_provenance; +mod sparse_shadow; +mod stream_pull_timing; +mod stream_writeback; +mod transfer_metrics; +mod writeback_coordinator; + +pub(crate) use artifact_role::*; +use codec_concurrency::*; +use decode_sinks::*; +pub use pull::*; +pub use replica_provenance::*; +pub use sparse_shadow::*; +use stream_pull_timing::*; +use stream_writeback::*; +use transfer_metrics::*; +pub use writeback_coordinator::*; + // zstd frame magic: first 4 bytes of any zstd-compressed stream. const ZSTD_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD]; @@ -55,8 +77,6 @@ const ZSTD_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD]; // binary files. Level 1 is ~2× faster with ~10% // lower ratio; level 9 is ~4× slower with ~5% better ratio. const ZSTD_LEVEL: i32 = 3; -const MEMORY_SPARSE_BLOCK_BYTES: usize = 4 * 1024; -const ROOTFS_SPARSE_BLOCK_BYTES: usize = 64 * 1024; // A decode whose frame pledges fewer decompressed bytes than this skips // codec admission: it costs at most about a second of CPU, and queueing @@ -94,400 +114,8 @@ async fn frame_content_size(path: &Path) -> Option { .and_then(|h| zstd::zstd_safe::get_frame_content_size(h).ok().flatten()) } -/// How many CPU-bound zstd passes (compress or decompress) may run at -/// once. `0` means auto-size to the CPU quota, or to the deployed CPU -/// request when the pod has no CPU limit. -/// -/// Each zstd pass is single-threaded, and `spawn_blocking`'s pool is ~512 -/// threads, so without admission a pull storm runs every decode -/// concurrently on the pod's CPU share — N jobs each at 1/N speed instead -/// of share-many at full speed. Bounding to the share keeps the same -/// aggregate throughput but finishes jobs in arrival order, which is what -/// per-job completion latency is made of. -fn resolved_codec_concurrency(configured: usize) -> usize { - if configured > 0 { - return configured; - } - // Kubernetes enforces CPU *limits* through the cgroup CFS quota, which - // `available_parallelism` (scheduler-affinity based) does not observe - // unless the pod is cpuset-pinned. Read the quota directly. Without a - // quota (the pod runs limit-less), the deployed whole-core CPU request - // arrives as WORKER_THREADS (Downward API on Kubernetes, Environment= - // in the direct systemd units) and is the intended share — affinity - // would size to every core of a metal host. Affinity remains the last - // resort for undeployed contexts. - detect_cgroup_cpu_quota_cores() - .or_else(cpu_request_cores_from_env) - .or_else(|| { - std::thread::available_parallelism() - .ok() - .map(std::num::NonZeroUsize::get) - }) - .unwrap_or(2) - .max(1) -} - -/// The whole-core CPU request the deployment stamped into `WORKER_THREADS`, -/// or `None` when unset or malformed. -fn cpu_request_cores_from_env() -> Option { - parse_cpu_request_cores(&std::env::var("WORKER_THREADS").ok()?) -} - -fn parse_cpu_request_cores(raw: &str) -> Option { - raw.trim().parse().ok().filter(|cores| *cores > 0) -} - -/// Cores implied by a CFS `quota`/`period` pair, rounded to the nearest -/// whole core and floored at 1. Integer arithmetic keeps this free of -/// float sign/precision casts. -fn cores_from_cfs_quota(quota: u64, period: u64) -> Option { - if period == 0 { - return None; - } - let rounded = (quota + period / 2) / period; - Some(usize::try_from(rounded).unwrap_or(usize::MAX).max(1)) -} - -/// Read the cgroup CPU quota (in whole cores) the process is limited to, -/// or `None` when unlimited or unreadable. Handles cgroup v2 (`cpu.max`) -/// and v1 (`cpu.cfs_quota_us` / `cpu.cfs_period_us`). Sync file reads at -/// startup only. -fn detect_cgroup_cpu_quota_cores() -> Option { - // cgroup v2: "$QUOTA $PERIOD" where QUOTA is "max" or microseconds. - if let Ok(contents) = std::fs::read_to_string("/sys/fs/cgroup/cpu.max") { - let mut fields = contents.split_whitespace(); - let quota = fields.next()?; - let period: u64 = fields.next()?.parse().ok()?; - if quota == "max" { - return None; - } - return cores_from_cfs_quota(quota.parse().ok()?, period); - } - // cgroup v1. - let quota: i64 = std::fs::read_to_string("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") - .ok()? - .trim() - .parse() - .ok()?; - // -1 (unlimited) and any non-positive value have no quota. - let quota = u64::try_from(quota).ok().filter(|&q| q > 0)?; - let period: u64 = std::fs::read_to_string("/sys/fs/cgroup/cpu/cpu.cfs_period_us") - .ok()? - .trim() - .parse() - .ok()?; - cores_from_cfs_quota(quota, period) -} - type DestinationPullLocks = Arc>>>; -fn transfer_artifact_name(local_path: &Path) -> Option<&str> { - let name = local_path.file_name()?.to_str()?; - let name = name.strip_prefix('.').unwrap_or(name); - Some(name.split(".restore-").next().unwrap_or(name)) -} - -pub(crate) fn snapshot_artifact_from_local_path(local_path: &Path) -> Option { - match transfer_artifact_name(local_path)? { - "snapshot.json" => Some(SnapshotArtifact::Metadata), - "vmstate" => Some(SnapshotArtifact::VmState), - "memory" => Some(SnapshotArtifact::Memory), - "rootfs.ext4" => Some(SnapshotArtifact::Rootfs), - "rootfs.btrfs-send" => Some(SnapshotArtifact::RootfsBtrfsSend), - "snapshot.luks" => Some(SnapshotArtifact::EncryptedSnapshotVolume), - "cow.luks" => Some(SnapshotArtifact::EncryptedRootfsCow), - _ => None, - } -} - -pub(crate) fn snapshot_artifact_from_entry_role(role: &str) -> Option { - [ - SnapshotArtifact::Rootfs, - SnapshotArtifact::RootfsBtrfsSend, - SnapshotArtifact::Memory, - SnapshotArtifact::VmState, - SnapshotArtifact::EncryptedSnapshotVolume, - SnapshotArtifact::EncryptedRootfsCow, - SnapshotArtifact::Metadata, - ] - .into_iter() - .find(|artifact| artifact.file_name() == role) -} - -fn transfer_artifact_tag(artifact: Option) -> &'static str { - match artifact { - Some(SnapshotArtifact::Metadata) => "metadata", - Some(SnapshotArtifact::VmState) => "vmstate", - Some(SnapshotArtifact::Memory) => "memory", - Some(SnapshotArtifact::Rootfs) => "rootfs_ext4", - Some(SnapshotArtifact::RootfsBtrfsSend) => "rootfs_btrfs_send", - Some(SnapshotArtifact::EncryptedSnapshotVolume) => "encrypted_snapshot_volume", - Some(SnapshotArtifact::EncryptedRootfsCow) => "encrypted_rootfs_cow", - None => "other", - } -} - -/// Closed vocabulary for the `phase` tag on transfer metrics. The string -/// form exists only at the emission boundary; everything in between -/// passes the enum so additions stay exhaustive and typos can't compile. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum TransferPhase { - Download, - DownloadDecode, - Decrypt, - DecryptDecompress, - DecryptWait, - Decompress, - DecompressWait, - CompressWait, - Promote, - ShadowReflink, -} - -impl TransferPhase { - fn as_str(self) -> &'static str { - match self { - Self::Download => "download", - Self::DownloadDecode => "download_decode", - Self::Decrypt => "decrypt", - Self::DecryptDecompress => "decrypt_decompress", - Self::DecryptWait => "decrypt_wait", - Self::Decompress => "decompress", - Self::DecompressWait => "decompress_wait", - Self::CompressWait => "compress_wait", - Self::Promote => "promote", - Self::ShadowReflink => "shadow_reflink", - } - } -} - -impl std::fmt::Display for TransferPhase { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str(self.as_str()) - } -} - -fn emit_transfer_phase_duration( - kind: &str, - artifact: &str, - phase: TransferPhase, - mode: MaterializationMode, - started_at: Instant, -) { - pplx_space_metrics::volume_manager::transfer_phase_duration( - kind, - artifact, - phase.as_str(), - mode, - started_at, - ); -} - -fn effective_materialization_mode( - artifact: Option, - requested: MaterializationMode, -) -> MaterializationMode { - if matches!( - artifact, - Some( - SnapshotArtifact::Memory - | SnapshotArtifact::Rootfs - | SnapshotArtifact::EncryptedSnapshotVolume - | SnapshotArtifact::EncryptedRootfsCow - ) - ) { - requested - } else { - MaterializationMode::Disabled - } -} - -fn execution_materialization_mode( - artifact: Option, - requested: MaterializationMode, -) -> MaterializationMode { - // LUKS backings contain literal zero regions outside their written encrypted - // sectors. Exact-zero materialization preserves the ciphertext and logical - // length, including when retrying jobs persisted before sparse support. - let requested = if matches!( - artifact, - Some( - SnapshotArtifact::Memory - | SnapshotArtifact::EncryptedSnapshotVolume - | SnapshotArtifact::EncryptedRootfsCow - ) - ) { - MaterializationMode::Enabled - } else { - requested - }; - effective_materialization_mode(artifact, requested) -} - -fn materialization_block_bytes(artifact: Option) -> usize { - match artifact { - Some(SnapshotArtifact::Rootfs) => ROOTFS_SPARSE_BLOCK_BYTES, - _ => MEMORY_SPARSE_BLOCK_BYTES, - } -} - -fn emit_transfer_throughput_mib_s( - kind: &str, - artifact: &str, - phase: TransferPhase, - mode: MaterializationMode, - bytes: u64, - started_at: Instant, -) { - let elapsed = started_at.elapsed().as_secs_f64(); - if elapsed > 0.0 { - let mib_s = (bytes as f64 / 1024.0 / 1024.0) / elapsed; - pplx_space_metrics::volume_manager::transfer_throughput_mib_s( - kind, - artifact, - phase.as_str(), - mode, - mib_s, - ); - } -} - -#[cfg(unix)] -async fn observe_decompression_allocation( - kind: &str, - artifact: &str, - job_id: &JobId, - subject_id: &str, - local_path: &Path, - report: DecompressionReport, -) -> Option { - use std::os::unix::fs::MetadataExt; - - let metadata = match tokio::fs::metadata(local_path).await { - Ok(metadata) => metadata, - Err(error) => { - tracing::debug!( - %error, - job_id = %job_id, - subject_id, - local_path = %local_path.display(), - "could not observe decompressed artifact allocation" - ); - return None; - } - }; - let logical_bytes = metadata.len(); - let allocated_bytes = metadata.blocks().saturating_mul(512); - let saved_bytes = logical_bytes.saturating_sub(allocated_bytes); - let allocation_ratio = if logical_bytes == 0 { - 0.0 - } else { - allocated_bytes as f64 / logical_bytes as f64 - }; - let layout = if report.mode == MaterializationMode::Enabled { - MaterializationLayout::Sparse - } else { - MaterializationLayout::Dense - }; - pplx_space_metrics::volume_manager::transfer_materialization_allocation( - kind, - artifact, - report.mode, - layout, - logical_bytes, - allocated_bytes, - report.candidate_hole_bytes, - ); - let projected_allocated_bytes = logical_bytes.saturating_sub(report.candidate_hole_bytes); - tracing::info!( - job_id = %job_id, - subject_id, - local_path = %local_path.display(), - artifact, - layout = %layout, - mode = %report.mode, - logical_bytes, - allocated_bytes, - candidate_hole_bytes = report.candidate_hole_bytes, - projected_allocated_bytes, - saved_bytes, - allocation_ratio, - "decompressed artifact allocation observed" - ); - Some(MaterializationObservation { - mode: report.mode, - layout, - logical_bytes, - allocated_bytes, - candidate_hole_bytes: report.candidate_hole_bytes, - }) -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum ShadowComparisonOutcome { - Match, - Mismatch, - Error, - SkippedBusy, - SkippedDisk, -} - -impl ShadowComparisonOutcome { - const fn as_str(self) -> &'static str { - match self { - Self::Match => "match", - Self::Mismatch => "mismatch", - Self::Error => "error", - Self::SkippedBusy => "skipped_busy", - Self::SkippedDisk => "skipped_disk", - } - } -} - -impl std::fmt::Display for ShadowComparisonOutcome { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(self.as_str()) - } -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -struct ShadowComparisonObservation { - logical_bytes: u64, - allocated_bytes: u64, - candidate_hole_bytes: u64, -} - -#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] -pub struct MaterializationObservation { - pub mode: MaterializationMode, - pub layout: MaterializationLayout, - pub logical_bytes: u64, - pub allocated_bytes: u64, - pub candidate_hole_bytes: u64, -} - -#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] -struct FileReplicaProvenance { - #[serde(flatten)] - replica: VerifiedReplicaProvenance, - #[serde(default, skip_serializing_if = "Option::is_none")] - materialization: Option, -} - -struct DestinationPullGuard { - locks: DestinationPullLocks, - path: PathBuf, - guard: Option>, -} - -impl Drop for DestinationPullGuard { - fn drop(&mut self) { - self.guard.take(); - self.locks - .remove_if(&self.path, |_path, lock| Arc::strong_count(lock) == 1); - } -} - #[derive(Debug, thiserror::Error)] pub enum FileModeError { #[error("file_mode must only contain rwx bits: {mode:o}")] @@ -516,33 +144,6 @@ enum FileModeApplyPolicy { ExistingFile, } -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PullExistingPolicy { - #[default] - LegacySkipPresent, - VerifiedReplica, -} - -impl PullExistingPolicy { - pub fn is_legacy_skip_present(policy: &Self) -> bool { - *policy == Self::LegacySkipPresent - } - - pub fn as_dedup_str(self) -> &'static str { - match self { - Self::LegacySkipPresent => "legacy_skip_present", - Self::VerifiedReplica => "verified_replica", - } - } -} - -struct PullExecutionOptions<'a> { - resolved_dek: Option<&'a [u8]>, - existing_policy: PullExistingPolicy, - cancellation: CancellationToken, -} - impl serde::Serialize for FileMode { fn serialize(&self, serializer: S) -> Result where @@ -911,91 +512,6 @@ fn partial_path(local_path: &Path) -> PathBuf { p } -fn requested_file_mode_bits(p: &TransferPayload) -> Option { - p.file_mode.map(FileMode::bits) -} - -fn file_replica_provenance( - p: &TransferPayload, - object: &ObjectMeta, - local_size: u64, -) -> VerifiedReplicaProvenance { - VerifiedReplicaProvenance { - version: VERIFIED_REPLICA_PROVENANCE_VERSION, - subject_id: p.subject_id.clone(), - s3_uri: p.s3_uri.clone(), - raw: p.raw, - file_mode: requested_file_mode_bits(p), - snapshot_artifact: p.snapshot_artifact, - object_size: object.size, - object_etag: object.etag.clone(), - local_size, - } -} - -impl FileReplicaProvenance { - fn new( - p: &TransferPayload, - object: &ObjectMeta, - local_size: u64, - materialization: Option, - ) -> Self { - Self { - replica: file_replica_provenance(p, object, local_size), - materialization, - } - } -} - -#[cfg(test)] -fn read_materialization_observation(local_path: &Path) -> Option { - let bytes = std::fs::read(verified_replica_provenance_path(local_path)).ok()?; - serde_json::from_slice::(&bytes) - .ok()? - .materialization -} - -fn materialization_is_compatible( - requested: MaterializationMode, - observation: Option<&MaterializationObservation>, -) -> bool { - match requested { - MaterializationMode::Disabled => true, - // The dense observation records the requested writer mode, not a - // successful background comparison. Shadow results are not persisted, - // so no completed replica can prove that postcondition yet. - MaterializationMode::Shadow => false, - MaterializationMode::Enabled => observation.is_some_and(|observation| { - observation.mode == MaterializationMode::Enabled - && observation.layout == MaterializationLayout::Sparse - }), - } -} - -/// Best-effort unlink of a `.partial` sibling. -/// -/// `NotFound` is the common expected case (no stale file to sweep), so we -/// silence it. Any other failure is logged so silent fs corruption — a -/// permission flip, a stuck mount — doesn't accumulate without diagnostics. -/// On a successful removal we emit a debug log so audits ("what was deleted -/// and why?") can trace each unlink to the code path that triggered it. -async fn unlink_partial(path: &Path, reason: &str) { - match tokio::fs::remove_file(path).await { - Ok(()) => tracing::debug!( - path = %path.display(), - reason, - "unlinked .partial" - ), - Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} - Err(err) => tracing::warn!( - error = %err, - path = %path.display(), - reason, - "unlink .partial failed" - ), - } -} - fn classify(e: &ObjectStoreError) -> JobError { let msg = e.to_string(); match e { @@ -1071,262 +587,6 @@ async fn is_zstd(path: &Path) -> bool { f.read_exact(&mut magic).await.is_ok() && magic == ZSTD_MAGIC } -fn sparse_shadow_path(path: &Path) -> PathBuf { - let mut shadow = path.to_path_buf(); - let mut name = shadow.file_name().unwrap_or_default().to_os_string(); - name.push(format!(".sparse-shadow-{}.tmp", uuid::Uuid::now_v7())); - shadow.set_file_name(name); - shadow -} - -fn sparse_shadow_source_path(staging_dir: &Path) -> PathBuf { - staging_dir.join(format!("sparse-shadow-source-{}.tmp", uuid::Uuid::now_v7())) -} - -#[derive(Debug, thiserror::Error)] -enum SparseShadowCmpError { - #[error("sparse shadow comparison cancelled")] - Cancelled, - #[error("sparse shadow bytes differ from dense output")] - Mismatch, - #[error("{operation}: {source}")] - Io { - operation: &'static str, - #[source] - source: std::io::Error, - }, -} - -impl SparseShadowCmpError { - fn io(operation: &'static str, source: std::io::Error) -> Self { - Self::Io { operation, source } - } -} - -fn compare_completed_files( - dense: &Path, - sparse: &Path, - shutdown_token: &CancellationToken, -) -> Result<(), SparseShadowCmpError> { - let mut dense_file = std::fs::File::open(dense) - .map_err(|error| SparseShadowCmpError::io("open dense comparison input", error))?; - let mut sparse_file = std::fs::File::open(sparse) - .map_err(|error| SparseShadowCmpError::io("open sparse comparison input", error))?; - let dense_len = dense_file - .metadata() - .map_err(|error| SparseShadowCmpError::io("stat dense comparison input", error))? - .len(); - let sparse_len = sparse_file - .metadata() - .map_err(|error| SparseShadowCmpError::io("stat sparse comparison input", error))? - .len(); - if dense_len != sparse_len { - return Err(SparseShadowCmpError::Mismatch); - } - - let mut dense_buffer = vec![0; 1024 * 1024]; - let mut sparse_buffer = vec![0; 1024 * 1024]; - let mut compared_through = 0_u64; - let mut advised_through = 0_u64; - while compared_through < dense_len { - if shutdown_token.is_cancelled() { - return Err(SparseShadowCmpError::Cancelled); - } - let remaining = dense_len.saturating_sub(compared_through); - let read_len = dense_buffer - .len() - .min(usize::try_from(remaining).unwrap_or(usize::MAX)); - let dense_chunk = dense_buffer.get_mut(..read_len).ok_or_else(|| { - SparseShadowCmpError::io( - "prepare dense comparison buffer", - std::io::Error::other("comparison read exceeds dense buffer"), - ) - })?; - let sparse_chunk = sparse_buffer.get_mut(..read_len).ok_or_else(|| { - SparseShadowCmpError::io( - "prepare sparse comparison buffer", - std::io::Error::other("comparison read exceeds sparse buffer"), - ) - })?; - dense_file - .read_exact(dense_chunk) - .map_err(|error| SparseShadowCmpError::io("read dense comparison input", error))?; - sparse_file - .read_exact(sparse_chunk) - .map_err(|error| SparseShadowCmpError::io("read sparse comparison input", error))?; - if dense_chunk != sparse_chunk { - return Err(SparseShadowCmpError::Mismatch); - } - compared_through = compared_through.saturating_add(read_len as u64); - if compared_through.saturating_sub(advised_through) >= STREAM_WRITEBACK_WINDOW_BYTES { - advise_shadow_read_range( - &dense_file, - dense, - advised_through, - compared_through.saturating_sub(advised_through), - "dense sparse-shadow comparison input", - ); - advise_shadow_read_range( - &sparse_file, - sparse, - advised_through, - compared_through.saturating_sub(advised_through), - "sparse shadow comparison input", - ); - advised_through = compared_through; - } - } - if advised_through < compared_through { - let length = compared_through.saturating_sub(advised_through); - advise_shadow_read_range( - &dense_file, - dense, - advised_through, - length, - "dense sparse-shadow comparison input", - ); - advise_shadow_read_range( - &sparse_file, - sparse, - advised_through, - length, - "sparse shadow comparison input", - ); - } - Ok(()) -} - -fn advise_shadow_read_range( - file: &std::fs::File, - path: &Path, - offset: u64, - length: u64, - operation: &'static str, -) { - if let Err(error) = page_cache::advise_range_dontneed(file, offset, length) { - tracing::warn!( - %error, - path_len = path.as_os_str().len(), - offset, - length, - operation, - "failed to evict consumed sparse-shadow read range" - ); - } -} - -fn compare_sparse_shadow( - path: &Path, - block_bytes: usize, - shutdown_token: &CancellationToken, - writeback_coordinator: Arc, -) -> Result { - let shadow = sparse_shadow_path(path); - let result = (|| { - if shutdown_token.is_cancelled() { - return Err(SparseShadowCmpError::Cancelled); - } - let mut reader = std::fs::File::open(path) - .map_err(|error| SparseShadowCmpError::io("open dense shadow source", error))?; - let output = std::fs::File::create(&shadow) - .map_err(|error| SparseShadowCmpError::io("create sparse shadow output", error))?; - let candidate_observer = Arc::new(std::sync::atomic::AtomicU64::new(0)); - let writeback_timing = Arc::new(StreamPullTiming::default()); - let mut writer = DecryptedFileSink::new( - output, - MaterializationMode::Enabled, - block_bytes, - Arc::clone(&candidate_observer), - Some(writeback_coordinator), - Some(writeback_timing), - JobPriority::default(), - ) - .map_err(|error| SparseShadowCmpError::io("create sparse shadow writer", error))?; - let mut buffer = vec![0; 1024 * 1024]; - let mut read_through = 0_u64; - let mut advised_through = 0_u64; - loop { - if shutdown_token.is_cancelled() { - return Err(SparseShadowCmpError::Cancelled); - } - let read = reader - .read(&mut buffer) - .map_err(|error| SparseShadowCmpError::io("read dense shadow source", error))?; - if read == 0 { - break; - } - let chunk = buffer.get(..read).ok_or_else(|| { - SparseShadowCmpError::io( - "read dense shadow source", - std::io::Error::other("read exceeded shadow buffer"), - ) - })?; - writer - .write_all(chunk) - .map_err(|error| SparseShadowCmpError::io("write sparse shadow output", error))?; - read_through = read_through.saturating_add(read as u64); - if read_through.saturating_sub(advised_through) >= STREAM_WRITEBACK_WINDOW_BYTES { - advise_shadow_read_range( - &reader, - path, - advised_through, - read_through.saturating_sub(advised_through), - "read dense sparse-shadow source", - ); - advised_through = read_through; - } - } - writer - .complete() - .map_err(|error| SparseShadowCmpError::io("finalize sparse shadow output", error))?; - if advised_through < read_through { - advise_shadow_read_range( - &reader, - path, - advised_through, - read_through.saturating_sub(advised_through), - "read dense sparse-shadow source", - ); - } - let candidate_hole_bytes = candidate_observer.load(std::sync::atomic::Ordering::Relaxed); - let output = std::fs::File::open(&shadow) - .map_err(|error| SparseShadowCmpError::io("open sparse shadow output", error))?; - page_cache::advise_dontneed(&reader, path, "read dense sparse-shadow source"); - page_cache::advise_dontneed(&output, &shadow, "write sparse shadow comparison file"); - drop(reader); - drop(output); - - compare_completed_files(path, &shadow, shutdown_token)?; - - let shadow_metadata = std::fs::metadata(&shadow) - .map_err(|error| SparseShadowCmpError::io("stat sparse shadow output", error))?; - let shadow_logical_bytes = shadow_metadata.len(); - #[cfg(unix)] - let shadow_allocated_bytes = { - use std::os::unix::fs::MetadataExt; - shadow_metadata.blocks().saturating_mul(512) - }; - #[cfg(not(unix))] - let shadow_allocated_bytes = shadow_logical_bytes; - Ok(ShadowComparisonObservation { - logical_bytes: shadow_logical_bytes, - allocated_bytes: shadow_allocated_bytes, - candidate_hole_bytes, - }) - })(); - - match std::fs::remove_file(&shadow) { - Ok(()) => {} - Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} - Err(error) => tracing::warn!( - path = %shadow.display(), - %error, - "failed to remove sparse shadow comparison file" - ), - } - result -} - /// Decompress a zstd-compressed file in-place: decompress to a /// `{path}..dec.tmp` scratch sibling, then atomically replace /// `path`. @@ -1466,1677 +726,28 @@ fn should_ignore_chmod_permission_denied( && current_rwx & desired_read_bits == desired_read_bits } -/// Buffered writer size between the (optional) zstd decoder and the -/// destination file when decrypt + decompress run back-to-back on the -/// pull side. Match the decoder output window to batch file writes and their -/// accounting. Both buffers remain bounded per admitted decoder. -const DECRYPT_PULL_WRITE_BUF_BYTES: usize = 1024 * 1024; - -// Submit decoded output to the process-wide writeback worker in the same -// increments used by ranged downloads. Writers share one byte budget, so -// aggregate dirty output stays bounded independently of pull concurrency. -const STREAM_WRITEBACK_WINDOW_BYTES: u64 = 32 * 1024 * 1024; -pub const DEFAULT_STREAM_WRITEBACK_BUDGET_BYTES: u64 = 2 * 1024 * 1024 * 1024; -const STREAM_WRITEBACK_IDLE_DRAIN: std::time::Duration = std::time::Duration::from_secs(1); -const STREAM_WRITEBACK_METRIC_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); - -#[derive(Debug, Default)] -struct StreamPullTiming { - compressed_input_bytes: std::sync::atomic::AtomicU64, - decoded_output_bytes: std::sync::atomic::AtomicU64, - decode_writer_us: std::sync::atomic::AtomicU64, - decode_writer_cpu_us: std::sync::atomic::AtomicU64, - output_writer_us: std::sync::atomic::AtomicU64, - output_writer_cpu_us: std::sync::atomic::AtomicU64, - writeback_reserve_us: std::sync::atomic::AtomicU64, - writeback_reserve_cpu_us: std::sync::atomic::AtomicU64, - writeback_start_us: std::sync::atomic::AtomicU64, - writeback_start_cpu_us: std::sync::atomic::AtomicU64, - writeback_wait_us: std::sync::atomic::AtomicU64, - writeback_start_calls: std::sync::atomic::AtomicU64, - writeback_submitted_logical_bytes: std::sync::atomic::AtomicU64, - writeback_submitted_dirty_bytes: std::sync::atomic::AtomicU64, - writeback_hole_bytes_released: std::sync::atomic::AtomicU64, - writeback_low_watermark_activations: std::sync::atomic::AtomicU64, -} - -fn current_thread_cpu_us() -> u64 { - let Ok(value) = nix::time::clock_gettime(nix::time::ClockId::CLOCK_THREAD_CPUTIME_ID) else { - return 0; - }; - let seconds = u64::try_from(value.tv_sec()).unwrap_or(0); - let nanoseconds = u64::try_from(value.tv_nsec()).unwrap_or(0); - seconds - .saturating_mul(1_000_000) - .saturating_add(nanoseconds / 1_000) -} - -impl StreamPullTiming { - fn add(counter: &std::sync::atomic::AtomicU64, value: u64) { - match counter.fetch_update( - std::sync::atomic::Ordering::Relaxed, - std::sync::atomic::Ordering::Relaxed, - |current| Some(current.saturating_add(value)), - ) { - Ok(_) | Err(_) => {} - } - } - - fn add_elapsed(counter: &std::sync::atomic::AtomicU64, started_at: Instant) { - Self::add( - counter, - u64::try_from(started_at.elapsed().as_micros()).unwrap_or(u64::MAX), - ); - } - - fn add_wall_and_thread_cpu( - wall_counter: &std::sync::atomic::AtomicU64, - cpu_counter: &std::sync::atomic::AtomicU64, - wall_started_at: Instant, - cpu_started_at_us: u64, - ) { - let cpu_elapsed_us = current_thread_cpu_us().saturating_sub(cpu_started_at_us); - let wall_elapsed_us = - u64::try_from(wall_started_at.elapsed().as_micros()).unwrap_or(u64::MAX); - Self::add(wall_counter, wall_elapsed_us); - Self::add(cpu_counter, cpu_elapsed_us); - } - - fn load(counter: &std::sync::atomic::AtomicU64) -> u64 { - counter.load(std::sync::atomic::Ordering::Relaxed) - } +/// Pull-side counterpart of [`stage_encrypted_temp`]. Reads the +/// downloaded ciphertext at `ciphertext_path`, unwraps the DEK, +/// AES-256-GCM-decrypts in 1 MiB chunks, and feeds the plaintext +/// straight into a [`MaybeZstdDecodingWriter`] writing the final +/// `.plaintext.tmp` next to the input. No `.plaintext.tmp` for the +/// decrypted-but-still-compressed intermediate — when the plaintext +/// is zstd-magic and the caller asked for transparent decompression, +/// it lands decompressed on disk in one pass. +/// +/// Returns the plaintext temp path and decompression report on success. A +/// [`DecodeGate::Capped`] pass that overflows its cap removes its own +/// partial output and returns [`CappedDecode::OverCap`] with the ciphertext +/// intact. On failure the partial `.plaintext.tmp` may exist — the caller is +/// responsible for cleaning it up, same contract as +/// [`stage_encrypted_temp`]. +#[derive(Clone, Copy)] +struct DecryptedPullOptions { + decompress_if_zstd: bool, + materialization_mode: MaterializationMode, } -fn log_streamed_pull_pipeline_timing( - job: &Job, - payload: &TransferPayload, - timing: &StreamPullTiming, - streamed_decoded: bool, - download_wall_ms: u64, -) { - let decode_writer_us = StreamPullTiming::load(&timing.decode_writer_us); - let decode_writer_cpu_us = StreamPullTiming::load(&timing.decode_writer_cpu_us); - let output_writer_us = StreamPullTiming::load(&timing.output_writer_us); - let output_writer_cpu_us = StreamPullTiming::load(&timing.output_writer_cpu_us); - let writeback_reserve_us = StreamPullTiming::load(&timing.writeback_reserve_us); - let writeback_reserve_cpu_us = StreamPullTiming::load(&timing.writeback_reserve_cpu_us); - tracing::info!( - job_id = %job.job_id, - subject_id = %payload.subject_id, - s3_uri_len = payload.s3_uri.len(), - local_path_len = payload.local_path.len(), - streamed_decoded, - download_decode_wall_ms = download_wall_ms, - compressed_input_bytes = StreamPullTiming::load(&timing.compressed_input_bytes), - decoded_output_bytes = StreamPullTiming::load(&timing.decoded_output_bytes), - decode_writer_ms = decode_writer_us / 1_000, - decode_stream_processing_ms = decode_writer_us - .saturating_sub(output_writer_us) - .saturating_sub(writeback_reserve_us) - / 1_000, - decode_stream_processing_cpu_ms = decode_writer_cpu_us - .saturating_sub(output_writer_cpu_us) - .saturating_sub(writeback_reserve_cpu_us) - / 1_000, - output_writer_ms = output_writer_us / 1_000, - output_writer_cpu_ms = output_writer_cpu_us / 1_000, - writeback_reserve_ms = writeback_reserve_us / 1_000, - writeback_reserve_cpu_ms = writeback_reserve_cpu_us / 1_000, - writeback_start_ms = StreamPullTiming::load(&timing.writeback_start_us) / 1_000, - writeback_start_cpu_ms = StreamPullTiming::load(&timing.writeback_start_cpu_us) / 1_000, - writeback_wait_observed_ms = StreamPullTiming::load(&timing.writeback_wait_us) / 1_000, - writeback_start_calls = StreamPullTiming::load(&timing.writeback_start_calls), - writeback_submitted_logical_bytes = - StreamPullTiming::load(&timing.writeback_submitted_logical_bytes), - writeback_submitted_dirty_bytes = - StreamPullTiming::load(&timing.writeback_submitted_dirty_bytes), - writeback_hole_bytes_released = - StreamPullTiming::load(&timing.writeback_hole_bytes_released), - writeback_low_watermark_activations = - StreamPullTiming::load(&timing.writeback_low_watermark_activations), - "streamed pull pipeline timing" - ); -} - -#[derive(Debug)] -struct StreamWritebackRange { - file: Arc, - offset: u64, - length: u64, - debt_bytes: u64, - started: bool, - timing: Arc, -} - -#[derive(Debug)] -enum StreamWritebackCommand { - Range(StreamWritebackRange), - ReapToLowWatermark, -} - -trait StreamPageCacheOperations: std::fmt::Debug + Send + Sync { - fn start(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()>; - fn wait(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()>; - fn advise_dontneed( - &self, - file: &std::fs::File, - offset: u64, - length: u64, - ) -> std::io::Result<()>; -} - -#[derive(Debug)] -struct SystemStreamPageCacheOperations; - -impl StreamPageCacheOperations for SystemStreamPageCacheOperations { - fn start(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()> { - page_cache::start_writeback_range(file, offset, length) - } - - fn wait(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()> { - page_cache::wait_writeback_range(file, offset, length) - } - - fn advise_dontneed( - &self, - file: &std::fs::File, - offset: u64, - length: u64, - ) -> std::io::Result<()> { - page_cache::advise_range_dontneed(file, offset, length) - } -} - -#[derive(Debug, Default)] -struct StreamWritebackBudget { - outstanding_bytes: std::sync::Mutex, - available: std::sync::Condvar, - failed: std::sync::atomic::AtomicBool, -} - -#[derive(Debug)] -struct StreamWritebackCoordinator { - sender: std::sync::mpsc::Sender, - budget: Arc, - high_watermark_bytes: u64, - priority_threshold: Option, - priority_reserve_bytes: u64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum StreamWritebackFailureOperation { - Start, - Wait, - Advice, - Queue, -} - -impl StreamWritebackFailureOperation { - fn as_str(self) -> &'static str { - match self { - Self::Start => "start", - Self::Wait => "wait", - Self::Advice => "advice", - Self::Queue => "queue", - } - } -} - -impl std::fmt::Display for StreamWritebackFailureOperation { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(self.as_str()) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum StreamWritebackPriorityClass { - Shared, - High, -} - -impl StreamWritebackPriorityClass { - fn as_str(self) -> &'static str { - match self { - Self::Shared => "shared", - Self::High => "high", - } - } -} - -impl std::fmt::Display for StreamWritebackPriorityClass { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(self.as_str()) - } -} - -fn fail_stream_writeback_budget( - budget: &StreamWritebackBudget, - operation: StreamWritebackFailureOperation, -) { - let was_failed = budget - .failed - .swap(true, std::sync::atomic::Ordering::AcqRel); - budget.available.notify_all(); - if !was_failed { - pplx_space_metrics::volume_manager::stream_writeback_failure(operation.as_str()); - } -} - -fn release_stream_writeback_budget(budget: &StreamWritebackBudget, bytes: u64) { - let mut outstanding = budget - .outstanding_bytes - .lock() - .expect("stream writeback budget lock poisoned"); - *outstanding = outstanding.saturating_sub(bytes); - budget.available.notify_all(); -} - -fn emit_stream_writeback_outstanding_if_due( - budget: &StreamWritebackBudget, - last_emitted_at: &mut Instant, - force: bool, -) { - let outstanding = stream_writeback_outstanding_bytes(budget); - if !should_emit_stream_writeback_outstanding(last_emitted_at.elapsed(), force) { - return; - } - pplx_space_metrics::volume_manager::stream_writeback_outstanding_bytes(outstanding); - *last_emitted_at = Instant::now(); -} - -fn should_emit_stream_writeback_outstanding(elapsed: std::time::Duration, force: bool) -> bool { - force || elapsed >= STREAM_WRITEBACK_METRIC_INTERVAL -} - -fn retire_stream_writeback_range( - range: StreamWritebackRange, - budget: &StreamWritebackBudget, - page_cache_operations: &dyn StreamPageCacheOperations, - writeback_enabled: &mut bool, - cache_advice_enabled: &mut bool, -) { - if *writeback_enabled && range.started { - let wait_started_at = Instant::now(); - let result = page_cache_operations.wait(&range.file, range.offset, range.length); - StreamPullTiming::add_elapsed(&range.timing.writeback_wait_us, wait_started_at); - if let Err(error) = result { - tracing::warn!( - %error, - offset = range.offset, - length = range.length, - "waiting for streamed output range writeback failed; closing streamed pull admission" - ); - *writeback_enabled = false; - fail_stream_writeback_budget(budget, StreamWritebackFailureOperation::Wait); - } - } - if *writeback_enabled - && range.started - && *cache_advice_enabled - && let Err(error) = - page_cache_operations.advise_dontneed(&range.file, range.offset, range.length) - { - tracing::warn!( - %error, - offset = range.offset, - length = range.length, - "streamed output range cache advice failed; disabling cache advice" - ); - pplx_space_metrics::volume_manager::stream_writeback_failure( - StreamWritebackFailureOperation::Advice.as_str(), - ); - *cache_advice_enabled = false; - } - release_stream_writeback_budget(budget, range.debt_bytes); -} - -fn stream_writeback_outstanding_bytes(budget: &StreamWritebackBudget) -> u64 { - *budget - .outstanding_bytes - .lock() - .expect("stream writeback budget lock poisoned") -} - -fn pending_stream_writeback_debt( - pending: &std::collections::VecDeque, -) -> u64 { - pending - .iter() - .fold(0_u64, |total, range| total.saturating_add(range.debt_bytes)) -} - -fn drain_stream_writeback_ranges_to( - pending: &mut std::collections::VecDeque, - budget: &StreamWritebackBudget, - target_bytes: u64, - page_cache_operations: &dyn StreamPageCacheOperations, - writeback_enabled: &mut bool, - cache_advice_enabled: &mut bool, -) -> bool { - while stream_writeback_outstanding_bytes(budget) > target_bytes { - let Some(range) = pending.pop_front() else { - return false; - }; - retire_stream_writeback_range( - range, - budget, - page_cache_operations, - writeback_enabled, - cache_advice_enabled, - ); - } - true -} - -fn drain_all_stream_writeback_ranges( - pending: &mut std::collections::VecDeque, - budget: &StreamWritebackBudget, - page_cache_operations: &dyn StreamPageCacheOperations, - writeback_enabled: &mut bool, - cache_advice_enabled: &mut bool, -) { - while let Some(range) = pending.pop_front() { - retire_stream_writeback_range( - range, - budget, - page_cache_operations, - writeback_enabled, - cache_advice_enabled, - ); - } -} - -fn queue_stream_writeback_range( - pending: &mut std::collections::VecDeque, - range: StreamWritebackRange, -) { - if let Some(previous) = pending.iter_mut().rev().find(|previous| { - Arc::ptr_eq(&previous.file, &range.file) && Arc::ptr_eq(&previous.timing, &range.timing) - }) { - let can_merge = !previous.started - && previous - .offset - .checked_add(previous.length) - .is_some_and(|end| end <= range.offset) - && previous - .debt_bytes - .checked_add(range.debt_bytes) - .is_some_and(|debt| debt <= STREAM_WRITEBACK_WINDOW_BYTES); - if can_merge - && let (Some(length), Some(debt_bytes)) = ( - range - .offset - .checked_add(range.length) - .and_then(|end| end.checked_sub(previous.offset)), - previous.debt_bytes.checked_add(range.debt_bytes), - ) - { - previous.length = length; - previous.debt_bytes = debt_bytes; - return; - } - } - pending.push_back(range); -} - -fn start_pending_stream_writeback_ranges( - pending: &mut std::collections::VecDeque, - budget: &StreamWritebackBudget, - page_cache_operations: &dyn StreamPageCacheOperations, - writeback_enabled: &mut bool, - minimum_debt_bytes: u64, -) { - for range in pending - .iter_mut() - .filter(|range| !range.started && range.debt_bytes >= minimum_debt_bytes) - { - if !*writeback_enabled { - break; - } - let start_started_at = Instant::now(); - let start_cpu_started_at = current_thread_cpu_us(); - let result = page_cache_operations.start(&range.file, range.offset, range.length); - StreamPullTiming::add_wall_and_thread_cpu( - &range.timing.writeback_start_us, - &range.timing.writeback_start_cpu_us, - start_started_at, - start_cpu_started_at, - ); - match result { - Ok(()) => { - range.started = true; - StreamPullTiming::add(&range.timing.writeback_start_calls, 1); - StreamPullTiming::add( - &range.timing.writeback_submitted_logical_bytes, - range.length, - ); - StreamPullTiming::add( - &range.timing.writeback_submitted_dirty_bytes, - range.debt_bytes, - ); - } - Err(error) => { - tracing::warn!( - %error, - offset = range.offset, - length = range.length, - debt_bytes = range.debt_bytes, - "starting streamed output range writeback failed; closing streamed pull admission" - ); - *writeback_enabled = false; - fail_stream_writeback_budget(budget, StreamWritebackFailureOperation::Start); - } - } - } -} - -impl StreamWritebackCoordinator { - fn new(high_watermark_bytes: u64, priority_threshold: Option) -> Arc { - Self::new_with_page_cache_operations( - high_watermark_bytes, - priority_threshold, - Arc::new(SystemStreamPageCacheOperations), - ) - } - - fn new_with_page_cache_operations( - high_watermark_bytes: u64, - priority_threshold: Option, - page_cache_operations: Arc, - ) -> Arc { - let (sender, receiver) = std::sync::mpsc::channel::(); - let budget = Arc::new(StreamWritebackBudget::default()); - let worker_budget = Arc::clone(&budget); - let low_watermark_bytes = high_watermark_bytes / 2; - let writeback_window_bytes = STREAM_WRITEBACK_WINDOW_BYTES.min(low_watermark_bytes); - std::thread::Builder::new() - .name("volmgr-writeback".to_string()) - .spawn(move || { - let mut writeback_enabled = true; - let mut cache_advice_enabled = true; - let mut pending = std::collections::VecDeque::new(); - let mut drain_target = None; - let mut pressure_active = false; - let mut last_outstanding_metric_at = Instant::now(); - - loop { - let mut force_outstanding_metric = false; - let writeback_enabled_at_start = writeback_enabled; - match receiver.recv_timeout(STREAM_WRITEBACK_IDLE_DRAIN) { - Ok(StreamWritebackCommand::Range(range)) if writeback_enabled => { - let activation_timing = Arc::clone(&range.timing); - queue_stream_writeback_range(&mut pending, range); - // Overlap complete dirty windows with decoding even - // below the pressure watermark. Keep their debt - // charged until the existing completion/drain path. - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - writeback_window_bytes, - ); - if pending_stream_writeback_debt(&pending) >= low_watermark_bytes { - let newly_activated = !pressure_active; - if !pressure_active { - pressure_active = true; - force_outstanding_metric = true; - StreamPullTiming::add( - &activation_timing.writeback_low_watermark_activations, - 1, - ); - } - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - if newly_activated { - 0 - } else { - writeback_window_bytes - }, - ); - } - } - Ok(StreamWritebackCommand::Range(range)) => { - release_stream_writeback_budget(&worker_budget, range.debt_bytes); - } - Ok(StreamWritebackCommand::ReapToLowWatermark) => { - force_outstanding_metric = true; - drain_target = Some(low_watermark_bytes); - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - 0, - ); - } - Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { - if pending_stream_writeback_debt(&pending) >= low_watermark_bytes { - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - 0, - ); - drain_target = Some(low_watermark_bytes); - } - } - Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - 0, - ); - drain_all_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - &mut cache_advice_enabled, - ); - emit_stream_writeback_outstanding_if_due( - &worker_budget, - &mut last_outstanding_metric_at, - true, - ); - break; - } - } - - if drain_target.is_some() { - start_pending_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - 0, - ); - } - if drain_target.is_some_and(|target| { - drain_stream_writeback_ranges_to( - &mut pending, - &worker_budget, - target, - page_cache_operations.as_ref(), - &mut writeback_enabled, - &mut cache_advice_enabled, - ) - }) { - drain_target = None; - pressure_active = stream_writeback_outstanding_bytes(&worker_budget) - > low_watermark_bytes; - force_outstanding_metric = true; - } - if !writeback_enabled { - drain_all_stream_writeback_ranges( - &mut pending, - &worker_budget, - page_cache_operations.as_ref(), - &mut writeback_enabled, - &mut cache_advice_enabled, - ); - drain_target = None; - pressure_active = false; - force_outstanding_metric |= writeback_enabled_at_start; - } - emit_stream_writeback_outstanding_if_due( - &worker_budget, - &mut last_outstanding_metric_at, - force_outstanding_metric, - ); - } - }) - .expect("spawn volume-manager writeback worker"); - Arc::new(Self { - sender, - budget, - high_watermark_bytes, - priority_threshold, - priority_reserve_bytes: if priority_threshold.is_some() { - STREAM_WRITEBACK_WINDOW_BYTES.min(high_watermark_bytes / 2) - } else { - 0 - }, - }) - } - - fn admission_limit(&self, priority: JobPriority) -> u64 { - if self - .priority_threshold - .is_some_and(|threshold| priority < threshold) - { - self.high_watermark_bytes - .saturating_sub(self.priority_reserve_bytes) - } else { - self.high_watermark_bytes - } - } - - fn try_reserve(&self, bytes: u64, priority: JobPriority) -> std::io::Result { - if self - .budget - .failed - .load(std::sync::atomic::Ordering::Acquire) - { - return Err(std::io::Error::other("stream writeback coordinator failed")); - } - let mut outstanding = self - .budget - .outstanding_bytes - .lock() - .map_err(|_| std::io::Error::other("stream writeback budget lock poisoned"))?; - if self - .budget - .failed - .load(std::sync::atomic::Ordering::Acquire) - { - return Err(std::io::Error::other("stream writeback coordinator failed")); - } - let admission_limit = self.admission_limit(priority); - if *outstanding != 0 && outstanding.saturating_add(bytes) > admission_limit { - return Ok(false); - } - *outstanding = outstanding.saturating_add(bytes); - Ok(true) - } - - fn reserve(&self, bytes: u64, priority: JobPriority) -> std::io::Result<()> { - let reserve_started_at = Instant::now(); - if self.try_reserve(bytes, priority)? { - return Ok(()); - } - let mut outstanding = self - .budget - .outstanding_bytes - .lock() - .map_err(|_| std::io::Error::other("stream writeback budget lock poisoned"))?; - if self - .budget - .failed - .load(std::sync::atomic::Ordering::Acquire) - { - return Err(std::io::Error::other("stream writeback coordinator failed")); - } - let admission_limit = self.admission_limit(priority); - let priority_class = if admission_limit == self.high_watermark_bytes { - StreamWritebackPriorityClass::High - } else { - StreamWritebackPriorityClass::Shared - }; - while *outstanding != 0 && outstanding.saturating_add(bytes) > admission_limit { - if self - .sender - .send(StreamWritebackCommand::ReapToLowWatermark) - .is_err() - { - tracing::warn!( - "stream writeback worker stopped while a producer was blocked; closing streamed pull admission" - ); - fail_stream_writeback_budget(&self.budget, StreamWritebackFailureOperation::Queue); - return Err(std::io::Error::other("stream writeback worker stopped")); - } - outstanding = self.budget.available.wait(outstanding).map_err(|_| { - std::io::Error::other("stream writeback budget lock poisoned while waiting") - })?; - if self - .budget - .failed - .load(std::sync::atomic::Ordering::Acquire) - { - return Err(std::io::Error::other("stream writeback coordinator failed")); - } - } - *outstanding = outstanding.saturating_add(bytes); - pplx_space_metrics::volume_manager::stream_writeback_producer_blocked_ms( - reserve_started_at.elapsed().as_secs_f64() * 1000.0, - priority_class.as_str(), - ); - Ok(()) - } - - fn release(&self, bytes: u64) { - let mut outstanding = self - .budget - .outstanding_bytes - .lock() - .expect("stream writeback budget lock poisoned"); - *outstanding = outstanding.saturating_sub(bytes); - self.budget.available.notify_all(); - } - - fn submit( - &self, - file: Arc, - offset: u64, - length: u64, - debt_bytes: u64, - timing: Arc, - ) { - let result = self - .sender - .send(StreamWritebackCommand::Range(StreamWritebackRange { - file, - offset, - length, - debt_bytes, - started: false, - timing, - })) - .map_err(|_| std::io::Error::other("stream writeback worker stopped")); - if let Err(error) = result { - tracing::warn!( - %error, - offset, - length, - debt_bytes, - "streamed output range could not be queued for writeback; closing streamed pull admission" - ); - fail_stream_writeback_budget(&self.budget, StreamWritebackFailureOperation::Queue); - self.release(debt_bytes); - } - } -} - -#[derive(Debug)] -struct StreamWritebackAccumulator { - offset: u64, - length: u64, - debt_bytes: u64, -} - -#[derive(Debug)] -struct StreamWriteback { - coordinator: Option>, - timing: Option>, - priority: JobPriority, - writeback_file: Option>, - unclassified_bytes: u64, - reserved_bytes: u64, - deferred_hole_bytes: u64, - pending: Option, -} - -impl StreamWriteback { - fn new( - coordinator: Option>, - timing: Option>, - priority: JobPriority, - ) -> Self { - Self { - coordinator, - timing, - priority, - writeback_file: None, - unclassified_bytes: 0, - reserved_bytes: 0, - deferred_hole_bytes: 0, - pending: None, - } - } - - fn pending_debt_bytes(&self) -> u64 { - self.pending - .as_ref() - .map_or(0, |pending| pending.debt_bytes) - } - - fn additional_reservation_bytes(&self, bytes: u64) -> u64 { - let Some(coordinator) = &self.coordinator else { - return 0; - }; - let owned_bytes = self - .unclassified_bytes - .saturating_add(self.pending_debt_bytes()) - .saturating_add(self.deferred_hole_bytes); - let required_bytes = owned_bytes.saturating_add(bytes); - if required_bytes <= self.reserved_bytes { - return 0; - } - let reservation_window = STREAM_WRITEBACK_WINDOW_BYTES - .min(coordinator.admission_limit(self.priority) / 2) - .max(bytes); - owned_bytes - .saturating_add(reservation_window) - .saturating_sub(self.reserved_bytes) - } - - fn reserve(&mut self, bytes: usize) -> std::io::Result<()> { - let bytes = bytes as u64; - if self.coordinator.is_some() { - let mut additional = self.additional_reservation_bytes(bytes); - if additional != 0 { - self.release_deferred_holes(); - additional = self.additional_reservation_bytes(bytes); - let coordinator = self - .coordinator - .as_ref() - .expect("coordinator checked above"); - let reserve_started_at = Instant::now(); - let reserve_cpu_started_at = current_thread_cpu_us(); - coordinator.reserve(additional, self.priority)?; - if let Some(timing) = &self.timing { - StreamPullTiming::add_wall_and_thread_cpu( - &timing.writeback_reserve_us, - &timing.writeback_reserve_cpu_us, - reserve_started_at, - reserve_cpu_started_at, - ); - } - self.reserved_bytes = self.reserved_bytes.saturating_add(additional); - } - } - Ok(()) - } - - fn try_reserve(&mut self, bytes: usize) -> std::io::Result { - let bytes = bytes as u64; - if self.coordinator.is_none() { - return Ok(true); - } - let mut additional = self.additional_reservation_bytes(bytes); - if additional == 0 { - return Ok(true); - } - self.release_deferred_holes(); - additional = self.additional_reservation_bytes(bytes); - let coordinator = self - .coordinator - .as_ref() - .expect("coordinator checked above"); - let reserve_started_at = Instant::now(); - let reserve_cpu_started_at = current_thread_cpu_us(); - if !coordinator.try_reserve(additional, self.priority)? { - return Ok(false); - } - if let Some(timing) = &self.timing { - StreamPullTiming::add_wall_and_thread_cpu( - &timing.writeback_reserve_us, - &timing.writeback_reserve_cpu_us, - reserve_started_at, - reserve_cpu_started_at, - ); - } - self.reserved_bytes = self.reserved_bytes.saturating_add(additional); - Ok(true) - } - - fn record_written(&mut self, bytes: usize) { - if self.coordinator.is_some() { - self.unclassified_bytes = self.unclassified_bytes.saturating_add(bytes as u64); - } - } - - fn release_unwritten(&mut self, bytes: usize) { - let bytes = bytes as u64; - if bytes == 0 { - return; - } - if let Some(coordinator) = &self.coordinator { - self.reserved_bytes = self.reserved_bytes.saturating_sub(bytes); - coordinator.release(bytes); - } - } - - fn writeback_file(&mut self, file: &std::fs::File) -> std::io::Result> { - if let Some(file) = &self.writeback_file { - return Ok(Arc::clone(file)); - } - let file = Arc::new(file.try_clone()?); - self.writeback_file = Some(Arc::clone(&file)); - Ok(file) - } - - fn release_deferred_holes(&mut self) { - let hole_bytes = std::mem::take(&mut self.deferred_hole_bytes); - if hole_bytes == 0 { - return; - } - self.reserved_bytes = self - .reserved_bytes - .checked_sub(hole_bytes) - .expect("deferred holes exceed producer reservation"); - let coordinator = self - .coordinator - .as_ref() - .expect("deferred holes require a coordinator"); - coordinator.release(hole_bytes); - if let Some(timing) = &self.timing { - StreamPullTiming::add(&timing.writeback_hole_bytes_released, hole_bytes); - } - } - - fn flush_pending(&mut self, file: &std::fs::File) -> std::io::Result<()> { - self.release_deferred_holes(); - let Some(pending) = self.pending.take() else { - return Ok(()); - }; - let writeback_file = self.writeback_file(file)?; - let coordinator = self - .coordinator - .as_ref() - .expect("pending writeback requires a coordinator"); - let timing = self - .timing - .clone() - .unwrap_or_else(|| Arc::new(StreamPullTiming::default())); - self.reserved_bytes = self - .reserved_bytes - .checked_sub(pending.debt_bytes) - .ok_or_else(|| { - std::io::Error::other("pending writeback exceeds producer reservation") - })?; - coordinator.submit( - writeback_file, - pending.offset, - pending.length, - pending.debt_bytes, - timing, - ); - Ok(()) - } - - fn accumulate_dirty_range( - &mut self, - file: &std::fs::File, - mut offset: u64, - mut length: u64, - ) -> std::io::Result<()> { - while length != 0 { - if self.pending_debt_bytes() == STREAM_WRITEBACK_WINDOW_BYTES { - self.flush_pending(file)?; - } - let available = STREAM_WRITEBACK_WINDOW_BYTES - self.pending_debt_bytes(); - let debt_bytes = length.min(available); - let end = offset - .checked_add(debt_bytes) - .ok_or_else(|| std::io::Error::other("streamed output range exceeds u64::MAX"))?; - if let Some(pending) = &mut self.pending { - let pending_end = pending.offset.checked_add(pending.length).ok_or_else(|| { - std::io::Error::other("pending streamed output range exceeds u64::MAX") - })?; - if offset < pending_end { - return Err(std::io::Error::other( - "streamed dirty output ranges are not ordered", - )); - } - pending.length = end.checked_sub(pending.offset).ok_or_else(|| { - std::io::Error::other("streamed output range precedes pending range") - })?; - pending.debt_bytes = - pending.debt_bytes.checked_add(debt_bytes).ok_or_else(|| { - std::io::Error::other("pending streamed dirty debt exceeds u64::MAX") - })?; - } else { - self.pending = Some(StreamWritebackAccumulator { - offset, - length: debt_bytes, - debt_bytes, - }); - } - offset = end; - length -= debt_bytes; - } - if self.pending_debt_bytes() == STREAM_WRITEBACK_WINDOW_BYTES { - self.flush_pending(file)?; - } - Ok(()) - } - - fn classify( - &mut self, - file: &std::fs::File, - materialization: SparseMaterialization, - ) -> std::io::Result<()> { - let (data_ranges, hole_bytes) = materialization.into_parts(); - let dirty_bytes = data_ranges - .iter() - .try_fold(0_u64, |total, range| total.checked_add(range.length())) - .ok_or_else(|| std::io::Error::other("streamed dirty output exceeds u64::MAX"))?; - let classified_bytes = dirty_bytes - .checked_add(hole_bytes) - .ok_or_else(|| std::io::Error::other("streamed classified output exceeds u64::MAX"))?; - if classified_bytes == 0 || self.coordinator.is_none() { - return Ok(()); - } - if classified_bytes > self.unclassified_bytes || classified_bytes > self.reserved_bytes { - return Err(std::io::Error::other( - "sparse writer classified more bytes than writeback admission reserved", - )); - } - self.unclassified_bytes -= classified_bytes; - if hole_bytes != 0 { - self.deferred_hole_bytes = self - .deferred_hole_bytes - .checked_add(hole_bytes) - .ok_or_else(|| std::io::Error::other("deferred sparse holes exceed u64::MAX"))?; - if self.deferred_hole_bytes >= STREAM_WRITEBACK_WINDOW_BYTES { - self.release_deferred_holes(); - } - } - for range in data_ranges { - self.accumulate_dirty_range(file, range.offset(), range.length())?; - } - Ok(()) - } - - fn classify_dense( - &mut self, - file: &std::fs::File, - offset: u64, - length: u64, - ) -> std::io::Result<()> { - self.classify( - file, - SparseMaterialization::from(SparseDataRange::new(offset, length)), - ) - } -} - -impl Drop for StreamWriteback { - fn drop(&mut self) { - self.release_deferred_holes(); - if self.reserved_bytes != 0 - && let Some(coordinator) = &self.coordinator - { - coordinator.release(self.reserved_bytes); - } - } -} - -/// A stream destination whose buffered state must be finalized after its -/// upstream decoder reaches end-of-stream. -trait CompletingWrite: Write { - fn complete(self) -> std::io::Result<()>; -} - -impl CompletingWrite for &mut W { - fn complete(self) -> std::io::Result<()> { - self.flush() - } -} - -/// Pull-side writer adapter that peeks the first four bytes written -/// into it, compares against the zstd magic, and either routes the -/// stream through a `zstd::stream::write::Decoder` (if the magic -/// matches AND `decompress_if_zstd` is true) or passes the bytes -/// through to the inner sink unchanged. -/// -/// Lets `decrypt_stream_blocking` feed straight into this writer so a -/// pull of an encrypted-and-compressed object decrypts + decompresses -/// in one streaming pass: no `.plaintext.tmp` intermediate file, no -/// re-read of the decrypted bytes for the magic sniff. The decision -/// is final after the first four bytes — empty or sub-4-byte writes -/// keep the bytes pending until enough has arrived to decide, and -/// `finish` flushes any still-pending bytes through the passthrough -/// path so a sub-4-byte plaintext (never zstd-magic) still round-trips -/// correctly. -struct MaybeZstdDecodingWriter -where - Input: Write, - Output: Write, - Select: FnOnce(Input, bool) -> std::io::Result, -{ - state: MaybeZstdState, - decompress_if_zstd: bool, - pending: Vec, - select_sink: Option, +} + +pub(super) enum MaybeZstdState { + /// First four bytes not yet seen; sink is parked here until we can + /// decide whether to wrap it with a zstd decoder. + Init(Input), + /// First four bytes were the zstd magic and the caller wanted + /// transparent decompress; subsequent writes flow through the + /// decoder. Built on the raw zio writer (not + /// `stream::write::Decoder`) because only the zio writer exposes the + /// fallible end-of-stream `finish` that rejects a truncated frame — + /// the higher-level decoder can only flush, which silently accepts + /// incomplete input. + Decoding( + zstd::stream::zio::Writer, zstd::stream::raw::Decoder<'static>>, + ), + /// First four bytes either weren't the zstd magic, or the caller + /// asked for raw (`raw = true`); subsequent writes go straight to + /// the buffered sink. + Passthrough(std::io::BufWriter), + /// Transition slot: `state` holds this value briefly while we + /// move the inner writer between variants. Should never be the + /// observed state at the start or end of a `Write` call. + Transient, +} + +/// A `BufWriter` flush that failed while finishing a +/// [`MaybeZstdDecodingWriter`]. The sink's error stays the source, so a +/// typed sink error such as a decode output-cap overflow survives the wrap. +#[derive(Debug, thiserror::Error)] +#[error("flush BufWriter at {stage}: {source}")] +pub(super) struct BufferedFlushError { + pub(super) stage: &'static str, + #[source] + pub(super) source: std::io::Error, +} + +impl BufferedFlushError { + pub(super) fn io(stage: &'static str, error: std::io::IntoInnerError) -> std::io::Error { + std::io::Error::other(Self { + stage, + source: error.into_error(), + }) + } +} + +#[cfg(test)] +pub(super) fn identity_decoding_sink( + sink: W, + _decompressing: bool, +) -> std::io::Result { + Ok(sink) +} + +#[cfg(test)] +impl MaybeZstdDecodingWriter std::io::Result> { + pub(super) fn new(sink: W, decompress_if_zstd: bool) -> Self { + Self::with_sink_selector(sink, decompress_if_zstd, identity_decoding_sink::) + } +} + +impl MaybeZstdDecodingWriter +where + Input: Write, + Output: Write, + Select: FnOnce(Input, bool) -> std::io::Result, +{ + pub(super) fn with_sink_selector( + sink: Input, + decompress_if_zstd: bool, + select_sink: Select, + ) -> Self { + Self { + state: MaybeZstdState::Init(sink), + decompress_if_zstd, + pending: Vec::with_capacity(ZSTD_MAGIC.len()), + select_sink: Some(select_sink), + } + } + + pub(super) fn select_sink( + &mut self, + sink: Input, + decompressing: bool, + ) -> std::io::Result { + self.select_sink + .take() + .ok_or_else(|| std::io::Error::other("zstd sink selector already consumed"))?( + sink, + decompressing, + ) + } + + /// Consume the writer. Returns `true` iff the stream was routed + /// through a zstd decoder (i.e. transparent decompression actually + /// happened). Flushes the trailing buffer and finalises the + /// decoder and destination when applicable. + pub(super) fn finish(mut self) -> std::io::Result + where + Output: CompletingWrite, + { + // If we never saw 4 bytes of plaintext (sub-magic-length input, + // or fully empty), flush the pending bytes through the + // passthrough path so the file at least contains whatever we + // did receive. + let pending = std::mem::take(&mut self.pending); + match std::mem::replace(&mut self.state, MaybeZstdState::Transient) { + MaybeZstdState::Init(sink) => { + let mut sink = self.select_sink(sink, false)?; + sink.write_all(&pending)?; + sink.complete()?; + Ok(false) + } + MaybeZstdState::Decoding(mut decoder) => { + // Fallible end-of-stream finalization FIRST: the raw + // decoder rejects a truncated or unterminated zstd frame + // with UnexpectedEof ("incomplete frame"). Flushing alone + // would silently accept partial input and commit corrupt + // decoded output. Then recover the BufWriter and + // explicitly flush it through to the destination file — + // without the explicit BufWriter::into_inner the trailing + // bytes would only land on disk via Drop, and a Drop + // error would be silently swallowed instead of surfaced + // as a JobError::Retriable. + decoder.finish()?; + let (buf_writer, _decoder) = decoder.into_inner(); + let sink = buf_writer + .into_inner() + .map_err(|e| BufferedFlushError::io("decoder finish", e))?; + sink.complete()?; + Ok(true) + } + MaybeZstdState::Passthrough(buf_writer) => { + let sink = buf_writer + .into_inner() + .map_err(|e| BufferedFlushError::io("stream finish", e))?; + sink.complete()?; + Ok(false) + } + MaybeZstdState::Transient => Err(std::io::Error::other( + "MaybeZstdDecodingWriter::finish observed transient state", + )), + } + } +} + +impl Write for MaybeZstdDecodingWriter +where + Input: Write, + Output: Write, + Select: FnOnce(Input, bool) -> std::io::Result, +{ + fn write(&mut self, buf: &[u8]) -> std::io::Result { + // Decide once: accumulate up to 4 bytes in `pending`, then + // transition state and dispatch the remaining bytes of this + // write call into the chosen sink. + if matches!(self.state, MaybeZstdState::Init(_)) { + let needed = ZSTD_MAGIC.len() - self.pending.len(); + let take = needed.min(buf.len()); + let (prefix_chunk, suffix) = buf.split_at(take); + self.pending.extend_from_slice(prefix_chunk); + if self.pending.len() < ZSTD_MAGIC.len() { + return Ok(take); + } + let is_zstd = self.pending == ZSTD_MAGIC; + let want_decompress = self.decompress_if_zstd && is_zstd; + let prefix = std::mem::take(&mut self.pending); + let sink = match std::mem::replace(&mut self.state, MaybeZstdState::Transient) { + MaybeZstdState::Init(sink) => sink, + _ => unreachable!("matches!(.., Init) above"), + }; + let sink = self.select_sink(sink, want_decompress)?; + if want_decompress { + let buf_writer = + std::io::BufWriter::with_capacity(DECRYPT_PULL_WRITE_BUF_BYTES, sink); + let raw_decoder = zstd::stream::raw::Decoder::new() + .map_err(|e| std::io::Error::other(format!("zstd decoder init: {e}")))?; + let mut decoder = zstd::stream::zio::Writer::new_with_capacity( + buf_writer, + raw_decoder, + DECRYPT_PULL_WRITE_BUF_BYTES, + ); + decoder.write_all(&prefix)?; + if !suffix.is_empty() { + decoder.write_all(suffix)?; + } + self.state = MaybeZstdState::Decoding(decoder); + } else { + let mut buf_writer = + std::io::BufWriter::with_capacity(DECRYPT_PULL_WRITE_BUF_BYTES, sink); + buf_writer.write_all(&prefix)?; + if !suffix.is_empty() { + buf_writer.write_all(suffix)?; + } + self.state = MaybeZstdState::Passthrough(buf_writer); + } + return Ok(buf.len()); + } + + match &mut self.state { + MaybeZstdState::Decoding(decoder) => decoder.write(buf), + MaybeZstdState::Passthrough(buf_writer) => buf_writer.write(buf), + MaybeZstdState::Init(_) | MaybeZstdState::Transient => unreachable!( + "Init handled above; Transient never observed outside the state-swap window" + ), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match &mut self.state { + MaybeZstdState::Init(sink) => sink.flush(), + MaybeZstdState::Decoding(decoder) => decoder.flush(), + MaybeZstdState::Passthrough(buf_writer) => buf_writer.flush(), + MaybeZstdState::Transient => Ok(()), + } + } +} + +pub(super) enum DecryptedFileSinkInner { + Dense(std::fs::File), + Classified(Box), +} + +pub(super) struct DecryptedFileSink { + pub(super) inner: DecryptedFileSinkInner, + pub(super) writeback: StreamWriteback, + pub(super) materialized_through: u64, +} + +impl DecryptedFileSink { + pub(super) fn new( + file: std::fs::File, + mode: MaterializationMode, + block_bytes: usize, + candidate_observer: Arc, + writeback_coordinator: Option>, + timing: Option>, + priority: JobPriority, + ) -> std::io::Result { + let inner = match mode { + MaterializationMode::Disabled | MaterializationMode::Shadow => { + DecryptedFileSinkInner::Dense(file) + } + MaterializationMode::Enabled => DecryptedFileSinkInner::Classified(Box::new( + SparseFileWriter::with_observer(file, block_bytes, candidate_observer)?, + )), + }; + Ok(Self { + inner, + writeback: StreamWriteback::new(writeback_coordinator, timing, priority), + materialized_through: 0, + }) + } + + pub(super) fn expose_classified_debt_before_blocking(&mut self) -> std::io::Result<()> { + let DecryptedFileSinkInner::Classified(writer) = &mut self.inner else { + return Ok(()); + }; + writer.materialize_complete_blocks_if_buffered_at_least(1)?; + self.materialized_through = writer.materialized_through(); + let materialization = writer.take_materialization(); + self.writeback.classify(writer.file(), materialization) + } +} + +impl CompletingWrite for DecryptedFileSink { + fn complete(self) -> std::io::Result<()> { + let Self { + inner, + mut writeback, + materialized_through: _, + } = self; + let flush_started_at = Instant::now(); + let flush_cpu_started_at = current_thread_cpu_us(); + let result: std::io::Result<(std::fs::File, Option)> = match inner { + DecryptedFileSinkInner::Dense(mut file) => { + file.flush()?; + Ok((file, None)) + } + DecryptedFileSinkInner::Classified(writer) => { + let (file, _, materialization) = writer.finish_with_stats_and_materialization()?; + Ok((file, Some(materialization))) + } + }; + if let Some(timing) = &writeback.timing { + StreamPullTiming::add_wall_and_thread_cpu( + &timing.output_writer_us, + &timing.output_writer_cpu_us, + flush_started_at, + flush_cpu_started_at, + ); + } + let (file, materialization) = result?; + if let Some(materialization) = materialization { + writeback.classify(&file, materialization)?; + } + writeback.flush_pending(&file)?; + Ok(()) + } +} + +impl Write for DecryptedFileSink { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + if !self.writeback.try_reserve(buf.len())? { + if matches!(&self.inner, DecryptedFileSinkInner::Classified(_)) { + self.expose_classified_debt_before_blocking()?; + } + let file = match &self.inner { + DecryptedFileSinkInner::Dense(file) => file, + DecryptedFileSinkInner::Classified(writer) => writer.file(), + }; + self.writeback.flush_pending(file)?; + self.writeback.reserve(buf.len())?; + } + let write_started_at = Instant::now(); + let write_cpu_started_at = current_thread_cpu_us(); + let write_offset = self.materialized_through; + let result = match &mut self.inner { + DecryptedFileSinkInner::Dense(file) => file.write(buf).and_then(|written| { + self.materialized_through = + self.materialized_through + .checked_add(written as u64) + .ok_or_else(|| std::io::Error::other("streamed output exceeds u64::MAX"))?; + Ok((written, None)) + }), + DecryptedFileSinkInner::Classified(writer) => writer.write(buf).and_then(|written| { + writer.materialize_complete_blocks_if_buffered_at_least( + STREAM_WRITEBACK_WINDOW_BYTES, + )?; + self.materialized_through = writer.materialized_through(); + Ok((written, Some(writer.take_materialization()))) + }), + }; + if let Some(timing) = &self.writeback.timing { + StreamPullTiming::add_wall_and_thread_cpu( + &timing.output_writer_us, + &timing.output_writer_cpu_us, + write_started_at, + write_cpu_started_at, + ); + } + let (written, materialization) = match result { + Ok(result) => result, + Err(error) => { + self.writeback.release_unwritten(buf.len()); + return Err(error); + } + }; + if let Some(timing) = &self.writeback.timing { + StreamPullTiming::add(&timing.decoded_output_bytes, written as u64); + } + self.writeback.record_written(written); + self.writeback + .release_unwritten(buf.len().saturating_sub(written)); + let file = match &self.inner { + DecryptedFileSinkInner::Dense(file) => file, + DecryptedFileSinkInner::Classified(writer) => writer.file(), + }; + if let Some(materialization) = materialization { + self.writeback.classify(file, materialization)?; + } else if written != 0 { + self.writeback + .classify_dense(file, write_offset, written as u64)?; + } + Ok(written) + } + + fn flush(&mut self) -> std::io::Result<()> { + let flush_started_at = Instant::now(); + let flush_cpu_started_at = current_thread_cpu_us(); + let result = match &mut self.inner { + DecryptedFileSinkInner::Dense(file) => file.flush().map(|()| None), + DecryptedFileSinkInner::Classified(writer) => { + writer.flush()?; + self.materialized_through = writer.materialized_through(); + Ok(Some(writer.take_materialization())) + } + }; + if let Some(timing) = &self.writeback.timing { + StreamPullTiming::add_wall_and_thread_cpu( + &timing.output_writer_us, + &timing.output_writer_cpu_us, + flush_started_at, + flush_cpu_started_at, + ); + } + let materialization = result?; + let file = match &self.inner { + DecryptedFileSinkInner::Dense(file) => file, + DecryptedFileSinkInner::Classified(writer) => writer.file(), + }; + if let Some(materialization) = materialization { + self.writeback.classify(file, materialization)?; + } + self.writeback.flush_pending(file)?; + Ok(()) + } +} + +pub(super) type StreamSinkSelector = + Box std::io::Result + Send>; + +/// Adapts [`MaybeZstdDecodingWriter`] to the object store's +/// [`pplx_space_object_store::FinishingWrite`] so the ranged download's +/// sequential writer decodes inline (one pass: fetch, decode, write). +/// Sparse selection happens only after the zstd-magic decision so a +/// passthrough stream remains dense. +pub(super) struct StreamDecodeSink { + pub(super) inner: MaybeZstdDecodingWriter, + pub(super) timing: Arc, +} + +impl StreamDecodeSink { + pub(super) fn new( + file: std::fs::File, + mode_if_decompressed: MaterializationMode, + block_bytes: usize, + candidate_observer: Arc, + writeback_coordinator: Arc, + timing: Arc, + priority: JobPriority, + ) -> Self { + let timing_for_sink = Arc::clone(&timing); + let select_sink: StreamSinkSelector = Box::new(move |file, decompressing| { + let mode = if decompressing { + mode_if_decompressed + } else { + MaterializationMode::Disabled + }; + DecryptedFileSink::new( + file, + mode, + block_bytes, + candidate_observer, + Some(writeback_coordinator), + Some(timing_for_sink), + priority, + ) + }); + Self { + inner: MaybeZstdDecodingWriter::with_sink_selector(file, true, select_sink), + timing, + } + } +} + +impl Write for StreamDecodeSink { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let write_started_at = Instant::now(); + let write_cpu_started_at = current_thread_cpu_us(); + let result = self.inner.write(buf); + StreamPullTiming::add_wall_and_thread_cpu( + &self.timing.decode_writer_us, + &self.timing.decode_writer_cpu_us, + write_started_at, + write_cpu_started_at, + ); + if let Ok(written) = &result { + StreamPullTiming::add(&self.timing.compressed_input_bytes, *written as u64); + } + result + } + + fn flush(&mut self) -> std::io::Result<()> { + let flush_started_at = Instant::now(); + let flush_cpu_started_at = current_thread_cpu_us(); + let result = self.inner.flush(); + StreamPullTiming::add_wall_and_thread_cpu( + &self.timing.decode_writer_us, + &self.timing.decode_writer_cpu_us, + flush_started_at, + flush_cpu_started_at, + ); + result + } +} + +impl pplx_space_object_store::FinishingWrite for StreamDecodeSink { + fn finish(self: Box) -> std::io::Result { + let StreamDecodeSink { inner, timing } = *self; + let finish_started_at = Instant::now(); + let finish_cpu_started_at = current_thread_cpu_us(); + let result = inner.finish(); + StreamPullTiming::add_wall_and_thread_cpu( + &timing.decode_writer_us, + &timing.decode_writer_cpu_us, + finish_started_at, + finish_cpu_started_at, + ); + result + } +} + +/// Codec admission bridge for streaming-decode pulls: one permit spans the +/// whole overlapped download+decode. The streamed path only runs for +/// ranged-size objects, whose decode cannot be proven small before the +/// bytes arrive (the frame header is not on disk yet), so admission is +/// unconditional — the conservative direction, and exactly what keeps the +/// highly-compressible multi-GiB snapshot class bounded. The artifact-class +/// bypass in `decode_admission` does not apply here: a stream needs more +/// than one download chunk of wire bytes (8 MiB minimum, 32 MiB by default), +/// far above any legitimate `metadata` or `vmstate` object, so such an object +/// is already suspect and keeps its permit. +pub(super) struct StreamedPullAdmission<'a> { + pub(super) executor: &'a TransferExecutor, + pub(super) kind: &'a str, + pub(super) artifact: &'static str, + pub(super) mode: MaterializationMode, + pub(super) priority: JobPriority, +} + +#[async_trait::async_trait] +impl pplx_space_object_store::DecodeAdmission for StreamedPullAdmission<'_> { + async fn admit( + &self, + _object_size: u64, + ) -> Result< + Option, + pplx_space_object_store::DecodeAdmissionError, + > { + let wait_started = Instant::now(); + // Typed so the failure crosses the object-store boundary as a source + // error rather than a string. + let permit = self + .executor + .codec + .acquire(self.executor.codec_priority(self.priority)) + .await + .map_err(pplx_space_object_store::DecodeAdmissionError::admission)?; + emit_transfer_phase_duration( + self.kind, + self.artifact, + TransferPhase::DecompressWait, + self.mode, + wait_started, + ); + Ok(Some(Box::new(permit))) + } +} diff --git a/crates/libs/vol/mgr/server/src/executor/pull.rs b/crates/libs/vol/mgr/server/src/executor/pull.rs new file mode 100644 index 000000000..b7f5d1dde --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/pull.rs @@ -0,0 +1,761 @@ +//! Pull jobs: destination locks, the existing-file policy, cancellation cleanup and the pull +//! itself. + +use super::*; + +pub(super) struct DestinationPullGuard { + pub(super) locks: DestinationPullLocks, + pub(super) path: PathBuf, + pub(super) guard: Option>, +} + +impl Drop for DestinationPullGuard { + fn drop(&mut self) { + self.guard.take(); + self.locks + .remove_if(&self.path, |_path, lock| Arc::strong_count(lock) == 1); + } +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PullExistingPolicy { + #[default] + LegacySkipPresent, + VerifiedReplica, +} + +impl PullExistingPolicy { + pub fn is_legacy_skip_present(policy: &Self) -> bool { + *policy == Self::LegacySkipPresent + } + + pub fn as_dedup_str(self) -> &'static str { + match self { + Self::LegacySkipPresent => "legacy_skip_present", + Self::VerifiedReplica => "verified_replica", + } + } +} + +pub(super) struct PullExecutionOptions<'a> { + pub(super) resolved_dek: Option<&'a [u8]>, + pub(super) existing_policy: PullExistingPolicy, + pub(super) cancellation: CancellationToken, +} + +/// Best-effort unlink of a `.partial` sibling. +/// +/// `NotFound` is the common expected case (no stale file to sweep), so we +/// silence it. Any other failure is logged so silent fs corruption — a +/// permission flip, a stuck mount — doesn't accumulate without diagnostics. +/// On a successful removal we emit a debug log so audits ("what was deleted +/// and why?") can trace each unlink to the code path that triggered it. +pub(super) async fn unlink_partial(path: &Path, reason: &str) { + match tokio::fs::remove_file(path).await { + Ok(()) => tracing::debug!( + path = %path.display(), + reason, + "unlinked .partial" + ), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} + Err(err) => tracing::warn!( + error = %err, + path = %path.display(), + reason, + "unlink .partial failed" + ), + } +} + +pub(super) async fn cleanup_cancelled_pull_temp_paths(local_path: &Path) -> Result<(), JobError> { + remove_cancelled_path(&partial_path(local_path)).await?; + remove_cancelled_path(&encrypted_sibling(local_path)).await?; + remove_cancelled_path(&decrypted_sibling(local_path)).await +} + +impl TransferExecutor { + pub(super) async fn lock_pull_destination(&self, local_path: &Path) -> DestinationPullGuard { + let path = local_path.to_path_buf(); + let lock = self + .destination_pull_locks + .entry(path.clone()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let guard = lock.lock_owned().await; + DestinationPullGuard { + locks: Arc::clone(&self.destination_pull_locks), + path, + guard: Some(guard), + } + } + + pub(super) async fn cleanup_cancelled_pull_path( + &self, + local_path: &Path, + ) -> Result<(), JobError> { + let _destination_guard = self.lock_pull_destination(local_path).await; + cleanup_cancelled_pull_temp_paths(local_path).await + } + + pub(super) async fn execute_pull_cancellable( + &self, + job: &Job, + p: &TransferPayload, + creds: &pplx_space_object_store::ObjectCredentials, + target: ObjectTarget<'_>, + options: PullExecutionOptions<'_>, + ) -> Result { + let local_path = Path::new(&p.local_path); + let _destination_guard = self.lock_pull_destination(local_path).await; + self.execute_pull_locked(job, p, creds, target, options) + .await + } + + #[cfg(test)] + pub(super) async fn execute_pull( + &self, + job: &Job, + p: &TransferPayload, + creds: &pplx_space_object_store::ObjectCredentials, + target: ObjectTarget<'_>, + resolved_dek: Option<&[u8]>, + existing_policy: PullExistingPolicy, + ) -> Result { + self.execute_pull_cancellable( + job, + p, + creds, + target, + PullExecutionOptions { + resolved_dek, + existing_policy, + cancellation: CancellationToken::new(), + }, + ) + .await + } + + pub(super) async fn execute_pull_locked( + &self, + job: &Job, + p: &TransferPayload, + creds: &pplx_space_object_store::ObjectCredentials, + target: ObjectTarget<'_>, + options: PullExecutionOptions<'_>, + ) -> Result { + let PullExecutionOptions { + resolved_dek, + existing_policy, + cancellation, + } = options; + let local_path = Path::new(&p.local_path); + let snapshot_artifact = p.effective_snapshot_artifact(); + let execution_materialization_mode = execution_materialization_mode( + snapshot_artifact, + p.requested_materialization_mode(snapshot_artifact), + ); + let artifact = transfer_artifact_tag(snapshot_artifact); + let tmp_path = partial_path(local_path); + + let source_meta = match existing_policy { + PullExistingPolicy::LegacySkipPresent => { + // Pre-check: if the local file already exists, looks the way + // our raw/non-raw contract says it should, and S3 still has + // the source object, we assume a previous download completed. + // + // Non-raw mode: the on-disk shape is the decompressed form, so + // a zstd-magic leftover means a previous attempt died between + // `get_file` and `decompress_in_place`. Fall through and + // re-download. + let half_decompressed = !p.raw && is_zstd(local_path).await; + if tokio::fs::metadata(local_path).await.is_ok() + && !half_decompressed + && let Ok(Some(_)) = self.store.head(creds, target).await + { + tracing::info!( + job_id = %job.job_id, + subject_id = %p.subject_id, + s3_uri = %p.s3_uri, + local_path = %p.local_path, + "pull skipped: local file already present" + ); + // Sweep any leftover .partial from an earlier aborted run + // so it doesn't accumulate. + unlink_partial(&tmp_path, "skip_present").await; + apply_file_mode(local_path, p.file_mode, FileModeApplyPolicy::ExistingFile) + .await?; + let size = tokio::fs::metadata(local_path) + .await + .map(|m| m.len()) + .unwrap_or(0); + return bytes_len_to_job_i64(size, "existing pull output size"); + } + None + } + PullExistingPolicy::VerifiedReplica => { + if let Some(size) = self + .validate_existing_file_replica( + p, + creds, + target, + local_path, + execution_materialization_mode, + ) + .await? + { + tracing::info!( + job_id = %job.job_id, + subject_id = %p.subject_id, + s3_uri = %p.s3_uri, + local_path = %p.local_path, + "file replica already present" + ); + unlink_partial(&tmp_path, "verified_replica").await; + apply_file_mode(local_path, p.file_mode, FileModeApplyPolicy::ExistingFile) + .await?; + return bytes_len_to_job_i64(size, "existing file replica size"); + } + Some(self.head_required(p, creds, target).await?) + } + }; + + if let Some(parent) = local_path.parent() { + tokio::fs::create_dir_all(parent).await.map_err(|e| { + JobError::permanent_with_reason( + JobFailureReason::Internal, + format!("create destination dir: {e}"), + ) + })?; + } + + // Clear any stale .partial left by a previous aborted attempt before + // streaming new bytes into it. + unlink_partial(&tmp_path, "pre_download").await; + + // Stream into the sibling .partial path; rename to the caller-visible + // local_path only on success so callers never observe a torn file. + // + // Plain compressed pulls decode inline with the ranged download when + // enabled: the reorder window already delivers sequential bytes, so + // routing them through the zstd decoder on the way to disk makes the + // job ≈ max(download, decode) instead of download + decode, and + // skips writing + re-reading the compressed intermediate entirely. + // Encrypted pulls keep their own chained decrypt+decode pass, and + // raw pulls must land wire bytes untouched. + let download_started_at = Instant::now(); + let stream_decode = self.config.stream_decode_pulls && p.encryption.is_none() && !p.raw; + let stream_pull_timing = Arc::new(StreamPullTiming::default()); + let (bytes, streamed, streamed_decoded, streamed_decompression_report) = if stream_decode { + let admission = StreamedPullAdmission { + executor: self, + kind: &job.kind, + artifact, + mode: execution_materialization_mode, + priority: job.priority, + }; + let mode_if_decompressed = + effective_materialization_mode(snapshot_artifact, execution_materialization_mode); + let block_bytes = materialization_block_bytes(snapshot_artifact); + let candidate_observer = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let observer_for_sink = Arc::clone(&candidate_observer); + let writeback_coordinator = Arc::clone(&self.stream_writeback); + let timing_for_sink = Arc::clone(&stream_pull_timing); + let stream_priority = job.priority; + let wrap: &pplx_space_object_store::WrapWriter = &move |file| { + Box::new(StreamDecodeSink::new( + file, + mode_if_decompressed, + block_bytes, + Arc::clone(&observer_for_sink), + Arc::clone(&writeback_coordinator), + Arc::clone(&timing_for_sink), + stream_priority, + )) + }; + let streamed_result = tokio::select! { + result = self.store.get_file_streamed( + creds, + target, + &tmp_path, + pplx_space_object_store::StreamDecode { + admission: &admission, + wrap_writer: wrap, + }, + ) => result, + _ = cancellation.cancelled() => { + unlink_partial(&tmp_path, "cancelled").await; + return Err(JobError::Permanent("pull cancelled".to_string())); + } + }; + match streamed_result { + Ok(pplx_space_object_store::StreamedGetFile::Streamed { bytes, decoded }) => { + let report = decoded.then(|| DecompressionReport { + mode: mode_if_decompressed, + candidate_hole_bytes: candidate_observer + .load(std::sync::atomic::Ordering::Relaxed), + }); + (bytes, true, decoded, report) + } + Ok(pplx_space_object_store::StreamedGetFile::Fallback { bytes }) => { + (bytes, false, false, None) + } + Err(e) => { + unlink_partial(&tmp_path, "download_failed").await; + return Err(classify(&e)); + } + } + } else { + let download_result = tokio::select! { + result = self.store.get_file(creds, target, &tmp_path) => result, + _ = cancellation.cancelled() => { + unlink_partial(&tmp_path, "cancelled").await; + return Err(JobError::Permanent("pull cancelled".to_string())); + } + }; + match download_result { + Ok(b) => (b, false, false, None), + Err(e) => { + unlink_partial(&tmp_path, "download_failed").await; + return Err(classify(&e)); + } + } + }; + let download_wall_ms = + u64::try_from(download_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + let download_phase = if streamed { + TransferPhase::DownloadDecode + } else { + TransferPhase::Download + }; + emit_transfer_phase_duration( + &job.kind, + artifact, + download_phase, + execution_materialization_mode, + download_started_at, + ); + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "downloaded", + bytes, + ); + emit_transfer_throughput_mib_s( + &job.kind, + artifact, + download_phase, + execution_materialization_mode, + bytes, + download_started_at, + ); + if streamed { + log_streamed_pull_pipeline_timing( + job, + p, + &stream_pull_timing, + streamed_decoded, + download_wall_ms, + ); + } + + // Decrypt + maybe-decompress in one pass when the request + // carries envelope-encryption. `.partial` holds the ciphertext + // we just downloaded; `stage_decrypted_temp` decrypts it + // streaming into a `MaybeZstdDecodingWriter` that peeks the + // first four plaintext bytes for the zstd magic and either + // pipes the rest through `zstd::stream::write::Decoder` or + // dumps it directly to `.plaintext.tmp`. No intermediate + // decrypted-but-still-compressed file on disk for the common + // BYOK + compressed case. + // + // After this block, `tmp_path` (the `.partial`) is gone and + // the bytes live at `plaintext_tmp` ready to be promoted to + // `local_path`. The variable `chained_decompressed` records + // whether the chain already ran the zstd decode so the + // separate `decompress_in_place` step further below can skip + // when it would otherwise be a no-op. + let mut chained_decompressed = false; + let mut decompression_report = streamed_decompression_report.unwrap_or_default(); + if let Some(enc) = &p.encryption { + let decompress_if_zstd = !p.raw; + // The chained decode runs inside the decrypt pass, so admit it + // to the codec limit up front. The frame size can't be read + // through ciphertext, so this admits conservatively on the + // compressed-input floor whenever a decode will run, unless the + // artifact class bypasses admission under an output cap. + let ciphertext_path: &Path = &tmp_path; + let options = DecryptedPullOptions { + decompress_if_zstd, + materialization_mode: execution_materialization_mode, + }; + let (decrypt_started_at, staged) = self + .run_admitted_decode( + job, + p, + execution_materialization_mode, + decompress_if_zstd.then_some(TransferPhase::DecryptWait), + ciphertext_path, + |gate| { + stage_decrypted_temp( + job, + p, + enc, + ciphertext_path, + options, + resolved_dek, + gate, + ) + }, + ) + .await?; + match staged { + Ok((plaintext_tmp, was_decompressed, report)) => { + chained_decompressed = was_decompressed; + decompression_report = report; + let phase = if was_decompressed { + TransferPhase::DecryptDecompress + } else { + TransferPhase::Decrypt + }; + emit_transfer_phase_duration( + &job.kind, + artifact, + phase, + execution_materialization_mode, + decrypt_started_at, + ); + let plaintext_size = tokio::fs::metadata(&plaintext_tmp) + .await + .map(|m| m.len()) + .unwrap_or(bytes); + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + if was_decompressed { + "decompressed" + } else { + "decrypted" + }, + plaintext_size, + ); + emit_transfer_throughput_mib_s( + &job.kind, + artifact, + phase, + execution_materialization_mode, + plaintext_size, + decrypt_started_at, + ); + if was_decompressed && bytes > 0 { + pplx_space_metrics::volume_manager::transfer_inflation_ratio( + &job.kind, + artifact, + plaintext_size as f64 / bytes as f64, + ); + } + // Replace `.partial` with the plaintext temp so the + // rename below targets a single canonical path + // regardless of whether encryption ran. + unlink_partial(&tmp_path, "post_decrypt").await; + if let Err(e) = tokio::fs::rename(&plaintext_tmp, &tmp_path).await { + unlink_crypto_tmp( + &plaintext_tmp, + &job.job_id, + &p.subject_id, + "rename_plaintext_failed", + ) + .await; + return Err(JobError::permanent_with_reason( + JobFailureReason::Internal, + format!("rename .plaintext.tmp -> .partial: {e}"), + )); + } + } + Err((partial_plaintext, err)) => { + unlink_partial(&tmp_path, "decrypt_failed").await; + unlink_crypto_tmp( + &partial_plaintext, + &job.job_id, + &p.subject_id, + "decrypt_failed", + ) + .await; + return Err(err); + } + } + } + + let promote_started_at = Instant::now(); + if let Err(err) = self + .invalidate_file_replica_provenance(local_path, "promote pull output") + .await + { + unlink_partial(&tmp_path, "provenance_invalidate_failed").await; + return Err(err); + } + + if let Err(e) = tokio::fs::rename(&tmp_path, local_path).await { + unlink_partial(&tmp_path, "rename_failed").await; + return Err(JobError::permanent_with_reason( + JobFailureReason::Internal, + format!("rename .partial -> final: {e}"), + )); + } + emit_transfer_phase_duration( + &job.kind, + artifact, + TransferPhase::Promote, + execution_materialization_mode, + promote_started_at, + ); + tracing::debug!( + job_id = %job.job_id, + subject_id = %p.subject_id, + from = %tmp_path.display(), + to = %p.local_path, + bytes, + chained_decompressed, + "renamed .partial to final path" + ); + + // The chained decrypt + maybe-decompress above already routes + // an encrypted-and-compressed object through the zstd decoder + // when applicable, so the standalone decompress step only + // matters for non-encrypted-but-compressed pulls. + // A streamed download already made the decode decision in-flight + // (decode or passthrough) — never re-sniff its output: decoded + // content that happens to start with the zstd magic must not be + // decoded a second time. + let file_is_zstd = !chained_decompressed && !streamed && is_zstd(local_path).await; + if p.should_transparent_decompress(file_is_zstd) { + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "compressed", + bytes, + ); + tracing::debug!( + job_id = %job.job_id, + subject_id = %p.subject_id, + local_path = %p.local_path, + compressed_bytes = bytes, + "decompressing downloaded artifact" + ); + let (decompress_started_at, report) = self + .run_admitted_decode( + job, + p, + execution_materialization_mode, + Some(TransferPhase::DecompressWait), + local_path, + |gate| { + decompress_in_place( + local_path, + snapshot_artifact, + execution_materialization_mode, + gate, + ) + }, + ) + .await?; + let report = report?; + apply_file_mode(local_path, p.file_mode, FileModeApplyPolicy::Required).await?; + let decompressed = tokio::fs::metadata(local_path) + .await + .map(|m| m.len()) + .unwrap_or(bytes); + emit_transfer_phase_duration( + &job.kind, + artifact, + TransferPhase::Decompress, + execution_materialization_mode, + decompress_started_at, + ); + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "decompressed", + decompressed, + ); + emit_transfer_throughput_mib_s( + &job.kind, + artifact, + TransferPhase::Decompress, + execution_materialization_mode, + decompressed, + decompress_started_at, + ); + if bytes > 0 { + pplx_space_metrics::volume_manager::transfer_inflation_ratio( + &job.kind, + artifact, + decompressed as f64 / bytes as f64, + ); + } + tracing::debug!( + job_id = %job.job_id, + subject_id = %p.subject_id, + local_path = %p.local_path, + decompressed_bytes = decompressed, + "artifact decompression complete" + ); + let materialization_mode = report.mode; + let materialization = { + #[cfg(unix)] + { + observe_decompression_allocation( + &job.kind, + artifact, + &job.job_id, + &p.subject_id, + local_path, + report, + ) + .await + } + #[cfg(not(unix))] + { + None + } + }; + if let Some(object_meta) = source_meta.as_ref() { + let object_meta = self + .confirm_source_meta_unchanged(p, creds, target, object_meta) + .await?; + self.write_file_replica_provenance( + job, + p, + local_path, + &object_meta, + decompressed, + materialization, + ) + .await?; + } + drop( + self.schedule_sparse_shadow_comparison( + &job.job_id, + &job.kind, + &p.subject_id, + local_path, + snapshot_artifact, + materialization_mode, + ) + .await, + ); + return bytes_len_to_job_i64(decompressed, "decompressed pull output size"); + } + + if streamed_decoded { + // The ranged download decoded inline; account the compressed and + // decoded sizes the way the legacy decompress pass would have. + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "compressed", + bytes, + ); + let decompressed = tokio::fs::metadata(local_path) + .await + .map(|m| m.len()) + .unwrap_or(bytes); + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "decompressed", + decompressed, + ); + if bytes > 0 { + pplx_space_metrics::volume_manager::transfer_inflation_ratio( + &job.kind, + artifact, + decompressed as f64 / bytes as f64, + ); + } + } + apply_file_mode(local_path, p.file_mode, FileModeApplyPolicy::Required).await?; + page_cache::advise_path_dontneed(local_path, "download volume output"); + // For any encrypted pull — chained decrypt+decompress *or* + // straight decrypt — the on-disk file is the decrypted plaintext + // payload, not the downloaded ciphertext: AEAD stripped the + // 17-byte header and a 16-byte tag per chunk during decrypt, and + // zstd-decompression (when chained) stripped that further. Report + // the on-disk size so `bytes_processed` and transfer metrics + // stay in parity across the BYOK toggle. Fall back to `bytes` + // (download/ciphertext size) only when the stat fails — a + // missing local path here is a downstream invariant we'd rather + // surface as metric noise than as a JobError, since the rename + // already succeeded. + let output_size = if chained_decompressed || streamed_decoded || p.encryption.is_some() { + tokio::fs::metadata(local_path) + .await + .map(|m| m.len()) + .unwrap_or(bytes) + } else { + bytes + }; + pplx_space_metrics::volume_manager::transfer_bytes( + &job.kind, + artifact, + "output", + output_size, + ); + let materialization_mode = decompression_report.mode; + let materialization = if chained_decompressed || streamed_decoded { + #[cfg(unix)] + { + observe_decompression_allocation( + &job.kind, + artifact, + &job.job_id, + &p.subject_id, + local_path, + decompression_report, + ) + .await + } + #[cfg(not(unix))] + { + None + } + } else { + None + }; + if let Some(object_meta) = source_meta.as_ref() { + let object_meta = self + .confirm_source_meta_unchanged(p, creds, target, object_meta) + .await?; + self.write_file_replica_provenance( + job, + p, + local_path, + &object_meta, + output_size, + materialization, + ) + .await?; + } + if chained_decompressed || streamed_decoded { + drop( + self.schedule_sparse_shadow_comparison( + &job.job_id, + &job.kind, + &p.subject_id, + local_path, + snapshot_artifact, + materialization_mode, + ) + .await, + ); + } + let size_context = if chained_decompressed || p.encryption.is_some() { + "decrypted pull output size" + } else if streamed_decoded { + "decompressed pull output size" + } else { + "pull download size" + }; + bytes_len_to_job_i64(output_size, size_context) + } +} diff --git a/crates/libs/vol/mgr/server/src/executor/replica_provenance.rs b/crates/libs/vol/mgr/server/src/executor/replica_provenance.rs new file mode 100644 index 000000000..9e1490520 --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/replica_provenance.rs @@ -0,0 +1,199 @@ +//! The provenance recorded next to a pulled file, and whether an existing replica can be reused. + +use super::*; + +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +pub struct MaterializationObservation { + pub mode: MaterializationMode, + pub layout: MaterializationLayout, + pub logical_bytes: u64, + pub allocated_bytes: u64, + pub candidate_hole_bytes: u64, +} + +#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +pub(super) struct FileReplicaProvenance { + #[serde(flatten)] + pub(super) replica: VerifiedReplicaProvenance, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) materialization: Option, +} + +pub(super) fn requested_file_mode_bits(p: &TransferPayload) -> Option { + p.file_mode.map(FileMode::bits) +} + +pub(super) fn file_replica_provenance( + p: &TransferPayload, + object: &ObjectMeta, + local_size: u64, +) -> VerifiedReplicaProvenance { + VerifiedReplicaProvenance { + version: VERIFIED_REPLICA_PROVENANCE_VERSION, + subject_id: p.subject_id.clone(), + s3_uri: p.s3_uri.clone(), + raw: p.raw, + file_mode: requested_file_mode_bits(p), + snapshot_artifact: p.snapshot_artifact, + object_size: object.size, + object_etag: object.etag.clone(), + local_size, + } +} + +impl FileReplicaProvenance { + pub(super) fn new( + p: &TransferPayload, + object: &ObjectMeta, + local_size: u64, + materialization: Option, + ) -> Self { + Self { + replica: file_replica_provenance(p, object, local_size), + materialization, + } + } +} + +#[cfg(test)] +pub(super) fn read_materialization_observation( + local_path: &Path, +) -> Option { + let bytes = std::fs::read(verified_replica_provenance_path(local_path)).ok()?; + serde_json::from_slice::(&bytes) + .ok()? + .materialization +} + +pub(super) fn materialization_is_compatible( + requested: MaterializationMode, + observation: Option<&MaterializationObservation>, +) -> bool { + match requested { + MaterializationMode::Disabled => true, + // The dense observation records the requested writer mode, not a + // successful background comparison. Shadow results are not persisted, + // so no completed replica can prove that postcondition yet. + MaterializationMode::Shadow => false, + MaterializationMode::Enabled => observation.is_some_and(|observation| { + observation.mode == MaterializationMode::Enabled + && observation.layout == MaterializationLayout::Sparse + }), + } +} + +impl TransferExecutor { + pub(super) async fn validate_existing_file_replica( + &self, + p: &TransferPayload, + creds: &pplx_space_object_store::ObjectCredentials, + target: ObjectTarget<'_>, + local_path: &Path, + requested_materialization: MaterializationMode, + ) -> Result, JobError> { + let local_size = match tokio::fs::metadata(local_path).await { + Ok(meta) if meta.is_file() => meta.len(), + Ok(_) => return Ok(None), + Err(_) => return Ok(None), + }; + + if !p.raw && is_zstd(local_path).await { + return Ok(None); + } + + let provenance_path = verified_replica_provenance_path(local_path); + let provenance_bytes = match tokio::fs::read(&provenance_path).await { + Ok(bytes) => bytes, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(_) => return Ok(None), + }; + let provenance = match serde_json::from_slice::(&provenance_bytes) { + Ok(provenance) => provenance, + Err(_) => return Ok(None), + }; + + let Some(object_meta) = self + .store + .head(creds, target) + .await + .map_err(|e| classify(&e))? + else { + return Ok(None); + }; + + // Sparse LUKS is a write optimization, not a new replica contract. + // Keep verified dense backings from older downloads: fetching them + // again only to punch holes would add disk and network work on retry. + let requested_materialization = if matches!( + p.effective_snapshot_artifact(), + Some(SnapshotArtifact::EncryptedSnapshotVolume | SnapshotArtifact::EncryptedRootfsCow) + ) { + MaterializationMode::Disabled + } else { + requested_materialization + }; + if provenance.replica == file_replica_provenance(p, &object_meta, local_size) + && materialization_is_compatible( + requested_materialization, + provenance.materialization.as_ref(), + ) + { + Ok(Some(local_size)) + } else { + Ok(None) + } + } + + pub(super) async fn write_file_replica_provenance( + &self, + job: &Job, + p: &TransferPayload, + local_path: &Path, + object_meta: &ObjectMeta, + local_size: u64, + materialization: Option, + ) -> Result<(), JobError> { + let provenance_path = verified_replica_provenance_path(local_path); + let mut tmp_path = provenance_path.clone(); + let mut tmp_name = provenance_path + .file_name() + .map(|name| name.to_os_string()) + .unwrap_or_else(|| OsString::from(".volmgr-replica.json")); + tmp_name.push(format!(".tmp-{}", job.job_id)); + tmp_path.set_file_name(tmp_name); + + let provenance = FileReplicaProvenance::new(p, object_meta, local_size, materialization); + let bytes = serde_json::to_vec(&provenance) + .map_err(|e| JobError::Permanent(format!("serialize file replica provenance: {e}")))?; + tokio::fs::write(&tmp_path, bytes) + .await + .map_err(|e| JobError::Permanent(format!("write file replica provenance: {e}")))?; + tokio::fs::rename(&tmp_path, &provenance_path) + .await + .map_err(|e| JobError::Permanent(format!("promote file replica provenance: {e}")))?; + Ok(()) + } + + pub(super) async fn invalidate_file_replica_provenance( + &self, + local_path: &Path, + reason: &'static str, + ) -> Result<(), JobError> { + let provenance_path = verified_replica_provenance_path(local_path); + match tokio::fs::remove_file(&provenance_path).await { + Ok(()) => { + tracing::debug!( + path = %provenance_path.display(), + reason, + "invalidated file replica provenance" + ); + Ok(()) + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err(JobError::Retriable(format!( + "remove file replica provenance {} before {reason}: {err}", + provenance_path.display() + ))), + } + } +} diff --git a/crates/libs/vol/mgr/server/src/executor/sparse_shadow.rs b/crates/libs/vol/mgr/server/src/executor/sparse_shadow.rs new file mode 100644 index 000000000..4af3ba60c --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/sparse_shadow.rs @@ -0,0 +1,640 @@ +//! The sparse-shadow rollout rung: compare a sparse materialization with a dense one, and schedule +//! those comparisons. + +use super::*; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum ShadowComparisonOutcome { + Match, + Mismatch, + Error, + SkippedBusy, + SkippedDisk, +} + +impl ShadowComparisonOutcome { + pub(super) const fn as_str(self) -> &'static str { + match self { + Self::Match => "match", + Self::Mismatch => "mismatch", + Self::Error => "error", + Self::SkippedBusy => "skipped_busy", + Self::SkippedDisk => "skipped_disk", + } + } +} + +impl std::fmt::Display for ShadowComparisonOutcome { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct ShadowComparisonObservation { + pub(super) logical_bytes: u64, + pub(super) allocated_bytes: u64, + pub(super) candidate_hole_bytes: u64, +} + +pub(super) fn sparse_shadow_path(path: &Path) -> PathBuf { + let mut shadow = path.to_path_buf(); + let mut name = shadow.file_name().unwrap_or_default().to_os_string(); + name.push(format!(".sparse-shadow-{}.tmp", uuid::Uuid::now_v7())); + shadow.set_file_name(name); + shadow +} + +pub(super) fn sparse_shadow_source_path(staging_dir: &Path) -> PathBuf { + staging_dir.join(format!("sparse-shadow-source-{}.tmp", uuid::Uuid::now_v7())) +} + +#[derive(Debug, thiserror::Error)] +pub(super) enum SparseShadowCmpError { + #[error("sparse shadow comparison cancelled")] + Cancelled, + #[error("sparse shadow bytes differ from dense output")] + Mismatch, + #[error("{operation}: {source}")] + Io { + operation: &'static str, + #[source] + source: std::io::Error, + }, +} + +impl SparseShadowCmpError { + pub(super) fn io(operation: &'static str, source: std::io::Error) -> Self { + Self::Io { operation, source } + } +} + +pub(super) fn compare_completed_files( + dense: &Path, + sparse: &Path, + shutdown_token: &CancellationToken, +) -> Result<(), SparseShadowCmpError> { + let mut dense_file = std::fs::File::open(dense) + .map_err(|error| SparseShadowCmpError::io("open dense comparison input", error))?; + let mut sparse_file = std::fs::File::open(sparse) + .map_err(|error| SparseShadowCmpError::io("open sparse comparison input", error))?; + let dense_len = dense_file + .metadata() + .map_err(|error| SparseShadowCmpError::io("stat dense comparison input", error))? + .len(); + let sparse_len = sparse_file + .metadata() + .map_err(|error| SparseShadowCmpError::io("stat sparse comparison input", error))? + .len(); + if dense_len != sparse_len { + return Err(SparseShadowCmpError::Mismatch); + } + + let mut dense_buffer = vec![0; 1024 * 1024]; + let mut sparse_buffer = vec![0; 1024 * 1024]; + let mut compared_through = 0_u64; + let mut advised_through = 0_u64; + while compared_through < dense_len { + if shutdown_token.is_cancelled() { + return Err(SparseShadowCmpError::Cancelled); + } + let remaining = dense_len.saturating_sub(compared_through); + let read_len = dense_buffer + .len() + .min(usize::try_from(remaining).unwrap_or(usize::MAX)); + let dense_chunk = dense_buffer.get_mut(..read_len).ok_or_else(|| { + SparseShadowCmpError::io( + "prepare dense comparison buffer", + std::io::Error::other("comparison read exceeds dense buffer"), + ) + })?; + let sparse_chunk = sparse_buffer.get_mut(..read_len).ok_or_else(|| { + SparseShadowCmpError::io( + "prepare sparse comparison buffer", + std::io::Error::other("comparison read exceeds sparse buffer"), + ) + })?; + dense_file + .read_exact(dense_chunk) + .map_err(|error| SparseShadowCmpError::io("read dense comparison input", error))?; + sparse_file + .read_exact(sparse_chunk) + .map_err(|error| SparseShadowCmpError::io("read sparse comparison input", error))?; + if dense_chunk != sparse_chunk { + return Err(SparseShadowCmpError::Mismatch); + } + compared_through = compared_through.saturating_add(read_len as u64); + if compared_through.saturating_sub(advised_through) >= STREAM_WRITEBACK_WINDOW_BYTES { + advise_shadow_read_range( + &dense_file, + dense, + advised_through, + compared_through.saturating_sub(advised_through), + "dense sparse-shadow comparison input", + ); + advise_shadow_read_range( + &sparse_file, + sparse, + advised_through, + compared_through.saturating_sub(advised_through), + "sparse shadow comparison input", + ); + advised_through = compared_through; + } + } + if advised_through < compared_through { + let length = compared_through.saturating_sub(advised_through); + advise_shadow_read_range( + &dense_file, + dense, + advised_through, + length, + "dense sparse-shadow comparison input", + ); + advise_shadow_read_range( + &sparse_file, + sparse, + advised_through, + length, + "sparse shadow comparison input", + ); + } + Ok(()) +} + +pub(super) fn advise_shadow_read_range( + file: &std::fs::File, + path: &Path, + offset: u64, + length: u64, + operation: &'static str, +) { + if let Err(error) = page_cache::advise_range_dontneed(file, offset, length) { + tracing::warn!( + %error, + path_len = path.as_os_str().len(), + offset, + length, + operation, + "failed to evict consumed sparse-shadow read range" + ); + } +} + +pub(super) fn compare_sparse_shadow( + path: &Path, + block_bytes: usize, + shutdown_token: &CancellationToken, + writeback_coordinator: Arc, +) -> Result { + let shadow = sparse_shadow_path(path); + let result = (|| { + if shutdown_token.is_cancelled() { + return Err(SparseShadowCmpError::Cancelled); + } + let mut reader = std::fs::File::open(path) + .map_err(|error| SparseShadowCmpError::io("open dense shadow source", error))?; + let output = std::fs::File::create(&shadow) + .map_err(|error| SparseShadowCmpError::io("create sparse shadow output", error))?; + let candidate_observer = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let writeback_timing = Arc::new(StreamPullTiming::default()); + let mut writer = DecryptedFileSink::new( + output, + MaterializationMode::Enabled, + block_bytes, + Arc::clone(&candidate_observer), + Some(writeback_coordinator), + Some(writeback_timing), + JobPriority::default(), + ) + .map_err(|error| SparseShadowCmpError::io("create sparse shadow writer", error))?; + let mut buffer = vec![0; 1024 * 1024]; + let mut read_through = 0_u64; + let mut advised_through = 0_u64; + loop { + if shutdown_token.is_cancelled() { + return Err(SparseShadowCmpError::Cancelled); + } + let read = reader + .read(&mut buffer) + .map_err(|error| SparseShadowCmpError::io("read dense shadow source", error))?; + if read == 0 { + break; + } + let chunk = buffer.get(..read).ok_or_else(|| { + SparseShadowCmpError::io( + "read dense shadow source", + std::io::Error::other("read exceeded shadow buffer"), + ) + })?; + writer + .write_all(chunk) + .map_err(|error| SparseShadowCmpError::io("write sparse shadow output", error))?; + read_through = read_through.saturating_add(read as u64); + if read_through.saturating_sub(advised_through) >= STREAM_WRITEBACK_WINDOW_BYTES { + advise_shadow_read_range( + &reader, + path, + advised_through, + read_through.saturating_sub(advised_through), + "read dense sparse-shadow source", + ); + advised_through = read_through; + } + } + writer + .complete() + .map_err(|error| SparseShadowCmpError::io("finalize sparse shadow output", error))?; + if advised_through < read_through { + advise_shadow_read_range( + &reader, + path, + advised_through, + read_through.saturating_sub(advised_through), + "read dense sparse-shadow source", + ); + } + let candidate_hole_bytes = candidate_observer.load(std::sync::atomic::Ordering::Relaxed); + let output = std::fs::File::open(&shadow) + .map_err(|error| SparseShadowCmpError::io("open sparse shadow output", error))?; + page_cache::advise_dontneed(&reader, path, "read dense sparse-shadow source"); + page_cache::advise_dontneed(&output, &shadow, "write sparse shadow comparison file"); + drop(reader); + drop(output); + + compare_completed_files(path, &shadow, shutdown_token)?; + + let shadow_metadata = std::fs::metadata(&shadow) + .map_err(|error| SparseShadowCmpError::io("stat sparse shadow output", error))?; + let shadow_logical_bytes = shadow_metadata.len(); + #[cfg(unix)] + let shadow_allocated_bytes = { + use std::os::unix::fs::MetadataExt; + shadow_metadata.blocks().saturating_mul(512) + }; + #[cfg(not(unix))] + let shadow_allocated_bytes = shadow_logical_bytes; + Ok(ShadowComparisonObservation { + logical_bytes: shadow_logical_bytes, + allocated_bytes: shadow_allocated_bytes, + candidate_hole_bytes, + }) + })(); + + match std::fs::remove_file(&shadow) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => tracing::warn!( + path = %shadow.display(), + %error, + "failed to remove sparse shadow comparison file" + ), + } + result +} + +pub(super) struct ShadowComparisonLifecycle { + pub(super) admission: std::sync::Mutex<()>, + pub(super) shutdown_token: CancellationToken, + pub(super) tasks: TaskTracker, +} + +#[derive(Clone)] +pub struct ShadowComparisonController { + pub(super) inner: Arc, +} + +impl ShadowComparisonController { + pub(super) fn new() -> Self { + Self { + inner: Arc::new(ShadowComparisonLifecycle { + admission: std::sync::Mutex::new(()), + shutdown_token: CancellationToken::new(), + tasks: TaskTracker::new(), + }), + } + } + + pub(super) fn spawn(&self, future: F) -> Option> + where + F: std::future::Future + Send + 'static, + F::Output: Send + 'static, + { + let _admission = self + .inner + .admission + .lock() + .expect("shadow comparison admission lock poisoned"); + (!self.inner.shutdown_token.is_cancelled()).then(|| self.inner.tasks.spawn(future)) + } + + pub(super) fn is_accepting(&self) -> bool { + !self.inner.shutdown_token.is_cancelled() + } + + pub(super) fn shutdown_token(&self) -> CancellationToken { + self.inner.shutdown_token.clone() + } + + pub async fn shutdown(&self, timeout: std::time::Duration) -> bool { + { + let _admission = self + .inner + .admission + .lock() + .expect("shadow comparison admission lock poisoned"); + self.inner.shutdown_token.cancel(); + self.inner.tasks.close(); + } + tokio::time::timeout(timeout, self.inner.tasks.wait()) + .await + .is_ok() + } +} + +pub(super) struct SparseShadowTask { + pub(super) job_id: JobId, + pub(super) subject_id: String, + pub(super) local_path: PathBuf, + pub(super) snapshot_path: PathBuf, + pub(super) block_bytes: usize, + pub(super) artifact: &'static str, +} + +pub(super) struct SparseShadowAdmission { + pub(super) _shadow_permit: tokio::sync::OwnedSemaphorePermit, + pub(super) _codec_permit: CodecPermit, +} + +pub(super) fn sparse_shadow_disk_floor_bytes( + status: pplx_space_common_core::host::disk::DiskStatus, +) -> u64 { + pplx_space_common_core::host::disk_admission::CacheDiskAdmissionConfig::default() + .hard_floor_bytes(status.total_bytes) +} + +pub(super) fn sparse_shadow_has_disk_headroom( + status: pplx_space_common_core::host::disk::DiskStatus, + logical_bytes: u64, +) -> bool { + status.available_bytes >= sparse_shadow_disk_floor_bytes(status).saturating_add(logical_bytes) +} + +pub(super) async fn remove_sparse_shadow_snapshot(path: &Path, reason: &'static str) { + if let Err(error) = tokio::fs::remove_file(path).await + && error.kind() != std::io::ErrorKind::NotFound + { + tracing::warn!( + shadow_source_path = %path.display(), + %error, + reason, + "failed to remove sparse shadow snapshot" + ); + } +} + +impl TransferExecutor { + pub(super) fn try_acquire_sparse_shadow_admission(&self) -> Option { + let shadow_permit = Arc::clone(&self.shadow_comparison_permits) + .try_acquire_owned() + .ok()?; + let codec_permit = self.codec.try_acquire_shared()?; + Some(SparseShadowAdmission { + _shadow_permit: shadow_permit, + _codec_permit: codec_permit, + }) + } + + pub(super) async fn schedule_sparse_shadow_comparison( + &self, + job_id: &JobId, + kind: &str, + subject_id: &str, + local_path: &Path, + snapshot_artifact: Option, + mode: MaterializationMode, + ) -> Option>> { + if mode != MaterializationMode::Shadow { + return None; + } + if !self.shadow_comparisons.is_accepting() { + return None; + } + + let snapshot_artifact = + snapshot_artifact.or_else(|| snapshot_artifact_from_local_path(local_path)); + let artifact = transfer_artifact_tag(snapshot_artifact); + let block_bytes = materialization_block_bytes(snapshot_artifact); + let Some(admission) = self.try_acquire_sparse_shadow_admission() else { + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::SkippedBusy.as_str(), + ); + tracing::info!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + "background sparse shadow comparison skipped because the node is busy" + ); + return None; + }; + let Ok(metadata) = tokio::fs::metadata(local_path).await else { + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::Error.as_str(), + ); + tracing::error!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + "could not inspect completed artifact before sparse shadow comparison" + ); + return None; + }; + let logical_bytes = metadata.len(); + #[cfg(unix)] + let disk_status = self + .shadow_staging_dir + .to_str() + .and_then(pplx_space_common_core::host::disk::statvfs_status); + #[cfg(not(unix))] + let disk_status: Option = None; + let disk_floor_bytes = disk_status.map_or(0, sparse_shadow_disk_floor_bytes); + let has_disk_headroom = disk_status + .is_some_and(|status| sparse_shadow_has_disk_headroom(status, logical_bytes)); + if !has_disk_headroom { + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::SkippedDisk.as_str(), + ); + tracing::warn!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + logical_bytes, + available_bytes = disk_status.map(|status| status.available_bytes), + disk_floor_bytes, + "background sparse shadow comparison skipped for disk headroom" + ); + return None; + } + let snapshot_path = sparse_shadow_source_path(&self.shadow_staging_dir); + let snapshot_started_at = Instant::now(); + if let Err(error) = crate::staging::reflink_only(local_path, &snapshot_path).await { + remove_sparse_shadow_snapshot(&snapshot_path, "reflink failed").await; + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::Error.as_str(), + ); + tracing::error!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + shadow_source_path = %snapshot_path.display(), + %error, + "could not snapshot completed artifact for background sparse shadow comparison" + ); + return None; + } + emit_transfer_phase_duration( + kind, + artifact, + TransferPhase::ShadowReflink, + MaterializationMode::Shadow, + snapshot_started_at, + ); + + tracing::info!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + shadow_source_path = %snapshot_path.display(), + snapshot_elapsed_ms = snapshot_started_at.elapsed().as_millis(), + "background sparse shadow comparison scheduled" + ); + + let cleanup_path = snapshot_path.clone(); + let handle = self.spawn_sparse_shadow_comparison( + SparseShadowTask { + job_id: *job_id, + subject_id: subject_id.to_string(), + local_path: local_path.to_path_buf(), + snapshot_path, + block_bytes, + artifact, + }, + admission, + ); + if handle.is_none() { + remove_sparse_shadow_snapshot(&cleanup_path, "task admission closed").await; + } + handle + } + + pub(super) fn spawn_sparse_shadow_comparison( + &self, + task: SparseShadowTask, + admission: SparseShadowAdmission, + ) -> Option>> { + let shutdown_token = self.shadow_comparisons.shutdown_token(); + let writeback_coordinator = Arc::clone(&self.stream_writeback); + self.shadow_comparisons.spawn(async move { + let SparseShadowTask { + job_id, + subject_id, + local_path, + snapshot_path, + block_bytes, + artifact, + } = task; + let _admission = admission; + if shutdown_token.is_cancelled() { + remove_sparse_shadow_snapshot(&snapshot_path, "shutdown before comparison").await; + return None; + } + let comparison_path = snapshot_path.clone(); + let blocking_shutdown_token = shutdown_token.clone(); + let result = tokio::task::spawn_blocking(move || { + compare_sparse_shadow( + &comparison_path, + block_bytes, + &blocking_shutdown_token, + writeback_coordinator, + ) + }) + .await; + remove_sparse_shadow_snapshot(&snapshot_path, "comparison finished").await; + + match result { + Ok(Ok(observation)) => { + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::Match.as_str(), + ); + pplx_space_metrics::volume_manager::transfer_materialization_allocation( + "pull", + artifact, + MaterializationMode::Shadow, + MaterializationLayout::Sparse, + observation.logical_bytes, + observation.allocated_bytes, + observation.candidate_hole_bytes, + ); + tracing::info!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + shadow_logical_bytes = observation.logical_bytes, + shadow_allocated_bytes = observation.allocated_bytes, + candidate_hole_bytes = observation.candidate_hole_bytes, + "background sparse shadow comparison completed" + ); + Some(observation) + } + Ok(Err(SparseShadowCmpError::Cancelled)) => { + tracing::info!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + "background sparse shadow comparison cancelled during shutdown" + ); + None + } + Ok(Err(error)) => { + let outcome = if matches!(error, SparseShadowCmpError::Mismatch) { + ShadowComparisonOutcome::Mismatch + } else { + ShadowComparisonOutcome::Error + }; + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + outcome.as_str(), + ); + tracing::error!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + ?error, + comparison = outcome.as_str(), + "background sparse shadow comparison failed" + ); + None + } + Err(error) => { + pplx_space_metrics::volume_manager::transfer_materialization_shadow_comparison( + artifact, + ShadowComparisonOutcome::Error.as_str(), + ); + tracing::error!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + %error, + "background sparse shadow comparison task panicked" + ); + None + } + } + }) + } +} diff --git a/crates/libs/vol/mgr/server/src/executor/stream_pull_timing.rs b/crates/libs/vol/mgr/server/src/executor/stream_pull_timing.rs new file mode 100644 index 000000000..d8e320070 --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/stream_pull_timing.rs @@ -0,0 +1,121 @@ +//! Per-stage timing of a streamed pull and the summary it logs. + +use super::*; + +#[derive(Debug, Default)] +pub(super) struct StreamPullTiming { + pub(super) compressed_input_bytes: std::sync::atomic::AtomicU64, + pub(super) decoded_output_bytes: std::sync::atomic::AtomicU64, + pub(super) decode_writer_us: std::sync::atomic::AtomicU64, + pub(super) decode_writer_cpu_us: std::sync::atomic::AtomicU64, + pub(super) output_writer_us: std::sync::atomic::AtomicU64, + pub(super) output_writer_cpu_us: std::sync::atomic::AtomicU64, + pub(super) writeback_reserve_us: std::sync::atomic::AtomicU64, + pub(super) writeback_reserve_cpu_us: std::sync::atomic::AtomicU64, + pub(super) writeback_start_us: std::sync::atomic::AtomicU64, + pub(super) writeback_start_cpu_us: std::sync::atomic::AtomicU64, + pub(super) writeback_wait_us: std::sync::atomic::AtomicU64, + pub(super) writeback_start_calls: std::sync::atomic::AtomicU64, + pub(super) writeback_submitted_logical_bytes: std::sync::atomic::AtomicU64, + pub(super) writeback_submitted_dirty_bytes: std::sync::atomic::AtomicU64, + pub(super) writeback_hole_bytes_released: std::sync::atomic::AtomicU64, + pub(super) writeback_low_watermark_activations: std::sync::atomic::AtomicU64, +} + +pub(super) fn current_thread_cpu_us() -> u64 { + let Ok(value) = nix::time::clock_gettime(nix::time::ClockId::CLOCK_THREAD_CPUTIME_ID) else { + return 0; + }; + let seconds = u64::try_from(value.tv_sec()).unwrap_or(0); + let nanoseconds = u64::try_from(value.tv_nsec()).unwrap_or(0); + seconds + .saturating_mul(1_000_000) + .saturating_add(nanoseconds / 1_000) +} + +impl StreamPullTiming { + pub(super) fn add(counter: &std::sync::atomic::AtomicU64, value: u64) { + match counter.fetch_update( + std::sync::atomic::Ordering::Relaxed, + std::sync::atomic::Ordering::Relaxed, + |current| Some(current.saturating_add(value)), + ) { + Ok(_) | Err(_) => {} + } + } + + pub(super) fn add_elapsed(counter: &std::sync::atomic::AtomicU64, started_at: Instant) { + Self::add( + counter, + u64::try_from(started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + ); + } + + pub(super) fn add_wall_and_thread_cpu( + wall_counter: &std::sync::atomic::AtomicU64, + cpu_counter: &std::sync::atomic::AtomicU64, + wall_started_at: Instant, + cpu_started_at_us: u64, + ) { + let cpu_elapsed_us = current_thread_cpu_us().saturating_sub(cpu_started_at_us); + let wall_elapsed_us = + u64::try_from(wall_started_at.elapsed().as_micros()).unwrap_or(u64::MAX); + Self::add(wall_counter, wall_elapsed_us); + Self::add(cpu_counter, cpu_elapsed_us); + } + + pub(super) fn load(counter: &std::sync::atomic::AtomicU64) -> u64 { + counter.load(std::sync::atomic::Ordering::Relaxed) + } +} + +pub(super) fn log_streamed_pull_pipeline_timing( + job: &Job, + payload: &TransferPayload, + timing: &StreamPullTiming, + streamed_decoded: bool, + download_wall_ms: u64, +) { + let decode_writer_us = StreamPullTiming::load(&timing.decode_writer_us); + let decode_writer_cpu_us = StreamPullTiming::load(&timing.decode_writer_cpu_us); + let output_writer_us = StreamPullTiming::load(&timing.output_writer_us); + let output_writer_cpu_us = StreamPullTiming::load(&timing.output_writer_cpu_us); + let writeback_reserve_us = StreamPullTiming::load(&timing.writeback_reserve_us); + let writeback_reserve_cpu_us = StreamPullTiming::load(&timing.writeback_reserve_cpu_us); + tracing::info!( + job_id = %job.job_id, + subject_id = %payload.subject_id, + s3_uri_len = payload.s3_uri.len(), + local_path_len = payload.local_path.len(), + streamed_decoded, + download_decode_wall_ms = download_wall_ms, + compressed_input_bytes = StreamPullTiming::load(&timing.compressed_input_bytes), + decoded_output_bytes = StreamPullTiming::load(&timing.decoded_output_bytes), + decode_writer_ms = decode_writer_us / 1_000, + decode_stream_processing_ms = decode_writer_us + .saturating_sub(output_writer_us) + .saturating_sub(writeback_reserve_us) + / 1_000, + decode_stream_processing_cpu_ms = decode_writer_cpu_us + .saturating_sub(output_writer_cpu_us) + .saturating_sub(writeback_reserve_cpu_us) + / 1_000, + output_writer_ms = output_writer_us / 1_000, + output_writer_cpu_ms = output_writer_cpu_us / 1_000, + writeback_reserve_ms = writeback_reserve_us / 1_000, + writeback_reserve_cpu_ms = writeback_reserve_cpu_us / 1_000, + writeback_start_ms = StreamPullTiming::load(&timing.writeback_start_us) / 1_000, + writeback_start_cpu_ms = StreamPullTiming::load(&timing.writeback_start_cpu_us) / 1_000, + writeback_wait_observed_ms = StreamPullTiming::load(&timing.writeback_wait_us) / 1_000, + writeback_start_calls = StreamPullTiming::load(&timing.writeback_start_calls), + writeback_submitted_logical_bytes = + StreamPullTiming::load(&timing.writeback_submitted_logical_bytes), + writeback_submitted_dirty_bytes = + StreamPullTiming::load(&timing.writeback_submitted_dirty_bytes), + writeback_hole_bytes_released = + StreamPullTiming::load(&timing.writeback_hole_bytes_released), + writeback_low_watermark_activations = + StreamPullTiming::load(&timing.writeback_low_watermark_activations), + "streamed pull pipeline timing" + ); +} diff --git a/crates/libs/vol/mgr/server/src/executor/stream_writeback.rs b/crates/libs/vol/mgr/server/src/executor/stream_writeback.rs new file mode 100644 index 000000000..8224e21d7 --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/stream_writeback.rs @@ -0,0 +1,312 @@ +//! One stream's writeback state: dirty-range accumulation, reservations and flushes. + +use super::*; + +#[derive(Debug)] +pub(super) struct StreamWritebackAccumulator { + pub(super) offset: u64, + pub(super) length: u64, + pub(super) debt_bytes: u64, +} + +#[derive(Debug)] +pub(super) struct StreamWriteback { + pub(super) coordinator: Option>, + pub(super) timing: Option>, + pub(super) priority: JobPriority, + pub(super) writeback_file: Option>, + pub(super) unclassified_bytes: u64, + pub(super) reserved_bytes: u64, + pub(super) deferred_hole_bytes: u64, + pub(super) pending: Option, +} + +impl StreamWriteback { + pub(super) fn new( + coordinator: Option>, + timing: Option>, + priority: JobPriority, + ) -> Self { + Self { + coordinator, + timing, + priority, + writeback_file: None, + unclassified_bytes: 0, + reserved_bytes: 0, + deferred_hole_bytes: 0, + pending: None, + } + } + + pub(super) fn pending_debt_bytes(&self) -> u64 { + self.pending + .as_ref() + .map_or(0, |pending| pending.debt_bytes) + } + + pub(super) fn additional_reservation_bytes(&self, bytes: u64) -> u64 { + let Some(coordinator) = &self.coordinator else { + return 0; + }; + let owned_bytes = self + .unclassified_bytes + .saturating_add(self.pending_debt_bytes()) + .saturating_add(self.deferred_hole_bytes); + let required_bytes = owned_bytes.saturating_add(bytes); + if required_bytes <= self.reserved_bytes { + return 0; + } + let reservation_window = STREAM_WRITEBACK_WINDOW_BYTES + .min(coordinator.admission_limit(self.priority) / 2) + .max(bytes); + owned_bytes + .saturating_add(reservation_window) + .saturating_sub(self.reserved_bytes) + } + + pub(super) fn reserve(&mut self, bytes: usize) -> std::io::Result<()> { + let bytes = bytes as u64; + if self.coordinator.is_some() { + let mut additional = self.additional_reservation_bytes(bytes); + if additional != 0 { + self.release_deferred_holes(); + additional = self.additional_reservation_bytes(bytes); + let coordinator = self + .coordinator + .as_ref() + .expect("coordinator checked above"); + let reserve_started_at = Instant::now(); + let reserve_cpu_started_at = current_thread_cpu_us(); + coordinator.reserve(additional, self.priority)?; + if let Some(timing) = &self.timing { + StreamPullTiming::add_wall_and_thread_cpu( + &timing.writeback_reserve_us, + &timing.writeback_reserve_cpu_us, + reserve_started_at, + reserve_cpu_started_at, + ); + } + self.reserved_bytes = self.reserved_bytes.saturating_add(additional); + } + } + Ok(()) + } + + pub(super) fn try_reserve(&mut self, bytes: usize) -> std::io::Result { + let bytes = bytes as u64; + if self.coordinator.is_none() { + return Ok(true); + } + let mut additional = self.additional_reservation_bytes(bytes); + if additional == 0 { + return Ok(true); + } + self.release_deferred_holes(); + additional = self.additional_reservation_bytes(bytes); + let coordinator = self + .coordinator + .as_ref() + .expect("coordinator checked above"); + let reserve_started_at = Instant::now(); + let reserve_cpu_started_at = current_thread_cpu_us(); + if !coordinator.try_reserve(additional, self.priority)? { + return Ok(false); + } + if let Some(timing) = &self.timing { + StreamPullTiming::add_wall_and_thread_cpu( + &timing.writeback_reserve_us, + &timing.writeback_reserve_cpu_us, + reserve_started_at, + reserve_cpu_started_at, + ); + } + self.reserved_bytes = self.reserved_bytes.saturating_add(additional); + Ok(true) + } + + pub(super) fn record_written(&mut self, bytes: usize) { + if self.coordinator.is_some() { + self.unclassified_bytes = self.unclassified_bytes.saturating_add(bytes as u64); + } + } + + pub(super) fn release_unwritten(&mut self, bytes: usize) { + let bytes = bytes as u64; + if bytes == 0 { + return; + } + if let Some(coordinator) = &self.coordinator { + self.reserved_bytes = self.reserved_bytes.saturating_sub(bytes); + coordinator.release(bytes); + } + } + + pub(super) fn writeback_file( + &mut self, + file: &std::fs::File, + ) -> std::io::Result> { + if let Some(file) = &self.writeback_file { + return Ok(Arc::clone(file)); + } + let file = Arc::new(file.try_clone()?); + self.writeback_file = Some(Arc::clone(&file)); + Ok(file) + } + + pub(super) fn release_deferred_holes(&mut self) { + let hole_bytes = std::mem::take(&mut self.deferred_hole_bytes); + if hole_bytes == 0 { + return; + } + self.reserved_bytes = self + .reserved_bytes + .checked_sub(hole_bytes) + .expect("deferred holes exceed producer reservation"); + let coordinator = self + .coordinator + .as_ref() + .expect("deferred holes require a coordinator"); + coordinator.release(hole_bytes); + if let Some(timing) = &self.timing { + StreamPullTiming::add(&timing.writeback_hole_bytes_released, hole_bytes); + } + } + + pub(super) fn flush_pending(&mut self, file: &std::fs::File) -> std::io::Result<()> { + self.release_deferred_holes(); + let Some(pending) = self.pending.take() else { + return Ok(()); + }; + let writeback_file = self.writeback_file(file)?; + let coordinator = self + .coordinator + .as_ref() + .expect("pending writeback requires a coordinator"); + let timing = self + .timing + .clone() + .unwrap_or_else(|| Arc::new(StreamPullTiming::default())); + self.reserved_bytes = self + .reserved_bytes + .checked_sub(pending.debt_bytes) + .ok_or_else(|| { + std::io::Error::other("pending writeback exceeds producer reservation") + })?; + coordinator.submit( + writeback_file, + pending.offset, + pending.length, + pending.debt_bytes, + timing, + ); + Ok(()) + } + + pub(super) fn accumulate_dirty_range( + &mut self, + file: &std::fs::File, + mut offset: u64, + mut length: u64, + ) -> std::io::Result<()> { + while length != 0 { + if self.pending_debt_bytes() == STREAM_WRITEBACK_WINDOW_BYTES { + self.flush_pending(file)?; + } + let available = STREAM_WRITEBACK_WINDOW_BYTES - self.pending_debt_bytes(); + let debt_bytes = length.min(available); + let end = offset + .checked_add(debt_bytes) + .ok_or_else(|| std::io::Error::other("streamed output range exceeds u64::MAX"))?; + if let Some(pending) = &mut self.pending { + let pending_end = pending.offset.checked_add(pending.length).ok_or_else(|| { + std::io::Error::other("pending streamed output range exceeds u64::MAX") + })?; + if offset < pending_end { + return Err(std::io::Error::other( + "streamed dirty output ranges are not ordered", + )); + } + pending.length = end.checked_sub(pending.offset).ok_or_else(|| { + std::io::Error::other("streamed output range precedes pending range") + })?; + pending.debt_bytes = + pending.debt_bytes.checked_add(debt_bytes).ok_or_else(|| { + std::io::Error::other("pending streamed dirty debt exceeds u64::MAX") + })?; + } else { + self.pending = Some(StreamWritebackAccumulator { + offset, + length: debt_bytes, + debt_bytes, + }); + } + offset = end; + length -= debt_bytes; + } + if self.pending_debt_bytes() == STREAM_WRITEBACK_WINDOW_BYTES { + self.flush_pending(file)?; + } + Ok(()) + } + + pub(super) fn classify( + &mut self, + file: &std::fs::File, + materialization: SparseMaterialization, + ) -> std::io::Result<()> { + let (data_ranges, hole_bytes) = materialization.into_parts(); + let dirty_bytes = data_ranges + .iter() + .try_fold(0_u64, |total, range| total.checked_add(range.length())) + .ok_or_else(|| std::io::Error::other("streamed dirty output exceeds u64::MAX"))?; + let classified_bytes = dirty_bytes + .checked_add(hole_bytes) + .ok_or_else(|| std::io::Error::other("streamed classified output exceeds u64::MAX"))?; + if classified_bytes == 0 || self.coordinator.is_none() { + return Ok(()); + } + if classified_bytes > self.unclassified_bytes || classified_bytes > self.reserved_bytes { + return Err(std::io::Error::other( + "sparse writer classified more bytes than writeback admission reserved", + )); + } + self.unclassified_bytes -= classified_bytes; + if hole_bytes != 0 { + self.deferred_hole_bytes = self + .deferred_hole_bytes + .checked_add(hole_bytes) + .ok_or_else(|| std::io::Error::other("deferred sparse holes exceed u64::MAX"))?; + if self.deferred_hole_bytes >= STREAM_WRITEBACK_WINDOW_BYTES { + self.release_deferred_holes(); + } + } + for range in data_ranges { + self.accumulate_dirty_range(file, range.offset(), range.length())?; + } + Ok(()) + } + + pub(super) fn classify_dense( + &mut self, + file: &std::fs::File, + offset: u64, + length: u64, + ) -> std::io::Result<()> { + self.classify( + file, + SparseMaterialization::from(SparseDataRange::new(offset, length)), + ) + } +} + +impl Drop for StreamWriteback { + fn drop(&mut self) { + self.release_deferred_holes(); + if self.reserved_bytes != 0 + && let Some(coordinator) = &self.coordinator + { + coordinator.release(self.reserved_bytes); + } + } +} diff --git a/crates/libs/vol/mgr/server/src/executor/transfer_metrics.rs b/crates/libs/vol/mgr/server/src/executor/transfer_metrics.rs new file mode 100644 index 000000000..6887100fe --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/transfer_metrics.rs @@ -0,0 +1,151 @@ +//! Transfer phases and the duration, throughput and decompression-allocation metrics they emit. + +use super::*; + +/// Closed vocabulary for the `phase` tag on transfer metrics. The string +/// form exists only at the emission boundary; everything in between +/// passes the enum so additions stay exhaustive and typos can't compile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum TransferPhase { + Download, + DownloadDecode, + Decrypt, + DecryptDecompress, + DecryptWait, + Decompress, + DecompressWait, + CompressWait, + Promote, + ShadowReflink, +} + +impl TransferPhase { + pub(super) fn as_str(self) -> &'static str { + match self { + Self::Download => "download", + Self::DownloadDecode => "download_decode", + Self::Decrypt => "decrypt", + Self::DecryptDecompress => "decrypt_decompress", + Self::DecryptWait => "decrypt_wait", + Self::Decompress => "decompress", + Self::DecompressWait => "decompress_wait", + Self::CompressWait => "compress_wait", + Self::Promote => "promote", + Self::ShadowReflink => "shadow_reflink", + } + } +} + +impl std::fmt::Display for TransferPhase { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +pub(super) fn emit_transfer_phase_duration( + kind: &str, + artifact: &str, + phase: TransferPhase, + mode: MaterializationMode, + started_at: Instant, +) { + pplx_space_metrics::volume_manager::transfer_phase_duration( + kind, + artifact, + phase.as_str(), + mode, + started_at, + ); +} + +pub(super) fn emit_transfer_throughput_mib_s( + kind: &str, + artifact: &str, + phase: TransferPhase, + mode: MaterializationMode, + bytes: u64, + started_at: Instant, +) { + let elapsed = started_at.elapsed().as_secs_f64(); + if elapsed > 0.0 { + let mib_s = (bytes as f64 / 1024.0 / 1024.0) / elapsed; + pplx_space_metrics::volume_manager::transfer_throughput_mib_s( + kind, + artifact, + phase.as_str(), + mode, + mib_s, + ); + } +} + +#[cfg(unix)] +pub(super) async fn observe_decompression_allocation( + kind: &str, + artifact: &str, + job_id: &JobId, + subject_id: &str, + local_path: &Path, + report: DecompressionReport, +) -> Option { + use std::os::unix::fs::MetadataExt; + + let metadata = match tokio::fs::metadata(local_path).await { + Ok(metadata) => metadata, + Err(error) => { + tracing::debug!( + %error, + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + "could not observe decompressed artifact allocation" + ); + return None; + } + }; + let logical_bytes = metadata.len(); + let allocated_bytes = metadata.blocks().saturating_mul(512); + let saved_bytes = logical_bytes.saturating_sub(allocated_bytes); + let allocation_ratio = if logical_bytes == 0 { + 0.0 + } else { + allocated_bytes as f64 / logical_bytes as f64 + }; + let layout = if report.mode == MaterializationMode::Enabled { + MaterializationLayout::Sparse + } else { + MaterializationLayout::Dense + }; + pplx_space_metrics::volume_manager::transfer_materialization_allocation( + kind, + artifact, + report.mode, + layout, + logical_bytes, + allocated_bytes, + report.candidate_hole_bytes, + ); + let projected_allocated_bytes = logical_bytes.saturating_sub(report.candidate_hole_bytes); + tracing::info!( + job_id = %job_id, + subject_id, + local_path = %local_path.display(), + artifact, + layout = %layout, + mode = %report.mode, + logical_bytes, + allocated_bytes, + candidate_hole_bytes = report.candidate_hole_bytes, + projected_allocated_bytes, + saved_bytes, + allocation_ratio, + "decompressed artifact allocation observed" + ); + Some(MaterializationObservation { + mode: report.mode, + layout, + logical_bytes, + allocated_bytes, + candidate_hole_bytes: report.candidate_hole_bytes, + }) +} diff --git a/crates/libs/vol/mgr/server/src/executor/writeback_coordinator.rs b/crates/libs/vol/mgr/server/src/executor/writeback_coordinator.rs new file mode 100644 index 000000000..9ace45247 --- /dev/null +++ b/crates/libs/vol/mgr/server/src/executor/writeback_coordinator.rs @@ -0,0 +1,668 @@ +//! Page-cache writeback pacing shared by every streamed pull: ranges, the byte budget and the +//! coordinator. + +use super::*; + +// Submit decoded output to the process-wide writeback worker in the same +// increments used by ranged downloads. Writers share one byte budget, so +// aggregate dirty output stays bounded independently of pull concurrency. +pub(super) const STREAM_WRITEBACK_WINDOW_BYTES: u64 = 32 * 1024 * 1024; + +pub const DEFAULT_STREAM_WRITEBACK_BUDGET_BYTES: u64 = 2 * 1024 * 1024 * 1024; + +pub(super) const STREAM_WRITEBACK_IDLE_DRAIN: std::time::Duration = + std::time::Duration::from_secs(1); + +pub(super) const STREAM_WRITEBACK_METRIC_INTERVAL: std::time::Duration = + std::time::Duration::from_secs(1); + +#[derive(Debug)] +pub(super) struct StreamWritebackRange { + pub(super) file: Arc, + pub(super) offset: u64, + pub(super) length: u64, + pub(super) debt_bytes: u64, + pub(super) started: bool, + pub(super) timing: Arc, +} + +#[derive(Debug)] +pub(super) enum StreamWritebackCommand { + Range(StreamWritebackRange), + ReapToLowWatermark, +} + +pub(super) trait StreamPageCacheOperations: std::fmt::Debug + Send + Sync { + fn start(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()>; + fn wait(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()>; + fn advise_dontneed( + &self, + file: &std::fs::File, + offset: u64, + length: u64, + ) -> std::io::Result<()>; +} + +#[derive(Debug)] +pub(super) struct SystemStreamPageCacheOperations; + +impl StreamPageCacheOperations for SystemStreamPageCacheOperations { + fn start(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()> { + page_cache::start_writeback_range(file, offset, length) + } + + fn wait(&self, file: &std::fs::File, offset: u64, length: u64) -> std::io::Result<()> { + page_cache::wait_writeback_range(file, offset, length) + } + + fn advise_dontneed( + &self, + file: &std::fs::File, + offset: u64, + length: u64, + ) -> std::io::Result<()> { + page_cache::advise_range_dontneed(file, offset, length) + } +} + +#[derive(Debug, Default)] +pub(super) struct StreamWritebackBudget { + pub(super) outstanding_bytes: std::sync::Mutex, + pub(super) available: std::sync::Condvar, + pub(super) failed: std::sync::atomic::AtomicBool, +} + +#[derive(Debug)] +pub(super) struct StreamWritebackCoordinator { + pub(super) sender: std::sync::mpsc::Sender, + pub(super) budget: Arc, + pub(super) high_watermark_bytes: u64, + pub(super) priority_threshold: Option, + pub(super) priority_reserve_bytes: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StreamWritebackFailureOperation { + Start, + Wait, + Advice, + Queue, +} + +impl StreamWritebackFailureOperation { + pub(super) fn as_str(self) -> &'static str { + match self { + Self::Start => "start", + Self::Wait => "wait", + Self::Advice => "advice", + Self::Queue => "queue", + } + } +} + +impl std::fmt::Display for StreamWritebackFailureOperation { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.as_str()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StreamWritebackPriorityClass { + Shared, + High, +} + +impl StreamWritebackPriorityClass { + pub(super) fn as_str(self) -> &'static str { + match self { + Self::Shared => "shared", + Self::High => "high", + } + } +} + +impl std::fmt::Display for StreamWritebackPriorityClass { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.as_str()) + } +} + +pub(super) fn fail_stream_writeback_budget( + budget: &StreamWritebackBudget, + operation: StreamWritebackFailureOperation, +) { + let was_failed = budget + .failed + .swap(true, std::sync::atomic::Ordering::AcqRel); + budget.available.notify_all(); + if !was_failed { + pplx_space_metrics::volume_manager::stream_writeback_failure(operation.as_str()); + } +} + +pub(super) fn release_stream_writeback_budget(budget: &StreamWritebackBudget, bytes: u64) { + let mut outstanding = budget + .outstanding_bytes + .lock() + .expect("stream writeback budget lock poisoned"); + *outstanding = outstanding.saturating_sub(bytes); + budget.available.notify_all(); +} + +pub(super) fn emit_stream_writeback_outstanding_if_due( + budget: &StreamWritebackBudget, + last_emitted_at: &mut Instant, + force: bool, +) { + let outstanding = stream_writeback_outstanding_bytes(budget); + if !should_emit_stream_writeback_outstanding(last_emitted_at.elapsed(), force) { + return; + } + pplx_space_metrics::volume_manager::stream_writeback_outstanding_bytes(outstanding); + *last_emitted_at = Instant::now(); +} + +pub(super) fn should_emit_stream_writeback_outstanding( + elapsed: std::time::Duration, + force: bool, +) -> bool { + force || elapsed >= STREAM_WRITEBACK_METRIC_INTERVAL +} + +pub(super) fn retire_stream_writeback_range( + range: StreamWritebackRange, + budget: &StreamWritebackBudget, + page_cache_operations: &dyn StreamPageCacheOperations, + writeback_enabled: &mut bool, + cache_advice_enabled: &mut bool, +) { + if *writeback_enabled && range.started { + let wait_started_at = Instant::now(); + let result = page_cache_operations.wait(&range.file, range.offset, range.length); + StreamPullTiming::add_elapsed(&range.timing.writeback_wait_us, wait_started_at); + if let Err(error) = result { + tracing::warn!( + %error, + offset = range.offset, + length = range.length, + "waiting for streamed output range writeback failed; closing streamed pull admission" + ); + *writeback_enabled = false; + fail_stream_writeback_budget(budget, StreamWritebackFailureOperation::Wait); + } + } + if *writeback_enabled + && range.started + && *cache_advice_enabled + && let Err(error) = + page_cache_operations.advise_dontneed(&range.file, range.offset, range.length) + { + tracing::warn!( + %error, + offset = range.offset, + length = range.length, + "streamed output range cache advice failed; disabling cache advice" + ); + pplx_space_metrics::volume_manager::stream_writeback_failure( + StreamWritebackFailureOperation::Advice.as_str(), + ); + *cache_advice_enabled = false; + } + release_stream_writeback_budget(budget, range.debt_bytes); +} + +pub(super) fn stream_writeback_outstanding_bytes(budget: &StreamWritebackBudget) -> u64 { + *budget + .outstanding_bytes + .lock() + .expect("stream writeback budget lock poisoned") +} + +pub(super) fn pending_stream_writeback_debt( + pending: &std::collections::VecDeque, +) -> u64 { + pending + .iter() + .fold(0_u64, |total, range| total.saturating_add(range.debt_bytes)) +} + +pub(super) fn drain_stream_writeback_ranges_to( + pending: &mut std::collections::VecDeque, + budget: &StreamWritebackBudget, + target_bytes: u64, + page_cache_operations: &dyn StreamPageCacheOperations, + writeback_enabled: &mut bool, + cache_advice_enabled: &mut bool, +) -> bool { + while stream_writeback_outstanding_bytes(budget) > target_bytes { + let Some(range) = pending.pop_front() else { + return false; + }; + retire_stream_writeback_range( + range, + budget, + page_cache_operations, + writeback_enabled, + cache_advice_enabled, + ); + } + true +} + +pub(super) fn drain_all_stream_writeback_ranges( + pending: &mut std::collections::VecDeque, + budget: &StreamWritebackBudget, + page_cache_operations: &dyn StreamPageCacheOperations, + writeback_enabled: &mut bool, + cache_advice_enabled: &mut bool, +) { + while let Some(range) = pending.pop_front() { + retire_stream_writeback_range( + range, + budget, + page_cache_operations, + writeback_enabled, + cache_advice_enabled, + ); + } +} + +pub(super) fn queue_stream_writeback_range( + pending: &mut std::collections::VecDeque, + range: StreamWritebackRange, +) { + if let Some(previous) = pending.iter_mut().rev().find(|previous| { + Arc::ptr_eq(&previous.file, &range.file) && Arc::ptr_eq(&previous.timing, &range.timing) + }) { + let can_merge = !previous.started + && previous + .offset + .checked_add(previous.length) + .is_some_and(|end| end <= range.offset) + && previous + .debt_bytes + .checked_add(range.debt_bytes) + .is_some_and(|debt| debt <= STREAM_WRITEBACK_WINDOW_BYTES); + if can_merge + && let (Some(length), Some(debt_bytes)) = ( + range + .offset + .checked_add(range.length) + .and_then(|end| end.checked_sub(previous.offset)), + previous.debt_bytes.checked_add(range.debt_bytes), + ) + { + previous.length = length; + previous.debt_bytes = debt_bytes; + return; + } + } + pending.push_back(range); +} + +pub(super) fn start_pending_stream_writeback_ranges( + pending: &mut std::collections::VecDeque, + budget: &StreamWritebackBudget, + page_cache_operations: &dyn StreamPageCacheOperations, + writeback_enabled: &mut bool, + minimum_debt_bytes: u64, +) { + for range in pending + .iter_mut() + .filter(|range| !range.started && range.debt_bytes >= minimum_debt_bytes) + { + if !*writeback_enabled { + break; + } + let start_started_at = Instant::now(); + let start_cpu_started_at = current_thread_cpu_us(); + let result = page_cache_operations.start(&range.file, range.offset, range.length); + StreamPullTiming::add_wall_and_thread_cpu( + &range.timing.writeback_start_us, + &range.timing.writeback_start_cpu_us, + start_started_at, + start_cpu_started_at, + ); + match result { + Ok(()) => { + range.started = true; + StreamPullTiming::add(&range.timing.writeback_start_calls, 1); + StreamPullTiming::add( + &range.timing.writeback_submitted_logical_bytes, + range.length, + ); + StreamPullTiming::add( + &range.timing.writeback_submitted_dirty_bytes, + range.debt_bytes, + ); + } + Err(error) => { + tracing::warn!( + %error, + offset = range.offset, + length = range.length, + debt_bytes = range.debt_bytes, + "starting streamed output range writeback failed; closing streamed pull admission" + ); + *writeback_enabled = false; + fail_stream_writeback_budget(budget, StreamWritebackFailureOperation::Start); + } + } + } +} + +impl StreamWritebackCoordinator { + pub(super) fn new( + high_watermark_bytes: u64, + priority_threshold: Option, + ) -> Arc { + Self::new_with_page_cache_operations( + high_watermark_bytes, + priority_threshold, + Arc::new(SystemStreamPageCacheOperations), + ) + } + + pub(super) fn new_with_page_cache_operations( + high_watermark_bytes: u64, + priority_threshold: Option, + page_cache_operations: Arc, + ) -> Arc { + let (sender, receiver) = std::sync::mpsc::channel::(); + let budget = Arc::new(StreamWritebackBudget::default()); + let worker_budget = Arc::clone(&budget); + let low_watermark_bytes = high_watermark_bytes / 2; + let writeback_window_bytes = STREAM_WRITEBACK_WINDOW_BYTES.min(low_watermark_bytes); + std::thread::Builder::new() + .name("volmgr-writeback".to_string()) + .spawn(move || { + let mut writeback_enabled = true; + let mut cache_advice_enabled = true; + let mut pending = std::collections::VecDeque::new(); + let mut drain_target = None; + let mut pressure_active = false; + let mut last_outstanding_metric_at = Instant::now(); + + loop { + let mut force_outstanding_metric = false; + let writeback_enabled_at_start = writeback_enabled; + match receiver.recv_timeout(STREAM_WRITEBACK_IDLE_DRAIN) { + Ok(StreamWritebackCommand::Range(range)) if writeback_enabled => { + let activation_timing = Arc::clone(&range.timing); + queue_stream_writeback_range(&mut pending, range); + // Overlap complete dirty windows with decoding even + // below the pressure watermark. Keep their debt + // charged until the existing completion/drain path. + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + writeback_window_bytes, + ); + if pending_stream_writeback_debt(&pending) >= low_watermark_bytes { + let newly_activated = !pressure_active; + if !pressure_active { + pressure_active = true; + force_outstanding_metric = true; + StreamPullTiming::add( + &activation_timing.writeback_low_watermark_activations, + 1, + ); + } + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + if newly_activated { + 0 + } else { + writeback_window_bytes + }, + ); + } + } + Ok(StreamWritebackCommand::Range(range)) => { + release_stream_writeback_budget(&worker_budget, range.debt_bytes); + } + Ok(StreamWritebackCommand::ReapToLowWatermark) => { + force_outstanding_metric = true; + drain_target = Some(low_watermark_bytes); + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + 0, + ); + } + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + if pending_stream_writeback_debt(&pending) >= low_watermark_bytes { + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + 0, + ); + drain_target = Some(low_watermark_bytes); + } + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + 0, + ); + drain_all_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + &mut cache_advice_enabled, + ); + emit_stream_writeback_outstanding_if_due( + &worker_budget, + &mut last_outstanding_metric_at, + true, + ); + break; + } + } + + if drain_target.is_some() { + start_pending_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + 0, + ); + } + if drain_target.is_some_and(|target| { + drain_stream_writeback_ranges_to( + &mut pending, + &worker_budget, + target, + page_cache_operations.as_ref(), + &mut writeback_enabled, + &mut cache_advice_enabled, + ) + }) { + drain_target = None; + pressure_active = stream_writeback_outstanding_bytes(&worker_budget) + > low_watermark_bytes; + force_outstanding_metric = true; + } + if !writeback_enabled { + drain_all_stream_writeback_ranges( + &mut pending, + &worker_budget, + page_cache_operations.as_ref(), + &mut writeback_enabled, + &mut cache_advice_enabled, + ); + drain_target = None; + pressure_active = false; + force_outstanding_metric |= writeback_enabled_at_start; + } + emit_stream_writeback_outstanding_if_due( + &worker_budget, + &mut last_outstanding_metric_at, + force_outstanding_metric, + ); + } + }) + .expect("spawn volume-manager writeback worker"); + Arc::new(Self { + sender, + budget, + high_watermark_bytes, + priority_threshold, + priority_reserve_bytes: if priority_threshold.is_some() { + STREAM_WRITEBACK_WINDOW_BYTES.min(high_watermark_bytes / 2) + } else { + 0 + }, + }) + } + + pub(super) fn admission_limit(&self, priority: JobPriority) -> u64 { + if self + .priority_threshold + .is_some_and(|threshold| priority < threshold) + { + self.high_watermark_bytes + .saturating_sub(self.priority_reserve_bytes) + } else { + self.high_watermark_bytes + } + } + + pub(super) fn try_reserve(&self, bytes: u64, priority: JobPriority) -> std::io::Result { + if self + .budget + .failed + .load(std::sync::atomic::Ordering::Acquire) + { + return Err(std::io::Error::other("stream writeback coordinator failed")); + } + let mut outstanding = self + .budget + .outstanding_bytes + .lock() + .map_err(|_| std::io::Error::other("stream writeback budget lock poisoned"))?; + if self + .budget + .failed + .load(std::sync::atomic::Ordering::Acquire) + { + return Err(std::io::Error::other("stream writeback coordinator failed")); + } + let admission_limit = self.admission_limit(priority); + if *outstanding != 0 && outstanding.saturating_add(bytes) > admission_limit { + return Ok(false); + } + *outstanding = outstanding.saturating_add(bytes); + Ok(true) + } + + pub(super) fn reserve(&self, bytes: u64, priority: JobPriority) -> std::io::Result<()> { + let reserve_started_at = Instant::now(); + if self.try_reserve(bytes, priority)? { + return Ok(()); + } + let mut outstanding = self + .budget + .outstanding_bytes + .lock() + .map_err(|_| std::io::Error::other("stream writeback budget lock poisoned"))?; + if self + .budget + .failed + .load(std::sync::atomic::Ordering::Acquire) + { + return Err(std::io::Error::other("stream writeback coordinator failed")); + } + let admission_limit = self.admission_limit(priority); + let priority_class = if admission_limit == self.high_watermark_bytes { + StreamWritebackPriorityClass::High + } else { + StreamWritebackPriorityClass::Shared + }; + while *outstanding != 0 && outstanding.saturating_add(bytes) > admission_limit { + if self + .sender + .send(StreamWritebackCommand::ReapToLowWatermark) + .is_err() + { + tracing::warn!( + "stream writeback worker stopped while a producer was blocked; closing streamed pull admission" + ); + fail_stream_writeback_budget(&self.budget, StreamWritebackFailureOperation::Queue); + return Err(std::io::Error::other("stream writeback worker stopped")); + } + outstanding = self.budget.available.wait(outstanding).map_err(|_| { + std::io::Error::other("stream writeback budget lock poisoned while waiting") + })?; + if self + .budget + .failed + .load(std::sync::atomic::Ordering::Acquire) + { + return Err(std::io::Error::other("stream writeback coordinator failed")); + } + } + *outstanding = outstanding.saturating_add(bytes); + pplx_space_metrics::volume_manager::stream_writeback_producer_blocked_ms( + reserve_started_at.elapsed().as_secs_f64() * 1000.0, + priority_class.as_str(), + ); + Ok(()) + } + + pub(super) fn release(&self, bytes: u64) { + let mut outstanding = self + .budget + .outstanding_bytes + .lock() + .expect("stream writeback budget lock poisoned"); + *outstanding = outstanding.saturating_sub(bytes); + self.budget.available.notify_all(); + } + + pub(super) fn submit( + &self, + file: Arc, + offset: u64, + length: u64, + debt_bytes: u64, + timing: Arc, + ) { + let result = self + .sender + .send(StreamWritebackCommand::Range(StreamWritebackRange { + file, + offset, + length, + debt_bytes, + started: false, + timing, + })) + .map_err(|_| std::io::Error::other("stream writeback worker stopped")); + if let Err(error) = result { + tracing::warn!( + %error, + offset, + length, + debt_bytes, + "streamed output range could not be queued for writeback; closing streamed pull admission" + ); + fail_stream_writeback_budget(&self.budget, StreamWritebackFailureOperation::Queue); + self.release(debt_bytes); + } + } +}