diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index b6e5b8201..979df2551 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -200,6 +200,7 @@ rules_scala [Tt]oolchainize [Tt]oolchainized [Vv]endored +[Vv]alidator dev_dependency Fumadocs Turborepo @@ -318,3 +319,24 @@ subtree hardlink multiplicatively SELinux +[Pp]assthrough +[Rr]ecompression +[Rr]ecompressions +[Dd]ecodable +libzstd +windowLog +[Ss]eekable +[Tt]ranscodes +[Ii]nterop +[Dd]esynchronization +max_recompression_size +EXDEV +[Oo]bservability +[Uu]nguessable +fsync +max_concurrent_staged_uploads +max_concurrent_identity_ops +max_inline_commit_size +stage_timeout_s +commit_timeout_s +compressed_upload_idle_timeout_s diff --git a/Cargo.lock b/Cargo.lock index 72fb165e8..74f5e7b2b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3246,6 +3246,7 @@ dependencies = [ "webpki-roots 1.0.3", "wincode", "zip", + "zstd", ] [[package]] diff --git a/deployment-examples/docker-compose/README.md b/deployment-examples/docker-compose/README.md index a48790c4e..741529e5e 100644 --- a/deployment-examples/docker-compose/README.md +++ b/deployment-examples/docker-compose/README.md @@ -121,6 +121,7 @@ services: ### Single Worker Setup - [`docker-compose.yml`](./docker-compose.yml) - Docker Compose file for single worker deployment - [`local-storage-cas.json5`](./local-storage-cas.json5) - Local storage CAS configuration for single worker +- [`local-storage-cas-zstd.json5`](./local-storage-cas-zstd.json5) - Same, but the CAS selects `compression_algorithm.zstd` to keep blobs as zstd at rest for byte-for-byte `--remote_cache_compression` passthrough (see the file's header comment for the dedicated-namespace and placement rules) - [`scheduler.json5`](./scheduler.json5) - Scheduler configuration for single worker deployment - [`worker.json5`](./worker.json5) - Worker configuration for single worker deployment diff --git a/deployment-examples/docker-compose/local-storage-cas-zstd.json5 b/deployment-examples/docker-compose/local-storage-cas-zstd.json5 new file mode 100644 index 000000000..9b4be03f2 --- /dev/null +++ b/deployment-examples/docker-compose/local-storage-cas-zstd.json5 @@ -0,0 +1,242 @@ +// Variant of `local-storage-cas.json5` that stores CAS blobs as zstd streams +// at rest via the `zstd` algorithm of `compression`, instead of raw bytes +// behind `compression` with its `lz4` algorithm. +// +// The payoff over plain `compression`: a Bazel client built with +// `--remote_cache_compression` (enabled below via `remote_cache_compression` +// on the `capabilities` service) gets served the stored zstd bytes +// byte-for-byte on a cache hit — no decompress-then-recompress round trip at +// the gRPC boundary. Identity clients (no `--remote_cache_compression`) are +// unaffected: they still get plain decompressed bytes. +// +// Read this whole comment block before adapting this file. `compression` with +// the `zstd` algorithm has deployment constraints that are not optional +// footnotes; getting them wrong +// either silently corrupts a mixed namespace or forces a cache flush later. +// Full explanation: +// https://github.com/TraceMachina/nativelink/blob/main/web/apps/docs/content/docs/reference/nativelink-config/store-overview.mdx +// +// 1. DEDICATED, EMPTY NAMESPACE ONLY. The `compression.backend` below, +// when `compression_algorithm.zstd` is selected, +// (`CAS_ZSTD_CONTENT_STORE`, and the paths it writes) must never be shared +// with a store, instance, or NativeLink version that reads or writes those +// same keys as raw bytes. The digest key is identical either way, raw CAS +// values are unstructured (a raw blob can legitimately start with the +// zstd magic number), so there is no safe way to sniff which encoding an +// entry is in. That's why this file uses brand-new `*-cas-zstd` paths +// instead of the `content_path-cas` paths from `local-storage-cas.json5` — +// reusing those would mix raw and zstd values under one namespace. +// +// 2. NO IN-PLACE MIGRATION. Selecting `compression_algorithm.zstd` for an existing +// populated CAS, or turning it back off, both require a cache flush (or +// cutting over to a new namespace as done here) — not an in-place +// conversion. The same applies to a rolling/mixed-version deploy against +// one shared namespace: unsupported, because some replicas would write +// raw and others zstd to the same keys. +// +// 3. PLACEMENT: the `compression` store with `compression_algorithm.zstd` must +// be the store an instance/service points at directly for the passthrough +// optimization above to apply. Wrappers such +// as `fast_slow`, `dedup`, `existence_cache`, `cache_metrics`, `shard`, +// `ref_store`, and `size_partitioning` are fine *inside* it (they then +// operate on the physical zstd stream). A wrapper placed *outside* +// the zstd compression store (e.g. `verify` around it) is still correct, +// but forces that wrapper to decode the stream, which disables passthrough at that +// boundary. +// +// 4. CAS-ONLY. `compression_algorithm.zstd` only supports digest-keyed CAS entries; never +// point an AC store at it (`AC_MAIN_STORE` below stays plain `filesystem`). +{ + stores: [ + { + name: "CAS_MAIN_STORE", + compression: { + compression_algorithm: { + zstd: { + // Staging directory for upload validation. MUST be on the same + // filesystem as the `filesystem` backend's `content_path` below: + // that backend commits a staged upload with rename(2), which fails + // with EXDEV across filesystems. Kept on the same + // `/root/.cache/nativelink` tree here for exactly that reason. + temp_path: "/root/.cache/nativelink/tmp_path-cas-zstd-stage", + + // Reject a compressed upload over this many wire bytes with + // RESOURCE_EXHAUSTED rather than staging it. + max_compressed_upload_size: "512MiB", + + // How many uploads may hold a staged file at once. Sizes worst-case + // temp disk usage together with max_compressed_upload_size: + // max_concurrent_staged_uploads * max_compressed_upload_size + // Recompression adds no staging file (it reuses the same + // descriptor), but does hold up to max_recompression_size in memory. + // Monitor `temp_path` against the disk bound. + max_concurrent_staged_uploads: 4, + + // How many uncompressed (identity) reads and writes to admit at + // once. Identity clients — anything without + // --remote_cache_compression — are encoded/decoded on a blocking + // thread held for the whole transfer, so this bound keeps a flood of + // slow identity clients from starving the process-wide blocking pool + // that filesystem I/O also uses. + max_concurrent_identity_ops: 256, + + // Level used to encode uploads this store compresses itself, and to + // re-encode incoming compressed uploads when + // max_recompression_size > 0. Omit for the default of 3. + // + // Valid range is 1..=19 — enforced at startup. This is *not* the + // full zstd range (up to 22): standard levels 1-19 cap the encoded + // frame's window size at windowLog <= 23 (<= 8MiB), which every + // Bazel/zstd-jni client can decode with plain libzstd defaults. + // Bazel's decoder never calls setLongMax, so its hard ceiling is + // windowLog <= 27 (128MiB) - 19 keeps a large margin under that and + // is what caps the field, not a smaller "safe zstd" limit. Long- + // distance matching and dictionaries are never enabled by this + // store. + // + // 9 is a general-purpose default here: it costs a fraction of level + // 19's CPU per byte. Raise it toward 19 for an archival cache whose + // blobs are read far more often than written. + compression_level: 9, + + // Recompression is only attempted for uploads whose *decoded* size + // is within this bound, and the smaller of the original and + // re-encoded stream is kept. 0 disables recompression entirely (an + // incoming compressed stream is always stored as-is), while still + // letting `compression_level` pick the level used for brand-new, + // non-precompressed uploads. A positive value requires + // `compression_level`; without it, startup fails rather than + // silently doing nothing. + max_recompression_size: "64MiB", + + // Concurrent recompressions admitted by this store instance. + // Recompression is best-effort: an upload that finds every slot busy + // commits its original stream rather than queueing, so a small pool + // here throttles recompression without throttling uploads. + max_concurrent_recompressions: 1, + + // Compressed uploads at or below this size are validated and + // committed straight from memory, with no staging file and no + // fsync. BatchUpdateBlobs payloads are small and numerous, so a + // per-blob disk round trip would dominate their cost. + max_inline_commit_size: "4MiB", + + // Total time one upload may spend being validated and staged, + // measured from admission to a staging slot. Unlike the ByteStream + // per-message idle timeout, continuous slow progress does NOT reset + // it. On expiry the blocking validator retains its slot and cleanup + // guard until its input closes; ByteStream closes that input when + // the store fails. Size this against max_compressed_upload_size and + // the slowest upload bandwidth worth serving. + stage_timeout_s: 600, + + // Time budget for a staged upload's optional recompression plus the + // inner-store commit, after the client stream is done. Bounds a + // stalled backend holding a staging slot. + commit_timeout_s: 300, + }, + }, + + // MUST be a brand-new/empty namespace — see note 1 above. Everything + // under `backend` here is a complete, dedicated CAS stack; nothing + // else in this file (or any other deployment) may read or write it + // as raw bytes. + backend: { + fast_slow: { + fast: { + memory: { + eviction_policy: { + // 500mb. Absorbs hot re-reads without touching disk. + max_bytes: 500000000, + }, + }, + }, + slow: { + filesystem: { + content_path: "/root/.cache/nativelink/content_path-cas-zstd", + temp_path: "/root/.cache/nativelink/tmp_path-cas-zstd", + eviction_policy: { + // 10gb. + max_bytes: 10000000000, + }, + }, + }, + }, + }, + }, + }, + { + // Holds blob-to-chunks layouts for the SplitBlob/SpliceBlob RPCs used + // by Bazel's --experimental_remote_cache_chunking. Must not verify + // digests and must not be the same store as the CAS. + name: "CHUNK_INDEX_STORE", + filesystem: { + content_path: "/root/.cache/nativelink/content_path-chunk-index", + temp_path: "/root/.cache/nativelink/tmp_path-chunk-index", + eviction_policy: { + // 100mb. + max_bytes: 100000000, + }, + }, + }, + { + // Plain filesystem store: zstd compression is CAS-only (note 4 above), so + // the AC never goes anywhere near it. + name: "AC_MAIN_STORE", + filesystem: { + content_path: "/root/.cache/nativelink/content_path-ac", + temp_path: "/root/.cache/nativelink/tmp_path-ac", + eviction_policy: { + // 500mb. + max_bytes: 500000000, + }, + }, + }, + ], + servers: [ + { + listener: { + http: { + socket_address: "0.0.0.0:50051", + }, + }, + services: { + cas: [ + { + cas_store: "CAS_MAIN_STORE", + + // Optional: enables content-defined chunking + // (SplitBlob/SpliceBlob) for Bazel clients running with + // --experimental_remote_cache_chunking. + experimental_chunking: { + index_store: "CHUNK_INDEX_STORE", + }, + }, + ], + ac: [ + { + ac_store: "AC_MAIN_STORE", + }, + ], + capabilities: [ + { + // Advertises zstd to clients so ByteStream/CAS RPCs accept and + // serve compressed-blobs/zstd/... payloads. This is what makes + // the byte-for-byte passthrough in the zstd compression store reachable — + // without it, every client uses identity transfers and + // zstd compression still stores zstd at rest, but every read pays + // a decompress. + remote_cache_compression: true, + }, + ], + bytestream: { + cas_stores: { + "": "CAS_MAIN_STORE", + }, + }, + fetch: {}, + push: {}, + }, + }, + ], +} diff --git a/nativelink-config/examples/stores-config.json5 b/nativelink-config/examples/stores-config.json5 index 0748ce406..fc6922369 100644 --- a/nativelink-config/examples/stores-config.json5 +++ b/nativelink-config/examples/stores-config.json5 @@ -172,7 +172,7 @@ "content_path": "/tmp/nativelink/data/content_path-cas", "temp_path": "/tmp/nativelink/data/tmp_path-cas", "eviction_policy": { - "max_bytes": "2gb", + "max_bytes": "2gb" } } } @@ -180,6 +180,30 @@ }, { name: "12", + "compression": { + "compression_algorithm": { + "zstd": { + "temp_path": "/var/tmp/nativelink-zstd", + "max_compressed_upload_size": "512MiB", + "max_concurrent_staged_uploads": 4, + "max_concurrent_identity_ops": 256, + "compression_level": 9, + "max_recompression_size": "64MiB", + "max_concurrent_recompressions": 1, + "max_inline_commit_size": "4MiB", + "stage_timeout_s": 600, + "commit_timeout_s": 300 + } + }, + "backend": { + "memory": { + "eviction_policy": { "max_bytes": "10GiB" } + } + } + } + }, + { + name: "13", "dedup": { "index_store": { "memory": { @@ -218,7 +242,7 @@ } }, { - name: "13", + name: "14", "existence_cache": { "backend": { "memory": { @@ -234,7 +258,7 @@ } }, { - name: "14", + name: "15", "fast_slow": { "fast": { "filesystem": { @@ -257,7 +281,7 @@ } }, { - name: "15", + name: "16", "shard": { "stores": [ { @@ -273,7 +297,7 @@ } }, { - name: "16", + name: "17", "filesystem": { "content_path": "/tmp/nativelink/data-worker-test/content_path-cas", "temp_path": "/tmp/nativelink/data-worker-test/tmp_path-cas", @@ -283,13 +307,13 @@ } }, { - name: "17", + name: "18", "ref_store": { "name": "FS_CONTENT_STORE" } }, { - name: "18", + name: "19", "size_partitioning": { "size": "128mib", "lower_store": { @@ -306,7 +330,7 @@ } }, { - name: "19", + name: "20", "grpc": { "instance_name": "main", "endpoints": [ @@ -327,7 +351,7 @@ } }, { - name: "20", + name: "21", "redis_store": { "addresses": [ "redis://127.0.0.1:6379/", @@ -336,11 +360,11 @@ } }, { - name: "21", + name: "22", "noop": {} }, { - name: "22", + name: "23", "experimental_mongo": { "connection_string": "mongodb://localhost:27017", "database": "nativelink", diff --git a/nativelink-config/src/backcompat.rs b/nativelink-config/src/backcompat.rs index 00b9e9b3e..18c9ea8f5 100644 --- a/nativelink-config/src/backcompat.rs +++ b/nativelink-config/src/backcompat.rs @@ -102,6 +102,8 @@ where max_bytes_per_stream: old_config.max_bytes_per_stream, persist_stream_on_disconnect_timeout_s: old_config .persist_stream_on_disconnect_timeout_s, + // Not expressible in the deprecated config; take the default. + compressed_upload_idle_timeout_s: 0, }, }) .collect(); diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 727830b59..d739c155f 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -352,6 +352,21 @@ pub struct ByteStreamConfig { alias = "persist_stream_on_disconnect_timeout" )] pub persist_stream_on_disconnect_timeout_s: usize, + + /// How long, in seconds, to wait for the *next* `WriteRequest` of a + /// compressed (`compressed-blobs/...`) upload before failing it with + /// `DEADLINE_EXCEEDED`. A client that keeps making progress resets the + /// deadline on every message, so this never rejects a large upload — it + /// bounds a client that opens a compressed write, takes a store slot, then + /// stalls or never sends `finish_write`. + /// + /// Default: 60 seconds + #[serde( + default, + deserialize_with = "convert_duration_with_shellexpand", + skip_serializing_if = "is_default" + )] + pub compressed_upload_idle_timeout_s: usize, } // Older bytestream config. All fields are as per the newer docs, but this requires diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index b8e827492..3e7b31b46 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -315,7 +315,7 @@ pub enum StoreSpec { /// is a concern it is often faster and more efficient to use this /// store before those stores. /// - /// **Example JSON Config:** + /// **LZ4 example:** /// ```json /// "compression": { /// "compression_algorithm": { @@ -326,13 +326,50 @@ pub enum StoreSpec { /// "content_path": "/tmp/nativelink/data/content_path-cas", /// "temp_path": "/tmp/nativelink/data/tmp_path-cas", /// "eviction_policy": { - /// "max_bytes": "2gb", + /// "max_bytes": "2gb" /// } /// } /// } /// } /// ``` /// + /// **Zstd example:** + /// ```json + /// "compression": { + /// "compression_algorithm": { + /// "zstd": { + /// "temp_path": "/var/tmp/nativelink-zstd", + /// "max_compressed_upload_size": "512MiB", + /// "max_concurrent_staged_uploads": 4, + /// "max_concurrent_identity_ops": 256, + /// "compression_level": 9, + /// "max_recompression_size": "64MiB", + /// "max_concurrent_recompressions": 1, + /// "max_inline_commit_size": "4MiB", + /// "stage_timeout_s": 600, + /// "commit_timeout_s": 300 + /// } + /// }, + /// "backend": { + /// "memory": { + /// "eviction_policy": { "max_bytes": "10GiB" } + /// } + /// } + /// } + /// ``` + /// + /// The `zstd` algorithm keeps CAS blobs as zstd streams at rest and serves them + /// byte-for-byte to `--remote_cache_compression` clients. Its `backend` MUST be + /// a new or empty dedicated namespace, never shared with processes that read or + /// write the same keys as raw bytes. Rollout and rollback require a cache flush + /// or new namespace — there is no in-place migration. + /// + /// For byte-for-byte passthrough, this compression store must be the store the + /// instance points at directly. `fast_slow`, `dedup`, `existence_cache`, + /// `cache_metrics`, `shard`, `ref`, and `size_partitioning` may appear **inside** + /// its `backend`. Any wrapper **outside** it is correct but disables passthrough + /// at that boundary. + /// Compression(Box), /// A dedup store will take the inputs and run a rolling hash @@ -1151,7 +1188,7 @@ pub struct Lz4Config { pub max_decode_block_size: u32, } -#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Copy)] +#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)] #[serde(rename_all = "snake_case")] #[cfg_attr(feature = "dev-schema", derive(JsonSchema))] pub enum CompressionAlgorithm { @@ -1163,6 +1200,11 @@ pub enum CompressionAlgorithm { /// /// see: Lz4(Lz4Config), + + /// Zstd compression keeps blobs as standard zstd streams at rest. When this + /// compression store is directly configured for an instance, zstd wire-compression + /// clients can receive the stored stream byte-for-byte. + Zstd(ZstdConfig), } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -1180,6 +1222,89 @@ pub struct CompressionSpec { pub compression_algorithm: CompressionAlgorithm, } +#[derive(Serialize, Deserialize, Debug, Default, PartialEq, Eq, Clone)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "dev-schema", derive(JsonSchema))] +pub struct ZstdConfig { + /// Operator-controlled staging directory for validation. Put it on the same + /// filesystem as the `content_path` of a `filesystem` backend: that backend + /// commits by `rename(2)`, which fails with `EXDEV` across filesystems. + #[serde(default, deserialize_with = "convert_string_with_shellexpand")] + pub temp_path: String, + + /// Max compressed wire bytes accepted by an upload. Exceeding it returns + /// `RESOURCE_EXHAUSTED`. Accepts "512MiB"-style strings. + #[serde(default, deserialize_with = "convert_data_size_with_shellexpand")] + pub max_compressed_upload_size: u64, + + /// Max uploads holding staging files at once. `0` means "use default 4" + /// at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_concurrent_staged_uploads: usize, + + /// Max concurrent uncompressed (identity) reads and writes this store will + /// admit. Each one occupies a blocking thread for the whole transfer, so + /// this bound keeps a flood of slow identity clients from starving the + /// process-wide blocking pool that filesystem I/O also uses. `0` means "use + /// default 256" at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_concurrent_identity_ops: usize, + + /// Level used to encode uploads this store compresses itself (identity + /// uploads) and, when `max_recompression_size > 0`, to re-encode incoming + /// compressed uploads. Omitted uses level 3. Validated `1..=19` at startup. + pub compression_level: Option, + + /// Max uncompressed blob size eligible for optional recompression of an + /// already-compressed upload. `0` disables recompression. Requires + /// `compression_level` to be set; a positive value without it is rejected + /// at startup rather than silently doing nothing. + #[serde(default, deserialize_with = "convert_data_size_with_shellexpand")] + pub max_recompression_size: u64, + + /// Concurrent recompressions admitted by this store. Recompression is + /// best-effort: an upload that finds every slot busy commits the client's + /// original stream instead of queueing behind them, so this never blocks a + /// staging slot. `0` means "use default 1" at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub max_concurrent_recompressions: usize, + + /// Maximum time, in seconds, one upload may spend being validated and + /// staged, measured from the moment it is admitted to a staging slot. Unlike + /// a per-message idle timeout, continuous slow progress does not reset it. + /// On expiry it fails with `DEADLINE_EXCEEDED`; the non-cancellable blocking + /// validator retains its slot and staged-file cleanup guard until its input + /// closes and it exits, so callers must promptly close the input after the + /// failure. Size it against `max_compressed_upload_size` and the slowest + /// upload bandwidth worth serving. + /// `0` means "use default 600" at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub stage_timeout_s: u64, + + /// Maximum time, in seconds, allowed for the optional recompression of a + /// staged upload plus its inner-store commit. The timer starts after the + /// client stream has finished validation and staging, so it does not reject + /// a steadily progressing large upload. On expiry it fails with + /// `DEADLINE_EXCEEDED`, removes the staged file, and releases its slot. + /// `0` means "use default 300" at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] + pub commit_timeout_s: u64, + + /// Compressed uploads at or below this many bytes are validated and + /// committed straight from memory, with no staging file and no `fsync`. + /// `BatchUpdateBlobs` payloads are small and numerous, so writing each one + /// to disk would dominate their cost. Larger uploads always stage to + /// `temp_path`. `0` means "use default 4MiB" at construction time. + /// Default: 0 + #[serde(default, deserialize_with = "convert_data_size_with_shellexpand")] + pub max_inline_commit_size: u64, +} + /// Eviction policy always works on LRU (Least Recently Used). Any time an entry /// is touched it updates the timestamp. Inserts and updates will execute the /// eviction policy removing any expired entries and/or the oldest entries @@ -2020,3 +2145,53 @@ impl Retry { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn zstd_compression_algorithm_parses_and_defaults() { + let cfg: StoreSpec = serde_json5::from_str( + r#"{ compression: { compression_algorithm: { zstd: { + temp_path: "/var/tmp/nl-zstd", max_compressed_upload_size: "512MiB", + compression_level: 19, max_recompression_size: "64MiB" } }, + backend: { memory: {} } } }"#, + ) + .unwrap(); + let StoreSpec::Compression(spec) = cfg else { + panic!("wrong variant") + }; + let CompressionAlgorithm::Zstd(spec) = spec.compression_algorithm else { + panic!("wrong compression algorithm") + }; + assert_eq!(spec.max_compressed_upload_size, 512 * 1024 * 1024); + assert_eq!(spec.compression_level, Some(19)); + assert_eq!(spec.max_recompression_size, 64 * 1024 * 1024); + assert_eq!(spec.max_concurrent_staged_uploads, 0); // 0 = "use default 4" at construction + assert_eq!(spec.max_concurrent_recompressions, 0); // 0 = "use default 1" + assert_eq!(spec.max_concurrent_identity_ops, 0); // 0 = "use default 256" + assert_eq!(spec.stage_timeout_s, 0); // 0 = "use default 600" + assert_eq!(spec.max_inline_commit_size, 0); // 0 = "use default 4MiB" + } + + #[test] + fn zstd_compression_level_may_be_omitted() { + // `compression_level` has no `#[serde(default)]`; confirm the documented + // "omitted means the default level" spelling actually deserializes. + let cfg: StoreSpec = serde_json5::from_str( + r#"{ compression: { compression_algorithm: { zstd: { + temp_path: "/var/tmp/nl-zstd", max_compressed_upload_size: "512MiB" } }, + backend: { memory: {} } } }"#, + ) + .unwrap(); + let StoreSpec::Compression(spec) = cfg else { + panic!("wrong variant") + }; + let CompressionAlgorithm::Zstd(spec) = spec.compression_algorithm else { + panic!("wrong compression algorithm") + }; + assert_eq!(spec.compression_level, None); + assert_eq!(spec.max_recompression_size, 0); + } +} diff --git a/nativelink-config/tests/json5_test.rs b/nativelink-config/tests/json5_test.rs index 24dbeaa3a..6ce5ed1b3 100644 --- a/nativelink-config/tests/json5_test.rs +++ b/nativelink-config/tests/json5_test.rs @@ -1,24 +1,14 @@ use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use nativelink_config::cas_server::CasConfig; -#[test] -fn test_example_parsing() { - let mut examples_path = Path::new(".") - .canonicalize() - .expect("Can canonicalize current dir"); - - if examples_path.join("nativelink-config").exists() { - // inside bazel - examples_path = examples_path.join("nativelink-config"); - } - examples_path = examples_path.join("examples"); - - let mut found_at_least_one_entry = false; - - for entry in fs::read_dir(&examples_path) - .unwrap_or_else(|e| panic!("Failed to read from {:?}: {}", &examples_path, e)) +/// Parses every `.json5` config directly under `dir`, returning how many it +/// found. Panics with the offending path on the first parse failure. +fn parse_all_configs_in(dir: &Path) -> usize { + let mut parsed = 0; + for entry in + fs::read_dir(dir).unwrap_or_else(|e| panic!("Failed to read from {}: {e}", dir.display())) { let config_file = entry.unwrap().path().display().to_string(); if !config_file.contains(".json5") { @@ -26,8 +16,60 @@ fn test_example_parsing() { } CasConfig::try_from_json5_file(&config_file) .unwrap_or_else(|e| panic!("Error while reading {config_file}: {e}")); - found_at_least_one_entry = true; + parsed += 1; + } + parsed +} + +/// The repository root, whether the test runs from the workspace root or from +/// inside the `nativelink-config` package (as it does under bazel). +fn repo_root() -> PathBuf { + let cwd = Path::new(".") + .canonicalize() + .expect("Can canonicalize current dir"); + if cwd.join("nativelink-config").exists() { + cwd + } else { + cwd.parent() + .expect("nativelink-config has a parent directory") + .to_path_buf() + } +} + +#[test] +fn test_example_parsing() { + let examples_path = repo_root().join("nativelink-config").join("examples"); + assert!( + parse_all_configs_in(&examples_path) > 0, + "expected at least one example config in {}", + examples_path.display() + ); +} + +/// The `deployment-examples` trees are what operators copy, so a config that +/// does not even parse is worse than a missing one. Nothing used to check them, +/// which let an unbalanced brace ship in `local-storage-cas-zstd.json5`. +#[test] +fn test_deployment_example_parsing() { + let deployment_examples = repo_root().join("deployment-examples"); + if !deployment_examples.exists() { + // Not present in the bazel runfiles for this target; the cargo run covers it. + return; + } + + let mut parsed = 0; + for entry in fs::read_dir(&deployment_examples) + .unwrap_or_else(|e| panic!("Failed to read from {}: {e}", deployment_examples.display())) + { + let dir = entry.unwrap().path(); + if dir.is_dir() { + parsed += parse_all_configs_in(&dir); + } } - assert!(found_at_least_one_entry); + assert!( + parsed > 0, + "expected at least one deployment example config under {}", + deployment_examples.display() + ); } diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index 179abe29b..e64b0f824 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -59,7 +59,7 @@ use tokio::time::sleep; use tonic::{Request, Response, Status, Streaming}; use tracing::{Instrument, Level, debug, error, error_span, info, instrument, trace, warn}; -use crate::wire_compression::RemoteCacheCompressionInstances; +use crate::wire_compression::{RemoteCacheCompressionInstances, wire_compressor_capability}; /// If this value changes update the documentation in the config definition. const DEFAULT_PERSIST_STREAM_ON_DISCONNECT_TIMEOUT: Duration = Duration::from_mins(1); @@ -67,6 +67,9 @@ const DEFAULT_PERSIST_STREAM_ON_DISCONNECT_TIMEOUT: Duration = Duration::from_mi /// If this value changes update the documentation in the config definition. const DEFAULT_MAX_BYTES_PER_STREAM: usize = 64 * 1024; +/// If this value changes update the documentation in the config definition. +const DEFAULT_COMPRESSED_UPLOAD_IDLE_TIMEOUT: Duration = Duration::from_mins(1); + /// Metrics for `ByteStream` server operations. /// Tracks upload/download activity, throughput, and latency. #[derive(Debug, Default)] @@ -254,6 +257,10 @@ pub struct InstanceInfo { active_uploads: Arc>>, /// How long to keep idle streams before timing them out. idle_stream_timeout: Duration, + /// How long to wait for the next `WriteRequest` of a compressed upload. + /// Distinct from `idle_stream_timeout`, which governs how long a + /// disconnected upload stays resumable. + compressed_upload_idle_timeout: Duration, metrics: Arc, /// Handle to the global sweeper task. Kept alive for the lifetime of the instance. _sweeper_handle: Arc>, @@ -268,6 +275,10 @@ impl Debug for InstanceInfo { .field("max_bytes_per_stream", &self.max_bytes_per_stream) .field("active_uploads", &self.active_uploads) .field("idle_stream_timeout", &self.idle_stream_timeout) + .field( + "compressed_upload_idle_timeout", + &self.compressed_upload_idle_timeout, + ) .field("metrics", &self.metrics) .field( "remote_cache_compression_enabled", @@ -337,6 +348,34 @@ impl InstanceInfo { } } +/// Returns the first client-attributable error among `results`: a bad request or +/// an expired deadline. Those are root causes worth reporting verbatim, whereas +/// the other halves of a joined pipeline typically fail with a consequential +/// "channel disconnected" error that would mask them. +fn first_client_attributable_error<'a>( + results: impl IntoIterator>, +) -> Option { + results + .into_iter() + .flatten() + .find(|err| { + matches!( + err.code, + Code::InvalidArgument | Code::DeadlineExceeded | Code::ResourceExhausted + ) + }) + .cloned() +} + +/// Folds `err` into an accumulated upload error, keeping the earlier error's +/// messages first. +fn merge_upload_error(accumulated: Option, err: Error) -> Error { + match accumulated { + Some(existing) => existing.merge(err), + None => err, + } +} + /// Pump compressed `ByteStream` upload chunks into the decoder. /// /// Compressed uploads intentionally do not support the identity upload resume @@ -345,9 +384,25 @@ async fn process_compressed_client_stream( mut stream: WriteRequestStreamWrapper> + Unpin>, mut tx: DropCloserWriteHalf, bytes_received: &Arc, + idle_timeout: Duration, ) -> Result<(), Error> { loop { - match stream.next().await { + // Bounds the wait for the *next* `WriteRequest` only, so continuous + // progress never trips it. A client that trickles bytes to hold a store + // slot open is bounded separately, by the store's own total staging + // deadline. Dropping `tx` on return disconnects the compressed channel, + // unwinding the store-side task through the caller's `join`. + let next = match tokio::time::timeout(idle_timeout, stream.next()).await { + Ok(next) => next, + Err(_elapsed) => { + return Err(make_err!( + Code::DeadlineExceeded, + "Compressed upload idle timeout ({}s) elapsed waiting for the next WriteRequest", + idle_timeout.as_secs() + )); + } + }; + match next { Some(Ok(write_request)) => { if write_request.write_offset < 0 { return Err(make_input_err!( @@ -541,6 +596,11 @@ impl ByteStreamServer { } else { config.max_bytes_per_stream }; + let compressed_upload_idle_timeout = if config.compressed_upload_idle_timeout_s == 0 { + DEFAULT_COMPRESSED_UPLOAD_IDLE_TIMEOUT + } else { + Duration::from_secs(config.compressed_upload_idle_timeout_s as u64) + }; let active_uploads: Arc>> = Arc::new(Mutex::new(HashMap::new())); @@ -602,6 +662,7 @@ impl ByteStreamServer { max_bytes_per_stream, active_uploads, idle_stream_timeout, + compressed_upload_idle_timeout, metrics, _sweeper_handle: Arc::new(sweeper_handle), remote_cache_compression_enabled, @@ -1101,6 +1162,69 @@ impl ByteStreamServer { let (bytes_received, _guard) = instance.track_compressed_upload(uuid_key); let (compressed_tx, compressed_rx) = make_buf_channel_pair(); + + // Fast path: when the immediate instance store can accept the negotiated + // wire representation directly, hand it the client's COMPRESSED stream to + // validate/stage/commit byte-for-byte, skipping the decode + re-encode + // round trip. The client stream still runs through + // `process_compressed_client_stream` so `bytes_received` (and therefore + // QueryWriteStatus and `committed_size`) still tracks compressed + // wire-byte progress. + let maybe_wire_store = wire_compressor_capability(wire_compressor) + .and_then(|capability| Some((instance.store.wire_compression_store()?, capability))); + if let Some((wire_store, capability)) = maybe_wire_store { + let update_fut = + wire_store.update_compressed(digest, digest_function, capability, compressed_rx); + let client_stream_fut = process_compressed_client_stream( + stream, + compressed_tx, + &bytes_received, + instance.compressed_upload_idle_timeout, + ); + tokio::pin!(client_stream_fut); + tokio::pin!(update_fut); + let (client_stream_result, update_result) = tokio::select! { + // Preserve a client-side error when both sides are already + // ready, but do not keep pumping an unbounded client stream + // after the store has failed or timed out. Dropping the pump + // closes `compressed_tx`, allowing a detached blocking + // validator to exit while retaining admission until it does. + biased; + client_stream_result = &mut client_stream_fut => { + let update_result = update_fut.await; + (client_stream_result, update_result) + } + update_result = &mut update_fut => match update_result { + Err(err) => return Err(err), + ok @ Ok(_) => { + let client_stream_result = client_stream_fut.await; + (client_stream_result, ok) + } + }, + }; + + if let Some(err) = first_client_attributable_error([ + client_stream_result.as_ref().err(), + update_result.as_ref().err(), + ]) { + return Err(err); + } + let mut upload_error = update_result.err(); + if let Err(err) = client_stream_result { + upload_error = Some(merge_upload_error(upload_error, err)); + } + if let Some(err) = upload_error { + return Err(err); + } + + // `committed_size` stays the compressed wire byte count tracked by + // the atomic, which should agree with the capability's returned wire + // byte count. + let committed_size = i64::try_from(bytes_received.load(Ordering::Acquire)) + .err_tip(|| "Compressed upload size was not convertible to i64")?; + return Ok(Response::new(WriteResponse { committed_size })); + } + let (decompressed_tx, decompressed_rx) = make_buf_channel_pair(); let store = instance.store.clone(); let store_update_context = make_ctx_for_hash_func(digest_function)?; @@ -1126,33 +1250,27 @@ impl ByteStreamServer { digest_function, decompressed_tx, ); - let client_stream_fut = - process_compressed_client_stream(stream, compressed_tx, &bytes_received); + let client_stream_fut = process_compressed_client_stream( + stream, + compressed_tx, + &bytes_received, + instance.compressed_upload_idle_timeout, + ); let (client_stream_result, decode_result, store_update_result) = tokio::join!(client_stream_fut, decode_fut, store_update_fut); - if let Err(err) = &client_stream_result - && err.code == Code::InvalidArgument - { - return Err(err.clone()); - } - if let Err(err) = &decode_result - && err.code == Code::InvalidArgument - { - return Err(err.clone()); + if let Some(err) = first_client_attributable_error([ + client_stream_result.as_ref().err(), + decode_result.as_ref().err(), + ]) { + return Err(err); } let mut upload_error = store_update_result.err(); if let Err(err) = decode_result { - upload_error = Some(match upload_error { - Some(existing) => existing.merge(err), - None => err, - }); + upload_error = Some(merge_upload_error(upload_error, err)); } if let Err(err) = client_stream_result { - upload_error = Some(match upload_error { - Some(existing) => existing.merge(err), - None => err, - }); + upload_error = Some(merge_upload_error(upload_error, err)); } if let Some(err) = upload_error { return Err(err); @@ -1211,6 +1329,8 @@ impl ByteStreamServer { } } + type BoxedResultFuture = Pin> + Send>>; + if read_request.read_limit != 0 { return Err(make_input_err!( "read_limit must be 0 when reading compressed blobs" @@ -1224,31 +1344,63 @@ impl ByteStreamServer { let read_offset = u64::try_from(read_request.read_offset) .err_tip(|| "Could not convert read_offset to u64")?; - let (raw_tx, raw_rx) = make_buf_channel_pair(); - let (compressed_tx, compressed_rx) = make_buf_channel_pair(); + // Fast path: at offset 0, when the immediate instance store already holds + // the negotiated wire representation, pipe its stored stream + // byte-for-byte into the same chunking pipeline, skipping the raw + // `get_part` + re-encode. REAPI defines `read_offset` against the + // *uncompressed* blob, so a resume at offset > 0 (and any instance + // without the capability) keeps the raw read + re-encode path. + let maybe_wire_store = if read_offset == 0 { + wire_compressor_capability(wire_compressor) + .and_then(|capability| Some((instance.store.wire_compression_store()?, capability))) + } else { + None + }; - let store = instance.store.clone(); - let get_part_fut = Box::pin(async move { - store - .get_part(digest, raw_tx, read_offset, None) - .await - .err_tip(|| "Failed to read blob for wire compression") - }); - // The encode runs as a plain async future: it must not occupy a - // blocking-pool thread for the stream's lifetime, because it only - // progresses at the client's drain rate. Dropping the returned - // stream drops this future, which tears the encode down exactly - // like the previous task-abort-on-drop did. - let encode_fut = Box::pin(crate::wire_compression::stream_encode_compressed_download( - raw_rx, - wire_compressor, - crate::wire_compression::ZSTD_COMPRESSION_LEVEL, - compressed_tx, - )); + let (rx, get_part_fut, encode_fut): ( + DropCloserReadHalf, + BoxedResultFuture, + BoxedResultFuture, + ) = if let Some((wire_store, capability)) = maybe_wire_store { + let (zstd_tx, zstd_rx) = make_buf_channel_pair(); + let get_part_fut: BoxedResultFuture = Box::pin(async move { + wire_store + .get_compressed(digest, capability, zstd_tx) + .await + .err_tip(|| "Failed to read stored compressed stream for passthrough") + }); + // No re-encode on the fast path; a ready Ok keeps the existing + // error-merge machinery in `ReaderState::finish` a no-op here. + let encode_fut: BoxedResultFuture = Box::pin(async { Ok(()) }); + (zstd_rx, get_part_fut, encode_fut) + } else { + let (raw_tx, raw_rx) = make_buf_channel_pair(); + let (compressed_tx, compressed_rx) = make_buf_channel_pair(); + let store = instance.store.clone(); + let get_part_fut: BoxedResultFuture = Box::pin(async move { + store + .get_part(digest, raw_tx, read_offset, None) + .await + .err_tip(|| "Failed to read blob for wire compression") + }); + // The encode runs as a plain async future: it must not occupy a + // blocking-pool thread for the stream's lifetime, because it only + // progresses at the client's drain rate. Dropping the returned + // stream drops this future, which tears the encode down exactly + // like the previous task-abort-on-drop did. + let encode_fut: BoxedResultFuture = + Box::pin(crate::wire_compression::stream_encode_compressed_download( + raw_rx, + wire_compressor, + crate::wire_compression::ZSTD_COMPRESSION_LEVEL, + compressed_tx, + )); + (compressed_rx, get_part_fut, encode_fut) + }; let state = Some(ReaderState { max_bytes_per_stream: instance.max_bytes_per_stream, - rx: compressed_rx, + rx, maybe_get_part_result: None, maybe_encode_result: None, get_part_fut, diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 953c2bc6e..6d7ad536e 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -13,6 +13,7 @@ // limitations under the License. use core::convert::Into; +use core::future::Future; use core::pin::{Pin, pin}; use core::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use core::time::Duration; @@ -43,16 +44,18 @@ use nativelink_store::grpc_store::GrpcStore; use nativelink_store::store_manager::StoreManager; use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; -use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc, make_ctx_for_hash_func}; +use nativelink_util::digest_hasher::{ + DigestHasher, DigestHasherFunc, digest_hasher_func_from_context, make_ctx_for_hash_func, +}; use nativelink_util::spawn_blocking; -use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo}; +use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo, WireCompressor}; use opentelemetry::context::FutureExt; use prost::Message; use tokio_util::io::StreamReader; use tonic::{Request, Response, Status}; use tracing::{Instrument, Level, debug, error_span, instrument, warn}; -use crate::wire_compression::RemoteCacheCompressionInstances; +use crate::wire_compression::{RemoteCacheCompressionInstances, wire_compressor_capability}; /// Metrics for the experimental `SplitBlob`/`SpliceBlob` chunking RPCs. /// The split hit rate (`split_hits` / `split_requests_total`) indicates how @@ -183,6 +186,25 @@ type GetTreeStream = Pin> /// Per-blob deadline applied inside `BatchReadBlobs` / `BatchUpdateBlobs`. const BATCH_PER_BLOB_TIMEOUT: Duration = Duration::from_secs(30); +/// Runs one BatchReadBlobs/BatchUpdateBlobs store operation under its individual +/// deadline, preserving the RPC-specific timeout and store-error context. +async fn run_batch_blob_with_timeout( + rpc_name: &str, + digest: DigestInfo, + store_error_context: &'static str, + operation: impl Future>, +) -> Result { + match tokio::time::timeout(BATCH_PER_BLOB_TIMEOUT, operation).await { + Ok(result) => result.err_tip(|| store_error_context), + Err(_elapsed) => Err(make_err!( + Code::DeadlineExceeded, + "{rpc_name} per-blob timeout ({} s) elapsed for digest {}", + BATCH_PER_BLOB_TIMEOUT.as_secs(), + digest, + )), + } +} + /// Maximum size of a single chunk accepted in a `SpliceBlob` request. /// Deliberately looser than the largest chunk the server ever advertises /// (4x the maximum allowed average = 4 MiB) so clients using their own @@ -341,6 +363,17 @@ impl CasServer { let remote_cache_compression_enabled = self .remote_cache_compression_instances .enabled_for(instance_name); + // The digest hasher is set on the ambient context by the + // `batch_update_blobs` wrapper (via `make_ctx_for_hash_func`); resolve + // it once for the whole request so the wire-compression fast path can + // validate the staged blob against the client's declared digest function. + let digest_function = digest_hasher_func_from_context(); + // Detect an immediate wire-compression capability at the instance + // boundary so a client-supplied zstd blob can be passed through + // byte-for-byte (validate + stage + commit) instead of decompressed then + // recompressed. Resolved once and cloned into each per-blob future. + let maybe_wire_store = store.wire_compression_store(); + let maybe_wire_store = &maybe_wire_store; let update_futures: FuturesUnordered<_> = request .requests .into_iter() @@ -351,31 +384,55 @@ impl CasServer { .clone() .err_tip(|| "Digest not found in request")?; let digest_info = DigestInfo::try_from(digest.clone())?; - let size_bytes = usize::try_from(digest_info.size_bytes()) - .err_tip(|| "Digest size_bytes was not convertible to usize")?; - - let store_data = crate::wire_compression::decompress_batch_update( - request.data, - request.compressor, - size_bytes, - remote_cache_compression_enabled, - ) - .await?; + + let request_capability = compressor::Value::try_from(request.compressor) + .ok() + .and_then(wire_compressor_capability); // Apply a per-blob deadline so one slow upload does not // make the whole batch hit the client's overall deadline. - let result = match tokio::time::timeout( - BATCH_PER_BLOB_TIMEOUT, - store_ref.update_oneshot(digest_info, store_data), - ) - .await + // The fast path is only taken when remote cache + // compression is enabled for this instance; otherwise a + // zstd-compressed blob falls through to + // `decompress_batch_update` below, which rejects it the + // same way it would without a wire-compression capability. + let result = if remote_cache_compression_enabled + && let (Some(wire_store), Some(capability)) = + (maybe_wire_store.clone(), request_capability) { - Ok(r) => r.err_tip(|| "Error writing to store"), - Err(_elapsed) => Err(make_err!( - Code::DeadlineExceeded, - "BatchUpdateBlobs per-blob timeout ({} s) elapsed for digest {}", - BATCH_PER_BLOB_TIMEOUT.as_secs(), + // Fast path: hand the compressed bytes straight to the + // wire-capable store, which validates them against the + // digest and stores them verbatim. A malformed stream + // surfaces here as this blob's `InvalidArgument`, isolated + // from siblings. + run_batch_blob_with_timeout( + "BatchUpdateBlobs", digest_info, - )), + "Error writing to store", + wire_store.update_compressed_oneshot( + digest_info, + digest_function, + capability, + request.data, + ), + ) + .await + } else { + let size_bytes = usize::try_from(digest_info.size_bytes()) + .err_tip(|| "Digest size_bytes was not convertible to usize")?; + let store_data = crate::wire_compression::decompress_batch_update( + request.data, + request.compressor, + size_bytes, + remote_cache_compression_enabled, + ) + .await?; + run_batch_blob_with_timeout( + "BatchUpdateBlobs", + digest_info, + "Error writing to store", + store_ref.update_oneshot(digest_info, store_data), + ) + .await }; Ok::<_, Error>(batch_update_blobs_response::Response { digest: Some(digest), @@ -419,6 +476,13 @@ impl CasServer { compressor::Value::try_from(*compressor_i32) .is_ok_and(|compressor| compressor == compressor::Value::Zstd) }); + // Detect an immediate wire-compression capability at the instance + // boundary so stored zstd bytes can be returned byte-for-byte (when the + // client accepts zstd and compression actually helped) with no + // decompress-then-recompress round trip. Resolved once and cloned into + // each per-blob future. + let maybe_wire_store = store.wire_compression_store(); + let maybe_wire_store = &maybe_wire_store; let read_futures: FuturesUnordered<_> = request .digests .into_iter() @@ -428,48 +492,68 @@ impl CasServer { // TODO(palfrey) There is a security risk here of someone taking all the memory on the instance. // Apply a per-blob deadline so one slow read does not // make the whole batch hit the client's overall deadline. - let result = match tokio::time::timeout( - BATCH_PER_BLOB_TIMEOUT, - store_ref.get_part_unchunked(digest_copy, 0, None), - ) - .await - { - Ok(r) => r.err_tip(|| "Error reading from store"), - Err(_elapsed) => Err(make_err!( - Code::DeadlineExceeded, - "BatchReadBlobs per-blob timeout ({} s) elapsed for digest {}", - BATCH_PER_BLOB_TIMEOUT.as_secs(), - digest_copy, - )), - }; - let (status, data, response_compressor) = match result { - Ok(raw_data) => { - let (output_data, chosen_compressor) = if client_accepts_zstd { - let data_for_compression = raw_data.clone(); - match spawn_blocking!("cas_encode_compressed_download", move || { - crate::wire_compression::compress_for_batch_read( - data_for_compression, - ) - }) - .await - { - Ok(compressed_result) => compressed_result, - Err(e) => { - warn!("Wire compression task failed for digest {:?}, falling back to identity: {}", digest, e); - (raw_data, compressor::Value::Identity) + let outcome: Result<(Bytes, compressor::Value), Error> = + if let Some(wire_store) = maybe_wire_store.clone() { + // Fast path: the wire-capable store chooses between its + // stored zstd bytes (when the client accepts zstd and they + // are smaller than the raw size) and the decoded raw bytes, + // avoiding a decompress-then-recompress round trip. + let acceptable_compressors = + client_accepts_zstd.then_some(WireCompressor::Zstd); + run_batch_blob_with_timeout( + "BatchReadBlobs", + digest_copy, + "Error reading from store", + wire_store + .get_for_batch(digest_copy, acceptable_compressors.as_slice()), + ) + .await + .map(|(data, maybe_compressor)| { + let chosen = match maybe_compressor { + Some(WireCompressor::Zstd) => compressor::Value::Zstd, + None => compressor::Value::Identity, + }; + (data, chosen) + }) + } else { + let raw = run_batch_blob_with_timeout( + "BatchReadBlobs", + digest_copy, + "Error reading from store", + store_ref.get_part_unchunked(digest_copy, 0, None), + ) + .await; + match raw { + Ok(raw_data) if client_accepts_zstd => { + let data_for_compression = raw_data.clone(); + match spawn_blocking!("cas_encode_compressed_download", move || { + crate::wire_compression::compress_for_batch_read( + data_for_compression, + ) + }) + .await + { + Ok(compressed_result) => Ok(compressed_result), + Err(e) => { + warn!("Wire compression task failed for digest {:?}, falling back to identity: {}", digest, e); + Ok((raw_data, compressor::Value::Identity)) + } } } - } else { - (raw_data, compressor::Value::Identity) - }; - - (GrpcStatus::default(), output_data, chosen_compressor) + Ok(raw_data) => Ok((raw_data, compressor::Value::Identity)), + Err(e) => Err(e), + } + }; + let (status, data, response_compressor) = match outcome { + Ok((data, chosen_compressor)) => { + (GrpcStatus::default(), data, chosen_compressor) } Err(mut e) => { if e.code == Code::NotFound { - // Trim the error code. Not Found is quite common and we don't want to send a large - // error (debug) message for something that is common. We resize to just the last - // message as it will be the most relevant. + // Trim the message chain. Not Found is quite common and + // we don't want to send a large error (debug) message + // for something that is common. We resize to just the + // last message as it will be the most relevant. e.messages.resize_with(1, String::new); } (e.into(), Bytes::new(), compressor::Value::Identity) diff --git a/nativelink-service/src/wire_compression.rs b/nativelink-service/src/wire_compression.rs index 12ccd79f3..a921c6513 100644 --- a/nativelink-service/src/wire_compression.rs +++ b/nativelink-service/src/wire_compression.rs @@ -16,7 +16,14 @@ //! //! This module handles compression/decompression of blob data on the gRPC wire //! between client and server, per the REAPI compressed-blobs specification. -//! This is orthogonal to at-rest compression (`CompressionStore` with LZ4). +//! +//! Wire compression is independent of at-rest compression, but not always +//! disjoint from it: a store whose physical representation already matches the +//! negotiated wire representation can serve and accept those bytes directly. +//! Such a store advertises itself through +//! [`WireCompressionStore`](nativelink_util::store_trait::WireCompressionStore), +//! and the `ByteStream`/CAS services use that capability to skip the codecs +//! here. See [`wire_compressor_capability`]. use std::collections::HashSet; @@ -29,7 +36,7 @@ use nativelink_util::spawn_blocking; // existing service callers keep their import paths. pub use nativelink_util::wire_compression::{ ZSTD_COMPRESSION_LEVEL, compress, decompress, stream_decode_compressed_upload, - stream_encode_compressed_download, + stream_encode_compressed_download, wire_compressor_capability, }; use tracing::warn; diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index 20da09e6b..cbc8b50ad 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -24,7 +24,9 @@ use hyper_util::rt::TokioIo; use hyper_util::server::conn::auto; use hyper_util::service::TowerToHyperService; use nativelink_config::cas_server::{ByteStreamConfig, HttpListener, WithInstanceName}; -use nativelink_config::stores::{MemorySpec, StoreSpec, VerifySpec}; +use nativelink_config::stores::{ + CompressionAlgorithm, CompressionSpec, MemorySpec, StoreSpec, VerifySpec, ZstdConfig, +}; use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient; @@ -37,9 +39,9 @@ use nativelink_service::wire_compression::RemoteCacheCompressionInstances; use nativelink_store::default_store_factory::store_factory; use nativelink_store::store_manager::StoreManager; use nativelink_util::channel_body_for_tests::ChannelBody; -use nativelink_util::common::{DigestInfo, encode_stream_proto}; +use nativelink_util::common::{DigestInfo, encode_stream_proto, make_temp_path}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; -use nativelink_util::store_trait::StoreLike; +use nativelink_util::store_trait::{StoreLike, WireCompressionStore, WireCompressor}; use nativelink_util::task::JoinHandleDropGuard; use nativelink_util::{background_spawn, spawn}; use pretty_assertions::assert_eq; @@ -91,6 +93,39 @@ async fn make_verify_store_manager() -> Result, Error> { Ok(store_manager) } +/// Builds a [`StoreManager`] whose `main_cas` store is a pure-passthrough +/// zstd compression over an in-memory backend, so the `ByteStream` zstd fast +/// paths are exercised end-to-end. +async fn make_zstd_store_manager() -> Result, Error> { + let store_manager = Arc::new(StoreManager::new()); + let temp_path = make_temp_path("bytestream-zstd-store"); + std::fs::create_dir_all(&temp_path) + .map_err(|e| make_err!(Code::Internal, "Failed to create zstd temp dir: {e}"))?; + let spec = StoreSpec::Compression(Box::new(CompressionSpec { + backend: StoreSpec::Memory(MemorySpec::default()), + compression_algorithm: CompressionAlgorithm::Zstd(ZstdConfig { + temp_path, + max_compressed_upload_size: 512 * 1024 * 1024, + ..ZstdConfig::default() + }), + })); + store_manager.add_store( + "main_cas", + store_factory(&spec, &store_manager, None).await?, + )?; + Ok(store_manager) +} + +/// Gets the configured store's immediate wire-compression capability for +/// direct fast-path setup without coupling this service test to a store type. +fn wire_compression_store_of(store_manager: &StoreManager) -> Arc { + store_manager + .get_store("main_cas") + .expect("main_cas store missing") + .wire_compression_store() + .expect("main_cas store was not wire-compression capable") +} + fn make_bytestream_server( store_manager: &StoreManager, config: Option>>, @@ -109,6 +144,7 @@ fn make_bytestream_server_with_remote_cache_compression( config: ByteStreamConfig { cas_store: "main_cas".to_string(), persist_stream_on_disconnect_timeout_s: 0, + compressed_upload_idle_timeout_s: 0, max_bytes_per_stream: 1024, }, }] @@ -2034,3 +2070,599 @@ async fn uuid_collision_does_not_deadlock() -> Result<(), Box Result<(), Box> +{ + let raw_data = Bytes::from("byte-for-byte zstd passthrough ".repeat(256)); + let hash = sha256_hex(raw_data.as_ref()); + // A known zstd stream. With compression_level=None + max_recompression_size=0 + // the store is a pure passthrough, so these exact bytes are what is stored. + let compressed = Bytes::from(zstd::bulk::compress(raw_data.as_ref(), 3)?); + assert_ne!( + compressed.as_ref(), + raw_data.as_ref(), + "test data must actually compress" + ); + + let store_manager = make_zstd_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server_with_remote_cache_compression(store_manager.as_ref(), None, true) + .expect("Failed to make server"), + ); + let wire_store = wire_compression_store_of(store_manager.as_ref()); + + let digest = DigestInfo::try_new(&hash, raw_data.len())?; + wire_store + .update_compressed_oneshot( + digest, + DigestHasherFunc::Sha256, + WireCompressor::Zstd, + compressed.clone(), + ) + .await?; + + let read_data = read_all_bytes( + bs_server.as_ref(), + ReadRequest { + resource_name: format!( + "{}/compressed-blobs/zstd/{}/{}", + INSTANCE_NAME, + hash, + raw_data.len() + ), + read_offset: 0, + read_limit: 0, + }, + ) + .await?; + + assert_eq!( + read_data.as_slice(), + compressed.as_ref(), + "offset-0 compressed read from a ZstdStore must be byte-for-byte identical to what was stored" + ); + let decoded = zstd::bulk::decompress(&read_data, raw_data.len())?; + assert_eq!( + decoded.as_slice(), + raw_data.as_ref(), + "returned zstd stream must decode to the original blob" + ); + + Ok(()) +} + +/// A compressed upload to a `ZstdStore` instance stores the client's zstd bytes +/// byte-for-byte, reports the compressed wire byte count as `committed_size`, and +/// a re-upload of an already-present blob still completes. +#[nativelink_test] +pub async fn zstd_store_compressed_write_round_trip() -> Result<(), Box> { + let raw_data = "zstd store compressed write round trip ".repeat(256); + let hash = sha256_hex(raw_data.as_bytes()); + let compressed = zstd::bulk::compress(raw_data.as_bytes(), 3)?; + assert!(compressed.len() < raw_data.len(), "test data must compress"); + + let store_manager = make_zstd_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server_with_remote_cache_compression(store_manager.as_ref(), None, true) + .expect("Failed to make server"), + ); + + // Upload the compressed blob. + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server.clone(), None); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_compressed_resource_name( + "4dcec57e-1389-4ab5-b188-4a59f22ceb60", + &hash, + raw_data.len(), + ), + write_offset: 0, + finish_write: true, + data: compressed.clone().into(), + })?)) + .await?; + let server_result = join_handle + .await + .expect("Failed to join") + .expect("Failed write"); + assert_eq!( + server_result.into_inner().committed_size, + i64::try_from(compressed.len()).unwrap(), + "compressed write to a ZstdStore must report the compressed wire byte count" + ); + + // A subsequent compressed read returns the exact stored zstd bytes. + let read_data = read_all_bytes( + bs_server.as_ref(), + ReadRequest { + resource_name: format!( + "{}/compressed-blobs/zstd/{}/{}", + INSTANCE_NAME, + hash, + raw_data.len() + ), + read_offset: 0, + read_limit: 0, + }, + ) + .await?; + assert_eq!( + read_data.as_slice(), + compressed.as_slice(), + "read-back must be byte-for-byte identical to the uploaded zstd stream" + ); + assert_eq!( + zstd::bulk::decompress(&read_data, raw_data.len())?.as_slice(), + raw_data.as_bytes() + ); + + // Re-upload an already-present blob; it must still complete (REAPI allows + // -1 or the byte count for an already-present blob). + let (tx2, join2) = make_stream_and_writer_spawn(bs_server.clone(), None); + tx2.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: make_compressed_resource_name( + "4dcec57e-1389-4ab5-b188-4a59f22ceb61", + &hash, + raw_data.len(), + ), + write_offset: 0, + finish_write: true, + data: compressed.clone().into(), + })?)) + .await?; + let committed_size = join2 + .await + .expect("Failed to join") + .expect("Failed re-upload write") + .into_inner() + .committed_size; + assert!( + committed_size == -1 || committed_size == i64::try_from(compressed.len()).unwrap(), + "re-upload committed_size must be -1 or the compressed byte count {}; got {committed_size}", + compressed.len() + ); + + Ok(()) +} + +/// A compressed read at offset > 0 from a `ZstdStore` instance takes the +/// fallback path (`StoreDriver::get_part` decompresses, service re-encodes) and +/// returns a valid zstd stream decoding to the uncompressed suffix. +#[nativelink_test] +pub async fn zstd_store_compressed_read_offset_uses_fallback() +-> Result<(), Box> { + let raw_data = Bytes::from( + (0usize..4096) + .map(|i| u8::try_from((i * 31 + i / 7) % 251).expect("modulo 251 fits in u8")) + .collect::>(), + ); + let hash = sha256_hex(raw_data.as_ref()); + let compressed = Bytes::from(zstd::bulk::compress(raw_data.as_ref(), 3)?); + + let store_manager = make_zstd_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server_with_remote_cache_compression(store_manager.as_ref(), None, true) + .expect("Failed to make server"), + ); + let wire_store = wire_compression_store_of(store_manager.as_ref()); + + let digest = DigestInfo::try_new(&hash, raw_data.len())?; + wire_store + .update_compressed_oneshot( + digest, + DigestHasherFunc::Sha256, + WireCompressor::Zstd, + compressed, + ) + .await?; + + let read_offset = 100usize; + let ranged_data = read_all_bytes( + bs_server.as_ref(), + ReadRequest { + resource_name: format!( + "{}/compressed-blobs/zstd/{}/{}", + INSTANCE_NAME, + hash, + raw_data.len() + ), + read_offset: i64::try_from(read_offset).unwrap(), + read_limit: 0, + }, + ) + .await?; + + let decoded = zstd::bulk::decompress(&ranged_data, raw_data.len() - read_offset)?; + assert_eq!( + decoded.as_slice(), + &raw_data.as_ref()[read_offset..], + "offset > 0 compressed read must decode to the uncompressed suffix" + ); + + Ok(()) +} + +/// An identity (non-compressed) write + read round trip through a `ZstdStore` +/// instance returns the correct raw bytes via the `StoreDriver` path. +#[nativelink_test] +pub async fn zstd_store_identity_write_and_read_round_trip() +-> Result<(), Box> { + let raw_data = Bytes::from("identity round trip through a zstd store ".repeat(64)); + let hash = sha256_hex(raw_data.as_ref()); + + let store_manager = make_zstd_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), + ); + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server.clone(), None); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: format!( + "{}/uploads/{}/blobs/{}/{}", + INSTANCE_NAME, + "4dcec57e-1389-4ab5-b188-4a59f22ceb62", + hash, + raw_data.len() + ), + write_offset: 0, + finish_write: true, + data: raw_data.clone(), + })?)) + .await?; + let server_result = join_handle + .await + .expect("Failed to join") + .expect("Failed write"); + assert_eq!( + server_result.into_inner().committed_size, + i64::try_from(raw_data.len()).unwrap() + ); + + let read_data = read_all_bytes( + bs_server.as_ref(), + ReadRequest { + resource_name: format!("{}/blobs/{}/{}", INSTANCE_NAME, hash, raw_data.len()), + read_offset: 0, + read_limit: i64::try_from(raw_data.len()).unwrap(), + }, + ) + .await?; + assert_eq!( + read_data.as_slice(), + raw_data.as_ref(), + "identity round trip through a ZstdStore must return the raw bytes" + ); + + Ok(()) +} + +/// A compressed read with a non-zero `read_limit` against a `ZstdStore` instance +/// must be rejected with `INVALID_ARGUMENT` (regression: the guard runs before the +/// fast path). +#[nativelink_test] +pub async fn zstd_store_compressed_read_rejects_nonzero_read_limit() +-> Result<(), Box> { + let raw_data = Bytes::from("read_limit rejection over a zstd store ".repeat(64)); + let hash = sha256_hex(raw_data.as_ref()); + let compressed = Bytes::from(zstd::bulk::compress(raw_data.as_ref(), 3)?); + + let store_manager = make_zstd_store_manager().await?; + let bs_server = + make_bytestream_server_with_remote_cache_compression(store_manager.as_ref(), None, true) + .expect("Failed to make server"); + let wire_store = wire_compression_store_of(store_manager.as_ref()); + + let digest = DigestInfo::try_new(&hash, raw_data.len())?; + wire_store + .update_compressed_oneshot( + digest, + DigestHasherFunc::Sha256, + WireCompressor::Zstd, + compressed, + ) + .await?; + + let Err(status) = bs_server + .read(Request::new(ReadRequest { + resource_name: format!( + "{}/compressed-blobs/zstd/{}/{}", + INSTANCE_NAME, + hash, + raw_data.len() + ), + read_offset: 0, + read_limit: 1, + })) + .await + else { + panic!("compressed read with read_limit should fail"); + }; + + assert_eq!(status.code(), Code::InvalidArgument); + assert!( + status.message().contains("read_limit must be 0"), + "unexpected error: {}", + status.message() + ); + + Ok(()) +} + +/// Builds a [`StoreManager`] whose `main_cas` store is a pure-passthrough +/// `ZstdStore` over memory with a configurable staged-upload bound. +async fn make_zstd_store_manager_with_staging( + max_concurrent_staged_uploads: usize, + stage_timeout_s: u64, +) -> Result, Error> { + let store_manager = Arc::new(StoreManager::new()); + let temp_path = make_temp_path("bytestream-zstd-idle"); + std::fs::create_dir_all(&temp_path) + .map_err(|e| make_err!(Code::Internal, "Failed to create zstd temp dir: {e}"))?; + let spec = StoreSpec::Compression(Box::new(CompressionSpec { + backend: StoreSpec::Memory(MemorySpec::default()), + compression_algorithm: CompressionAlgorithm::Zstd(ZstdConfig { + temp_path, + max_compressed_upload_size: 512 * 1024 * 1024, + max_concurrent_staged_uploads, + stage_timeout_s, + ..ZstdConfig::default() + }), + })); + store_manager.add_store( + "main_cas", + store_factory(&spec, &store_manager, None).await?, + )?; + Ok(store_manager) +} + +/// A compressed `ByteStream` client that acquires the single staging slot and then +/// stalls (never sends `finish_write`, never disconnects) is timed out with +/// `DeadlineExceeded`, its staging permit is released, and a subsequent valid +/// upload behind the `max_concurrent_staged_uploads = 1` bound completes. +/// +/// Uses a short real-time (1s) idle timeout rather than a paused clock: the +/// stalled upload parks a `spawn_blocking` recv thread inside the store, and +/// tokio's `start_paused` auto-advance will not fire the idle timer while a +/// blocking task is outstanding — so a paused clock would deadlock here. +#[nativelink_test] +async fn zstd_compressed_upload_idle_timeout_frees_staging_slot() +-> Result<(), Box> { + let store_manager = make_zstd_store_manager_with_staging(1, 0).await?; + // Short idle timeout; under start_paused the virtual clock advances past it. + let config = vec![WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + persist_stream_on_disconnect_timeout_s: 1, + // The knob under test: how long to wait for the next WriteRequest + // of a compressed upload. + compressed_upload_idle_timeout_s: 1, + max_bytes_per_stream: 1024, + }, + }]; + let bs_server = Arc::new(make_bytestream_server_with_remote_cache_compression( + store_manager.as_ref(), + Some(config), + true, + )?); + + let raw = "idle timeout regression payload ".repeat(64); + let compressed = zstd::bulk::compress(raw.as_bytes(), 3)?; + let hash = sha256_hex(raw.as_bytes()); + + // --- Upload A: send a partial chunk, then stall (keep tx alive). --- + let resource_a = + make_compressed_resource_name("aaaaaaaa-1111-2222-3333-444444444444", &hash, raw.len()); + let (tx_a, join_a) = make_stream_and_writer_spawn(bs_server.clone(), None); + tx_a.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: resource_a, + write_offset: 0, + finish_write: false, + data: Bytes::copy_from_slice(&compressed[..compressed.len() / 2]), + })?)) + .await?; + + // Do not send more and do not drop tx_a: the server is left waiting for the + // next WriteRequest. The idle timeout must fire (DeadlineExceeded), not an + // InvalidArgument "stream ended without finish_write". + let status_a = join_a + .await + .expect("join A") + .expect_err("stalled upload A must fail"); + assert_eq!( + status_a.code(), + Code::DeadlineExceeded, + "an idle-stalled compressed upload must time out with DeadlineExceeded, got: {status_a:?}" + ); + drop(tx_a); + + // --- Upload B: a full valid upload must complete (slot was released). --- + let resource_b = + make_compressed_resource_name("bbbbbbbb-1111-2222-3333-444444444444", &hash, raw.len()); + let (tx_b, join_b) = make_stream_and_writer_spawn(bs_server.clone(), None); + tx_b.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: resource_b, + write_offset: 0, + finish_write: true, + data: compressed.clone().into(), + })?)) + .await?; + join_b + .await + .expect("join B") + .expect("valid upload B behind the freed slot must succeed"); + + let store = store_manager.get_store("main_cas").unwrap(); + let digest = DigestInfo::try_new(&hash, raw.len())?; + assert_eq!( + store.get_part_unchunked(digest, 0, None).await?.as_ref(), + raw.as_bytes(), + "the recovered upload must store the decompressed bytes" + ); + Ok(()) +} + +/// A client that keeps making progress is NOT idle-timed-out even when the whole +/// upload takes longer than the idle timeout: the deadline is per-WriteRequest +/// and resets on every chunk. Real-time test with sub-timeout inter-chunk gaps +/// whose sum exceeds the timeout. +#[nativelink_test] +async fn zstd_compressed_upload_making_progress_is_not_idle_timed_out() +-> Result<(), Box> { + let store_manager = make_zstd_store_manager_with_staging(1, 0).await?; + let config = vec![WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + persist_stream_on_disconnect_timeout_s: 1, + // The knob under test: how long to wait for the next WriteRequest + // of a compressed upload. + compressed_upload_idle_timeout_s: 1, + max_bytes_per_stream: 1024, + }, + }]; + let bs_server = Arc::new(make_bytestream_server_with_remote_cache_compression( + store_manager.as_ref(), + Some(config), + true, + )?); + + let raw = "progress keeps the stream alive ".repeat(256); + let compressed = zstd::bulk::compress(raw.as_bytes(), 3)?; + let hash = sha256_hex(raw.as_bytes()); + let resource_name = + make_compressed_resource_name("cccccccc-1111-2222-3333-444444444444", &hash, raw.len()); + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server.clone(), None); + + // Five chunks, ~300ms apart: each inter-chunk gap is well under the 1s idle + // timeout, but their sum (~1.2s) exceeds it. A per-message deadline that + // resets on progress lets this succeed; a whole-upload deadline would not. + let chunk_count = 5usize; + let chunk_len = compressed.len().div_ceil(chunk_count); + let mut offset = 0usize; + while offset < compressed.len() { + let end = (offset + chunk_len).min(compressed.len()); + tx.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: resource_name.clone(), + write_offset: i64::try_from(offset).unwrap(), + finish_write: end == compressed.len(), + data: Bytes::copy_from_slice(&compressed[offset..end]), + })?)) + .await?; + offset = end; + if offset < compressed.len() { + tokio::time::sleep(core::time::Duration::from_millis(300)).await; + } + } + + join_handle + .await + .expect("join") + .expect("a continuously-progressing upload must not be idle-timed-out"); + + let store = store_manager.get_store("main_cas").unwrap(); + let digest = DigestInfo::try_new(&hash, raw.len())?; + assert_eq!( + store.get_part_unchunked(digest, 0, None).await?.as_ref(), + raw.as_bytes(), + "the progressive upload must store the decompressed bytes" + ); + Ok(()) +} + +/// A total `ZstdStore` staging deadline must terminate the `ByteStream` RPC even +/// while the client stays connected and continues sending within the service's +/// per-message idle deadline. Terminating the pump closes the store channel so +/// its detached validator can release the staging slot. +#[nativelink_test] +async fn zstd_stage_timeout_stops_a_progressing_client_and_frees_the_slot() +-> Result<(), Box> { + let store_manager = make_zstd_store_manager_with_staging(1, 1).await?; + let config = vec![WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: ByteStreamConfig { + cas_store: "main_cas".to_string(), + persist_stream_on_disconnect_timeout_s: 1, + compressed_upload_idle_timeout_s: 5, + max_bytes_per_stream: 1024, + }, + }]; + let bs_server = Arc::new(make_bytestream_server_with_remote_cache_compression( + store_manager.as_ref(), + Some(config), + true, + )?); + + // A long valid frame ensures the byte-at-a-time prefix remains incomplete + // until the store's total staging deadline fires. + let mut state = 0x9E37_79B9_u32; + let raw = (0..256 * 1024) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 24) as u8 + }) + .collect::>(); + let compressed = zstd::bulk::compress(&raw, 3)?; + assert!(compressed.len() > 100); + let hash = sha256_hex(&raw); + + let resource_a = + make_compressed_resource_name("dddddddd-1111-2222-3333-444444444444", &hash, raw.len()); + let (tx_a, join_a) = make_stream_and_writer_spawn(bs_server.clone(), None); + let trickle = Bytes::from(compressed.clone()); + let sender = spawn!("bytestream_zstd_stage_timeout_trickle", async move { + let mut offset = 0usize; + while offset < trickle.len() { + if tx_a + .send(Frame::data( + encode_stream_proto(&WriteRequest { + resource_name: resource_a.clone(), + write_offset: i64::try_from(offset).unwrap(), + finish_write: false, + data: trickle.slice(offset..=offset), + }) + .unwrap(), + )) + .await + .is_err() + { + return; + } + offset += 1; + tokio::time::sleep(core::time::Duration::from_millis(100)).await; + } + }); + + let status_a = tokio::time::timeout(core::time::Duration::from_secs(3), join_a) + .await + .expect("the store timeout must stop the client pump before the idle timeout") + .expect("join A") + .expect_err("the progressively trickling upload must hit the total staging deadline"); + assert_eq!(status_a.code(), Code::DeadlineExceeded, "got: {status_a:?}"); + drop(sender); + + // The failed RPC dropped the pump, so the blocking validator can finish + // cleanup and a valid upload behind the single staging slot can proceed. + let resource_b = + make_compressed_resource_name("eeeeeeee-1111-2222-3333-444444444444", &hash, raw.len()); + let (tx_b, join_b) = make_stream_and_writer_spawn(bs_server, None); + tx_b.send(Frame::data(encode_stream_proto(&WriteRequest { + resource_name: resource_b, + write_offset: 0, + finish_write: true, + data: compressed.into(), + })?)) + .await?; + tokio::time::timeout(core::time::Duration::from_secs(3), join_b) + .await + .expect("the valid upload must acquire the slot after cleanup") + .expect("join B") + .expect("valid upload B must succeed"); + Ok(()) +} diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index 055bfe9a8..b68b634d9 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -21,7 +21,7 @@ use async_trait::async_trait; use bytes::Bytes; use futures::StreamExt; use nativelink_config::cas_server::WithInstanceName; -use nativelink_config::stores::{MemorySpec, StoreSpec}; +use nativelink_config::stores::{MemorySpec, StoreSpec, ZstdConfig}; use nativelink_error::Error; use nativelink_macro::nativelink_test; use nativelink_metric::MetricsComponent; @@ -38,9 +38,11 @@ use nativelink_service::cas_server::CasServer; use nativelink_service::wire_compression::RemoteCacheCompressionInstances; use nativelink_store::ac_utils::serialize_and_upload_message; use nativelink_store::default_store_factory::store_factory; +use nativelink_store::memory_store::MemoryStore; use nativelink_store::store_manager::StoreManager; +use nativelink_store::zstd_store::ZstdStore; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; -use nativelink_util::common::DigestInfo; +use nativelink_util::common::{DigestInfo, make_temp_path}; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; use nativelink_util::store_trait::{ @@ -1031,6 +1033,387 @@ async fn batch_read_blobs_zstd_falls_back_to_identity_when_not_smaller() Ok(()) } +/// A passthrough zstd configuration over an in-memory backend with no +/// recompression. +fn zstd_instance_spec(temp_path: String) -> ZstdConfig { + ZstdConfig { + temp_path, + max_compressed_upload_size: 512 * 1024 * 1024, + ..ZstdConfig::default() + } +} + +/// Builds a zstd-compressing `CasServer` over memory with wire compression +/// enabled, returning the concrete store so tests can seed it directly. +fn make_zstd_instance_cas_server() -> Result<(CasServer, Arc), Error> { + let temp_path = make_temp_path("cas_server_zstd_instance"); + std::fs::create_dir_all(&temp_path).expect("create temp dir"); + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let zstd_store = ZstdStore::new(&zstd_instance_spec(temp_path), inner)?; + let store_manager = Arc::new(StoreManager::new()); + store_manager.add_store("main_cas", Store::new(zstd_store.clone()))?; + let cas_server = make_cas_server_with_zstd(&store_manager)?; + Ok((cas_server, zstd_store)) +} + +fn sha256_digest(data: &[u8]) -> (DigestInfo, Digest) { + let mut hasher = DigestHasherFunc::Sha256.hasher(); + DigestHasher::update(&mut hasher, data); + let digest_info = hasher.finalize_digest(); + let digest = Digest::from(&digest_info); + (digest_info, digest) +} + +/// High-entropy (incompressible) bytes generated by a deterministic xorshift. +fn incompressible_bytes(len: usize) -> Vec { + let mut state: u64 = 0x2545_F491_4F6C_DD1D; + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state.to_le_bytes()[0] + }) + .collect() +} + +#[nativelink_test] +async fn batch_read_zstd_instance_passthrough_and_identity() +-> Result<(), Box> { + let (cas_server, zstd_store) = make_zstd_instance_cas_server()?; + + // Compressible blob: seed the store with the client's zstd bytes verbatim + // (level 19, which the passthrough store keeps as-is). + let raw: Vec = "hello world ".repeat(100).into_bytes(); + let (raw_di, raw_digest) = sha256_digest(&raw); + let stored_zstd = zstd::bulk::compress(&raw, 19)?; + zstd_store + .update_zstd_oneshot( + raw_di, + DigestHasherFunc::Sha256, + Bytes::from(stored_zstd.clone()), + ) + .await?; + + // Incompressible blob: store it via the identity path so the physical zstd + // bytes are larger than the raw size. + let incompressible = incompressible_bytes(8192); + let (inc_di, inc_digest) = sha256_digest(&incompressible); + Store::new(zstd_store.clone()) + .update_oneshot(inc_di, Bytes::from(incompressible.clone())) + .await?; + + let response = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![raw_digest.clone(), inc_digest.clone()], + acceptable_compressors: vec![compressor::Value::Zstd.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + + assert_eq!(response.responses.len(), 2); + let compressible = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&raw_digest)) + .expect("compressible response present"); + assert_eq!(compressible.compressor, compressor::Value::Zstd as i32); + // Byte-for-byte passthrough of the stored zstd stream, no recompression. + assert_eq!(compressible.data.as_ref(), stored_zstd.as_slice()); + assert_eq!( + zstd::bulk::decompress(&compressible.data, raw.len())?, + raw, + "stored zstd must decode to the original bytes" + ); + + let incompressible_resp = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&inc_digest)) + .expect("incompressible response present"); + assert_eq!( + incompressible_resp.compressor, + compressor::Value::Identity as i32, + "incompressible blob must not regress to a larger zstd response" + ); + assert_eq!(incompressible_resp.data.as_ref(), incompressible.as_slice()); + Ok(()) +} + +#[nativelink_test] +async fn batch_read_zstd_instance_error_isolation() -> Result<(), Box> { + let (cas_server, zstd_store) = make_zstd_instance_cas_server()?; + + let raw: Vec = "compress me ".repeat(64).into_bytes(); + let (raw_di, present_digest) = sha256_digest(&raw); + let stored_zstd = zstd::bulk::compress(&raw, 3)?; + zstd_store + .update_zstd_oneshot(raw_di, DigestHasherFunc::Sha256, Bytes::from(stored_zstd)) + .await?; + + let absent_digest = Digest { + hash: HASH3.to_string(), + size_bytes: 5, + }; + + let response = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![present_digest.clone(), absent_digest.clone()], + acceptable_compressors: vec![compressor::Value::Zstd.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + + assert_eq!(response.responses.len(), 2); + let present = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&present_digest)) + .expect("present response"); + assert_eq!(present.status.as_ref().map(|s| s.code), Some(0)); + assert_eq!(present.compressor, compressor::Value::Zstd as i32); + + let absent = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&absent_digest)) + .expect("absent response"); + let status = absent.status.as_ref().expect("status set"); + assert_eq!(status.code, Code::NotFound as i32); + assert_eq!( + status.message, + format!( + "Key {:?} not found", + StoreKey::from(DigestInfo::try_from(absent_digest)?) + ), + "NotFound message should be trimmed to a single message" + ); + assert!(absent.data.is_empty()); + Ok(()) +} + +#[nativelink_test] +async fn batch_update_zstd_instance_roundtrip_and_corrupt_isolation() +-> Result<(), Box> { + let (cas_server, _zstd_store) = make_zstd_instance_cas_server()?; + + // Valid zstd entry. + let good_raw: Vec = "round trip ".repeat(80).into_bytes(); + let (_good_di, good_digest) = sha256_digest(&good_raw); + let good_compressed = zstd::bulk::compress(&good_raw, 3)?; + + // Corrupt entry: valid zstd of some bytes but with a digest whose hash does + // not match the content (decodes to the wrong hash). + let bad_raw = b"totally different content".to_vec(); + let bad_compressed = zstd::bulk::compress(&bad_raw, 3)?; + let bad_digest = Digest { + hash: HASH1.to_string(), + size_bytes: i64::try_from(bad_raw.len()).unwrap(), + }; + + let response = cas_server + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + requests: vec![ + batch_update_blobs_request::Request { + digest: Some(good_digest.clone()), + data: good_compressed.into(), + compressor: compressor::Value::Zstd.into(), + }, + batch_update_blobs_request::Request { + digest: Some(bad_digest.clone()), + data: bad_compressed.into(), + compressor: compressor::Value::Zstd.into(), + }, + ], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + + assert_eq!(response.responses.len(), 2); + let good = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&good_digest)) + .expect("good response"); + assert_eq!( + good.status.as_ref().map(|s| s.code), + Some(0), + "valid zstd entry should succeed: {:?}", + good.status + ); + let bad = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&bad_digest)) + .expect("bad response"); + assert_eq!( + bad.status.as_ref().map(|s| s.code), + Some(Code::InvalidArgument as i32), + "corrupt entry should be InvalidArgument, got {:?}", + bad.status + ); + + // The valid entry round-trips byte-for-byte on read. + let read = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![good_digest.clone()], + acceptable_compressors: vec![compressor::Value::Identity.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + assert_eq!(read.responses.len(), 1); + assert_eq!(read.responses[0].data.as_ref(), good_raw.as_slice()); + Ok(()) +} + +#[nativelink_test] +async fn batch_update_zstd_instance_identity_entry_roundtrips() +-> Result<(), Box> { + let (cas_server, _zstd_store) = make_zstd_instance_cas_server()?; + + let raw: Vec = "identity entry ".repeat(50).into_bytes(); + let (_di, digest) = sha256_digest(&raw); + + let response = cas_server + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + requests: vec![batch_update_blobs_request::Request { + digest: Some(digest.clone()), + data: raw.clone().into(), + compressor: compressor::Value::Identity.into(), + }], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + assert_eq!(response.responses.len(), 1); + assert_eq!( + response.responses[0].status.as_ref().map(|s| s.code), + Some(0) + ); + + let read = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![digest.clone()], + acceptable_compressors: vec![compressor::Value::Identity.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + assert_eq!(read.responses.len(), 1); + assert_eq!(read.responses[0].data.as_ref(), raw.as_slice()); + Ok(()) +} + +/// Builds the same zstd-compressing server with wire compression disabled. +fn make_zstd_instance_cas_server_compression_disabled() -> Result<(CasServer, Arc), Error> +{ + let temp_path = make_temp_path("cas_server_zstd_instance_disabled"); + std::fs::create_dir_all(&temp_path).expect("create temp dir"); + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let zstd_store = ZstdStore::new(&zstd_instance_spec(temp_path), inner)?; + let store_manager = Arc::new(StoreManager::new()); + store_manager.add_store("main_cas", Store::new(zstd_store.clone()))?; + let cas_server = make_cas_server(&store_manager)?; + Ok((cas_server, zstd_store)) +} + +#[nativelink_test] +async fn batch_update_zstd_instance_rejected_when_remote_cache_compression_disabled() +-> Result<(), Box> { + // The wire-compression capability must not bypass the instance-level gate. + let (cas_server, _zstd_store) = make_zstd_instance_cas_server_compression_disabled()?; + + let raw_data = b"zstd disabled batch update to zstd-backed instance"; + let compressed_data = zstd::bulk::compress(raw_data, 3)?; + let digest = Digest { + hash: HASH1.to_string(), + size_bytes: i64::try_from(raw_data.len()).unwrap(), + }; + + let Err(status) = cas_server + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + requests: vec![batch_update_blobs_request::Request { + digest: Some(digest), + data: compressed_data.into(), + compressor: compressor::Value::Zstd.into(), + }], + digest_function: digest_function::Value::Sha256.into(), + })) + .await + else { + panic!( + "zstd BatchUpdateBlobs to a ZstdStore-backed instance should fail when remote cache compression is disabled" + ); + }; + + assert_eq!(status.code(), Code::InvalidArgument); + assert!( + status + .message() + .contains("Remote cache compression is not supported"), + "unexpected error: {}", + status.message() + ); + + Ok(()) +} + +#[nativelink_test] +async fn batch_update_zstd_instance_identity_stores_when_remote_cache_compression_disabled() +-> Result<(), Box> { + // Identity entries must be unaffected by the gate: they still store and + // round-trip through the same ZstdStore-backed, compression-disabled + // instance. + let (cas_server, _zstd_store) = make_zstd_instance_cas_server_compression_disabled()?; + + let raw: Vec = "identity entry to disabled zstd instance " + .repeat(20) + .into_bytes(); + let (_di, digest) = sha256_digest(&raw); + + let response = cas_server + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + requests: vec![batch_update_blobs_request::Request { + digest: Some(digest.clone()), + data: raw.clone().into(), + compressor: compressor::Value::Identity.into(), + }], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + assert_eq!(response.responses.len(), 1); + assert_eq!( + response.responses[0].status.as_ref().map(|s| s.code), + Some(0) + ); + + let read = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![digest.clone()], + acceptable_compressors: vec![compressor::Value::Identity.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await? + .into_inner(); + assert_eq!(read.responses.len(), 1); + assert_eq!(read.responses[0].data.as_ref(), raw.as_slice()); + Ok(()) +} + const CHUNK1_VALUE: &str = "hello "; const CHUNK2_VALUE: &str = "world"; @@ -1746,3 +2129,161 @@ async fn chunking_infers_blake3_when_digest_function_unset() assert_eq!(split_response.chunk_digests, vec![other_digest]); Ok(()) } + +/// Inner store that stalls indefinitely for one specific digest while serving +/// every other digest from an inner `MemoryStore`. Used to prove a batch's +/// per-blob timeout isolates a stalled sibling from a healthy one. +#[derive(MetricsComponent)] +struct SelectiveStallStore { + stall_digest: DigestInfo, + delay: Duration, + #[metric(group = "inner")] + inner: Store, +} + +impl SelectiveStallStore { + fn stalls(&self, key: &StoreKey<'_>) -> bool { + matches!(key, StoreKey::Digest(d) + if d.packed_hash() == self.stall_digest.packed_hash() + && d.size_bytes() == self.stall_digest.size_bytes()) + } +} + +#[async_trait] +impl StoreDriver for SelectiveStallStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner.has_with_results(keys, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + if self.stalls(&key) { + tokio::time::sleep(self.delay).await; + return Ok(0); + } + self.inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + if self.stalls(&key) { + tokio::time::sleep(self.delay).await; + return Ok(()); + } + self.inner + .as_store_driver_pin() + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(SelectiveStallStore); + +/// A batch update whose siblings are one healthy blob and one blob that stalls +/// past `BATCH_PER_BLOB_TIMEOUT` isolates the two: the healthy blob commits +/// (`Ok`), the stalled one fails with its own `DeadlineExceeded` status, and +/// neither outcome affects the other. +#[nativelink_test(start_paused = true)] +async fn batch_update_per_blob_timeout_isolates_siblings() -> Result<(), Box> +{ + let good_data = b"healthy sibling blob".to_vec(); + let (good_di, good_digest) = sha256_digest(&good_data); + let stall_data = b"stalling sibling blob".to_vec(); + let (stall_di, stall_digest) = sha256_digest(&stall_data); + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(Arc::new(SelectiveStallStore { + stall_digest: stall_di, + delay: Duration::from_mins(2), // Longer than BATCH_PER_BLOB_TIMEOUT (30s). + inner: inner.clone(), + })); + let store_manager = Arc::new(StoreManager::new()); + store_manager.add_store("main_cas", store)?; + let cas_server = make_cas_server(&store_manager)?; + + let response = cas_server + .batch_update_blobs(Request::new(BatchUpdateBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + requests: vec![ + batch_update_blobs_request::Request { + digest: Some(good_digest.clone()), + data: good_data.clone().into(), + compressor: compressor::Value::Identity.into(), + }, + batch_update_blobs_request::Request { + digest: Some(stall_digest.clone()), + data: stall_data.into(), + compressor: compressor::Value::Identity.into(), + }, + ], + digest_function: digest_function::Value::Sha256.into(), + })) + .await + .unwrap() + .into_inner(); + + assert_eq!(response.responses.len(), 2); + let good = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&good_digest)) + .expect("healthy sibling response present"); + assert_eq!( + good.status.as_ref().map(|s| s.code), + Some(Code::Ok as i32), + "healthy sibling must commit despite the stalled one: {:?}", + good.status + ); + let stalled = response + .responses + .iter() + .find(|r| r.digest.as_ref() == Some(&stall_digest)) + .expect("stalled sibling response present"); + assert_eq!( + stalled.status.as_ref().map(|s| s.code), + Some(Code::DeadlineExceeded as i32), + "stalled sibling must surface its own DeadlineExceeded: {:?}", + stalled.status + ); + + // The healthy blob actually committed; the stalled one did not. + assert_eq!(inner.has(good_di).await?, Some(good_data.len() as u64)); + assert_eq!(inner.has(stall_di).await?, None); + Ok(()) +} diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index c3bc47a1a..db88b70ea 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -49,6 +49,7 @@ rust_library( "src/size_partitioning_store.rs", "src/store_manager.rs", "src/verify_store.rs", + "src/zstd_store.rs", ], proc_macro_deps = [ "@crates//:async-trait", @@ -111,6 +112,7 @@ rust_library( "@crates//:uuid", "@crates//:webpki-roots", "@crates//:wincode", + "@crates//:zstd", ] + select({ "@platforms//os:macos": ["@crates//:libc"], "//conditions:default": [], @@ -151,6 +153,7 @@ rust_test_suite( "tests/size_partitioning_store_test.rs", "tests/store_manager_test.rs", "tests/verify_store_test.rs", + "tests/zstd_store_test.rs", ], compile_data = [ "tests/mongo_runner/mod.rs", @@ -212,6 +215,7 @@ rust_test_suite( "@crates//:uuid", "@crates//:wincode", "@crates//:zip", + "@crates//:zstd", ], ) @@ -235,6 +239,7 @@ rust_test( "@crates//:redis", "@crates//:serde_json", "@crates//:sha2", + "@crates//:tracing-test", ], ) diff --git a/nativelink-store/Cargo.toml b/nativelink-store/Cargo.toml index 6d96bb385..9612b4310 100644 --- a/nativelink-store/Cargo.toml +++ b/nativelink-store/Cargo.toml @@ -130,6 +130,7 @@ wincode = { version = "0.5.4", default-features = false, features = [ "alloc", "derive", ] } +zstd = { version = "0.13.3", default-features = false } [target.'cfg(target_os = "macos")'.dependencies] libc = { version = "0.2.177", default-features = false } diff --git a/nativelink-store/src/compression_store.rs b/nativelink-store/src/compression_store.rs index 9e6e52327..b291949dc 100644 --- a/nativelink-store/src/compression_store.rs +++ b/nativelink-store/src/compression_store.rs @@ -22,7 +22,7 @@ use bytes::{Buf, BufMut, BytesMut}; use futures::future::FutureExt; use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size}; use nativelink_config::stores::CompressionSpec; -use nativelink_error::{Code, Error, ResultExt, error_if, make_err}; +use nativelink_error::{Code, Error, ResultExt, error_if, make_err, make_input_err}; use nativelink_metric::MetricsComponent; use nativelink_util::buf_channel::{ DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, @@ -257,8 +257,9 @@ impl core::fmt::Debug for CompressionStore { impl CompressionStore { pub fn new(spec: &CompressionSpec, inner_store: Store) -> Result, Error> { - let lz4_config = match spec.compression_algorithm { - nativelink_config::stores::CompressionAlgorithm::Lz4(mut lz4_config) => { + let lz4_config = match &spec.compression_algorithm { + nativelink_config::stores::CompressionAlgorithm::Lz4(lz4_config) => { + let mut lz4_config = *lz4_config; if lz4_config.block_size == 0 { lz4_config.block_size = DEFAULT_BLOCK_SIZE; } @@ -267,6 +268,11 @@ impl CompressionStore { } lz4_config } + nativelink_config::stores::CompressionAlgorithm::Zstd(_) => { + return Err(make_input_err!( + "CompressionStore only implements the lz4 compression algorithm" + )); + } }; Ok(Arc::new(Self { inner_store, diff --git a/nativelink-store/src/default_store_factory.rs b/nativelink-store/src/default_store_factory.rs index 01ae9d064..ec92a5069 100644 --- a/nativelink-store/src/default_store_factory.rs +++ b/nativelink-store/src/default_store_factory.rs @@ -18,7 +18,9 @@ use std::time::SystemTime; use futures::stream::FuturesOrdered; use futures::{Future, TryStreamExt}; -use nativelink_config::stores::{ExperimentalCloudObjectSpec, RedisMode, StoreSpec}; +use nativelink_config::stores::{ + CompressionAlgorithm, ExperimentalCloudObjectSpec, RedisMode, StoreSpec, +}; use nativelink_error::Error; use nativelink_util::health_utils::HealthRegistryBuilder; use nativelink_util::store_trait::{Store, StoreDriver}; @@ -47,6 +49,7 @@ use crate::shard_store::ShardStore; use crate::size_partitioning_store::SizePartitioningStore; use crate::store_manager::StoreManager; use crate::verify_store::VerifyStore; +use crate::zstd_store::ZstdStore; type FutureMaybeStore<'a> = Box> + Send + 'a>; @@ -93,10 +96,13 @@ pub fn store_factory<'a>( spec, store_factory(&spec.backend, store_manager, None).await?, ), - StoreSpec::Compression(spec) => CompressionStore::new( - &spec.clone(), - store_factory(&spec.backend, store_manager, None).await?, - )?, + StoreSpec::Compression(spec) => { + let inner = store_factory(&spec.backend, store_manager, None).await?; + match &spec.compression_algorithm { + CompressionAlgorithm::Lz4(_) => CompressionStore::new(spec, inner)?, + CompressionAlgorithm::Zstd(config) => ZstdStore::new(config, inner)?, + } + } StoreSpec::Dedup(spec) => DedupStore::new( spec, store_factory(&spec.index_store, store_manager, None).await?, diff --git a/nativelink-store/src/lib.rs b/nativelink-store/src/lib.rs index 1ceea189a..b75dd3ca8 100644 --- a/nativelink-store/src/lib.rs +++ b/nativelink-store/src/lib.rs @@ -43,3 +43,4 @@ pub mod shard_store; pub mod size_partitioning_store; pub mod store_manager; pub mod verify_store; +pub mod zstd_store; diff --git a/nativelink-store/src/zstd_store.rs b/nativelink-store/src/zstd_store.rs new file mode 100644 index 000000000..4c572e006 --- /dev/null +++ b/nativelink-store/src/zstd_store.rs @@ -0,0 +1,1845 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use core::cmp; +use core::future::Future; +use core::pin::Pin; +use core::sync::atomic::{AtomicU64, Ordering}; +use core::time::Duration; +use std::ffi::OsString; +use std::io::{Read, Seek, SeekFrom, Write}; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; +use std::sync::{Arc, OnceLock}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::future::FutureExt; +use nativelink_config::stores::ZstdConfig; +use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; +use nativelink_metric::MetricsComponent; +use nativelink_util::buf_channel::{ + BufChannelReader, BufChannelWriter, DropCloserReadHalf, DropCloserWriteHalf, + make_buf_channel_pair, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::{ + DigestHasher, DigestHasherFunc, DigestHasherImpl, digest_hasher_func_from_context, +}; +use nativelink_util::fs::{self, FileSlot}; +use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; +use nativelink_util::store_trait::{ + RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, WireCompressionStore, + WireCompressor, +}; +use nativelink_util::{spawn, spawn_blocking}; +use tokio::io::{AsyncSeekExt, AsyncWriteExt}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +use crate::cas_utils::is_zero_digest; + +/// Default number of concurrent staged uploads when the config value is `0`. +const DEFAULT_MAX_CONCURRENT_STAGED_UPLOADS: usize = 4; +/// Default number of concurrent identity (uncompressed) reads/writes when the +/// config value is `0`. +const DEFAULT_MAX_CONCURRENT_IDENTITY_OPS: usize = 256; +/// Default number of concurrent recompressions when the config value is `0`. +const DEFAULT_MAX_CONCURRENT_RECOMPRESSIONS: usize = 1; +/// Default total validate-and-stage deadline (seconds) when the config value is +/// `0`. Bounds how long the caller waits for a trickling upload's validation. +const DEFAULT_STAGE_TIMEOUT_S: u64 = 600; +/// Default commit timeout (seconds) when the config value is `0`. Bounds how +/// long a stalled inner-store commit (or recompression) may hold a staging +/// permit before the upload fails with `DeadlineExceeded` and cleans up. +const DEFAULT_COMMIT_TIMEOUT_S: u64 = 300; +/// Default ceiling for committing a compressed upload straight from memory +/// (no staging file, no `fsync`) when the config value is `0`. +const DEFAULT_MAX_INLINE_COMMIT_SIZE: u64 = 4 * 1024 * 1024; +/// Maximum decoder history window accepted from a zstd frame. This matches the +/// largest window emitted by the configured `1..=19` compression-level range +/// while preventing a tiny frame from reserving zstd's 128 MiB default limit. +const MAX_ZSTD_WINDOW_LOG: u32 = 23; +/// Level used when no `compression_level` is configured, matching the service +/// wire codec (`nativelink-service` `wire_compression::ZSTD_COMPRESSION_LEVEL`). +const DEFAULT_ENCODE_LEVEL: i32 = 3; + +/// The canonical zstd encoding of empty input, computed once. Emitted verbatim +/// for zero-digest `get_zstd` responses so the client always receives a valid +/// zstd stream that decodes to nothing. +fn empty_zstd_frame() -> Bytes { + static EMPTY: OnceLock = OnceLock::new(); + EMPTY + .get_or_init(|| { + Bytes::from( + zstd::bulk::compress(&[], DEFAULT_ENCODE_LEVEL) + .expect("zstd compression of empty input cannot fail"), + ) + }) + .clone() +} + +/// Flattens a `spawn!`/`spawn_blocking!` join result into the task's own result. +fn flatten_join>( + joined: Result, E>, + tip: &'static str, +) -> Result { + joined.err_tip(|| tip)? +} + +/// Maps an `io::Error` from a staging-file operation onto an internal error. +fn stage_file_err(action: &'static str) -> impl FnOnce(std::io::Error) -> Error { + move |e| make_err!(Code::Internal, "Failed to {action} zstd staging file: {e}") +} + +/// Increments a gauge for as long as it is held. +struct Inflight<'a>(&'a AtomicU64); + +impl<'a> Inflight<'a> { + fn enter(gauge: &'a AtomicU64) -> Self { + gauge.fetch_add(1, Ordering::Relaxed); + Self(gauge) + } +} + +impl Drop for Inflight<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + +/// Owned variant used when an admission gauge must outlive the async future +/// that spawned non-cancellable blocking work. +struct OwnedInflight(Arc); + +impl OwnedInflight { + fn enter(gauge: Arc) -> Self { + gauge.fetch_add(1, Ordering::Relaxed); + Self(gauge) + } +} + +impl Drop for OwnedInflight { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + +/// Upper bound on the number of compressed bytes a zstd frame can produce for +/// `src_size` uncompressed bytes. Mirrors zstd's `ZSTD_COMPRESSBOUND` with a +/// little extra slack for the frame header/epilogue. Used to bound the inner +/// store upload since the physical size is not known ahead of time. +const fn zstd_compress_bound(src_size: u64) -> u64 { + let low_bound_slack = if src_size < 128 * 1024 { + (128 * 1024 - src_size) >> 11 + } else { + 0 + }; + src_size + .saturating_add(src_size >> 8) + .saturating_add(low_bound_slack) + .saturating_add(64) +} + +/// A blocking source of compressed chunks. An empty chunk signals EOF. +/// +/// This is what lets the streaming and whole-buffer entry points share one +/// bounded decoder: `DropCloserReadHalf` feeds it from a client stream, while +/// [`OneChunk`] feeds it a payload already in memory. +trait ChunkSource { + fn next_chunk(&mut self) -> Result; +} + +impl ChunkSource for DropCloserReadHalf { + fn next_chunk(&mut self) -> Result { + self.blocking_recv() + } +} + +/// Yields one in-memory buffer, then EOF. +struct OneChunk(Option); + +impl ChunkSource for OneChunk { + fn next_chunk(&mut self) -> Result { + Ok(self.0.take().unwrap_or_default()) + } +} + +/// Streaming zstd encode of an identity (raw) upload. +/// +/// The inner store EOF is withheld until the finalized hash and uncompressed +/// size match `digest`; only then is the encoder finished and EOF sent, so a +/// mismatch can never commit the inner upload. +fn encode_identity( + mut reader: DropCloserReadHalf, + tx: DropCloserWriteHalf, + level: i32, + hasher_func: DigestHasherFunc, + digest: DigestInfo, +) -> Result { + let expected_size = digest.size_bytes(); + let mut hasher = hasher_func.hasher(); + let mut uncompressed_size: u64 = 0; + + let mut encoder = + zstd::stream::write::Encoder::new(BufChannelWriter::new(tx), level).map_err(|e| { + make_err!( + Code::Internal, + "Zstd encoder init failed in zstd store: {e}" + ) + })?; + + loop { + let chunk = reader + .blocking_recv() + .err_tip(|| "Failed to read chunk in zstd store update")?; + if chunk.is_empty() { + break; // EOF. + } + + uncompressed_size = uncompressed_size + .checked_add(chunk.len() as u64) + .ok_or_else(|| make_input_err!("Uncompressed size overflow in zstd store update"))?; + if uncompressed_size > expected_size { + return Err(make_input_err!( + "Received more data than digest size in zstd store update, got at least {} but digest says {}", + uncompressed_size, + expected_size + )); + } + + hasher.update(&chunk); + encoder + .write_all(&chunk) + .map_err(|e| make_err!(Code::Internal, "Zstd encode failed in zstd store: {e}"))?; + } + + if uncompressed_size != expected_size { + return Err(make_input_err!( + "Expected size {} but got size {} in zstd store update", + expected_size, + uncompressed_size + )); + } + let actual_digest = hasher.finalize_digest(); + if actual_digest.packed_hash() != digest.packed_hash() { + return Err(make_input_err!( + "Hashes do not match in zstd store update, expected {} but got {}", + digest.packed_hash(), + actual_digest.packed_hash() + )); + } + + let mut writer = encoder.finish().map_err(|e| { + make_err!( + Code::Internal, + "Zstd encoder finish failed in zstd store: {e}" + ) + })?; + writer + .send_eof() + .err_tip(|| "Failed to send EOF in zstd store update")?; + Ok(expected_size) +} + +/// Streaming zstd decode of stored (physical) data back to raw bytes, +/// forwarding at most `length` bytes starting at `offset`. A decode failure of +/// already-stored data is `DataLoss`: it was validated at upload time. +fn decode_identity( + physical_rx: DropCloserReadHalf, + mut raw_tx: DropCloserWriteHalf, + offset: u64, + length: Option, +) -> Result<(), Error> { + let reader = BufChannelReader::new(physical_rx); + let mut decoder = zstd::stream::read::Decoder::new(reader).map_err(|e| { + make_err!( + Code::DataLoss, + "Zstd decoder init failed in zstd store: {e}" + ) + })?; + decoder.window_log_max(MAX_ZSTD_WINDOW_LOG).map_err(|e| { + make_err!( + Code::DataLoss, + "Failed to cap zstd decoder window in zstd store: {e}" + ) + })?; + + let mut to_skip = offset; + let mut remaining = length.unwrap_or(u64::MAX); + let mut buffer = vec![0u8; zstd::zstd_safe::DCtx::out_size()]; + + // The decoder must be drained to EOF even once the byte budget is met: + // dropping `physical_rx` early would fail the still-running inner `get` + // with "receiver disconnected". Mirrors `compression_store::get_part`. + loop { + let read = decoder + .read(&mut buffer) + .map_err(|e| make_err!(Code::DataLoss, "Zstd decode failed in zstd store: {e}"))?; + if read == 0 { + break; // EOF: physical reader fully drained. + } + if remaining == 0 { + continue; // Budget met; keep draining but stop forwarding. + } + + let mut slice = &buffer[..read]; + if to_skip > 0 { + // `slice.len()` bounds the min, so the result always fits in usize. + let skip = usize::try_from(cmp::min(to_skip, slice.len() as u64)).unwrap_or(usize::MAX); + slice = &slice[skip..]; + to_skip -= skip as u64; + } + if slice.is_empty() { + continue; + } + + // `slice.len()` bounds the min, so the result always fits in usize. + let take = usize::try_from(cmp::min(remaining, slice.len() as u64)).unwrap_or(usize::MAX); + remaining -= take as u64; + raw_tx + .blocking_send(Bytes::copy_from_slice(&slice[..take])) + .err_tip(|| "Failed to send decoded chunk in zstd store get_part")?; + } + + raw_tx + .send_eof() + .err_tip(|| "Failed to send decoded EOF in zstd store get_part")?; + Ok(()) +} + +/// Best-effort removal of a staged temp file when dropped. Staged files are pure +/// scratch — the inner store takes the bytes at commit — so the path is removed +/// on success, error, and cancellation alike. +#[derive(Debug)] +struct TempFileGuard { + path: Option, +} + +impl TempFileGuard { + const fn arm(path: String) -> Self { + Self { path: Some(path) } + } +} + +impl Drop for TempFileGuard { + fn drop(&mut self) { + if let Some(path) = &self.path { + // Ignore failures: a backend that commits by moving the file has + // already renamed this path away. + drop(std::fs::remove_file(path)); + } + } +} + +/// A `std::io::Write` sink that hashes and counts decoded output instead of +/// storing it, optionally buffering it for recompression up to `collect_limit`. +/// +/// This is the enforcement point for the decoded-output bound: a `write` that +/// would push `decoded_len` past `max_decoded_size` is rejected *before* any of +/// the offending bytes are hashed or collected. Because the zstd streaming +/// decoder emits output in bounded blocks, a small "zstd bomb" is stopped at the +/// first over-limit block rather than being fully materialized. +struct DecodeSink { + hasher: DigestHasherImpl, + decoded_len: u64, + /// Hard ceiling on total decoded bytes; equals the digest's uncompressed + /// size. + max_decoded_size: u64, + collect: Option>, + collect_limit: u64, +} + +impl DecodeSink { + fn new( + hasher_func: DigestHasherFunc, + max_decoded_size: u64, + collect_limit: Option, + ) -> Self { + Self { + hasher: hasher_func.hasher(), + decoded_len: 0, + max_decoded_size, + collect: collect_limit.map(|_| Vec::new()), + collect_limit: collect_limit.unwrap_or(0), + } + } +} + +impl Write for DecodeSink { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + // `checked_add` (never `saturating_add`) so a length overflow is an + // error, not a silently clamped value. + let new_len = self + .decoded_len + .checked_add(buf.len() as u64) + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "decoded length overflow in zstd store", + ) + })?; + if new_len > self.max_decoded_size { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "decoded output exceeded digest size {} in zstd store", + self.max_decoded_size + ), + )); + } + self.hasher.update(buf); + self.decoded_len = new_len; + if let Some(collected) = self.collect.as_mut() { + if new_len <= self.collect_limit { + collected.extend_from_slice(buf); + } else { + // Too large to re-compress in-memory; stop collecting. + self.collect = None; + } + } + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +/// Final check of a bounded decode against the requested digest. The sink +/// already guarantees `decoded_len <= digest.size_bytes()`; this catches the +/// short case and the hash. +fn verify_decoded_digest(sink: &mut DecodeSink, digest: DigestInfo) -> Result<(), Error> { + let expected_size = digest.size_bytes(); + if sink.decoded_len != expected_size { + return Err(make_input_err!( + "Decoded size {} does not match digest size {expected_size} in zstd store", + sink.decoded_len + )); + } + let actual_digest = sink.hasher.finalize_digest(); + if actual_digest.packed_hash() != digest.packed_hash() { + return Err(make_input_err!( + "Hashes do not match in zstd store update, expected {} but got {}", + digest.packed_hash(), + actual_digest.packed_hash() + )); + } + Ok(()) +} + +/// Exclusively create a staged temp file at `path`, returning its open +/// descriptor. `create_new(true)` (`O_CREAT | O_EXCL`) fails if the path already +/// exists or is a symlink, defeating an observe-and-replace race; on unix the +/// owner-only `0o600` mode is applied atomically at creation rather than via a +/// later `chmod` by path. Runs inside `spawn_blocking!`. +fn create_temp_exclusive(path: &str) -> Result { + let mut options = std::fs::OpenOptions::new(); + options.read(true).write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + options.open(path).map_err(|e| { + make_err!( + Code::Internal, + "Failed to exclusively create zstd staging file {path}: {e}" + ) + }) +} + +/// Decode a complete concatenation of zstd frames from `source` with bounded +/// wire and decoded sizes. `write_wire` receives every compressed chunk before +/// it is decoded (the staging path uses it to persist the exact client bytes). +/// +/// [`zio::Writer::finish`](zstd::stream::zio::Writer::finish) is deliberate: +/// unlike `flush`, it drives the decoder with EOF and rejects an incomplete +/// final frame, so a valid frame followed by a truncated second frame is +/// refused rather than committed. +fn decode_bounded_zstd_stream( + mut source: S, + sink: DecodeSink, + max_compressed_upload_size: u64, + read_context: &'static str, + decode_context: &'static str, + mut write_wire: F, +) -> Result<(DecodeSink, u64), Error> +where + S: ChunkSource, + F: FnMut(&[u8]) -> Result<(), Error>, +{ + let mut raw_decoder = zstd::stream::raw::Decoder::new().map_err(|e| { + make_err!( + Code::Internal, + "Zstd decoder init failed in zstd store: {e}" + ) + })?; + raw_decoder + .set_parameter(zstd::stream::raw::DParameter::WindowLogMax( + MAX_ZSTD_WINDOW_LOG, + )) + .map_err(|e| { + make_err!( + Code::Internal, + "Failed to cap zstd decoder window in zstd store: {e}" + ) + })?; + let mut decoder = zstd::stream::zio::Writer::new(sink, raw_decoder); + let mut wire_bytes_consumed: u64 = 0; + loop { + let chunk = source.next_chunk().err_tip(|| read_context)?; + if chunk.is_empty() { + break; + } + wire_bytes_consumed = wire_bytes_consumed + .checked_add(chunk.len() as u64) + .ok_or_else(|| make_err!(Code::Internal, "Wire byte count overflow in zstd store"))?; + if wire_bytes_consumed > max_compressed_upload_size { + return Err(make_err!( + Code::ResourceExhausted, + "Compressed upload exceeded max_compressed_upload_size ({max_compressed_upload_size} bytes) in zstd store" + )); + } + write_wire(&chunk)?; + decoder + .write_all(&chunk) + .map_err(|e| make_input_err!("{decode_context}: {e}"))?; + } + decoder + .finish() + .map_err(|e| make_input_err!("{decode_context}: {e}"))?; + let (sink, _raw_decoder) = decoder.into_inner(); + Ok((sink, wire_bytes_consumed)) +} + +/// What the blocking staging pass hands back on success. +struct StageOutput { + /// Compressed wire bytes consumed from the client, which is also the size of + /// the staged file. + wire_bytes_consumed: u64, + /// Decoded content, present only when recompression is eligible (decoded + /// length stayed within the configured `max_recompression_size`). + collected: Option>, + /// The descriptor that was validated, rewound to offset 0. + file: FileSlot, + /// The cleanup guard, owned by the blocking task until it returns. + guard: TempFileGuard, + /// Admission remains owned by the blocking task after an async timeout and + /// returns to the commit path only after validation has actually stopped. + permit: OwnedSemaphorePermit, + inflight: OwnedInflight, +} + +/// Resources transferred to the detached staging task. In particular, the +/// armed cleanup guard stays owned by the task until it returns. +struct StageCompressedInput { + reader: DropCloserReadHalf, + path: String, + digest: DigestInfo, + hasher_func: DigestHasherFunc, + max_compressed_upload_size: u64, + collect_limit: Option, + fs_permit: tokio::sync::SemaphorePermit<'static>, + guard: TempFileGuard, + permit: OwnedSemaphorePermit, + inflight: OwnedInflight, +} + +/// Blocking validation pass: exclusively create the temp file at `path`, stream +/// compressed chunks from `reader` into it while decoding them to recompute the +/// uncompressed length and hash, then check both against `digest`. Enforces the +/// compressed-size cap (`ResourceExhausted`) and, via [`DecodeSink`], the +/// decoded-output cap (`InvalidArgument`) inside the decoder write. +/// +/// The cleanup guard is owned by this detached blocking task until it returns, +/// so cancellation of the async caller cannot leak a file created after the +/// caller stopped awaiting. Only runs from within `spawn_blocking!`. +fn stage_compressed_blocking( + StageCompressedInput { + reader, + path, + digest, + hasher_func, + max_compressed_upload_size, + collect_limit, + fs_permit, + guard, + permit, + inflight, + }: StageCompressedInput, +) -> Result { + let mut file = create_temp_exclusive(&path)?; + + let sink = DecodeSink::new(hasher_func, digest.size_bytes(), collect_limit); + let (mut sink, wire_bytes_consumed) = decode_bounded_zstd_stream( + reader, + sink, + max_compressed_upload_size, + "Failed to read compressed chunk in zstd store staging", + "Zstd decode failed in zstd store staging", + |chunk| file.write_all(chunk).map_err(stage_file_err("write")), + )?; + let collected = sink.collect.take(); + verify_decoded_digest(&mut sink, digest)?; + + file.flush().map_err(stage_file_err("flush"))?; + file.sync_all().map_err(stage_file_err("sync"))?; + // Rewind the *same* descriptor so the commit streams it from the start. + file.seek(SeekFrom::Start(0)) + .map_err(stage_file_err("rewind"))?; + + Ok(StageOutput { + wire_bytes_consumed, + collected, + file: FileSlot::from_std(fs_permit, file), + guard, + permit, + inflight, + }) +} + +/// Blocking bounded decode used to validate that a stream decodes to exactly +/// `expected_size` bytes without buffering the decoded output or touching disk. +/// Returns the compressed wire bytes consumed. +fn validate_bounded_zstd_blocking( + source: S, + hasher_func: DigestHasherFunc, + expected_size: u64, + max_compressed_upload_size: u64, +) -> Result { + let sink = DecodeSink::new(hasher_func, expected_size, None); + let (sink, wire_bytes_consumed) = decode_bounded_zstd_stream( + source, + sink, + max_compressed_upload_size, + "Failed to read chunk validating zero-digest zstd", + "Zstd decode failed for zero digest", + |_| Ok(()), + )?; + if sink.decoded_len != expected_size { + return Err(make_input_err!( + "Zero-digest zstd upload did not decode to empty in zstd store" + )); + } + Ok(wire_bytes_consumed) +} + +/// Blocking whole-buffer validation of a compressed upload: bounded decode, no +/// disk involvement. Returns the wire bytes consumed and, when recompression is +/// eligible, the decoded content. +fn validate_zstd_buffer( + data: Bytes, + digest: DigestInfo, + hasher_func: DigestHasherFunc, + max_compressed_upload_size: u64, + collect_limit: Option, +) -> Result<(u64, Option>), Error> { + let sink = DecodeSink::new(hasher_func, digest.size_bytes(), collect_limit); + let (mut sink, wire_bytes_consumed) = decode_bounded_zstd_stream( + OneChunk(Some(data)), + sink, + max_compressed_upload_size, + "Failed to read compressed buffer in zstd store", + "Zstd decode failed in zstd store inline validation", + |_| Ok(()), + )?; + let collected = sink.collect.take(); + verify_decoded_digest(&mut sink, digest)?; + Ok((wire_bytes_consumed, collected)) +} + +/// A validated staged upload, ready to commit. Dropping it removes the temp file +/// and releases the staging permit. +struct StagedStream { + /// The descriptor validated by staging, rewound to offset 0, taken by the + /// commit. Declared before `_guard` deliberately: Rust drops struct fields + /// in declaration order, closing the descriptor before best-effort path + /// removal on every cancellation/error path. + file: Option, + /// Path of the temp file. Backends that commit by moving the file + /// (`filesystem`) act on this path rather than on the descriptor; see + /// [`ZstdStore::commit_staged`]. + path: String, + /// Physical (compressed) size of what will be committed. + compressed_size: u64, + /// Compressed wire bytes actually consumed from the client. + wire_bytes_consumed: u64, + /// Decoded bytes retained only for optional recompression. Staging bounds + /// this by `max_recompression_size`; it is consumed before commit. + collected: Option>, + _permit: OwnedSemaphorePermit, + _inflight: OwnedInflight, + _guard: TempFileGuard, +} + +/// A pass-through store that stores blobs zstd-compressed in the inner CAS +/// store while presenting the raw (uncompressed) view to clients. +/// +/// This is CAS-only: only digest keys are supported. Zero-byte digests are +/// never forwarded to the inner store. +/// +/// Zstd implementation selected by `CompressionAlgorithm::Zstd`. +#[derive(MetricsComponent)] +pub struct ZstdStore { + #[metric(group = "inner_store")] + inner_store: Store, + /// Operator-controlled staging directory. Created/permission-checked in + /// `post_init`; used by the zstd fast path. + #[metric(help = "Staging directory for validated compressed uploads")] + temp_path: String, + /// Configured encode level, already validated to `1..=19`; `None` uses + /// [`DEFAULT_ENCODE_LEVEL`]. + compression_level: Option, + /// Hard cap on the number of compressed wire bytes accepted for a single + /// upload before it is rejected with `ResourceExhausted`. + #[metric(help = "Max compressed wire bytes accepted for one upload")] + max_compressed_upload_size: u64, + /// Uncompressed-size ceiling below which an upload is eligible for + /// re-compression at `compression_level`. `0` disables re-compression. + #[metric(help = "Max uncompressed size eligible for recompression")] + max_recompression_size: u64, + /// Compressed uploads at or below this size skip the staging file entirely. + #[metric(help = "Max compressed upload size committed inline from memory")] + max_inline_commit_size: u64, + /// Bounds concurrent staged uploads (temp files being validated at once). + staged_upload_semaphore: Arc, + /// Bounds concurrent identity (uncompressed) reads and writes, each of which + /// occupies a blocking thread for the whole transfer. + identity_semaphore: Arc, + /// Bounds concurrent re-compression passes. Acquired with `try_acquire` so a + /// busy pool skips recompression instead of holding a staging slot. + recompression_semaphore: Arc, + /// Total time one upload may spend being validated and staged, from the + /// moment it is admitted. Unlike a per-message idle timeout, slow continuous + /// progress does not reset it. The admitted blocking validator retains its + /// slot until it actually exits, including after the caller times out. + stage_timeout: Duration, + /// Upper bound on how long the inner-store commit (and any recompression) + /// of a staged upload may take before it fails with `DeadlineExceeded`, + /// releasing the staging permit and removing the temp file. + commit_timeout: Duration, + + #[metric(help = "Compressed uploads accepted through the wire fast path")] + wire_uploads: AtomicU64, + #[metric(help = "Compressed wire bytes accepted through the wire fast path")] + wire_upload_bytes: AtomicU64, + #[metric(help = "Stored zstd streams served byte-for-byte to clients")] + wire_downloads: AtomicU64, + #[metric(help = "Batch reads answered with stored zstd bytes")] + batch_zstd_passthroughs: AtomicU64, + #[metric(help = "Batch reads answered with decoded identity bytes")] + batch_identity_decodes: AtomicU64, + #[metric(help = "Compressed uploads committed inline with no staging file")] + inline_commits: AtomicU64, + #[metric(help = "Compressed uploads committed from a staging file")] + staged_commits: AtomicU64, + #[metric(help = "Compressed uploads currently validating or staging")] + staged_uploads_inflight: Arc, + #[metric(help = "Identity reads and writes currently in flight")] + identity_ops_inflight: Arc, + #[metric(help = "Recompressions that produced a smaller stream and were kept")] + recompressions_applied: AtomicU64, + #[metric(help = "Recompressions that were not smaller and were discarded")] + recompressions_rejected: AtomicU64, + #[metric(help = "Recompressions skipped because every slot was busy")] + recompressions_skipped_busy: AtomicU64, + #[metric(help = "Uploads failed by the validate-and-stage deadline")] + stage_timeouts: AtomicU64, + #[metric(help = "Uploads failed by the recompress-and-commit deadline")] + commit_timeouts: AtomicU64, +} + +impl core::fmt::Debug for ZstdStore { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("ZstdStore") + .field("inner_store", &self.inner_store) + .field("compression_level", &self.compression_level) + .finish_non_exhaustive() + } +} + +impl ZstdStore { + pub fn new(spec: &ZstdConfig, inner_store: Store) -> Result, Error> { + if let Some(level) = spec.compression_level + && !(1..=19).contains(&level) + { + return Err(make_input_err!( + "ZstdStore compression_level must be in [1, 19], got {level}" + )); + } + if spec.temp_path.is_empty() { + return Err(make_input_err!("ZstdStore requires a non-empty temp_path")); + } + if spec.max_compressed_upload_size == 0 { + return Err(make_input_err!( + "ZstdStore requires a positive max_compressed_upload_size" + )); + } + // Recompression re-encodes at `compression_level`, so a ceiling without + // a level would silently do nothing. Reject it instead. + if spec.max_recompression_size > 0 && spec.compression_level.is_none() { + return Err(make_input_err!( + "ZstdStore max_recompression_size requires compression_level to be set" + )); + } + let nonzero_or = |configured: usize, default: usize| match configured { + 0 => default, + n => n, + }; + Ok(Arc::new(Self { + inner_store, + temp_path: spec.temp_path.clone(), + compression_level: spec.compression_level, + max_compressed_upload_size: spec.max_compressed_upload_size, + max_recompression_size: spec.max_recompression_size, + max_inline_commit_size: match spec.max_inline_commit_size { + 0 => DEFAULT_MAX_INLINE_COMMIT_SIZE, + n => n, + }, + staged_upload_semaphore: Arc::new(Semaphore::new(nonzero_or( + spec.max_concurrent_staged_uploads, + DEFAULT_MAX_CONCURRENT_STAGED_UPLOADS, + ))), + identity_semaphore: Arc::new(Semaphore::new(nonzero_or( + spec.max_concurrent_identity_ops, + DEFAULT_MAX_CONCURRENT_IDENTITY_OPS, + ))), + recompression_semaphore: Arc::new(Semaphore::new(nonzero_or( + spec.max_concurrent_recompressions, + DEFAULT_MAX_CONCURRENT_RECOMPRESSIONS, + ))), + stage_timeout: Duration::from_secs(match spec.stage_timeout_s { + 0 => DEFAULT_STAGE_TIMEOUT_S, + n => n, + }), + commit_timeout: Duration::from_secs(match spec.commit_timeout_s { + 0 => DEFAULT_COMMIT_TIMEOUT_S, + n => n, + }), + wire_uploads: AtomicU64::new(0), + wire_upload_bytes: AtomicU64::new(0), + wire_downloads: AtomicU64::new(0), + batch_zstd_passthroughs: AtomicU64::new(0), + batch_identity_decodes: AtomicU64::new(0), + inline_commits: AtomicU64::new(0), + staged_commits: AtomicU64::new(0), + staged_uploads_inflight: Arc::new(AtomicU64::new(0)), + identity_ops_inflight: Arc::new(AtomicU64::new(0)), + recompressions_applied: AtomicU64::new(0), + recompressions_rejected: AtomicU64::new(0), + recompressions_skipped_busy: AtomicU64::new(0), + stage_timeouts: AtomicU64::new(0), + commit_timeouts: AtomicU64::new(0), + })) + } + + #[inline] + fn encode_level(&self) -> i32 { + self.compression_level.unwrap_or(DEFAULT_ENCODE_LEVEL) + } + + /// `Some(limit)` when an upload's decoded bytes should be buffered for a + /// recompression attempt. `new` guarantees a level is configured whenever + /// the ceiling is positive. + #[inline] + const fn recompression_limit(&self) -> Option { + if self.max_recompression_size > 0 { + Some(self.max_recompression_size) + } else { + None + } + } + + async fn acquire_staging_permit(&self) -> Result { + self.staged_upload_semaphore + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "Staged upload semaphore closed: {e}")) + } + + async fn acquire_identity_permit(&self) -> Result { + self.identity_semaphore + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "Identity semaphore closed: {e}")) + } + + async fn update_identity( + self: Pin<&Self>, + key: StoreKey<'_>, + mut reader: DropCloserReadHalf, + ) -> Result { + let StoreKey::Digest(digest) = key else { + return Err(make_input_err!( + "ZstdStore only supports digest keys, got {key:?}" + )); + }; + + if is_zero_digest(digest) { + return reader.recv().await.and_then(|chunk| { + if chunk.is_empty() { + Ok(0) + } else { + Err(make_err!(Code::Internal, "Zero byte hash was not empty")) + } + }); + } + + // The encode below parks a blocking thread for the whole upload, so it + // must be admitted rather than allowed to exhaust the shared pool. + let _permit = self.acquire_identity_permit().await?; + let _inflight = Inflight::enter(&self.identity_ops_inflight); + + let (tx, rx) = make_buf_channel_pair(); + let inner_store = self.inner_store.clone(); + let max_output_size = zstd_compress_bound(digest.size_bytes()); + let update_fut = spawn!("zstd_store_update_spawn", async move { + inner_store + .update(digest, rx, UploadSizeInfo::MaxSize(max_output_size)) + .await + .err_tip(|| "Inner store update in zstd store failed") + }) + .map(|joined| flatten_join(joined, "Failed to run zstd store update spawn")); + + let level = self.encode_level(); + let hasher_func = digest_hasher_func_from_context(); + let encode_fut = spawn_blocking!("zstd_store_encode_identity", move || { + encode_identity(reader, tx, level, hasher_func, digest) + }) + .map(|joined| flatten_join(joined, "Failed to run zstd store encode task")); + + let (encode_res, update_res) = tokio::join!(encode_fut, update_fut); + match (encode_res, update_res) { + (Ok(size), Ok(_)) => Ok(size), + (Err(e), _) | (_, Err(e)) => Err(e), + } + } + + async fn get_part_identity( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + let StoreKey::Digest(digest) = key else { + return Err(make_input_err!( + "ZstdStore only supports digest keys, got {key:?}" + )); + }; + + if is_zero_digest(digest) { + writer + .send_eof() + .err_tip(|| "Failed to send zero-digest EOF in zstd store get_part")?; + return Ok(()); + } + + // The decode below parks a blocking thread until the client has consumed + // the whole blob, so it must be admitted like an upload. + let _permit = self.acquire_identity_permit().await?; + let _inflight = Inflight::enter(&self.identity_ops_inflight); + + let (physical_tx, physical_rx) = make_buf_channel_pair(); + let inner_store = self.inner_store.clone(); + let get_fut = spawn!("zstd_store_get_part_spawn", async move { + inner_store + .get_part(digest, physical_tx, 0, None) + .await + .err_tip(|| "Inner store get in zstd store failed") + }) + .map(|joined| flatten_join(joined, "Failed to run zstd store get spawn")); + + let (raw_tx, mut raw_rx) = make_buf_channel_pair(); + let decode_fut = spawn_blocking!("zstd_store_decode_identity", move || { + decode_identity(physical_rx, raw_tx, offset, length) + }) + .map(|joined| flatten_join(joined, "Failed to run zstd store decode task")); + + let pump_fut = async move { + loop { + let chunk = raw_rx + .recv() + .await + .err_tip(|| "Failed to read decoded chunk in zstd store get_part")?; + if chunk.is_empty() { + break; + } + writer + .send(chunk) + .await + .err_tip(|| "Failed to send chunk in zstd store get_part")?; + } + writer + .send_eof() + .err_tip(|| "Failed to send EOF in zstd store get_part")?; + Result::<(), Error>::Ok(()) + }; + + let (get_res, decode_res, pump_res) = tokio::join!(get_fut, decode_fut, pump_fut); + // Prioritize the inner get error (e.g. NotFound) as the root cause, then + // the decode error (DataLoss on corrupt data), then the pump error. + get_res?; + decode_res?; + pump_res?; + Ok(()) + } + + /// Byte-for-byte passthrough of the stored zstd stream. Always emits zstd: + /// zero digests are answered with the canonical empty encoding. + pub async fn get_zstd( + &self, + digest: DigestInfo, + mut writer: DropCloserWriteHalf, + ) -> Result<(), Error> { + if is_zero_digest(digest) { + writer + .send(empty_zstd_frame()) + .await + .err_tip(|| "Failed to send empty zstd in zstd store get_zstd")?; + writer + .send_eof() + .err_tip(|| "Failed to send EOF in zstd store get_zstd")?; + return Ok(()); + } + // The physical bytes are already zstd; pipe them straight through. + self.inner_store + .get_part(digest, writer, 0, None) + .await + .err_tip(|| "Inner store get_part in zstd store get_zstd failed")?; + self.wire_downloads.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + + /// Accept a client-supplied compressed (zstd) stream: validate, stage, + /// optionally re-compress, then commit. Returns the number of compressed + /// wire bytes consumed. + pub async fn update_zstd( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + reader: DropCloserReadHalf, + ) -> Result { + if is_zero_digest(digest) { + return self.validate_empty_zstd(digest_function, reader).await; + } + let wire_bytes = self + .update_staged_zstd(digest, digest_function, reader) + .await?; + self.record_wire_upload(wire_bytes); + Ok(wire_bytes) + } + + /// Whole-`Bytes` variant of [`ZstdStore::update_zstd`] for `BatchUpdate`. + pub async fn update_zstd_oneshot( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + data: Bytes, + ) -> Result { + // Zero digests must not be whole-buffer decoded by an allocating + // decoder, so they go through the same bounded decoder as everything + // else — just fed from memory instead of a channel. + if is_zero_digest(digest) { + return self + .validate_empty_zstd(digest_function, OneChunk(Some(data))) + .await; + } + + let wire_bytes = if data.len() as u64 <= self.max_inline_commit_size { + self.update_inline_zstd(digest, digest_function, data) + .await? + } else { + self.update_large_oneshot_zstd(digest, digest_function, data) + .await? + }; + self.record_wire_upload(wire_bytes); + Ok(wire_bytes) + } + + fn record_wire_upload(&self, wire_bytes: u64) { + self.wire_uploads.fetch_add(1, Ordering::Relaxed); + self.wire_upload_bytes + .fetch_add(wire_bytes, Ordering::Relaxed); + } + + /// Batch read selection: return `(data, is_zstd)`. Prefers the physical + /// zstd bytes when the client accepts zstd and compression actually helped; + /// otherwise decodes to raw. + pub async fn get_for_batch( + &self, + digest: DigestInfo, + client_accepts_zstd: bool, + ) -> Result<(Bytes, bool), Error> { + if is_zero_digest(digest) { + return Ok((Bytes::new(), false)); + } + let physical = self + .inner_store + .get_part_unchunked(digest, 0, None) + .await + .err_tip(|| "Inner store get in zstd store get_for_batch failed")?; + if client_accepts_zstd && (physical.len() as u64) < digest.size_bytes() { + self.batch_zstd_passthroughs.fetch_add(1, Ordering::Relaxed); + return Ok((physical, true)); + } + let permit = self.acquire_identity_permit().await?; + let inflight = OwnedInflight::enter(self.identity_ops_inflight.clone()); + let expected_size = usize::try_from(digest.size_bytes()) + .map_err(|_| make_err!(Code::Internal, "Digest size too large for this platform"))?; + let raw = flatten_join( + spawn_blocking!("zstd_store_batch_decode", move || { + let _permit = permit; + let _inflight = inflight; + // Stored data was validated at upload time, and + // the bulk decoder bounds its output by `expected_size`, so a + // failure here is corruption of already-stored bytes. + let mut decompressor = zstd::bulk::Decompressor::new().map_err(|e| { + make_err!( + Code::DataLoss, + "Zstd decoder init failed in zstd store: {e}" + ) + })?; + decompressor + .set_parameter(zstd::zstd_safe::DParameter::WindowLogMax( + MAX_ZSTD_WINDOW_LOG, + )) + .map_err(|e| { + make_err!( + Code::DataLoss, + "Failed to cap zstd decoder window in zstd store: {e}" + ) + })?; + decompressor + .decompress(&physical, expected_size) + .map_err(|e| make_err!(Code::DataLoss, "Zstd decode failed in zstd store: {e}")) + }) + .await, + "Failed to run zstd store batch decode task", + )?; + self.batch_identity_decodes.fetch_add(1, Ordering::Relaxed); + Ok((Bytes::from(raw), false)) + } + + /// Validate that a zero-digest client stream decodes to empty without + /// touching the inner store or allocating the decoded output. Takes a + /// staging permit and the staging deadline, so a flood of zero-digest + /// streams cannot spawn unbounded blocking decode jobs or park them. + /// Returns the compressed wire bytes consumed. + async fn validate_empty_zstd( + &self, + digest_function: DigestHasherFunc, + source: S, + ) -> Result { + let permit = self.acquire_staging_permit().await?; + let inflight = OwnedInflight::enter(self.staged_uploads_inflight.clone()); + let max_compressed_upload_size = self.max_compressed_upload_size; + let validate_fut = spawn_blocking!("zstd_store_validate_empty", move || { + let _permit = permit; + let _inflight = inflight; + validate_bounded_zstd_blocking(source, digest_function, 0, max_compressed_upload_size) + }); + match tokio::time::timeout(self.stage_timeout, validate_fut).await { + Ok(joined) => flatten_join(joined, "Failed to run zstd store empty validation task"), + Err(_elapsed) => Err(self.stage_timeout_err()), + } + } + + fn stage_timeout_err(&self) -> Error { + self.stage_timeouts.fetch_add(1, Ordering::Relaxed); + make_err!( + Code::DeadlineExceeded, + "zstd store validation and staging exceeded stage_timeout_s ({}s); cleanup remains admitted until validation exits", + self.stage_timeout.as_secs() + ) + } + + /// Bound the post-validation recompression and commit of an already-staged + /// or already-validated upload. + async fn with_commit_timeout(&self, commit: F) -> Result + where + F: Future>, + { + match tokio::time::timeout(self.commit_timeout, commit).await { + Ok(result) => result, + Err(_elapsed) => { + self.commit_timeouts.fetch_add(1, Ordering::Relaxed); + Err(make_err!( + Code::DeadlineExceeded, + "zstd store recompression or commit timed out after {}s; cleaning up staged files", + self.commit_timeout.as_secs() + )) + } + } + } + + /// Validate and stage a non-zero-digest compressed stream, then commit it. + /// The blocking staging task owns its cleanup guard, so cancellation while + /// it is detached still removes its temp file once it stops. + async fn update_staged_zstd( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + reader: DropCloserReadHalf, + ) -> Result { + let mut staged = self + .stage_compressed(digest, digest_function, reader) + .await?; + let wire_bytes_consumed = staged.wire_bytes_consumed; + self.with_commit_timeout(async { + self.recompress_staged(&mut staged).await?; + self.commit_staged(digest, staged).await?; + Ok(wire_bytes_consumed) + }) + .await + } + + /// Oneshot payload too large to hold a second copy of in memory: feed it + /// through the same bounded streaming staging path as a client stream. + async fn update_large_oneshot_zstd( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + data: Bytes, + ) -> Result { + let (mut tx, rx) = make_buf_channel_pair(); + let feed_fut = spawn!("zstd_store_oneshot_feed", async move { + if !data.is_empty() { + tx.send(data) + .await + .err_tip(|| "Failed to feed oneshot data in zstd store")?; + } + tx.send_eof() + .err_tip(|| "Failed to send oneshot EOF in zstd store")?; + Result::<(), Error>::Ok(()) + }); + + let stage_res = self.update_staged_zstd(digest, digest_function, rx).await; + let feed_res = flatten_join(feed_fut.await, "Failed to run zstd store oneshot feed task"); + // Staging owns the meaningful error: rejecting an upload drops the + // reader, which makes the feeder fail too. Only surface the feeder's + // error when staging itself succeeded. + let wire_bytes = stage_res?; + feed_res?; + Ok(wire_bytes) + } + + /// Validate and commit a small compressed upload straight from memory: no + /// staging file and no `fsync`. `BatchUpdateBlobs` payloads are small and + /// numerous, so a per-blob disk round trip would dominate their cost. + async fn update_inline_zstd( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + data: Bytes, + ) -> Result { + let permit = self.acquire_staging_permit().await?; + let inflight = OwnedInflight::enter(self.staged_uploads_inflight.clone()); + + let max_compressed_upload_size = self.max_compressed_upload_size; + let collect_limit = self.recompression_limit(); + // `Bytes` is refcounted, so the validation copy is free. + let to_validate = data.clone(); + let validate_fut = spawn_blocking!("zstd_store_validate_inline", move || { + let result = validate_zstd_buffer( + to_validate, + digest, + digest_function, + max_compressed_upload_size, + collect_limit, + ); + (result, permit, inflight) + }); + let ((wire_bytes_consumed, collected), _permit, _inflight) = + match tokio::time::timeout(self.stage_timeout, validate_fut).await { + Ok(joined) => { + let (result, permit, inflight) = + joined.err_tip(|| "Failed to run zstd store inline validation")?; + (result?, permit, inflight) + } + Err(_elapsed) => return Err(self.stage_timeout_err()), + }; + + self.with_commit_timeout(async { + let mut payload = data; + if let Some(collected) = collected + && let Some(smaller) = self + .maybe_recompress(collected, payload.len() as u64) + .await? + { + payload = Bytes::from(smaller); + } + self.inner_store + .update_oneshot(digest, payload) + .await + .err_tip(|| "Failed to commit inline zstd upload to inner store")?; + self.inline_commits.fetch_add(1, Ordering::Relaxed); + Ok(wire_bytes_consumed) + }) + .await + } + + /// Stage a validated compressed upload to a temp file. The returned + /// [`StagedStream`] owns the cleanup guard and staging permit. + async fn stage_compressed( + &self, + digest: DigestInfo, + digest_function: DigestHasherFunc, + reader: DropCloserReadHalf, + ) -> Result { + let permit = self.acquire_staging_permit().await?; + let inflight = OwnedInflight::enter(self.staged_uploads_inflight.clone()); + + let stage_path = format!("{}/zstd-stage-{}", self.temp_path, uuid::Uuid::new_v4()); + // Transfer an already-armed guard to the detached blocking task. If this + // async future is cancelled (including by `stage_timeout`) while that + // task is still creating or validating the file, the task keeps owning + // cleanup until it finishes; the guard comes back only on success. + let guard = TempFileGuard::arm(stage_path.clone()); + + let fs_permit = fs::get_permit().await?; + let max_compressed_upload_size = self.max_compressed_upload_size; + let collect_limit = self.recompression_limit(); + let blocking_path = stage_path.clone(); + let stage_fut = spawn_blocking!("zstd_store_stage", move || { + stage_compressed_blocking(StageCompressedInput { + reader, + path: blocking_path, + digest, + hasher_func: digest_function, + max_compressed_upload_size, + collect_limit, + fs_permit, + guard, + permit, + inflight, + }) + }); + let output = match tokio::time::timeout(self.stage_timeout, stage_fut).await { + Ok(joined) => flatten_join(joined, "Failed to run zstd store staging task")?, + Err(_elapsed) => return Err(self.stage_timeout_err()), + }; + + Ok(StagedStream { + file: Some(output.file), + path: stage_path, + compressed_size: output.wire_bytes_consumed, + wire_bytes_consumed: output.wire_bytes_consumed, + collected: output.collected, + _permit: output.permit, + _inflight: output.inflight, + _guard: output.guard, + }) + } + + /// Re-encode `collected` at the configured level, returning it only if it is + /// smaller than `current_size`. + /// + /// Best-effort by design: an upload that finds every recompression slot busy + /// returns `None` rather than queueing, because it is holding a staging + /// permit and queueing here would let a small recompression pool throttle + /// the whole upload path. + async fn maybe_recompress( + &self, + collected: Vec, + current_size: u64, + ) -> Result>, Error> { + let Ok(rec_permit) = self.recompression_semaphore.clone().try_acquire_owned() else { + self.recompressions_skipped_busy + .fetch_add(1, Ordering::Relaxed); + return Ok(None); + }; + let level = self.encode_level(); + let recompressed = flatten_join( + spawn_blocking!("zstd_store_recompress", move || { + let _rec_permit = rec_permit; + zstd::bulk::compress(&collected, level).map_err(|e| { + make_err!( + Code::Internal, + "Zstd re-compression failed in zstd store: {e}" + ) + }) + }) + .await, + "Failed to run zstd store recompress task", + )?; + if (recompressed.len() as u64) < current_size { + self.recompressions_applied.fetch_add(1, Ordering::Relaxed); + Ok(Some(recompressed)) + } else { + self.recompressions_rejected.fetch_add(1, Ordering::Relaxed); + Ok(None) + } + } + + /// Overwrite the staged file with a smaller re-encoding when one is + /// available. The validated descriptor stays open and is reused, so this + /// needs neither a second file permit (which could deadlock under + /// exhaustion) nor a reopen of the staged pathname. + async fn recompress_staged(&self, staged: &mut StagedStream) -> Result<(), Error> { + let Some(collected) = staged.collected.take() else { + return Ok(()); + }; + let Some(recompressed) = self + .maybe_recompress(collected, staged.compressed_size) + .await? + else { + return Ok(()); + }; + + let file = staged + .file + .as_mut() + .ok_or_else(|| make_err!(Code::Internal, "Staged zstd file already consumed"))?; + file.as_ref() + .set_len(0) + .await + .map_err(stage_file_err("truncate"))?; + file.seek(SeekFrom::Start(0)) + .await + .map_err(stage_file_err("rewind"))?; + file.write_all(&recompressed) + .await + .map_err(stage_file_err("write"))?; + file.flush().await.map_err(stage_file_err("flush"))?; + file.as_ref() + .sync_all() + .await + .map_err(stage_file_err("sync"))?; + file.seek(SeekFrom::Start(0)) + .await + .map_err(stage_file_err("rewind"))?; + staged.compressed_size = recompressed.len() as u64; + Ok(()) + } + + /// Commit a validated staged upload to the inner store. Only called after + /// validation succeeded, so the inner store never sees corrupt data. + /// Consumes `staged`, so the staging permit and temp file are released and + /// removed once the commit resolves. + /// + /// Both the validated descriptor and its pathname are handed over, and which + /// one the backend uses is backend-specific: stores that stream the bytes + /// (the default `StoreDriver::update_with_whole_file`, so `memory`, S3, …) + /// read the exact descriptor that was validated, while `filesystem` drops it + /// and commits by `rename(2)` on the path. For the latter, the guarantee + /// against an observe-and-replace race comes from `O_EXCL` creation of an + /// unguessable name inside an operator-private `temp_path` — enforced by + /// `post_init` — not from descriptor pinning. That backend also requires + /// `temp_path` and its `content_path` to share a filesystem; a cross-device + /// `rename` fails with `EXDEV` and the upload is rejected. + async fn commit_staged( + &self, + digest: DigestInfo, + mut staged: StagedStream, + ) -> Result<(), Error> { + let file = staged + .file + .take() + .ok_or_else(|| make_err!(Code::Internal, "Staged zstd file already consumed"))?; + let compressed_size = staged.compressed_size; + let path = OsString::from(staged.path.clone()); + self.inner_store + .update_with_whole_file( + digest, + path, + file, + UploadSizeInfo::ExactSize(compressed_size), + ) + .await + .err_tip(|| "Failed to commit staged zstd upload to inner store")?; + self.staged_commits.fetch_add(1, Ordering::Relaxed); + Ok(()) + // `staged` (permit + cleanup guard) drops here on every path. + } +} + +#[async_trait] +impl WireCompressionStore for ZstdStore { + async fn update_compressed( + self: Arc, + digest: DigestInfo, + digest_function: DigestHasherFunc, + compressor: WireCompressor, + reader: DropCloserReadHalf, + ) -> Result { + match compressor { + WireCompressor::Zstd => self.update_zstd(digest, digest_function, reader).await, + } + } + + async fn get_compressed( + self: Arc, + digest: DigestInfo, + compressor: WireCompressor, + writer: DropCloserWriteHalf, + ) -> Result<(), Error> { + match compressor { + WireCompressor::Zstd => self.get_zstd(digest, writer).await, + } + } + + async fn update_compressed_oneshot( + self: Arc, + digest: DigestInfo, + digest_function: DigestHasherFunc, + compressor: WireCompressor, + data: Bytes, + ) -> Result<(), Error> { + match compressor { + WireCompressor::Zstd => self + .update_zstd_oneshot(digest, digest_function, data) + .await + .map(|_| ()), + } + } + + async fn get_for_batch( + self: Arc, + digest: DigestInfo, + acceptable_compressors: &[WireCompressor], + ) -> Result<(Bytes, Option), Error> { + let accepts_zstd = acceptable_compressors.contains(&WireCompressor::Zstd); + let (data, is_zstd) = Self::get_for_batch(self.as_ref(), digest, accepts_zstd).await?; + Ok((data, is_zstd.then_some(WireCompressor::Zstd))) + } +} + +#[async_trait] +impl StoreDriver for ZstdStore { + async fn post_init(self: Arc) -> Result<(), Error> { + tokio::fs::create_dir_all(&self.temp_path) + .await + .map_err(|e| { + make_err!( + Code::Internal, + "Failed to create ZstdStore temp_path {}: {e}", + self.temp_path + ) + })?; + let meta = tokio::fs::metadata(&self.temp_path).await.map_err(|e| { + make_err!( + Code::Internal, + "Failed to stat ZstdStore temp_path {}: {e}", + self.temp_path + ) + })?; + if !meta.is_dir() { + return Err(make_input_err!( + "ZstdStore temp_path {} is not a directory", + self.temp_path + )); + } + #[cfg(unix)] + { + // Staged files hold validated-but-uncommitted blob contents, and a + // `filesystem` backend commits them by pathname, so the directory + // itself must not be writable by untrusted local users. A sticky + // world-writable dir (like /tmp) restricts deletes/renames to the + // owner, so it is tolerated. + let mode = meta.permissions().mode(); + if mode & 0o002 != 0 && mode & 0o1000 == 0 { + return Err(make_input_err!( + "ZstdStore temp_path {} is world-writable without the sticky bit (mode {:o}); \ + use an operator-private directory not writable by untrusted users", + self.temp_path, + mode & 0o7777 + )); + } + } + // Probe writability at startup rather than discovering it on the first + // upload. + let probe_path = format!( + "{}/.zstd-store-probe-{}", + self.temp_path, + uuid::Uuid::new_v4() + ); + tokio::fs::write(&probe_path, [0u8]).await.map_err(|e| { + make_err!( + Code::Internal, + "ZstdStore temp_path {} is not writable: {e}", + self.temp_path + ) + })?; + tokio::fs::remove_file(&probe_path).await.map_err(|e| { + make_err!( + Code::Internal, + "Failed to remove ZstdStore write-probe file {probe_path}: {e}" + ) + })?; + self.inner_store.clone().into_inner().post_init().await + } + + fn wire_compression_store(self: Arc) -> Option> { + Some(self) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + // ZstdStore is CAS-only. + for key in keys { + if let StoreKey::Str(_) = key { + return Err(make_input_err!("ZstdStore only supports digest keys")); + } + } + + // Invariant: zero digests never touch the inner store. Satisfy them + // directly and only forward the non-zero keys to the inner store. + for (key, result) in keys.iter().zip(results.iter_mut()) { + if is_zero_digest(key.borrow()) { + *result = Some(0); + } + } + + let nonzero_keys = keys + .iter() + .filter(|key| !is_zero_digest(key.borrow())) + .map(StoreKey::borrow) + .collect::>(); + if nonzero_keys.is_empty() { + return Ok(()); + } + + let mut nonzero_results = vec![None; nonzero_keys.len()]; + self.inner_store + .as_store_driver_pin() + .has_with_results(&nonzero_keys, &mut nonzero_results) + .await?; + + // Presence comes from the inner store, but the reported size is always + // the digest's uncompressed size (the physical zstd size is meaningless + // to clients). + let nonzero_slots = keys + .iter() + .zip(results.iter_mut()) + .filter_map(|(key, result)| (!is_zero_digest(key.borrow())).then_some((key, result))); + for ((key, result), inner_result) in nonzero_slots.zip(nonzero_results) { + *result = match (inner_result, key) { + (Some(_), StoreKey::Digest(digest)) => Some(digest.size_bytes()), + (other, _) => other, + }; + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + _upload_size: UploadSizeInfo, + ) -> Result { + self.update_identity(key, reader).await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.get_part_identity(key, writer, offset, length).await + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + // Representation-changing store => terminal. + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { + self.inner_store.register_remove_callback(callback) + } +} + +default_health_status_indicator!(ZstdStore); + +#[cfg(test)] +mod tests { + use std::io::Write; + + use bytes::Bytes; + use nativelink_config::stores::{MemorySpec, ZstdConfig}; + use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; + use nativelink_util::store_trait::{Store, StoreLike}; + + use super::{DecodeSink, ZstdStore, create_temp_exclusive}; + use crate::cas_utils::is_zero_digest; + use crate::memory_store::MemoryStore; + + fn unique_temp_path(tag: &str) -> std::path::PathBuf { + // Process- and call-unique name under the system temp dir. `Math`-style + // randomness is unavailable in some sandboxes, so combine the pid with a + // monotonic counter instead of a clock. + use core::sync::atomic::{AtomicU64, Ordering}; + static COUNTER: AtomicU64 = AtomicU64::new(0); + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "nativelink-zstd-unit-{tag}-{}-{n}", + std::process::id() + )) + } + + #[test] + fn create_temp_exclusive_rejects_existing_path() { + let path = unique_temp_path("exists"); + let path_str = path.to_str().unwrap(); + let first = create_temp_exclusive(path_str).expect("first exclusive create must succeed"); + drop(first); + let err = create_temp_exclusive(path_str) + .expect_err("second exclusive create on an existing path must fail"); + assert!( + err.to_string().contains("exclusively create"), + "unexpected error: {err}" + ); + drop(std::fs::remove_file(path_str)); + } + + #[cfg(unix)] + #[test] + fn create_temp_exclusive_rejects_symlink() { + let target = unique_temp_path("symlink-target"); + let link = unique_temp_path("symlink-link"); + let link_str = link.to_str().unwrap(); + // A dangling symlink at the path: O_EXCL must refuse to create through it. + std::os::unix::fs::symlink(&target, &link).expect("symlink must be created"); + let err = create_temp_exclusive(link_str) + .expect_err("exclusive create through a symlink must fail"); + assert!( + err.to_string().contains("exclusively create"), + "unexpected error: {err}" + ); + // The symlink target must never have been created (no follow). + assert!(!target.exists(), "O_EXCL must not follow the symlink"); + drop(std::fs::remove_file(link_str)); + } + + #[cfg(unix)] + #[test] + fn create_temp_exclusive_sets_owner_only_mode() { + use std::os::unix::fs::PermissionsExt; + let path = unique_temp_path("mode"); + let path_str = path.to_str().unwrap(); + let file = create_temp_exclusive(path_str).expect("exclusive create must succeed"); + let mode = file.metadata().unwrap().permissions().mode() & 0o777; + assert_eq!( + mode, 0o600, + "staging file must be created 0o600, got {mode:o}" + ); + drop(file); + drop(std::fs::remove_file(path_str)); + } + + #[test] + fn decode_sink_rejects_output_past_max_and_never_overcounts() { + // Cap at 4 bytes; a write that would exceed it is rejected wholesale and + // must not advance the counter or hash any of the offending bytes. + let mut sink = DecodeSink::new(DigestHasherFunc::Sha256, 4, None); + assert_eq!( + sink.write(b"abcd").unwrap(), + 4, + "exact-fit write is accepted" + ); + assert_eq!(sink.decoded_len, 4); + let err = sink + .write(b"e") + .expect_err("a byte past the cap must be rejected at the sink"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + assert_eq!(sink.decoded_len, 4, "rejected bytes must not be counted"); + } + + #[test] + fn decode_sink_zero_cap_rejects_first_byte() { + // The zero-digest configuration: any decoded byte is rejected. + let mut sink = DecodeSink::new(DigestHasherFunc::Sha256, 0, None); + assert!( + sink.write(b"x").is_err(), + "a zero-capacity sink must reject any output" + ); + assert_eq!(sink.decoded_len, 0); + // An empty write is a no-op success (decoders may flush zero bytes). + assert_eq!(sink.write(b"").unwrap(), 0); + } + + #[nativelink_macro::nativelink_test] + async fn batch_identity_decode_waits_for_identity_admission() { + let data = Bytes::from_static(b"batch identity admission regression payload"); + let mut hasher = DigestHasherFunc::Sha256.hasher(); + hasher.update(&data); + let digest = hasher.finalize_digest(); + let physical = zstd::bulk::compress(&data, 3).unwrap(); + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + inner + .update_oneshot(digest, Bytes::from(physical)) + .await + .unwrap(); + let zstd = ZstdStore::new( + &ZstdConfig { + temp_path: unique_temp_path("batch-admission") + .to_string_lossy() + .into_owned(), + max_compressed_upload_size: 1024 * 1024, + max_concurrent_identity_ops: 1, + ..ZstdConfig::default() + }, + inner, + ) + .unwrap(); + + let held_permit = zstd + .identity_semaphore + .clone() + .acquire_owned() + .await + .unwrap(); + let batch_read = zstd.get_for_batch(digest, false); + tokio::pin!(batch_read); + assert!( + tokio::time::timeout(core::time::Duration::from_millis(100), &mut batch_read) + .await + .is_err(), + "identity batch decode must wait for max_concurrent_identity_ops admission" + ); + + drop(held_permit); + let (decoded, is_zstd) = + tokio::time::timeout(core::time::Duration::from_secs(1), batch_read) + .await + .expect("batch decode must proceed after admission is released") + .unwrap(); + assert!(!is_zstd); + assert_eq!(decoded, data); + } + + #[test] + fn zero_byte_digests_are_recognized() { + // Guards the invariant the zero-digest fast paths rely on. + for digest in crate::cas_utils::ZERO_BYTE_DIGESTS { + assert!(is_zero_digest(digest)); + } + } +} diff --git a/nativelink-store/tests/zstd_store_test.rs b/nativelink-store/tests/zstd_store_test.rs new file mode 100644 index 000000000..abbb2b17e --- /dev/null +++ b/nativelink-store/tests/zstd_store_test.rs @@ -0,0 +1,2352 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use core::pin::Pin; +use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use core::time::Duration; +use std::ffi::OsString; +use std::io::{Read, Write}; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use nativelink_config::stores::{ + CacheMetricsSpec, CompressionAlgorithm, CompressionSpec, DedupSpec, ExistenceCacheSpec, + FastSlowSpec, FilesystemSpec, MemorySpec, NoopSpec, RefSpec, ShardConfig, ShardSpec, + SizePartitioningSpec, StoreDirection, StoreSpec, ZstdConfig, +}; +use nativelink_error::{Code, Error, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_metric::MetricsComponent; +use nativelink_store::cache_metrics_store::CacheMetricsStore; +use nativelink_store::cas_utils::{ZERO_BYTE_DIGESTS, is_zero_digest}; +use nativelink_store::dedup_store::DedupStore; +use nativelink_store::default_store_factory::store_factory; +use nativelink_store::existence_cache_store::ExistenceCacheStore; +use nativelink_store::fast_slow_store::FastSlowStore; +use nativelink_store::filesystem_store::{FileEntryImpl, FilesystemStore}; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::ref_store::RefStore; +use nativelink_store::shard_store::ShardStore; +use nativelink_store::size_partitioning_store::SizePartitioningStore; +use nativelink_store::store_manager::StoreManager; +use nativelink_store::zstd_store::ZstdStore; +use nativelink_util::buf_channel::{ + DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, +}; +use nativelink_util::common::{DigestInfo, make_temp_path}; +use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc, make_ctx_for_hash_func}; +use nativelink_util::fs::FileSlot; +use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; +use nativelink_util::store_trait::{ + RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, WireCompressor, +}; +use nativelink_util::{background_spawn, spawn}; +use opentelemetry::context::FutureExt; +use pretty_assertions::assert_eq; +use sha2::{Digest, Sha256}; +use tokio::io::{AsyncReadExt, AsyncSeekExt}; + +#[derive(Default, MetricsComponent)] +struct RecordingStore { + update_count: AtomicUsize, + /// Total number of keys the inner store observed via `has_with_results`. + has_key_count: AtomicUsize, + /// Set if any zero-byte digest was ever forwarded to the inner store. + has_saw_zero: AtomicBool, +} + +#[async_trait] +impl StoreDriver for RecordingStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + self.has_key_count.fetch_add(keys.len(), Ordering::Relaxed); + if keys.iter().any(|key| is_zero_digest(key.borrow())) { + self.has_saw_zero.store(true, Ordering::Relaxed); + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + self.update_count.fetch_add(1, Ordering::Relaxed); + Ok(0) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(Code::NotFound, "Not found")) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(RecordingStore); + +/// A store wrapper that stores data in an inner [`MemoryStore`] but, on +/// `get_part`, replays the physical bytes in many small `send` calls. This forces +/// the `ZstdStore` physical channel (a 2-slot buffer) to back up so the drain +/// behavior of the decode path is actually exercised. A real streaming backend +/// (S3, filesystem) chunks its output the same way. +#[derive(MetricsComponent)] +struct ChunkingStore { + #[metric(group = "inner")] + inner: Store, + chunk_size: usize, +} + +#[async_trait] +impl StoreDriver for ChunkingStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner.has_with_results(keys, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + self.inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + // Ignore offset/length: ZstdStore always requests the full physical blob. + let _ = (offset, length); + let full = self.inner.get_part_unchunked(key.borrow(), 0, None).await?; + for chunk in full.chunks(self.chunk_size) { + writer.send(Bytes::copy_from_slice(chunk)).await?; + } + writer.send_eof()?; + Ok(()) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(ChunkingStore); + +const TEMP_PATH: &str = "/tmp/nativelink-zstd-store-test"; + +fn spec() -> ZstdConfig { + ZstdConfig { + temp_path: TEMP_PATH.to_string(), + max_compressed_upload_size: 512 * 1024 * 1024, + // Exercise the staging path; the inline path has dedicated tests. + max_inline_commit_size: 1, + ..ZstdConfig::default() + } +} + +fn digest_for(data: &[u8]) -> DigestInfo { + let hash: [u8; 32] = Sha256::digest(data).into(); + DigestInfo::new(hash, data.len() as u64) +} + +#[nativelink_test] +async fn identity_round_trip() -> Result<(), Error> { + const DATA: &[u8] = b"hello zstd store, this compresses a bit aaaaaaaaaaaaaaaaaaaa"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + let digest = digest_for(DATA); + + store.update_oneshot(digest, DATA.into()).await?; + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(&got[..], DATA, "Expected round-tripped data to match"); + Ok(()) +} + +#[nativelink_test] +async fn identity_round_trip_partial_read() -> Result<(), Error> { + const DATA: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyz0123456789"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + let digest = digest_for(DATA); + + store.update_oneshot(digest, DATA.into()).await?; + let got = store.get_part_unchunked(digest, 5, Some(10)).await?; + assert_eq!(&got[..], &DATA[5..15], "Expected partial read to match"); + Ok(()) +} + +#[nativelink_test] +async fn zero_digest_read_returns_empty() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + for digest in ZERO_BYTE_DIGESTS { + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(got.len(), 0, "Expected zero-digest read to be empty"); + } + Ok(()) +} + +#[nativelink_test] +async fn zero_digest_write_empty_skips_inner() -> Result<(), Error> { + let recording = Arc::new(RecordingStore::default()); + let store = Store::new(ZstdStore::new(&spec(), Store::new(recording.clone()))?); + + for digest in ZERO_BYTE_DIGESTS { + store.update_oneshot(digest, Bytes::new()).await?; + } + // The inner store's update must never be called for zero digests. + assert_eq!( + recording.update_count.load(Ordering::Relaxed), + 0, + "Zero digest must not reach the inner store" + ); + Ok(()) +} + +#[nativelink_test] +async fn str_key_is_rejected() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + let keys = [StoreKey::new_str("some-string-key")]; + let mut results = [None]; + let err = store + .has_with_results(&keys, &mut results) + .await + .expect_err("Expected a Str key to be rejected"); + assert!( + err.to_string().contains("only supports digest keys"), + "Unexpected error: {err}" + ); + Ok(()) +} + +#[nativelink_test] +async fn has_reports_uncompressed_digest_size() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + // Highly compressible payload so the physical zstd size differs from the + // uncompressed digest size. + let data = vec![0u8; 4096]; + let digest = digest_for(&data); + + store.update_oneshot(digest, data.clone().into()).await?; + let reported = store.has(digest).await?; + assert_eq!( + reported, + Some(data.len() as u64), + "has must report the uncompressed digest size, not the physical zstd size" + ); + Ok(()) +} + +#[nativelink_test] +async fn ranged_read_large_blob_drains_inner_stream() -> Result<(), Error> { + // Regression for the partial-read drain bug: a ranged read on a blob whose + // physical (compressed) stream is larger than the 2-slot channel buffer must + // still return the requested bytes. Before the fix the decode loop returned + // as soon as the requested `length` was produced, dropping the physical + // reader while the spawned inner `get` was still streaming compressed bytes; + // its next `send` then failed with "receiver disconnected" and the whole + // get_part surfaced that error instead of the requested bytes. + // + // Use a repeating-but-nontrivial pattern (not all-zero) so the compressed + // stream is large enough to span multiple channel sends. + let mut data = Vec::with_capacity(1024 * 1024); + let mut counter: u32 = 0x1234_5678; + while data.len() < 1024 * 1024 { + counter = counter.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + data.extend_from_slice(&counter.to_le_bytes()); + } + + // A chunking inner store forces the physical stream to span many channel + // sends (16 KiB each) so the 2-slot channel backs up while the decoder is + // still consuming; this is what makes an early return drop the reader and + // fail the inner get. + let inner = Store::new(Arc::new(ChunkingStore { + inner: Store::new(MemoryStore::new(&MemorySpec::default())), + chunk_size: 16 * 1024, + })); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + + let offset = 300 * 1024; + let length = 100 * 1024; + let got = store + .get_part_unchunked(digest, offset as u64, Some(length as u64)) + .await?; + assert_eq!( + &got[..], + &data[offset..offset + length], + "Ranged read of a large blob must return the requested sub-range" + ); + Ok(()) +} + +#[nativelink_test] +async fn corrupt_upload_hash_mismatch_is_rejected_and_not_committed() -> Result<(), Error> { + const DATA: &[u8] = b"the quick brown fox jumps over the lazy dog"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner.clone())?); + + // Claim a digest whose hash does NOT match the data (but the size matches). + let mut wrong = DATA.to_vec(); + wrong[0] ^= 0xFF; + let bad_digest = digest_for(&wrong); + assert_eq!(bad_digest.size_bytes(), DATA.len() as u64); + + let err = store + .update_oneshot(bad_digest, DATA.into()) + .await + .expect_err("Expected a hash mismatch to be rejected"); + assert_eq!( + err.code, + Code::InvalidArgument, + "Hash mismatch must be InvalidArgument, got: {err}" + ); + + // The blob must NOT have committed to the inner store. + assert_eq!( + store.has(bad_digest).await?, + None, + "A rejected upload must not be visible via the zstd store" + ); + assert_eq!( + inner.has(bad_digest).await?, + None, + "A rejected upload must not commit to the inner store" + ); + Ok(()) +} + +#[nativelink_test] +async fn corrupt_upload_size_mismatch_is_rejected_and_not_committed() -> Result<(), Error> { + const DATA: &[u8] = b"the quick brown fox jumps over the lazy dog"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner.clone())?); + + // Correct hash, but claim a size larger than the actual data length. + let hash: [u8; 32] = Sha256::digest(DATA).into(); + let bad_digest = DigestInfo::new(hash, DATA.len() as u64 + 10); + + let err = store + .update_oneshot(bad_digest, DATA.into()) + .await + .expect_err("Expected a size mismatch to be rejected"); + assert_eq!( + err.code, + Code::InvalidArgument, + "Size mismatch must be InvalidArgument, got: {err}" + ); + + assert_eq!( + store.has(bad_digest).await?, + None, + "A rejected upload must not be visible via the zstd store" + ); + assert_eq!( + inner.has(bad_digest).await?, + None, + "A rejected upload must not commit to the inner store" + ); + Ok(()) +} + +#[nativelink_test] +async fn new_rejects_out_of_range_compression_levels() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + + for level in [0, 20] { + let mut bad_spec = spec(); + bad_spec.compression_level = Some(level); + assert!( + ZstdStore::new(&bad_spec, inner.clone()).is_err(), + "compression_level {level} must be rejected" + ); + } + + for level in [Some(1), Some(19), None] { + let mut ok_spec = spec(); + ok_spec.compression_level = level; + assert!( + ZstdStore::new(&ok_spec, inner.clone()).is_ok(), + "compression_level {level:?} must be accepted" + ); + } + Ok(()) +} + +#[nativelink_test] +async fn identity_round_trip_blake3_context() -> Result<(), Error> { + const DATA: &[u8] = b"blake3 round trip payload aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + // Compute the BLAKE3 digest of the payload using the context hasher. + let mut hasher = DigestHasherFunc::Blake3.hasher(); + hasher.update(DATA); + let digest = hasher.finalize_digest(); + + // Run both the update and get under a BLAKE3 request-digest context, proving + // the context-hasher mechanism (not the SHA-256 default) is honored. + let ctx = make_ctx_for_hash_func(DigestHasherFunc::Blake3)?; + store + .update_oneshot(digest, DATA.into()) + .with_context(ctx.clone()) + .await?; + let got = store + .get_part_unchunked(digest, 0, None) + .with_context(ctx) + .await?; + assert_eq!(&got[..], DATA, "Expected BLAKE3 round-trip to match"); + Ok(()) +} + +#[nativelink_test] +async fn has_zero_digest_never_touches_inner_store() -> Result<(), Error> { + let recording = Arc::new(RecordingStore::default()); + let store = Store::new(ZstdStore::new(&spec(), Store::new(recording.clone()))?); + + // Mixed batch: two zero digests and one non-zero (absent) digest. + let nonzero = digest_for(b"not present"); + let keys = [ + StoreKey::Digest(ZERO_BYTE_DIGESTS[0]), + StoreKey::Digest(nonzero), + StoreKey::Digest(ZERO_BYTE_DIGESTS[1]), + ]; + let mut results = [None, None, None]; + store.has_with_results(&keys, &mut results).await?; + + assert_eq!( + results, + [Some(0), None, Some(0)], + "Zero digests must report Some(0); absent non-zero must be None" + ); + assert!( + !recording.has_saw_zero.load(Ordering::Relaxed), + "Zero digests must never be forwarded to the inner store" + ); + assert_eq!( + recording.has_key_count.load(Ordering::Relaxed), + 1, + "Only the single non-zero key may reach the inner store" + ); + Ok(()) +} + +#[nativelink_test] +async fn has_mixed_batch_reports_correct_sizes() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner)?); + + let present = vec![7u8; 4096]; + let present_digest = digest_for(&present); + store + .update_oneshot(present_digest, present.clone().into()) + .await?; + let absent_digest = digest_for(b"definitely absent"); + + let keys = [ + StoreKey::Digest(ZERO_BYTE_DIGESTS[0]), + StoreKey::Digest(present_digest), + StoreKey::Digest(absent_digest), + ]; + let mut results = [None, None, None]; + store.has_with_results(&keys, &mut results).await?; + assert_eq!( + results, + [Some(0), Some(present.len() as u64), None], + "Mixed batch must report zero=Some(0), present=uncompressed size, absent=None" + ); + Ok(()) +} + +#[nativelink_test] +async fn factory_selects_zstd_wire_capability() -> Result<(), Error> { + let store_spec = StoreSpec::Compression(Box::new(CompressionSpec { + backend: StoreSpec::Memory(MemorySpec::default()), + compression_algorithm: CompressionAlgorithm::Zstd(spec()), + })); + let store_manager = Arc::new(StoreManager::new()); + let store = store_factory(&store_spec, &store_manager, None).await?; + assert!( + store.wire_compression_store().is_some(), + "Expected CompressionAlgorithm::Zstd to expose the wire-compression capability" + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// zstd fast-path (get_zstd / update_zstd / update_zstd_oneshot / get_for_batch) +// --------------------------------------------------------------------------- + +/// A spec pointing at a specific temp dir (for tests that assert temp-dir +/// emptiness and therefore must not share the global staging directory). +/// +/// `max_inline_commit_size: 1` forces even tiny oneshot payloads through the +/// staging path, which is what most tests below are actually asserting about. +/// The inline path has its own tests; see [`spec_for_inline`]. +fn spec_for(temp_path: String) -> ZstdConfig { + ZstdConfig { + temp_path, + max_compressed_upload_size: 512 * 1024 * 1024, + max_inline_commit_size: 1, + ..ZstdConfig::default() + } +} + +/// Like [`spec_for`], but leaves `max_inline_commit_size` at its default so +/// small compressed oneshot payloads take the in-memory commit path. +fn spec_for_inline(temp_path: String) -> ZstdConfig { + ZstdConfig { + temp_path, + max_compressed_upload_size: 512 * 1024 * 1024, + ..ZstdConfig::default() + } +} + +/// Builds a `ZstdStore` over a fresh `MemoryStore`, ensuring the staging dir +/// exists. Returns the concrete store (for the fast-path methods), a `Store` +/// wrapper (for the identity round-trip view), and the inner store (for +/// physical inspection). +async fn build(spec: &ZstdConfig) -> Result<(Arc, Store, Store), Error> { + std::fs::create_dir_all(&spec.temp_path) + .map_err(|e| make_err!(Code::Internal, "Failed to create test temp dir: {e}"))?; + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let zstd = ZstdStore::new(spec, inner.clone())?; + let store = Store::new(zstd.clone()); + Ok((zstd, store, inner)) +} + +/// Drive `get_zstd` and collect the full emitted (physical zstd) stream. +async fn collect_zstd(store: &ZstdStore, digest: DigestInfo) -> Result { + let (tx, mut rx) = make_buf_channel_pair(); + let (get_res, collected) = tokio::join!(store.get_zstd(digest, tx), rx.consume(None)); + get_res?; + collected +} + +/// A `DropCloserReadHalf` that yields the given chunks then EOF, fed by a task. +fn reader_from(chunks: Vec) -> DropCloserReadHalf { + let (mut tx, rx) = make_buf_channel_pair(); + background_spawn!("zstd_test_reader_feed", async move { + for chunk in chunks { + if tx.send(chunk).await.is_err() { + return; + } + } + drop(tx.send_eof()); + }); + rx +} + +/// Highly compressible data (a repeated pseudo-random block) whose compressed +/// size is sensitive to the zstd level. +fn compressible_data(len: usize) -> Vec { + let mut block = [0u8; 251]; + let mut state: u32 = 0x9E37_79B1; + for b in &mut block { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + *b = (state >> 24) as u8; + } + let mut out = Vec::with_capacity(len); + while out.len() < len { + let take = (len - out.len()).min(block.len()); + out.extend_from_slice(&block[..take]); + } + out +} + +fn random_bytes(len: usize) -> Vec { + use rand::RngCore; + let mut out = vec![0u8; len]; + rand::rng().fill_bytes(&mut out); + out +} + +/// Data whose compressed size is strongly level-sensitive: two copies of a +/// large incompressible block. Only a large-window (high level) encoder can +/// dedup the second copy, so a low level produces a much larger stream. +fn level_sensitive_data() -> Vec { + let block = random_bytes(512 * 1024); + let mut out = Vec::with_capacity(block.len() * 2); + out.extend_from_slice(&block); + out.extend_from_slice(&block); + out +} + +fn dir_entry_count(path: &str) -> usize { + std::fs::read_dir(path).map_or(0, Iterator::count) +} + +#[nativelink_test] +async fn get_zstd_is_byte_for_byte_passthrough() -> Result<(), Error> { + const DATA: &[u8] = b"passthrough payload aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + + let (zstd, _store, _inner) = build(&spec()).await?; + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + + let wire = zstd + .update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await?; + assert_eq!(wire, compressed.len() as u64, "wire bytes must match input"); + + let got = collect_zstd(&zstd, digest).await?; + assert_eq!( + &got[..], + &compressed[..], + "get_zstd must return the stored zstd stream byte-for-byte" + ); + let decoded = zstd::stream::decode_all(&got[..]).unwrap(); + assert_eq!( + &decoded[..], + DATA, + "the passthrough stream must decode back" + ); + Ok(()) +} + +#[nativelink_test] +async fn get_zstd_preserves_concatenated_frames() -> Result<(), Error> { + // Two independently-compressed zstd frames concatenated into one stream. + const PART_A: &[u8] = b"first frame content aaaaaaaaaaaaaaaaaaaaaaaa"; + const PART_B: &[u8] = b"second frame content bbbbbbbbbbbbbbbbbbbbbbbb"; + + let frame_a = zstd::bulk::compress(PART_A, 3).unwrap(); + let frame_b = zstd::bulk::compress(PART_B, 3).unwrap(); + let mut concatenated = frame_a.clone(); + concatenated.extend_from_slice(&frame_b); + + let mut raw = PART_A.to_vec(); + raw.extend_from_slice(PART_B); + let digest = digest_for(&raw); + + let (zstd, _store, _inner) = build(&spec()).await?; + // Upload the two frames as separate stream chunks to exercise the streaming + // path; passthrough must not re-frame them. + let wire = zstd + .update_zstd( + digest, + DigestHasherFunc::Sha256, + reader_from(vec![Bytes::from(frame_a), Bytes::from(frame_b)]), + ) + .await?; + assert_eq!(wire, concatenated.len() as u64); + + let got = collect_zstd(&zstd, digest).await?; + assert_eq!( + &got[..], + &concatenated[..], + "passthrough must preserve the exact concatenated-frame bytes" + ); + assert_eq!( + zstd::stream::decode_all(&got[..]).unwrap(), + raw, + "concatenated frames must decode to the concatenated raw content" + ); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_round_trips_through_identity() -> Result<(), Error> { + const DATA: &[u8] = b"round trip payload cccccccccccccccccccccccccccccccc"; + + let (zstd, store, _inner) = build(&spec()).await?; + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + + let wire = zstd + .update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await?; + assert_eq!(wire, compressed.len() as u64); + + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!( + &got[..], + DATA, + "identity read must return the original raw bytes" + ); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_round_trips_with_blake3() -> Result<(), Error> { + const DATA: &[u8] = b"blake3 compressed upload dddddddddddddddddddddddddddd"; + + let (zstd, store, _inner) = build(&spec()).await?; + let mut hasher = DigestHasherFunc::Blake3.hasher(); + hasher.update(DATA); + let digest = hasher.finalize_digest(); + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + + zstd.update_zstd_oneshot(digest, DigestHasherFunc::Blake3, Bytes::from(compressed)) + .await?; + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(&got[..], DATA, "BLAKE3 compressed round-trip must match"); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_rejects_hash_mismatch() -> Result<(), Error> { + const DATA: &[u8] = b"the quick brown fox jumps over the lazy dog"; + + let temp = make_temp_path("zstd-hash-mismatch"); + let (zstd, _store, inner) = build(&spec_for(temp.clone())).await?; + + // Compress DATA but claim the digest of a different blob of the same size. + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let mut other = DATA.to_vec(); + other[0] ^= 0xFF; + let bad_digest = digest_for(&other); + assert_eq!(bad_digest.size_bytes(), DATA.len() as u64); + + let err = zstd + .update_zstd_oneshot( + bad_digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed), + ) + .await + .expect_err("hash mismatch must be rejected"); + assert_eq!( + err.code, + Code::InvalidArgument, + "hash mismatch must be InvalidArgument, got: {err}" + ); + assert_eq!( + inner.has(bad_digest).await?, + None, + "a rejected compressed upload must not commit to the inner store" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "a rejected upload must not leave staged temp files" + ); + Ok(()) +} + +/// A complete frame followed by the beginning of a second frame is not a +/// complete concatenated zstd stream. `flush` alone accepted this shape, so it +/// is important that both the staged and zero-digest validation paths finalize +/// the decoder at EOF. +#[nativelink_test] +async fn update_zstd_rejects_incomplete_trailing_frame() -> Result<(), Error> { + const DATA: &[u8] = b"complete first frame followed by a truncated second frame"; + + let temp = make_temp_path("zstd-incomplete-trailing-frame"); + let (zstd, _store, inner) = build(&spec_for(temp.clone())).await?; + let digest = digest_for(DATA); + let mut trailing = zstd::bulk::compress(DATA, 3).unwrap(); + let next_frame = zstd::bulk::compress(b"second frame", 3).unwrap(); + trailing.extend_from_slice(&next_frame[..4]); // zstd magic, but no full frame. + + let err = zstd + .update_zstd_oneshot(digest, DigestHasherFunc::Sha256, Bytes::from(trailing)) + .await + .expect_err("a partial trailing zstd frame must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + assert_eq!( + inner.has(digest).await?, + None, + "a stream with an incomplete trailing frame must not commit" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "a rejected trailing frame must not leave a staging file" + ); + + let zero = ZERO_BYTE_DIGESTS[0]; + let mut empty_then_partial = zstd::bulk::compress(&[], 3).unwrap(); + empty_then_partial.extend_from_slice(&next_frame[..4]); + let err = zstd + .update_zstd_oneshot( + zero, + DigestHasherFunc::Sha256, + Bytes::from(empty_then_partial), + ) + .await + .expect_err("zero-digest validation must reject a partial trailing frame too"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + Ok(()) +} + +/// The decoded-size bound alone does not limit zstd's history allocation: a +/// small frame header may request a much larger window before producing output. +/// Reject frames beyond the store's 8 MiB interoperability ceiling. +#[nativelink_test] +async fn update_zstd_rejects_a_window_larger_than_eight_mib() -> Result<(), Error> { + let temp = make_temp_path("zstd-window-limit"); + let (zstd, _store, inner) = build(&spec_for(temp.clone())).await?; + let data = random_bytes(9 * 1024 * 1024); + let digest = digest_for(&data); + + let mut encoder = zstd::stream::Encoder::new(Vec::new(), 3).unwrap(); + encoder.window_log(24).unwrap(); + encoder.write_all(&data).unwrap(); + let compressed = encoder.finish().unwrap(); + + // Prove the fixture is otherwise valid and specifically needs a window + // above the store limit. + let mut decoder = zstd::stream::read::Decoder::new(compressed.as_slice()).unwrap(); + decoder.window_log_max(24).unwrap(); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).unwrap(); + assert_eq!(decoded, data); + let mut limited_decoder = zstd::stream::read::Decoder::new(compressed.as_slice()).unwrap(); + limited_decoder.window_log_max(23).unwrap(); + assert!( + limited_decoder.read_to_end(&mut Vec::new()).is_err(), + "the fixture must require a window above the store limit" + ); + + let err = zstd + .update_zstd_oneshot(digest, DigestHasherFunc::Sha256, Bytes::from(compressed)) + .await + .expect_err("a zstd frame requiring more than an 8 MiB window must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + assert_eq!(inner.has(digest).await?, None); + assert_eq!(dir_entry_count(&temp), 0); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_rejects_oversize_and_cleans_up() -> Result<(), Error> { + let temp = make_temp_path("zstd-oversize"); + let mut spec = spec_for(temp.clone()); + spec.max_compressed_upload_size = 32; // Very small cap. + let (zstd, _store, inner) = build(&spec).await?; + + // Incompressible data so the compressed stream comfortably exceeds 32 bytes. + let data = random_bytes(4096); + let compressed = zstd::bulk::compress(&data, 3).unwrap(); + assert!(compressed.len() as u64 > spec.max_compressed_upload_size); + let digest = digest_for(&data); + + let err = zstd + .update_zstd_oneshot(digest, DigestHasherFunc::Sha256, Bytes::from(compressed)) + .await + .expect_err("oversize upload must be rejected"); + assert_eq!( + err.code, + Code::ResourceExhausted, + "oversize upload must be ResourceExhausted, got: {err}" + ); + assert_eq!( + inner.has(digest).await?, + None, + "an oversize upload must not commit to the inner store" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "an oversize upload must not leave staged temp files" + ); + Ok(()) +} + +/// Cancelling the async request after the blocking stage creates its file must +/// not leak it. The guard is moved into the detached blocking task before that +/// task can create the path; after EOF lets it finish, dropping its unobserved +/// result closes the descriptor before removing the file. +#[nativelink_test] +async fn cancelled_staging_task_eventually_cleans_its_temp_file() -> Result<(), Error> { + const DATA: &[u8] = b"cancelled staging cleanup payload"; + + let temp = make_temp_path("zstd-cancelled-stage"); + let (zstd, _store, _inner) = build(&spec_for(temp.clone())).await?; + let digest = digest_for(DATA); + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let (mut tx, rx) = make_buf_channel_pair(); + let task_store = zstd.clone(); + let task = tokio::spawn(async move { + task_store + .update_zstd(digest, DigestHasherFunc::Sha256, rx) + .await + }); + + tx.send(Bytes::from(compressed)) + .await + .map_err(|e| make_err!(Code::Internal, "failed to feed staging task: {e}"))?; + tokio::time::timeout(Duration::from_secs(1), async { + while dir_entry_count(&temp) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the blocking stage must create its temp file before EOF"); + + task.abort(); + assert!(task.await.is_err(), "the request task must be cancelled"); + tx.send_eof() + .map_err(|e| make_err!(Code::Internal, "failed to finish staging input: {e}"))?; + + tokio::time::timeout(Duration::from_secs(1), async { + while dir_entry_count(&temp) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the detached blocking stage must eventually remove its temp file"); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_recompresses_to_smaller_stream() -> Result<(), Error> { + let data = level_sensitive_data(); + let digest = digest_for(&data); + + // Upload a poorly-compressed (level 1) stream; store re-compresses at 19. + let mut spec = spec_for(make_temp_path("zstd-recompress")); + spec.compression_level = Some(19); + spec.max_recompression_size = 16 * 1024 * 1024; + let (zstd, _store, inner) = build(&spec).await?; + + let poorly = zstd::bulk::compress(&data, 1).unwrap(); + zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(poorly.clone()), + ) + .await?; + let physical = inner.get_part_unchunked(digest, 0, None).await?; + assert!( + physical.len() < poorly.len(), + "re-compression must shrink the stored stream ({} !< {})", + physical.len(), + poorly.len() + ); + assert_eq!( + zstd::stream::decode_all(&physical[..]).unwrap(), + data, + "re-compressed stream must still decode to the original content" + ); + + // Control: an already-well-compressed upload is kept as-is (never enlarged). + let (zstd2, _store2, inner2) = build(&spec).await?; + let well = zstd::bulk::compress(&data, 19).unwrap(); + zstd2 + .update_zstd_oneshot(digest, DigestHasherFunc::Sha256, Bytes::from(well.clone())) + .await?; + let physical2 = inner2.get_part_unchunked(digest, 0, None).await?; + assert_eq!( + &physical2[..], + &well[..], + "a well-compressed upload must be kept byte-for-byte, not enlarged" + ); + Ok(()) +} + +#[nativelink_test] +async fn update_zstd_skips_recompression_above_threshold() -> Result<(), Error> { + let data = compressible_data(64 * 1024); + let digest = digest_for(&data); + + let mut spec = spec_for(make_temp_path("zstd-above-threshold")); + spec.compression_level = Some(19); + spec.max_recompression_size = 16; // Uncompressed size far exceeds this. + let (zstd, _store, inner) = build(&spec).await?; + + let poorly = zstd::bulk::compress(&data, 1).unwrap(); + zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(poorly.clone()), + ) + .await?; + let physical = inner.get_part_unchunked(digest, 0, None).await?; + assert_eq!( + &physical[..], + &poorly[..], + "above max_recompression_size the stored stream must equal the upload byte-for-byte" + ); + Ok(()) +} + +#[nativelink_test] +async fn get_for_batch_selects_zstd_or_raw() -> Result<(), Error> { + let (zstd, store, _inner) = build(&spec()).await?; + let wire_store = Store::new(zstd.clone()) + .wire_compression_store() + .expect("ZstdStore must expose the wire-compression capability"); + + // Compressible blob (stored via the identity path so inner holds zstd). + let compressible = vec![0u8; 4096]; + let comp_digest = digest_for(&compressible); + store + .update_oneshot(comp_digest, compressible.clone().into()) + .await?; + + let (payload, compressor) = wire_store + .clone() + .get_for_batch(comp_digest, &[WireCompressor::Zstd]) + .await?; + assert_eq!( + compressor, + Some(WireCompressor::Zstd), + "compressible blob with accepts_zstd must return zstd" + ); + assert!( + (payload.len() as u64) < comp_digest.size_bytes(), + "physical zstd must be smaller than the uncompressed size" + ); + assert_eq!( + zstd::stream::decode_all(&payload[..]).unwrap(), + compressible, + "returned zstd must decode to the original content" + ); + + // Same blob, but the client does not accept zstd => raw. + let (raw, compressor) = wire_store.clone().get_for_batch(comp_digest, &[]).await?; + assert_eq!(compressor, None); + assert_eq!( + &raw[..], + &compressible[..], + "must return raw when zstd not accepted" + ); + + // Incompressible blob: physical zstd is not smaller => always raw. + let incompressible = random_bytes(2048); + let inc_digest = digest_for(&incompressible); + store + .update_oneshot(inc_digest, incompressible.clone().into()) + .await?; + let (raw, compressor) = wire_store + .get_for_batch(inc_digest, &[WireCompressor::Zstd]) + .await?; + assert_eq!( + compressor, None, + "incompressible blob must not be served as zstd" + ); + assert_eq!(&raw[..], &incompressible[..]); + Ok(()) +} + +#[nativelink_test] +async fn get_for_batch_decodes_concatenated_frames() -> Result<(), Error> { + // Regression for `get_for_batch`'s raw-decode path: the physical stream + // stored for a blob can be two independently-compressed zstd frames + // concatenated together (see `get_zstd_preserves_concatenated_frames`). + // `zstd::bulk::decompress` must handle that multi-frame input, not just a + // single-frame stream, when the client does not accept zstd. + const PART_A: &[u8] = b"first frame content aaaaaaaaaaaaaaaaaaaaaaaa"; + const PART_B: &[u8] = b"second frame content bbbbbbbbbbbbbbbbbbbbbbbb"; + + let frame_a = zstd::bulk::compress(PART_A, 3).unwrap(); + let frame_b = zstd::bulk::compress(PART_B, 3).unwrap(); + + let mut raw = PART_A.to_vec(); + raw.extend_from_slice(PART_B); + let digest = digest_for(&raw); + + let (zstd, _store, _inner) = build(&spec()).await?; + zstd.update_zstd( + digest, + DigestHasherFunc::Sha256, + reader_from(vec![Bytes::from(frame_a), Bytes::from(frame_b)]), + ) + .await?; + + // Client does not accept zstd => `get_for_batch` must decode the stored + // (concatenated-frame) physical bytes to the full raw content. + let (payload, is_zstd) = zstd.get_for_batch(digest, false).await?; + assert!( + !is_zstd, + "client that does not accept zstd must get raw bytes" + ); + assert_eq!( + &payload[..], + &raw[..], + "get_for_batch must decode a concatenated-frame physical stream to the full raw content" + ); + Ok(()) +} + +#[nativelink_test] +async fn zero_digest_zstd_fast_path() -> Result<(), Error> { + let recording = Arc::new(RecordingStore::default()); + std::fs::create_dir_all(TEMP_PATH).map_err(|e| make_err!(Code::Internal, "temp dir: {e}"))?; + let zstd = ZstdStore::new(&spec(), Store::new(recording.clone()))?; + + let zero = ZERO_BYTE_DIGESTS[0]; + + // get_zstd yields a valid zstd stream that decodes to empty. + let got = collect_zstd(&zstd, zero).await?; + assert!( + !got.is_empty(), + "zero-digest get_zstd must emit a real zstd frame" + ); + assert_eq!( + zstd::stream::decode_all(&got[..]).unwrap().len(), + 0, + "zero-digest zstd stream must decode to empty" + ); + + // update_zstd of an empty-decoding stream succeeds without touching inner. + let empty = zstd::bulk::compress(&[], 3).unwrap(); + let wire = zstd + .update_zstd_oneshot(zero, DigestHasherFunc::Sha256, Bytes::from(empty.clone())) + .await?; + assert_eq!(wire, empty.len() as u64); + assert_eq!( + recording.update_count.load(Ordering::Relaxed), + 0, + "zero digest must never reach the inner store" + ); + + // A stream that decodes to non-empty content is rejected for a zero digest. + let non_empty = zstd::bulk::compress(b"not empty", 3).unwrap(); + let err = zstd + .update_zstd_oneshot(zero, DigestHasherFunc::Sha256, Bytes::from(non_empty)) + .await + .expect_err("non-empty content under a zero digest must be rejected"); + assert_eq!(err.code, Code::InvalidArgument); + Ok(()) +} + +#[nativelink_test] +async fn cancelled_update_zstd_leaves_no_temp_files() -> Result<(), Error> { + let temp = make_temp_path("zstd-cancel"); + let (zstd, _store, _inner) = build(&spec_for(temp.clone())).await?; + + let data = compressible_data(256 * 1024); + let compressed = zstd::bulk::compress(&data, 3).unwrap(); + let digest = digest_for(&data); + + let (mut tx, rx) = make_buf_channel_pair(); + let zstd_clone = zstd.clone(); + // `spawn!` returns a guard that aborts the task when it is dropped. + let handle = spawn!("zstd_test_cancel", async move { + zstd_clone + .update_zstd(digest, DigestHasherFunc::Sha256, rx) + .await + }); + + // Send only the first half of the compressed stream, never EOF. + let half = compressed.len() / 2; + tx.send(Bytes::copy_from_slice(&compressed[..half])) + .await + .ok(); + tokio::time::sleep(Duration::from_millis(50)).await; + + // Cancel the upload mid-stream (dropping the guard aborts it), then + // release the sender. + drop(handle); + drop(tx); + tokio::time::sleep(Duration::from_millis(150)).await; + + assert_eq!( + dir_entry_count(&temp), + 0, + "a cancelled upload must not leave staged temp files" + ); + Ok(()) +} + +#[nativelink_test] +async fn concurrent_uploads_respect_staging_semaphore() -> Result<(), Error> { + const DATA_A: &[u8] = b"concurrent upload alpha aaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + const DATA_B: &[u8] = b"concurrent upload beta bbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + + let mut spec = spec_for(make_temp_path("zstd-semaphore")); + spec.max_concurrent_staged_uploads = 1; // Serialize staging. + let (zstd, store, _inner) = build(&spec).await?; + + let digest_a = digest_for(DATA_A); + let digest_b = digest_for(DATA_B); + let comp_a = zstd::bulk::compress(DATA_A, 3).unwrap(); + let comp_b = zstd::bulk::compress(DATA_B, 3).unwrap(); + + let za = zstd.clone(); + let zb = zstd.clone(); + let (ra, rb) = tokio::join!( + za.update_zstd_oneshot(digest_a, DigestHasherFunc::Sha256, Bytes::from(comp_a)), + zb.update_zstd_oneshot(digest_b, DigestHasherFunc::Sha256, Bytes::from(comp_b)), + ); + ra?; + rb?; + + assert_eq!( + &store.get_part_unchunked(digest_a, 0, None).await?[..], + DATA_A, + "first concurrent upload must round-trip" + ); + assert_eq!( + &store.get_part_unchunked(digest_b, 0, None).await?[..], + DATA_B, + "second concurrent upload must round-trip" + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Cross-wrapper integration: ZstdStore composed with other stores. +// +// These prove ZstdStore is a well-behaved participant in a store stack. Every +// test does a full identity round-trip (write raw via `update_oneshot`, read raw +// back via `get_part_unchunked`) through the composed stack. Where a wrapper +// lives INSIDE ZstdStore (so it observes the physical zstd stream) the test also +// asserts that wrapper's side effect still fires. +// +// Identity-path tests do not touch `temp_path` (only the zstd fast path stages +// files), so they reuse the shared `spec()` without creating a staging dir. +// --------------------------------------------------------------------------- + +// 1. ZstdStore over `fast_slow` (fast=memory, slow=memory): round-trip, then +// prove a slow-hit read repopulates the fast tier with the physical zstd blob. +#[nativelink_test] +async fn over_fast_slow_populates_fast_tier() -> Result<(), Error> { + let fast_mem = MemoryStore::new(&MemorySpec::default()); + let slow_mem = MemoryStore::new(&MemorySpec::default()); + let fast_slow = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Memory(MemorySpec::default()), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(fast_mem.clone()), + Store::new(slow_mem.clone()), + )); + let store = Store::new(ZstdStore::new(&spec(), fast_slow)?); + + let data = compressible_data(64 * 1024); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-fast_slow must return the raw bytes" + ); + + // The physical (compressed) bytes landed in both tiers; capture the slow copy. + let physical = slow_mem.get_part_unchunked(digest, 0, None).await?; + assert!( + physical.len() < data.len(), + "inner store must hold the compressed physical stream ({} !< {})", + physical.len(), + data.len() + ); + + // Evict the fast tier, leaving only the slow tier populated. + assert!( + fast_mem.remove_entry(digest.into()).await, + "fast tier should have held the blob before eviction" + ); + assert_eq!( + fast_mem.has(digest).await?, + None, + "fast tier must be empty after eviction" + ); + + // A read through the stack is a slow-hit that must repopulate the fast tier. + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "slow-hit read-through must still return the raw bytes" + ); + assert_eq!( + fast_mem.has(digest).await?, + Some(physical.len() as u64), + "fast tier must be repopulated with the physical zstd blob" + ); + assert_eq!( + &fast_mem.get_part_unchunked(digest, 0, None).await?[..], + &physical[..], + "the repopulated fast-tier bytes must be the physical zstd stream" + ); + Ok(()) +} + +// 2. ZstdStore over `dedup`: a blob whose physical stream exceeds the dedup block +// size must split across multiple content chunks yet round-trip byte-for-byte. +#[nativelink_test] +async fn over_dedup_splits_and_round_trips() -> Result<(), Error> { + let content_mem = MemoryStore::new(&MemorySpec::default()); + let dedup = Store::new(DedupStore::new( + &DedupSpec { + index_store: StoreSpec::Memory(MemorySpec::default()), + content_store: StoreSpec::Memory(MemorySpec::default()), + min_size: 8 * 1024, + normal_size: 32 * 1024, + max_size: 128 * 1024, + max_concurrent_fetch_per_get: 10, + }, + Store::new(MemoryStore::new(&MemorySpec::default())), + Store::new(content_mem.clone()), + )?); + let store = Store::new(ZstdStore::new(&spec(), dedup)?); + + // Incompressible data so the physical zstd stream stays ~256 KiB, comfortably + // above the 128 KiB max block size and therefore split into several chunks. + let data = random_bytes(256 * 1024); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-dedup must return the raw bytes" + ); + assert!( + content_mem.len_for_test() > 1, + "dedup must split the physical zstd stream into multiple content chunks, got {}", + content_mem.len_for_test() + ); + Ok(()) +} + +// 3. ZstdStore over `existence_cache`: round-trip and prove the existence cache +// side effect (population on write) still fires for the forwarded digest. +#[nativelink_test] +async fn over_existence_cache_fires_side_effect() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let existence = ExistenceCacheStore::new( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), // Unused: inner is passed directly. + eviction_policy: None, + }, + inner, + ); + let store = Store::new(ZstdStore::new(&spec(), Store::new(existence.clone()))?); + + let data = compressible_data(4096); + let digest = digest_for(&data); + assert!( + !existence.exists_in_cache(&digest).await, + "digest must not be cached before the write" + ); + + store.update_oneshot(digest, data.clone().into()).await?; + assert!( + existence.exists_in_cache(&digest).await, + "write through zstd must populate the inner existence cache" + ); + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-existence_cache must return the raw bytes" + ); + // `has` is served from the now-populated existence cache and reports the + // uncompressed digest size (not the physical zstd size). + assert_eq!( + store.has(digest).await?, + Some(data.len() as u64), + "has must report the uncompressed size via the existence cache" + ); + Ok(()) +} + +// 4. ZstdStore over `cache_metrics`: round-trip through the metrics wrapper. The +// metrics themselves are process-global OpenTelemetry counters (not unit +// assertable without a metrics-reader harness, mirroring the existing +// `cache_metrics_store_test.rs`), so we assert that operations route through +// the wrapper and that zstd compression is still applied behind it. +#[nativelink_test] +async fn over_cache_metrics_round_trips() -> Result<(), Error> { + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let metrics = CacheMetricsStore::new( + &CacheMetricsSpec { + cache_type: "cas".to_string(), + backend: StoreSpec::Memory(MemorySpec::default()), // Unused: inner is passed directly. + }, + inner.clone(), + ); + let store = Store::new(ZstdStore::new(&spec(), Store::new(metrics))?); + + let data = compressible_data(4096); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-cache_metrics must return the raw bytes" + ); + assert_eq!(store.has(digest).await?, Some(data.len() as u64)); + let physical = inner.get_part_unchunked(digest, 0, None).await?; + assert!( + physical.len() < data.len(), + "zstd compression must still be applied behind the metrics wrapper" + ); + Ok(()) +} + +// 5a. ZstdStore over `size_partitioning`: identity round-trip; the physical zstd +// stream is routed by the (uncompressed) digest size. +#[nativelink_test] +async fn over_size_partitioning_round_trips() -> Result<(), Error> { + let lower = MemoryStore::new(&MemorySpec::default()); + let upper = MemoryStore::new(&MemorySpec::default()); + let size_part = Store::new(SizePartitioningStore::new( + &SizePartitioningSpec { + size: 100, + lower_store: StoreSpec::Memory(MemorySpec::default()), + upper_store: StoreSpec::Memory(MemorySpec::default()), + }, + Store::new(lower.clone()), + Store::new(upper.clone()), + )); + let store = Store::new(ZstdStore::new(&spec(), size_part)?); + + // Uncompressed size (4096) >= partition threshold (100) => routed to `upper`. + let data = compressible_data(4096); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-size_partitioning must return the raw bytes" + ); + assert!( + upper.has(digest).await?.is_some(), + "a blob above the partition threshold must be stored in the upper partition" + ); + assert_eq!( + lower.has(digest).await?, + None, + "the lower partition must not hold the blob" + ); + Ok(()) +} + +// 5b. ZstdStore over `shard`: identity round-trip through a two-way shard. +#[nativelink_test] +async fn over_shard_round_trips() -> Result<(), Error> { + let shard = Store::new(ShardStore::new( + &ShardSpec { + stores: vec![ + ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }, + ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }, + ], + }, + vec![ + Store::new(MemoryStore::new(&MemorySpec::default())), + Store::new(MemoryStore::new(&MemorySpec::default())), + ], + )?); + let store = Store::new(ZstdStore::new(&spec(), shard)?); + + let data = compressible_data(4096); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-shard must return the raw bytes" + ); + assert_eq!(store.has(digest).await?, Some(data.len() as u64)); + Ok(()) +} + +// 5c. ZstdStore over `ref`: identity round-trip through a ref store that resolves +// to a named memory backend via the StoreManager. +#[nativelink_test] +async fn over_ref_round_trips() -> Result<(), Error> { + let store_manager = Arc::new(StoreManager::new()); + let backing = Store::new(MemoryStore::new(&MemorySpec::default())); + store_manager.add_store("backing", backing.clone())?; + let ref_store = Store::new(RefStore::new( + &RefSpec { + name: "backing".to_string(), + }, + Arc::downgrade(&store_manager), + )); + store_manager.add_store("ref", ref_store.clone())?; + store_manager.run_post_init().await.unwrap(); + + let store = Store::new(ZstdStore::new(&spec(), ref_store)?); + + let data = compressible_data(4096); + let digest = digest_for(&data); + store.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "round-trip through zstd-over-ref must return the raw bytes" + ); + // The resolved backing store must physically hold the compressed stream. + let physical = backing.get_part_unchunked(digest, 0, None).await?; + assert!( + physical.len() < data.len(), + "the ref-resolved backend must hold the compressed physical stream" + ); + Ok(()) +} + +// 6. A wrapper OUTSIDE ZstdStore (cache_metrics -> zstd_store -> memory). This +// exercises only the StoreDriver identity path (the service zstd fast path is +// unreachable through an outer wrapper), proving that placing a wrapper +// outside ZstdStore preserves identity correctness and does not bypass the +// outer wrapper's operations while zstd compression stays active underneath. +#[nativelink_test] +async fn wrapper_outside_zstd_preserves_identity() -> Result<(), Error> { + let backing = Store::new(MemoryStore::new(&MemorySpec::default())); + let zstd = Store::new(ZstdStore::new(&spec(), backing.clone())?); + let outer = Store::new(CacheMetricsStore::new( + &CacheMetricsSpec { + cache_type: "cas".to_string(), + backend: StoreSpec::Memory(MemorySpec::default()), // Unused: zstd is passed directly. + }, + zstd, + )); + + let data = compressible_data(4096); + let digest = digest_for(&data); + outer.update_oneshot(digest, data.clone().into()).await?; + assert_eq!( + &outer.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "identity round-trip through an outer wrapper must return the raw bytes" + ); + assert_eq!( + outer.has(digest).await?, + Some(data.len() as u64), + "the outer wrapper must report the uncompressed size" + ); + let physical = backing.get_part_unchunked(digest, 0, None).await?; + assert!( + physical.len() < data.len(), + "zstd compression must still be active behind the outer wrapper" + ); + Ok(()) +} + +// 7. Concatenated-frame integrity through a wrapper: upload two concatenated zstd +// frames (fast path) into a ZstdStore-over-fast_slow, read back via `get_zstd` +// byte-for-byte, and via the identity `get_part` decode to the full content. +#[nativelink_test] +async fn concatenated_frames_through_fast_slow() -> Result<(), Error> { + const PART_A: &[u8] = b"first frame content aaaaaaaaaaaaaaaaaaaaaaaa"; + const PART_B: &[u8] = b"second frame content bbbbbbbbbbbbbbbbbbbbbbbb"; + + let frame_a = zstd::bulk::compress(PART_A, 3).unwrap(); + let frame_b = zstd::bulk::compress(PART_B, 3).unwrap(); + let mut concatenated = frame_a.clone(); + concatenated.extend_from_slice(&frame_b); + + let mut raw = PART_A.to_vec(); + raw.extend_from_slice(PART_B); + let digest = digest_for(&raw); + + // The fast path stages files, so a real writable temp dir is required. + let temp = make_temp_path("zstd-concat-fast-slow"); + std::fs::create_dir_all(&temp) + .map_err(|e| make_err!(Code::Internal, "Failed to create test temp dir: {e}"))?; + let fast_slow = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Memory(MemorySpec::default()), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + Store::new(MemoryStore::new(&MemorySpec::default())), + Store::new(MemoryStore::new(&MemorySpec::default())), + )); + let zstd = ZstdStore::new(&spec_for(temp), fast_slow)?; + let store = Store::new(zstd.clone()); + + let wire = zstd + .update_zstd( + digest, + DigestHasherFunc::Sha256, + reader_from(vec![Bytes::from(frame_a), Bytes::from(frame_b)]), + ) + .await?; + assert_eq!(wire, concatenated.len() as u64); + + let got = collect_zstd(&zstd, digest).await?; + assert_eq!( + &got[..], + &concatenated[..], + "passthrough through fast_slow must preserve the exact concatenated-frame bytes" + ); + + // The identity view decodes the concatenated frames to the full content. + let decoded = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!( + &decoded[..], + &raw[..], + "the identity read must decode the concatenated frames to the full raw content" + ); + Ok(()) +} + +// 8. Rollout from an empty namespace: a fresh memory backend, a full write->read +// cycle, confirming the digest is absent before the write and present after. +#[nativelink_test] +async fn rollout_from_empty_namespace() -> Result<(), Error> { + const DATA: &[u8] = b"rollout payload from an empty dedicated namespace aaaaaaaaaa"; + + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let store = Store::new(ZstdStore::new(&spec(), inner.clone())?); + let digest = digest_for(DATA); + + assert_eq!( + store.has(digest).await?, + None, + "an empty namespace must report the digest as absent" + ); + assert_eq!( + inner.has(digest).await?, + None, + "the inner backend must start empty" + ); + + store.update_oneshot(digest, DATA.into()).await?; + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + DATA, + "the rollout write->read cycle must return the raw bytes" + ); + assert_eq!( + store.has(digest).await?, + Some(DATA.len() as u64), + "after the write the digest must be present at its uncompressed size" + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Security & starvation hardening (bounded decode, descriptor pinning, commit +// deadline). See the PR's issues 1-3. +// --------------------------------------------------------------------------- + +/// Inner store that, on `update_with_whole_file`, first *replaces* the staged +/// pathname's contents (simulating an observe-and-replace attacker) and then +/// reads the retained descriptor it was handed, storing whatever the descriptor +/// yields into an inner `MemoryStore`. If the zstd store commits from the +/// validated descriptor (as it must), the stored bytes are the validated stream +/// regardless of what happened to the pathname. +#[derive(MetricsComponent)] +struct ClobberingStore { + #[metric(group = "inner")] + inner: Store, +} + +#[async_trait] +impl StoreDriver for ClobberingStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner.has_with_results(keys, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + self.inner + .as_store_driver_pin() + .update(key, reader, size_info) + .await + } + + async fn update_with_whole_file( + self: Pin<&Self>, + key: StoreKey<'_>, + path: OsString, + mut file: FileSlot, + _upload_size: UploadSizeInfo, + ) -> Result<(u64, Option), Error> { + // Replace the pathname with attacker-controlled content *after* the zstd + // store validated the descriptor. Unlink first so the retained fd points + // at a now-orphaned inode, then create a fresh file at the same path. + drop(std::fs::remove_file(&path)); + std::fs::write(&path, b"CLOBBERED-BY-ATTACKER").expect("clobber write"); + + // Read the *descriptor* the store handed us (the validated bytes). + file.rewind() + .await + .map_err(|e| make_err!(Code::Internal, "rewind clobber fd: {e}"))?; + let mut buf = Vec::new(); + file.read_to_end(&mut buf) + .await + .map_err(|e| make_err!(Code::Internal, "read clobber fd: {e}"))?; + let size = buf.len() as u64; + self.inner + .update_oneshot(key.into_digest(), Bytes::from(buf)) + .await?; + Ok((size, None)) + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.inner + .as_store_driver_pin() + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(ClobberingStore); + +/// Inner store whose whole-file commit never resolves, used to prove the zstd +/// store bounds a stalled backend commit with its `commit_timeout` and releases +/// the staging permit afterwards. +#[derive(Default, MetricsComponent)] +struct StallStore {} + +#[async_trait] +impl StoreDriver for StallStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + // Never completes. + core::future::pending::<()>().await; + unreachable!("StallStore::update never resolves") + } + + async fn update_with_whole_file( + self: Pin<&Self>, + _key: StoreKey<'_>, + _path: OsString, + _file: FileSlot, + _upload_size: UploadSizeInfo, + ) -> Result<(u64, Option), Error> { + // Simulate a permanently stalled backend commit. + core::future::pending::<()>().await; + unreachable!("StallStore::update_with_whole_file never resolves") + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(Code::NotFound, "Not found")) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(StallStore); + +/// A compressed stream that decodes to non-empty content is rejected under a +/// zero digest — via both the streaming and one-shot entry points — and stops +/// at the output sink without materializing the (large) decoded output. +#[nativelink_test] +async fn zero_digest_decoding_nonempty_is_rejected_at_sink() -> Result<(), Error> { + let temp = make_temp_path("zstd-zero-bomb"); + let (zstd, _store, _inner) = build(&spec_for(temp.clone())).await?; + + // 16 MiB of zeros compresses to a few KiB. Under a zero digest the decoded + // output cap is 0, so the first decoded block is rejected immediately — the + // 16 MiB is never materialized (the test would be far slower / OOM if it + // were). + let bomb = zstd::bulk::compress(&vec![0u8; 16 * 1024 * 1024], 3).unwrap(); + assert!( + bomb.len() < 64 * 1024, + "bomb should be tiny compressed ({} bytes)", + bomb.len() + ); + let zero = ZERO_BYTE_DIGESTS[0]; + + // One-shot entry point. + let err = zstd + .update_zstd_oneshot(zero, DigestHasherFunc::Sha256, Bytes::from(bomb.clone())) + .await + .expect_err("a non-empty-decoding zero-digest upload must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + + // Streaming entry point. + let err = zstd + .update_zstd( + zero, + DigestHasherFunc::Sha256, + reader_from(vec![Bytes::from(bomb)]), + ) + .await + .expect_err("a non-empty-decoding zero-digest stream must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + + assert_eq!( + dir_entry_count(&temp), + 0, + "zero-digest validation must never stage a temp file" + ); + Ok(()) +} + +/// A non-zero upload whose decoded output exceeds the declared digest size is +/// stopped at the decoder sink (`InvalidArgument`), commits nothing, and leaves +/// no staging file. The decoded output (256 KiB) dwarfs the declared size (100 +/// bytes), so the sink rejects the very first over-limit block. +#[nativelink_test] +async fn decoded_output_exceeding_digest_size_is_stopped_at_sink() -> Result<(), Error> { + let temp = make_temp_path("zstd-decode-overflow"); + let (zstd, _store, inner) = build(&spec_for(temp.clone())).await?; + + let data = compressible_data(256 * 1024); + let compressed = zstd::bulk::compress(&data, 3).unwrap(); + // Claim a digest for only the first 100 decoded bytes: the actual decoded + // stream is far larger than the declared size. + let bad_digest = digest_for(&data[..100]); + assert_eq!(bad_digest.size_bytes(), 100); + + let err = zstd + .update_zstd_oneshot( + bad_digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed), + ) + .await + .expect_err("decoded output larger than the digest size must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + + assert_eq!( + inner.has(bad_digest).await?, + None, + "an over-decoding upload must not commit to the inner store" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "an over-decoding upload must not leave staged temp files" + ); + Ok(()) +} + +/// The commit streams the exact validated descriptor, not a reopened pathname: +/// replacing the pathname's contents after validation cannot change what is +/// committed. +#[nativelink_test] +async fn commit_uses_validated_descriptor_not_reopened_path() -> Result<(), Error> { + const DATA: &[u8] = b"descriptor-pinned commit payload aaaaaaaaaaaaaaaaaaaaaaaa"; + + let temp = make_temp_path("zstd-fd-pin"); + std::fs::create_dir_all(&temp).map_err(|e| make_err!(Code::Internal, "temp dir: {e}"))?; + let inner_mem = Store::new(MemoryStore::new(&MemorySpec::default())); + let clobber = Store::new(Arc::new(ClobberingStore { + inner: inner_mem.clone(), + })); + let zstd = ZstdStore::new(&spec_for(temp.clone()), clobber)?; + + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + + zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await?; + + // The bytes the inner store committed came from the retained descriptor, so + // they equal the validated stream — NOT the "CLOBBERED-BY-ATTACKER" content + // that replaced the pathname during the commit. + let committed = inner_mem.get_part_unchunked(digest, 0, None).await?; + assert_eq!( + &committed[..], + &compressed[..], + "committed bytes must be the validated descriptor's contents, not the replaced pathname's" + ); + assert_eq!( + zstd::stream::decode_all(&committed[..]).unwrap(), + DATA, + "the descriptor-committed stream must still decode to the original content" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "a successful commit must leave no staged temp files" + ); + Ok(()) +} + +/// A permanently stalled backend commit is bounded by `commit_timeout`: the +/// upload fails with `DeadlineExceeded`, the staged file is removed, and the +/// staging permit is released — so a second upload behind a `max_concurrent +/// staged_uploads = 1` bound is admitted rather than blocked forever. +// Uses a short real-time `commit_timeout_s = 1` rather than a paused clock: +// the `nativelink-store` test crate does not enable tokio's `test-util` +// feature, so `start_paused` is unavailable here. +#[nativelink_test] +async fn stalled_commit_times_out_and_releases_staging_slot() -> Result<(), Error> { + const DATA: &[u8] = b"payload whose commit stalls forever bbbbbbbbbbbbbbbbbbbb"; + + let temp = make_temp_path("zstd-commit-stall"); + let mut spec = spec_for(temp.clone()); + spec.max_concurrent_staged_uploads = 1; + spec.commit_timeout_s = 1; + std::fs::create_dir_all(&temp).map_err(|e| make_err!(Code::Internal, "temp dir: {e}"))?; + let zstd = ZstdStore::new(&spec, Store::new(Arc::new(StallStore {})))?; + + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + + let err = zstd + .update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await + .expect_err("a stalled commit must fail rather than hang"); + assert_eq!( + err.code, + Code::DeadlineExceeded, + "a stalled commit must surface DeadlineExceeded, got: {err}" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "a timed-out commit must remove its staged temp file" + ); + + // If the first upload had not released the staging permit, this second + // upload (staging bound = 1) would block on the semaphore forever and the + // test would hang. It resolving at all proves the slot was freed. + let err2 = zstd + .update_zstd_oneshot(digest, DigestHasherFunc::Sha256, Bytes::from(compressed)) + .await + .expect_err("second stalled commit must also time out, not hang"); + assert_eq!(err2.code, Code::DeadlineExceeded, "got: {err2}"); + assert_eq!( + dir_entry_count(&temp), + 0, + "the second timed-out commit must also clean up" + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Inline commit path (small compressed oneshot payloads) +// --------------------------------------------------------------------------- + +/// A compressed oneshot payload at or below `max_inline_commit_size` is +/// validated and committed straight from memory: it round-trips correctly and +/// never creates a staging file. +#[nativelink_test] +async fn inline_commit_skips_the_staging_file() -> Result<(), Error> { + const DATA: &[u8] = b"batch-sized payload committed without touching disk aaaaaaaa"; + + let temp = make_temp_path("zstd-inline"); + let spec = spec_for_inline(temp.clone()); + let (zstd, store, inner) = build(&spec).await?; + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + assert!( + (compressed.len() as u64) <= 4 * 1024 * 1024, + "the payload must be inline-eligible for this test to mean anything" + ); + + zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await?; + + assert_eq!( + &inner.get_part_unchunked(digest, 0, None).await?[..], + &compressed[..], + "the inline commit must store the client's bytes verbatim" + ); + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + DATA, + "the inline-committed blob must still decode through the identity view" + ); + assert_eq!( + dir_entry_count(&temp), + 0, + "an inline commit must never create a staging file" + ); + Ok(()) +} + +/// The inline path validates against the declared digest just like staging does, +/// and rejects a mismatch with `InvalidArgument` without storing anything. +#[nativelink_test] +async fn inline_commit_rejects_hash_mismatch() -> Result<(), Error> { + let temp = make_temp_path("zstd-inline-mismatch"); + let (zstd, _store, inner) = build(&spec_for_inline(temp.clone())).await?; + + let compressed = zstd::bulk::compress(b"inline payload", 3).unwrap(); + let wrong_digest = digest_for(b"a completely different payload"); + + let err = zstd + .update_zstd_oneshot( + wrong_digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed), + ) + .await + .expect_err("an inline upload whose digest does not match must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + assert!( + inner + .get_part_unchunked(wrong_digest, 0, None) + .await + .is_err(), + "a rejected inline upload must not be committed" + ); + assert_eq!(dir_entry_count(&temp), 0); + Ok(()) +} + +/// A oneshot payload larger than `max_inline_commit_size` is fed through the +/// streaming staging path. When staging rejects it, the reported error must be +/// staging's own `InvalidArgument` — not the `Internal` "failed to feed" error +/// the feeder task produces as a *consequence* of the reader being dropped. +#[nativelink_test] +async fn large_oneshot_reports_the_validation_error_not_the_feed_error() -> Result<(), Error> { + let temp = make_temp_path("zstd-oneshot-error-order"); + let mut spec = spec_for_inline(temp.clone()); + spec.max_inline_commit_size = 1; // Force the streaming staging path. + let (zstd, _store, inner) = build(&spec).await?; + + // Large enough that the feeder cannot hand the whole payload off before + // staging notices the mismatch and drops the reader. + let data = compressible_data(2 * 1024 * 1024); + let compressed = zstd::bulk::compress(&data, 3).unwrap(); + let wrong_digest = digest_for(b"not the uploaded content"); + + let err = zstd + .update_zstd_oneshot( + wrong_digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed), + ) + .await + .expect_err("a digest mismatch must be rejected"); + assert_eq!( + err.code, + Code::InvalidArgument, + "the staging error must win over the feeder's channel error, got: {err}" + ); + assert!( + inner + .get_part_unchunked(wrong_digest, 0, None) + .await + .is_err(), + "a rejected upload must not be committed" + ); + assert_eq!(dir_entry_count(&temp), 0); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Admission, deadlines, and config validation +// --------------------------------------------------------------------------- + +/// `max_recompression_size` re-encodes at `compression_level`, so a positive +/// ceiling with no level configured would silently do nothing. It is rejected at +/// construction instead. +#[nativelink_test] +async fn new_rejects_recompression_size_without_a_level() -> Result<(), Error> { + let mut spec = spec_for(make_temp_path("zstd-recompress-no-level")); + spec.max_recompression_size = 1024 * 1024; + spec.compression_level = None; + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let err = ZstdStore::new(&spec, inner.clone()) + .expect_err("max_recompression_size without compression_level must be rejected"); + assert_eq!(err.code, Code::InvalidArgument, "got: {err}"); + + // With a level it constructs fine. + spec.compression_level = Some(9); + ZstdStore::new(&spec, inner)?; + Ok(()) +} + +/// A client that trickles bytes forever resets any per-message idle timeout, so +/// the store applies a *total* validate-and-stage deadline. On expiry the upload +/// fails with `DeadlineExceeded`. The detached validator must retain its staging +/// slot until its input closes; a second upload can proceed only after that +/// validator exits. +#[nativelink_test] +async fn stage_timeout_retains_the_slot_until_the_validator_exits() -> Result<(), Error> { + const DATA: &[u8] = b"payload delivered one byte at a time aaaaaaaaaaaaaaaaaaaa"; + + let temp = make_temp_path("zstd-stage-timeout"); + let mut spec = spec_for(temp.clone()); + spec.max_concurrent_staged_uploads = 1; + spec.stage_timeout_s = 1; + let (zstd, _store, _inner) = build(&spec).await?; + + let compressed = zstd::bulk::compress(DATA, 3).unwrap(); + let digest = digest_for(DATA); + + // A reader that hands over one byte every 100ms and never reaches EOF: every + // individual wait is short, but the upload never finishes. + let (mut tx, rx) = make_buf_channel_pair(); + let stop = Arc::new(AtomicBool::new(false)); + let trickle_stop = stop.clone(); + let trickle = Bytes::from(compressed.clone()); + background_spawn!("zstd_test_trickle", async move { + let mut index = 0usize; + while !trickle_stop.load(Ordering::Relaxed) { + let byte = trickle.slice(index % trickle.len()..).slice(..1); + if tx.send(byte).await.is_err() { + return; + } + index += 1; + tokio::time::sleep(Duration::from_millis(100)).await; + } + }); + + let err = zstd + .update_zstd(digest, DigestHasherFunc::Sha256, rx) + .await + .expect_err("a trickling upload must be bounded by the staging deadline"); + assert_eq!( + err.code, + Code::DeadlineExceeded, + "a trickling upload must surface DeadlineExceeded, got: {err}" + ); + // The async timeout must not release admission while the non-cancellable + // validator is still consuming the live stream. Keep this second future so + // the short timeout only observes it rather than cancelling it. + let second_upload = zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ); + tokio::pin!(second_upload); + assert!( + tokio::time::timeout(Duration::from_millis(250), &mut second_upload) + .await + .is_err(), + "the second upload must wait while the timed-out validator still owns the only slot" + ); + + // Closing the first validator's input lets it unwind, remove its file, and + // release admission. The already-waiting valid upload can then complete. + stop.store(true, Ordering::Relaxed); + tokio::time::timeout(Duration::from_secs(3), second_upload) + .await + .expect("the second upload must acquire the slot after validator cleanup")?; + while dir_entry_count(&temp) != 0 { + tokio::time::sleep(Duration::from_millis(20)).await; + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Commit against a `filesystem` backend (path-based commit, not descriptor) +// --------------------------------------------------------------------------- + +/// A `filesystem` backend drops the staged descriptor and commits by +/// `rename(2)` on the pathname, so the descriptor-pinning property that holds +/// for streaming backends does *not* apply to it. What must still hold: the +/// validated bytes land in the CAS, the staging directory is left empty, and the +/// blob reads back correctly through both the physical and identity views. +/// +/// `temp_path` deliberately sits next to the backend's `content_path` on one +/// filesystem: a cross-device staging directory would make the commit fail with +/// `EXDEV`. +#[nativelink_test] +async fn commit_to_filesystem_backend_round_trips() -> Result<(), Error> { + let data = compressible_data(96 * 1024); + let digest = digest_for(&data); + + let root = make_temp_path("zstd-fs-backend"); + let stage_path = format!("{root}/stage"); + let content_path = format!("{root}/content"); + let fs_temp_path = format!("{root}/fs-temp"); + for path in [&stage_path, &content_path, &fs_temp_path] { + std::fs::create_dir_all(path) + .map_err(|e| make_err!(Code::Internal, "Failed to create {path}: {e}"))?; + } + + let inner = Store::new( + FilesystemStore::::new(&FilesystemSpec { + content_path: content_path.clone(), + temp_path: fs_temp_path, + eviction_policy: None, + read_buffer_size: 0, + block_size: 0, + max_concurrent_writes: 0, + evict_page_cache: false, + }) + .await?, + ); + let mut spec = spec_for(stage_path.clone()); + spec.max_inline_commit_size = 1; // Force the staging + commit path. + let zstd = ZstdStore::new(&spec, inner.clone())?; + let store = Store::new(zstd.clone()); + + let compressed = zstd::bulk::compress(&data, 3).unwrap(); + zstd.update_zstd_oneshot( + digest, + DigestHasherFunc::Sha256, + Bytes::from(compressed.clone()), + ) + .await?; + + assert_eq!( + &inner.get_part_unchunked(digest, 0, None).await?[..], + &compressed[..], + "the filesystem backend must hold the validated zstd stream" + ); + assert_eq!( + &store.get_part_unchunked(digest, 0, None).await?[..], + &data[..], + "the committed blob must decode to the original content" + ); + assert_eq!( + collect_zstd(&zstd, digest).await?, + &compressed[..], + "passthrough reads must serve the stored stream byte-for-byte" + ); + assert_eq!( + dir_entry_count(&stage_path), + 0, + "a committed upload must leave the staging directory empty" + ); + Ok(()) +} diff --git a/nativelink-util/src/buf_channel.rs b/nativelink-util/src/buf_channel.rs index b685ca623..f198c74ac 100644 --- a/nativelink-util/src/buf_channel.rs +++ b/nativelink-util/src/buf_channel.rs @@ -429,6 +429,91 @@ impl DropCloserReadHalf { } } +/// A `std::io::Read` adapter over a [`DropCloserReadHalf`] that blocks the +/// current (blocking) thread while waiting for the next chunk. Only safe to use +/// from within `spawn_blocking!`. +#[derive(Debug)] +pub struct BufChannelReader { + rx: DropCloserReadHalf, + chunk: Bytes, + chunk_offset: usize, +} + +impl BufChannelReader { + #[must_use] + pub const fn new(rx: DropCloserReadHalf) -> Self { + Self { + rx, + chunk: Bytes::new(), + chunk_offset: 0, + } + } + + fn refill_chunk(&mut self) -> std::io::Result { + while self.chunk_offset == self.chunk.len() { + self.chunk = self.rx.blocking_recv().map_err(Error::to_std_err)?; + self.chunk_offset = 0; + if self.chunk.is_empty() { + return Ok(false); + } + } + Ok(true) + } +} + +impl std::io::Read for BufChannelReader { + fn read(&mut self, output: &mut [u8]) -> std::io::Result { + if output.is_empty() { + return Ok(0); + } + if !self.refill_chunk()? { + return Ok(0); + } + let chunk_remaining = &self.chunk[self.chunk_offset..]; + let bytes_to_copy = output.len().min(chunk_remaining.len()); + output[..bytes_to_copy].copy_from_slice(&chunk_remaining[..bytes_to_copy]); + self.chunk_offset += bytes_to_copy; + Ok(bytes_to_copy) + } +} + +/// A `std::io::Write` adapter over a [`DropCloserWriteHalf`] that blocks the +/// current (blocking) thread while forwarding chunks. Only safe to use from +/// within `spawn_blocking!`. EOF is intentionally NOT sent on drop; the caller +/// must call [`BufChannelWriter::send_eof`] explicitly once the write is +/// validated so a failed write never commits downstream. +#[derive(Debug)] +pub struct BufChannelWriter { + tx: DropCloserWriteHalf, +} + +impl BufChannelWriter { + #[must_use] + pub const fn new(tx: DropCloserWriteHalf) -> Self { + Self { tx } + } + + pub fn send_eof(&mut self) -> Result<(), Error> { + self.tx.send_eof() + } +} + +impl std::io::Write for BufChannelWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + if buf.is_empty() { + return Ok(0); + } + self.tx + .blocking_send(Bytes::copy_from_slice(buf)) + .map_err(Error::to_std_err)?; + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + impl Stream for DropCloserReadHalf { type Item = Result; diff --git a/nativelink-util/src/fs.rs b/nativelink-util/src/fs.rs index 47c5ccf3f..b271d5932 100644 --- a/nativelink-util/src/fs.rs +++ b/nativelink-util/src/fs.rs @@ -44,6 +44,23 @@ pub struct FileSlot { } impl FileSlot { + /// Wrap an already-open [`std::fs::File`] that is accounted for by `permit` + /// (a permit from [`OPEN_FILE_SEMAPHORE`], e.g. via [`get_permit`]). + /// + /// This is the escape hatch for callers that open/create a descriptor inside + /// their own `spawn_blocking!` closure (for example, to write synchronously + /// with the [`std::io`] API) and later need to hand the *same* descriptor to + /// an async consumer such as [`crate::store_trait::StoreLike::update_with_whole_file`]. + /// The caller is responsible for having acquired `permit` before opening the + /// file so the descriptor stays accounted for by the global semaphore. + #[must_use] + pub fn from_std(permit: SemaphorePermit<'static>, file: std::fs::File) -> Self { + Self { + _permit: permit, + inner: tokio::fs::File::from_std(file), + } + } + /// Advise the kernel to drop page cache for this file's contents. /// Only available on Linux; #[cfg(target_os = "linux")] diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index 5fb650e2b..1524a8891 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -353,6 +353,63 @@ impl Display for StoreKey<'_> { } } +/// A wire representation understood by a [`WireCompressionStore`]. +/// +/// This is intentionally separate from the REAPI protobuf enum: stores own +/// their physical representation, while protocol adapters translate their +/// negotiated compressor to this util-level capability. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum WireCompressor { + /// A standard zstd frame. + Zstd, +} + +/// Optional immediate-store capability for serving and accepting a compressed +/// wire representation without a decode/re-encode round trip. +/// +/// A `StoreDriver` exposes this capability only when it is itself the store +/// that owns the physical representation. Wrappers must not forward it: their +/// own transformation may make the wrapped bytes unsuitable for wire +/// passthrough. +#[async_trait] +pub trait WireCompressionStore: Send + Sync + 'static { + /// Stores a client-supplied compressed stream after validating it against + /// the declared digest. Returns the number of compressed wire bytes stored. + async fn update_compressed( + self: Arc, + digest: DigestInfo, + digest_function: DigestHasherFunc, + compressor: WireCompressor, + reader: DropCloserReadHalf, + ) -> Result; + + /// Streams the stored compressed representation for a blob. + async fn get_compressed( + self: Arc, + digest: DigestInfo, + compressor: WireCompressor, + writer: DropCloserWriteHalf, + ) -> Result<(), Error>; + + /// Stores a complete client-supplied compressed blob after validation. + async fn update_compressed_oneshot( + self: Arc, + digest: DigestInfo, + digest_function: DigestHasherFunc, + compressor: WireCompressor, + data: Bytes, + ) -> Result<(), Error>; + + /// Returns either a stored compressed representation accepted by the + /// client, or the decoded identity representation. The optional compressor + /// describes `data`; `None` means identity. + async fn get_for_batch( + self: Arc, + digest: DigestInfo, + acceptable_compressors: &[WireCompressor], + ) -> Result<(Bytes, Option), Error>; +} + #[derive(Clone, MetricsComponent)] #[repr(transparent)] pub struct Store { @@ -394,6 +451,16 @@ impl Store { self.inner.inner_store(maybe_digest).as_any().downcast_ref() } + /// Returns the wire-compression capability of the immediate store driver. + /// + /// This deliberately does not follow [`Self::inner_store`]. A wrapper + /// outside a representation-changing store must disable passthrough unless + /// it explicitly exposes its own valid wire representation. + #[inline] + pub fn wire_compression_store(&self) -> Option> { + self.inner.clone().wire_compression_store() + } + /// Register health checks used to monitor the store. #[inline] pub fn register_health(&self, registry: &mut HealthRegistryBuilder) { @@ -624,6 +691,12 @@ pub trait StoreDriver: // for ref stores async fn post_init(self: Arc) -> Result<(), Error>; + /// Returns an optional capability for the driver's immediate compressed + /// wire representation. The default intentionally disables passthrough. + fn wire_compression_store(self: Arc) -> Option> { + None + } + /// See: [`StoreLike::has`] for details. #[inline] async fn has(self: Pin<&Self>, key: StoreKey<'_>) -> Result, Error> { diff --git a/nativelink-util/src/wire_compression.rs b/nativelink-util/src/wire_compression.rs index c9760c7ea..d301cf65d 100644 --- a/nativelink-util/src/wire_compression.rs +++ b/nativelink-util/src/wire_compression.rs @@ -27,6 +27,7 @@ use nativelink_proto::build::bazel::remote::execution::v2::compressor; use crate::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use crate::common::DigestInfo; use crate::digest_hasher::{DigestHasher, DigestHasherFunc}; +use crate::store_trait::WireCompressor; /// Zstd compression level for wire compression. /// Level 0 in the zstd crate means "use default" (currently 3). @@ -347,3 +348,17 @@ pub async fn stream_encode_compressed_download_from_reader( .err_tip(|| "Failed to send compressed download EOF")?; Ok(()) } + +/// Translate a negotiated REAPI compressor into the store-level wire +/// representation, or `None` when no store can serve that compressor's bytes +/// directly. Keeps the passthrough fast paths from assuming a compressor the +/// client did not actually negotiate. +#[must_use] +pub const fn wire_compressor_capability( + compressor_value: compressor::Value, +) -> Option { + match compressor_value { + compressor::Value::Zstd => Some(WireCompressor::Zstd), + _ => None, + } +} diff --git a/nativelink-util/tests/store_trait_test.rs b/nativelink-util/tests/store_trait_test.rs index 0cecc7ee3..83dba428e 100644 --- a/nativelink-util/tests/store_trait_test.rs +++ b/nativelink-util/tests/store_trait_test.rs @@ -1,14 +1,18 @@ use core::pin::Pin; use std::sync::Arc; -use nativelink_error::Error; +use bytes::Bytes; +use nativelink_error::{Code, Error, make_err}; use nativelink_macro::nativelink_test; use nativelink_metric::MetricsComponent; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; +use nativelink_util::common::DigestInfo; use nativelink_util::default_health_status_indicator; +use nativelink_util::digest_hasher::DigestHasherFunc; use nativelink_util::health_utils::HealthStatusIndicator; use nativelink_util::store_trait::{ - RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, + RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, WireCompressionStore, + WireCompressor, }; use tonic::async_trait; @@ -22,6 +26,10 @@ impl StoreDriver for FakeStore { Ok(()) } + fn wire_compression_store(self: Arc) -> Option> { + Some(self) + } + async fn has_with_results( self: Pin<&Self>, _keys: &[StoreKey<'_>], @@ -68,6 +76,124 @@ impl StoreDriver for FakeStore { default_health_status_indicator!(FakeStore); +#[async_trait] +impl WireCompressionStore for FakeStore { + async fn update_compressed( + self: Arc, + _digest: DigestInfo, + _digest_function: DigestHasherFunc, + _compressor: WireCompressor, + _reader: DropCloserReadHalf, + ) -> Result { + Err(make_err!(Code::Unimplemented, "fake wire store")) + } + + async fn get_compressed( + self: Arc, + _digest: DigestInfo, + _compressor: WireCompressor, + _writer: DropCloserWriteHalf, + ) -> Result<(), Error> { + Err(make_err!(Code::Unimplemented, "fake wire store")) + } + + async fn update_compressed_oneshot( + self: Arc, + _digest: DigestInfo, + _digest_function: DigestHasherFunc, + _compressor: WireCompressor, + _data: Bytes, + ) -> Result<(), Error> { + Err(make_err!(Code::Unimplemented, "fake wire store")) + } + + async fn get_for_batch( + self: Arc, + _digest: DigestInfo, + _acceptable_compressors: &[WireCompressor], + ) -> Result<(Bytes, Option), Error> { + Err(make_err!(Code::Unimplemented, "fake wire store")) + } +} + +/// A wrapper `StoreDriver` whose `inner_store()` forwards to its wrapped +/// `Store`, unlike most production wrappers (e.g. `CompressionStore`, +/// `VerifyStore`) which return `self`. It proves wire-compression capability +/// discovery remains deliberately immediate. +#[derive(Debug, MetricsComponent)] +struct ForwardingWrapperStore { + inner: Store, +} + +#[async_trait] +#[allow(clippy::todo)] +impl StoreDriver for ForwardingWrapperStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + todo!(); + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + todo!(); + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + todo!(); + } + + fn inner_store(&self, digest: Option) -> &dyn StoreDriver { + self.inner.inner_store(digest) + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + todo!(); + } +} + +default_health_status_indicator!(ForwardingWrapperStore); + +#[nativelink_test] +async fn wire_compression_capability_only_matches_outer_driver() -> Result<(), Error> { + let immediate_store = Store::new(Arc::new(FakeStore {})); + assert!(immediate_store.wire_compression_store().is_some()); + + let wrapped_store = Store::new(Arc::new(ForwardingWrapperStore { + inner: immediate_store, + })); + assert!( + wrapped_store.wire_compression_store().is_none(), + "wrappers must not recursively expose a representation-changing store's wire capability" + ); + + Ok(()) +} + #[nativelink_test] async fn fast_has_with_results() -> Result<(), Error> { let store = Store::new(Arc::new(FakeStore {})); diff --git a/web/apps/docs/content/docs/configuration/compression.mdx b/web/apps/docs/content/docs/configuration/compression.mdx index 2422d0a6c..fa264ef65 100644 --- a/web/apps/docs/content/docs/configuration/compression.mdx +++ b/web/apps/docs/content/docs/configuration/compression.mdx @@ -194,6 +194,23 @@ To also shrink bytes at rest, wrap the backing store in a [`compression` store](/reference/nativelink-config/main#compressionspec); the two compose, since one covers the wire and the other covers storage. +## Skipping the decompress/recompress round trip: `compression_algorithm.zstd` + +`compression_algorithm: { lz4: ... }` always decompresses to raw bytes at +the store boundary, so a wire-compressed request still pays a +decompress-on-write and recompress-on-read even though the bytes were zstd on +the wire the whole time. `compression_algorithm: { zstd: ... }` closes that +gap: it keeps blobs as zstd **at rest** and, when it's the outermost store an +instance points at, serves them +byte-for-byte to `--remote_cache_compression` clients — no recompress on a +hit. Identity clients still get plain decompressed bytes. + +It has real deployment constraints — a dedicated, empty backend namespace, +no in-place migration, and a capped `compression_level` — that don't apply +to `compression_algorithm: { lz4: ... }`. See +[Store overview → Pass-through compression](/reference/nativelink-config/store-overview#pass-through-compression-compression_algorithmzstd) +for the full picture before adopting it. + ## FAQ diff --git a/web/apps/docs/content/docs/reference/nativelink-config/store-overview.mdx b/web/apps/docs/content/docs/reference/nativelink-config/store-overview.mdx index 81637daf1..464f69924 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/store-overview.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/store-overview.mdx @@ -69,7 +69,7 @@ store name the server exposes. | Store key | Wraps | What it adds | | --- | --- | --- | -| `compression` | one store | LZ4-compresses on write, decompresses on read. | +| `compression` | one store | Select `lz4` for NativeLink's LZ4 framing, or `zstd` for standard zstd frames and optional wire-compression passthrough. | | `dedup` | two stores (`index_store` + `content_store`) | Rolling-hash chunking so only changed slices upload. | | `fast_slow` | two stores (`fast` + `slow`) | Reads try `fast` first, fall back to `slow`, and backfill `fast`. Writes mirror to both. | | `shard` | N stores | Routes by digest hash. The standard shape for scaling CAS past one backend. | @@ -91,8 +91,239 @@ store name the server exposes. - `existence_cache` and `size_partitioning` are CAS-only; `completeness_checking` is AC-only. Using them on the other store type produces confusing correctness bugs, not a config error. + - `compression` with `compression_algorithm: { zstd: ... }` is the one + algorithm choice that isn't position-agnostic: it must be the store an + instance points at directly, or its main benefit doesn't apply. See the + dedicated section below before deploying it. +## Pass-through compression: `compression_algorithm.zstd` + +The `zstd` algorithm of `compression` keeps CAS blobs as zstd streams **at +rest** instead of raw bytes. Unlike `compression_algorithm: { lz4: ... }` +(which uses NativeLink's LZ4 framing and decompresses back to raw bytes at +the `StoreDriver` boundary), `compression_algorithm: { zstd: ... }` can serve its stored bytes +**byte-for-byte** to a Bazel client running `--remote_cache_compression` on +a cache hit — no decompress-then-recompress round trip at the gRPC +boundary. Identity clients (no `--remote_cache_compression`) are unaffected: +they still receive plain decompressed bytes. See +[Remote cache compression](/configuration/compression) for the wire-level +feature this pairs with. + +It is **CAS-only** — every operation requires a digest key. Never point an +AC store at it. + +**Example JSON5 config:** + +```json5 +{ + name: "CAS_MAIN_STORE", + compression: { + compression_algorithm: { + zstd: { + temp_path: "/var/tmp/nativelink-zstd", + max_compressed_upload_size: "512MiB", + max_concurrent_staged_uploads: 4, + max_concurrent_identity_ops: 256, + compression_level: 9, + max_recompression_size: "64MiB", + max_concurrent_recompressions: 1, + max_inline_commit_size: "4MiB", + stage_timeout_s: 600, + commit_timeout_s: 300, + }, + }, + backend: { memory: { eviction_policy: { max_bytes: "10GiB" } } }, + }, +} +``` + +### Placement: it must be the outermost store + +Every wrapper in the table above can nest at any depth. `compression` with +`compression_algorithm.zstd` is different: **it must be the store the +instance/CAS/ByteStream service points at directly** for the byte-for-byte +passthrough to apply. `fast_slow`, +`dedup`, `existence_cache`, `cache_metrics`, `shard`, `ref_store`, and +`size_partitioning` are all fine **inside** it — they then operate on the +physical zstd stream exactly as they would on raw bytes. + +A wrapper placed **outside** the zstd compression store (for example, +`verify` wrapping it) is still correct, but that wrapper has to decode the +stream to do its job, which disables passthrough at that boundary. + + + The `backend` of `compression_algorithm.zstd` stores digest → zstd bytes, + not digest → raw bytes. The digest key is identical either way, and raw CAS payloads are + unstructured — an arbitrary raw blob can legitimately start with the zstd + magic number — so there is no safe way to sniff which encoding a given + entry is in. Version 1 has no in-place migration. + + - `backend` must point at a brand-new or empty store namespace, never one + a raw-bytes reader or writer touches — including a plain + `filesystem`/`experimental_cloud_object_store` CAS, a different + instance, or a pre-zstd NativeLink version pointed at the same + path/bucket/prefix. + - Rollout **and** rollback both require a cache flush (or cutting over to + a new namespace) — there is no supported way to select or remove zstd + compression in place against a populated namespace. + - A rolling/mixed-version deployment against **one shared namespace** is + unsupported: while some replicas write raw bytes and others write zstd + to the same keys, every reader misinterprets roughly half the entries. + + +### `compression_level`: capped at 19, not 22 + +`compression_level` is optional and sets the level used for **two** things: the +uploads this store compresses itself (identity clients), and the re-encoding of +incoming compressed uploads when `max_recompression_size > 0`. Omitted, it uses +level 3 and never re-encodes an already-compressed upload. When set, it must be +in `1..=19` — NativeLink rejects anything outside that range at +config-validation time (startup). + +Level 9 is a reasonable general-purpose choice; it costs a fraction of level +19's CPU per byte. Reach for the high end only when blobs are read far more often +than written. + +The zstd level isn't recorded in the frame; decoders are level-agnostic. The +actual cross-client constraint is **window size**. Standard levels 1–19 cap +the encoded frame at `windowLog ≤ 23` (≤ 8 MiB), decodable by every +Bazel/zstd-jni client using plain `libzstd` defaults. Bazel's decoder never +calls `setLongMax`, so its hard ceiling is `windowLog ≤ 27` (128 MiB) — +level 19 keeps a wide margin under that ceiling rather than chasing the +absolute cap. Zstd compression never enables long-distance matching or external +dictionaries (REAPI has no dictionary channel, and long-distance matching +alone can push `windowLog` past 23). + +### Bounds and admission: stopping a zstd bomb + +Every compressed input is bounded from both ends. `max_compressed_upload_size` +caps the compressed bytes a client may push; exceeding it fails the upload with +`RESOURCE_EXHAUSTED`. Independently, the decoder's output sink is capped at the +digest's declared uncompressed size and rejects with `INVALID_ARGUMENT` the +moment decoded bytes *would* exceed it — so a small "zstd bomb" that expands to +a huge buffer is stopped at the first over-limit block rather than being fully +materialized. The decoder also rejects frames requesting more than an 8 MiB +history window, before their header can cause zstd's larger default allocation. +Zero-digest validation streams through the same bounded decoder with an output +cap of zero: it never allocates the decoded output, and it takes an admission +permit like any other staged upload, so it cannot spawn unbounded decode jobs. + +`max_concurrent_staged_uploads` admits a bounded number of compressed uploads at +a time, but admission alone can't stop a stalled client or backend from parking +on a slot forever. The deadlines below fail the RPC with `DEADLINE_EXCEEDED`. +Validation runs in a non-cancellable blocking task, so a staging timeout does +not release admission early: the task retains its slot and cleanup guard until +its input closes and it exits. The ByteStream fast path stops its client pump on +a store failure, closing that input promptly. Direct store callers must do the +same. Cancellation, validation failure, and success all clean up once the +validator has stopped. + +- **`stage_timeout_s`** — total seconds one upload may spend being validated and + staged, measured from the moment it takes a staging slot. Unlike an idle + timeout, continuous slow progress does not reset it. Size it against + `max_compressed_upload_size` and the slowest upload bandwidth worth serving. + `0` uses the default of 600. +- **`commit_timeout_s`** — seconds to wait, after validation and staging, for + optional recompression and the inner-store commit. Bounds a stalled backend. + `0` uses the default of 300. +- **`compressed_upload_idle_timeout_s`** — a ByteStream instance setting, not a + store setting: how long to wait for the *next* `WriteRequest` of a compressed + upload. A client making continuous progress is never timed out, only an idle + one. `0` uses the default of 60. (Earlier builds reused + `persist_stream_on_disconnect_timeout_s` for this; it is now its own knob, + because how long a disconnected upload stays resumable has nothing to do with + how long an attached client may stall.) + +Two admission bounds sit alongside those deadlines: + +- **`max_concurrent_staged_uploads`** — compressed uploads validating or staging + at once. `0` uses the default of 4. +- **`max_concurrent_identity_ops`** — uncompressed (identity) reads and writes at + once. Each one holds a blocking thread for the whole transfer, so without a + bound a flood of slow identity clients can starve the process-wide blocking + pool that filesystem I/O also uses. Raising it trades that risk for more + queuing on this store; `0` uses the default of 256. + +Small compressed uploads skip staging altogether. **`max_inline_commit_size`** +(default 4 MiB) is the ceiling below which an upload is validated and committed +straight from memory — no staging file, no `fsync`. `BatchUpdateBlobs` payloads +are small and numerous, so a per-blob disk round trip would dominate their cost. +Inline uploads are validated exactly as staged ones are and take the same +staging admission permit and `stage_timeout_s`. + +### `temp_path`: staging directory, size and place it deliberately + +Zstd compression stages every compressed upload larger than +`max_inline_commit_size` as a temp file under `temp_path` before committing it to +`backend`. Recompression keeps the decoded bytes in bounded memory and overwrites +the same validated descriptor only when the result is smaller, so it never adds a +second staging file. + +`temp_path` must be a private directory owned by the operator and **not** writable +by local users you do not trust. At startup the store verifies it is a directory, rejects +a world-writable directory that lacks the sticky bit, and probes that it is +writable. Each staging file is created exclusively (`O_EXCL` / `create_new`) under +an unguessable name, with `0o600` applied atomically at creation rather than by a +later `chmod`. Every staged file is removed on success, error, and cancellation. + +- **What defends against an observe-and-replace race depends on the backend.** + Both the validated descriptor and its pathname are handed to `backend` at + commit. Backends that stream the bytes — the default, so `memory`, + `experimental_cloud_object_store`, and friends — read the exact descriptor that + was validated, and never reopen the path. A `filesystem` backend instead drops + the descriptor and commits with `rename(2)` on the pathname, so for that backend + the guarantee rests on exclusive creation of an unguessable name inside an + operator-private directory, which is why the directory checks above are not + optional. +- **Worst-case temp disk usage** is roughly `max_concurrent_staged_uploads + * max_compressed_upload_size`. Recompression uses bounded memory, capped by + `max_recompression_size`, rather than an additional staging file. Size and + monitor `temp_path` against the disk bound, not just steady-state usage. +- **Put `temp_path` on the same filesystem as the `content_path` of a + `filesystem` backend.** That backend commits with the `rename(2)` above, and a + cross-filesystem rename fails with `EXDEV` — the upload fails, the blob is + never stored, and the temp file is still cleaned up — instead of silently + falling back to a copy. + +### Recompression: keep whichever encoding is smaller + +With `compression_level` set and `max_recompression_size > 0`, an incoming +upload whose *decoded* size is within `max_recompression_size` gets +re-encoded at `compression_level`, and zstd compression keeps whichever of the +original or re-encoded stream is smaller. `max_recompression_size: 0` +disables recompression entirely — `compression_level` still picks the level +used to encode brand-new, not-already-compressed uploads; incoming compressed +streams are left unchanged. + +Because re-encoding needs a level, `max_recompression_size > 0` **requires** +`compression_level`. Setting the ceiling without a level is rejected at startup +rather than silently disabling recompression. + +`max_concurrent_recompressions` bounds how many re-encodes run at once, and it is +**best-effort**: an upload that finds every slot busy commits its original stream +instead of queuing. That matters because the upload is holding a staging slot +while it waits, so a small recompression pool would otherwise throttle the entire +upload path. Skipped recompressions are counted in the store's metrics +(`recompressions_skipped_busy`); a persistently high rate means the pool is too +small for the offered load, not that anything failed. Recompression also runs +under the same `commit_timeout_s` deadline as the inner-store commit. + +### Observability + +The store publishes its own counters, so the bounds above can be tuned against +measurements rather than guesses: + +| Metric | What it tells you | +| --- | --- | +| `wire_uploads`, `wire_upload_bytes` | How much traffic actually takes the compressed fast path. | +| `wire_downloads` | Cache hits served byte-for-byte, with no re-encode. | +| `batch_zstd_passthroughs`, `batch_identity_decodes` | The same split for `BatchReadBlobs`. | +| `inline_commits`, `staged_commits` | How much of the upload mix stays in memory vs. hits `temp_path`. Tune `max_inline_commit_size` from this. | +| `staged_uploads_inflight`, `identity_ops_inflight` | Live occupancy against `max_concurrent_staged_uploads` / `max_concurrent_identity_ops`. Sustained saturation means requests are queuing. | +| `recompressions_applied`, `recompressions_rejected`, `recompressions_skipped_busy` | Whether recompression is earning its CPU, and whether its pool is too small. | +| `stage_timeouts`, `commit_timeouts` | Deadlines that fired. A nonzero `stage_timeouts` is either a trickling client or a `stage_timeout_s` set below real upload times. | + ## Choosing a terminal backend | Backend | Durability | Shared across instances | Typical role | @@ -189,3 +420,6 @@ instead of declaring the same backend twice: fit alongside servers, schedulers, and workers. - [Configuration → Production](/configuration/production) — sharded, multi-region store topologies end to end. +- [`deployment-examples/docker-compose/local-storage-cas-zstd.json5`](https://github.com/TraceMachina/nativelink/blob/main/deployment-examples/docker-compose/local-storage-cas-zstd.json5) + — a complete zstd compression deployment, annotated with the namespace and + placement rules above.