/build/source/nativelink-service/src/wire_compression.rs

Source
// 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.

//! Wire compression utilities for REAPI compressed-blobs.
//!
//! 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).

use std::collections::HashSet;

use bytes::Bytes;
use nativelink_config::cas_server::{CapabilitiesConfig, InstanceName, WithInstanceName};
use nativelink_error::{Code, Error, make_err, make_input_err};
use nativelink_proto::build::bazel::remote::execution::v2::compressor;
use nativelink_util::spawn_blocking;
// The codecs are shared with client-side GrpcStore transfers; re-export so
// 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,
};
use tracing::warn;

/// Which instances accept and advertise REAPI compressed-blobs (zstd).
///
/// Derived from `CapabilitiesConfig.remote_cache_compression` so that
/// advertisement (capabilities) and acceptance (ByteStream/CAS) cannot drift.
#[derive(Debug, Clone, Default)]
pub struct RemoteCacheCompressionInstances(HashSet<InstanceName>);

impl RemoteCacheCompressionInstances {
    #[must_use]
    pub fn from_capabilities_configs(configs: &[WithInstanceName<CapabilitiesConfig>]) -> Self {
        Self(
            configs
                .iter()
                .filter(|config| config.remote_cache_compression)
                .map(|config| config.instance_name2.clone2())
                .collect(),
        )
    }

    #[must_use]
    pub fn from_enabled_instance_names(
        instance_names: impl IntoIterator<Item = InstanceName>,
    ) -> Self {
        Self(instance_names.into_iter().collect())
    }

    #[must_use]
    pub fn enabled_for(&self, instance_name: &str) -> bool {
        self.0.contains(instance_name)
    }
}

/// Resolve a wire compressor from a URI compressor string (as it appears in
/// the `compressed-blobs/{compressor}/...` resource name) and validate it
/// against whether the instance supports remote cache compression.
///
/// Returns `compressor::Value::Identity` when `compressor_str` is `None` or
/// `"identity"`. This is the single source of truth for URI-to-compressor
/// parsing so that `ByteStream` and any other callers stay consistent.
pub fn resolve_wire_compressor(
    compressor_str: Option<&str>,
    remote_cache_compression_enabled: bool,
) -> Result<compressor::Value, Error> {
    match compressor_str {
        None | Some("identity"21) => Ok(compressor::Value::Identity),
        Some("zstd") => {
            if remote_cache_compression_enabled {
                Ok(compressor::Value::Zstd)
            } else {
                Err(make_input_err!(
                    "Remote cache compression is not supported by this instance"
                ))
            }
        }
        Some(other) => Err(make_input_err!("Unsupported wire compressor: '{}'", other)),
    }
}

#[must_use]
pub fn compress_for_batch_read(data: Bytes) -> (Bytes, compressor::Value) {
    match compress(data.clone(), compressor::Value::Zstd) {
        Ok(compressed) => {
            if compressed.len() < data.len() {
                (compressed, compressor::Value::Zstd)
            } else {
                (data, compressor::Value::Identity)
            }
        }
        Err(err) => {
            warn!("Wire compression failed, falling back to identity: {}", err);
            (data, compressor::Value::Identity)
        }
    }
}

fn validate_identity_batch_update(data: Bytes, expected_size: usize) -> Result<Bytes, Error> {
    if data.len() != expected_size {
        return Err(make_err!(
            Code::InvalidArgument,
            "Identity data size {} does not match expected size {}",
            data.len(),
            expected_size
        ));
    }
    Ok(data)
}

pub async fn decompress_batch_update(
    data: Bytes,
    compressor_i32: i32,
    expected_size: usize,
    remote_cache_compression_enabled: bool,
) -> Result<Bytes, Error> {
    let request_compressor = compressor::Value::try_from(compressor_i32)
        .map_err(|_| make_input_err!0("Unknown compressor value: {}", compressor_i32))?0;

    match request_compressor {
        compressor::Value::Identity => validate_identity_batch_update(data, expected_size),
        compressor::Value::Zstd => {
            if !remote_cache_compression_enabled {
                return Err(make_input_err!(
                    "Remote cache compression is not supported by this instance"
                ));
            }
            spawn_blocking!("cas_decode_compressed_upload", move || {
                decompress(&data, compressor::Value::Zstd, expected_size)
            })
            .await
            .map_err(|e| make_err!0(Code::Internal0, "Decompression task failed: {}", e))?0
        }
        other => Err(make_input_err!("Unsupported wire compressor: {:?}", other)),
    }
}

Line	Count	Source
1		// Copyright 2026 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		//! Wire compression utilities for REAPI compressed-blobs.
16		//!
17		//! This module handles compression/decompression of blob data on the gRPC wire
18		//! between client and server, per the REAPI compressed-blobs specification.
19		//! This is orthogonal to at-rest compression (`CompressionStore` with LZ4).
20
21		use std::collections::HashSet;
22
23		use bytes::Bytes;
24		use nativelink_config::cas_server::{CapabilitiesConfig, InstanceName, WithInstanceName};
25		use nativelink_error::{Code, Error, make_err, make_input_err};
26		use nativelink_proto::build::bazel::remote::execution::v2::compressor;
27		use nativelink_util::spawn_blocking;
28		// The codecs are shared with client-side GrpcStore transfers; re-export so
29		// existing service callers keep their import paths.
30		pub use nativelink_util::wire_compression::{
31		ZSTD_COMPRESSION_LEVEL, compress, decompress, stream_decode_compressed_upload,
32		stream_encode_compressed_download,
33		};
34		use tracing::warn;
35
36		/// Which instances accept and advertise REAPI compressed-blobs (zstd).
37		///
38		/// Derived from `CapabilitiesConfig.remote_cache_compression` so that
39		/// advertisement (capabilities) and acceptance (ByteStream/CAS) cannot drift.
40		#[derive(Debug, Clone, Default)]
41		pub struct RemoteCacheCompressionInstances(HashSet<InstanceName>);
42
43		impl RemoteCacheCompressionInstances {
44		#[must_use]
45	3	pub fn from_capabilities_configs(configs: &[WithInstanceName<CapabilitiesConfig>]) -> Self {
46		Self(
47	3	configs
48	3	.iter()
49	3	.filter(\|config\| config.remote_cache_compression)
50	3	.map(\|config\| config.instance_name2 . clone2 ())
51	3	.collect(),
52		)
53	3	}
54
55		#[must_use]
56	16	pub fn from_enabled_instance_names(
57	16	instance_names: impl IntoIterator<Item = InstanceName>,
58	16	) -> Self {
59	16	Self(instance_names.into_iter().collect())
60	16	}
61
62		#[must_use]
63	52	pub fn enabled_for(&self, instance_name: &str) -> bool {
64	52	self.0.contains(instance_name)
65	52	}
66		}
67
68		/// Resolve a wire compressor from a URI compressor string (as it appears in
69		/// the `compressed-blobs/{compressor}/...` resource name) and validate it
70		/// against whether the instance supports remote cache compression.
71		///
72		/// Returns `compressor::Value::Identity` when `compressor_str` is `None` or
73		/// `"identity"`. This is the single source of truth for URI-to-compressor
74		/// parsing so that `ByteStream` and any other callers stay consistent.
75	52	pub fn resolve_wire_compressor(
76	52	compressor_str: Option<&str>,
77	52	remote_cache_compression_enabled: bool,
78	52	) -> Result<compressor::Value, Error> {
79	52	match compressor_str {
80	34	None \| Some( "identity"21 ) => Ok(compressor::Value::Identity),
81	18	Some("zstd") => {
82	17	if remote_cache_compression_enabled {
83	13	Ok(compressor::Value::Zstd)
84		} else {
85	4	Err(make_input_err!(
86	4	"Remote cache compression is not supported by this instance"
87	4	))
88		}
89		}
90	1	Some(other) => Err(make_input_err!("Unsupported wire compressor: '{}'", other)),
91		}
92	52	}
93
94		#[must_use]
95	2	pub fn compress_for_batch_read(data: Bytes) -> (Bytes, compressor::Value) {
96	2	match compress(data.clone(), compressor::Value::Zstd) {
97	2	Ok(compressed) => {
98	2	if compressed.len() < data.len() {
99	1	(compressed, compressor::Value::Zstd)
100		} else {
101	1	(data, compressor::Value::Identity)
102		}
103		}
104	0	Err(err) => {
105	0	warn!("Wire compression failed, falling back to identity: {}", err);
106	0	(data, compressor::Value::Identity)
107		}
108		}
109	2	}
110
111	6	fn validate_identity_batch_update(data: Bytes, expected_size: usize) -> Result<Bytes, Error> {
112	6	if data.len() != expected_size {
113	0	return Err(make_err!(
114	0	Code::InvalidArgument,
115	0	"Identity data size {} does not match expected size {}",
116	0	data.len(),
117	0	expected_size
118	0	));
119	6	}
120	6	Ok(data)
121	6	}
122
123	8	pub async fn decompress_batch_update(
124	8	data: Bytes,
125	8	compressor_i32: i32,
126	8	expected_size: usize,
127	8	remote_cache_compression_enabled: bool,
128	8	) -> Result<Bytes, Error> {
129	8	let request_compressor = compressor::Value::try_from(compressor_i32)
130	8	.map_err(\|_\| make_input_err!0 ("Unknown compressor value: {}", compressor_i32)) ?0 ;
131
132	8	match request_compressor {
133	6	compressor::Value::Identity => validate_identity_batch_update(data, expected_size),
134		compressor::Value::Zstd => {
135	2	if !remote_cache_compression_enabled {
136	1	return Err(make_input_err!(
137	1	"Remote cache compression is not supported by this instance"
138	1	));
139	1	}
140	1	spawn_blocking!("cas_decode_compressed_upload", move \|\| {
141	1	decompress(&data, compressor::Value::Zstd, expected_size)
142	1	})
143	1	.await
144	1	.map_err(\|e\| make_err!0 ( Code::Internal0 , "Decompression task failed: {}", e)) ?0
145		}
146	0	other => Err(make_input_err!("Unsupported wire compressor: {:?}", other)),
147		}
148	8	}

Coverage Report

Created: 2026-09-18 20:40