/build/source/nativelink-store/src/cache_metrics_store.rs

Source
// Copyright 2026 Trace Machina, Inc. 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::borrow::BorrowMut;
use core::ops::Bound;
use core::pin::Pin;
use std::ffi::OsString;
use std::sync::Arc;
use std::time::Instant;

use async_trait::async_trait;
use bytes::Bytes;
use nativelink_config::stores::CacheMetricsSpec;
use nativelink_error::{Code, Error};
use nativelink_metric::{
    MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, group, publish,
};
use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
use nativelink_util::fs;
use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatusIndicator};
use nativelink_util::metrics::{CACHE_METRICS, CACHE_TYPE, CacheMetricAttrs};
use nativelink_util::store_trait::{
    RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, StoreOptimizations, UploadSizeInfo,
};
use opentelemetry::KeyValue;

#[derive(Debug)]
pub struct CacheMetricsStore {
    backend: Store,
    cache_type: String,
    allow_remove: bool,
    attrs: CacheMetricAttrs,
}

impl CacheMetricsStore {
    pub fn new(spec: &CacheMetricsSpec, backend: Store) -> Arc<Self> {
        let base_attrs = [KeyValue::new(CACHE_TYPE, spec.cache_type.clone())];
        // A store that tracks its own size reports it under the same cache
        // type as the operation metrics. Other stores report no size.
        backend
            .inner_store(None::<StoreKey<'_>>)
            .enable_cache_size_metrics(&base_attrs);
        Arc::new(Self {
            backend,
            cache_type: spec.cache_type.clone(),
            allow_remove: spec.allow_remove,
            attrs: CacheMetricAttrs::new(&base_attrs),
        })
    }

    fn duration_ms(start: Instant) -> f64 {
        start.elapsed().as_secs_f64() * 1000.0
    }

    const fn size_info_bytes(size_info: UploadSizeInfo) -> Option<u64> {
        match size_info {
            UploadSizeInfo::ExactSize(size) => Some(size),
            UploadSizeInfo::MaxSize(_) => None,
        }
    }

    fn record_duration(&self, start: Instant, attrs: &[KeyValue]) {
        CACHE_METRICS
            .cache_operation_duration
            .record(Self::duration_ms(start), attrs);
    }

    fn record_write_io(&self, bytes: Option<u64>) {
        if let Some(bytes) = bytes {
            CACHE_METRICS
                .cache_io
                .add(bytes, self.attrs.write_success());
            CACHE_METRICS
                .cache_entry_size
                .record(bytes, self.attrs.write_success());
        }0
    }
}

impl MetricsComponent for CacheMetricsStore {
    fn publish(
        &self,
        _kind: MetricKind,
        _field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        publish!(
            "cache_type",
            &self.cache_type,
            MetricKind::String,
            "Low-cardinality cache type label emitted on cache metrics"
        );
        let _enter = group!("backend").entered();
        self.backend
            .publish(MetricKind::Component, MetricFieldData::default())?;
        Ok(MetricPublishKnownKindData::Component)
    }
}

#[async_trait]
impl StoreDriver for CacheMetricsStore {
    async fn post_init(self: Arc<Self>) -> Result<(), Error> {
        self.backend.clone().into_inner().post_init().await?;
        Ok(())
    }

    async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
        let start = Instant::now();
        // Backends such as memory and filesystem cannot infer whether their
        // copy is disposable. Require an explicit opt-in on this tier.
        let result = if self.allow_remove {
            self.backend.as_store_driver_pin().remove(key).await
        } else {
            Ok(false)
        };
        let attrs = match &result {
            Ok(true) => self.attrs.delete_success(),
            Ok(false) => self.attrs.delete_miss(),
            Err(_) => self.attrs.delete_error(),
        };
        CACHE_METRICS.cache_operations.add(1, attrs);
        self.record_duration(start, attrs);
        result
    }

    /// Not counted as a delete: this tier's own copy is never dropped.
    async fn remove_cached(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
        if self.allow_remove {
            self.backend.as_store_driver_pin().remove_cached(key).await
        } else {
            Ok(false)
        }
    }

    async fn has_with_results(
        self: Pin<&Self>,
        keys: &[StoreKey<'_>],
        results: &mut [Option<u64>],
    ) -> Result<(), Error> {
        let start = Instant::now();
        let result = self.backend.has_with_results(keys, results).await;
        if result.is_ok() {
            let hits = results.iter().filter(|result| result.is_some()).count();
            let misses = results.len().saturating_sub(hits);
            if hits > 0 {
                CACHE_METRICS
                    .cache_operations
                    .add(hits as u64, self.attrs.read_hit());
            }
            if misses > 0 {
                CACHE_METRICS
                    .cache_operations
                    .add(misses as u64, self.attrs.read_miss());
            }
            let duration_attrs = if hits > 0 {
                self.attrs.read_hit()
            } else {
                self.attrs.read_miss()
            };
            self.record_duration(start, duration_attrs);
        } else {
            CACHE_METRICS
                .cache_operations
                .add(keys.len() as u64, self.attrs.read_error());
            self.record_duration(start, self.attrs.read_error());
        }
        result
    }

    async fn list(
        self: Pin<&Self>,
        range: (Bound<StoreKey<'_>>, Bound<StoreKey<'_>>),
        handler: &mut (dyn for<'a> FnMut(&'a StoreKey) -> bool + Send + Sync + '_),
    ) -> Result<u64, Error> {
        self.backend.list(range, handler).await
    }

    async fn update(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        reader: DropCloserReadHalf,
        upload_size: UploadSizeInfo,
    ) -> Result<u64, Error> {
        let start = Instant::now();
        let result = self.backend.update(key, reader, upload_size).await;
        if let Ok(size) = &result {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_success());
            self.record_write_io(Some(*size));
            self.record_duration(start, self.attrs.write_success());
        } else {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_error());
            self.record_duration(start, self.attrs.write_error());
        }
        result
    }

    fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
        self.backend.optimized_for(optimization)
    }

    async fn update_with_whole_file(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        path: OsString,
        file: fs::FileSlot,
        upload_size: UploadSizeInfo,
    ) -> Result<(u64, Option<fs::FileSlot>), Error> {
        let start = Instant::now();
        let bytes = Self::size_info_bytes(upload_size);
        let result = self
            .backend
            .update_with_whole_file(key, path, file, upload_size)
            .await;
        if result.is_ok() {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_success());
            self.record_write_io(bytes);
            self.record_duration(start, self.attrs.write_success());
        } else {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_error());
            self.record_duration(start, self.attrs.write_error());
        }
        result
    }

    async fn update_oneshot(self: Pin<&Self>, key: StoreKey<'_>, data: Bytes) -> Result<(), Error> {
        let start = Instant::now();
        let bytes = data.len() as u64;
        let result = self.backend.update_oneshot(key, data).await;
        if result.is_ok() {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_success());
            self.record_write_io(Some(bytes));
            self.record_duration(start, self.attrs.write_success());
        } else {
            CACHE_METRICS
                .cache_operations
                .add(1, self.attrs.write_error());
            self.record_duration(start, self.attrs.write_error());
        }
        result
    }

    async fn get_part(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        writer: &mut DropCloserWriteHalf,
        offset: u64,
        length: Option<u64>,
    ) -> Result<(), Error> {
        let start = Instant::now();
        let result = self
            .backend
            .get_part(key, writer.borrow_mut(), offset, length)
            .await;
        match &result {
            Ok(()) => {
                CACHE_METRICS.cache_operations.add(1, self.attrs.read_hit());
                CACHE_METRICS
                    .cache_io
                    .add(writer.get_bytes_written(), self.attrs.read_hit());
                self.record_duration(start, self.attrs.read_hit());
            }
            Err(err) if err.code == Code::NotFound => {
                CACHE_METRICS
                    .cache_operations
                    .add(1, self.attrs.read_miss());
                self.record_duration(start, self.attrs.read_miss());
            }
            Err(_) => {
                CACHE_METRICS
                    .cache_operations
                    .add(1, self.attrs.read_error());
                self.record_duration(start, self.attrs.read_error());
            }
        }
        result
    }

    fn inner_store(&self, _key: Option<StoreKey>) -> &dyn StoreDriver {
        self
    }

    fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) {
        self
    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
        self
    }

    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
        self.backend.clone().register_health(registry);
    }

    fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
        self.backend.register_remove_callback(callback)
    }
}

#[async_trait]
impl HealthStatusIndicator for CacheMetricsStore {
    fn get_name(&self) -> &'static str {
        "CacheMetricsStore"
    }

    async fn check_health(
        &self,
        namespace: std::borrow::Cow<'static, str>,
    ) -> nativelink_util::health_utils::HealthStatus {
        self.backend
            .as_store_driver_pin()
            .check_health(namespace)
            .await
    }
}

Coverage Report

Created: 2026-11-13 23:29

Line	Count	Source
1		// Copyright 2026 Trace Machina, Inc. 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		use core::borrow::BorrowMut;
16		use core::ops::Bound;
17		use core::pin::Pin;
18		use std::ffi::OsString;
19		use std::sync::Arc;
20		use std::time::Instant;
21
22		use async_trait::async_trait;
23		use bytes::Bytes;
24		use nativelink_config::stores::CacheMetricsSpec;
25		use nativelink_error::{Code, Error};
26		use nativelink_metric::{
27		MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, group, publish,
28		};
29		use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
30		use nativelink_util::fs;
31		use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatusIndicator};
32		use nativelink_util::metrics::{CACHE_METRICS, CACHE_TYPE, CacheMetricAttrs};
33		use nativelink_util::store_trait::{
34		RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, StoreOptimizations, UploadSizeInfo,
35		};
36		use opentelemetry::KeyValue;
37
38		#[derive(Debug)]
39		pub struct CacheMetricsStore {
40		backend: Store,
41		cache_type: String,
42		allow_remove: bool,
43		attrs: CacheMetricAttrs,
44		}
45
46		impl CacheMetricsStore {
47	17	pub fn new(spec: &CacheMetricsSpec, backend: Store) -> Arc<Self> {
48	17	let base_attrs = [KeyValue::new(CACHE_TYPE, spec.cache_type.clone())];
49		// A store that tracks its own size reports it under the same cache
50		// type as the operation metrics. Other stores report no size.
51	17	backend
52	17	.inner_store(None::<StoreKey<'_>>)
53	17	.enable_cache_size_metrics(&base_attrs);
54	17	Arc::new(Self {
55	17	backend,
56	17	cache_type: spec.cache_type.clone(),
57	17	allow_remove: spec.allow_remove,
58	17	attrs: CacheMetricAttrs::new(&base_attrs),
59	17	})
60	17	}
61
62	95	fn duration_ms(start: Instant) -> f64 {
63	95	start.elapsed().as_secs_f64() * 1000.0
64	95	}
65
66	0	const fn size_info_bytes(size_info: UploadSizeInfo) -> Option<u64> {
67	0	match size_info {
68	0	UploadSizeInfo::ExactSize(size) => Some(size),
69	0	UploadSizeInfo::MaxSize(_) => None,
70		}
71	0	}
72
73	95	fn record_duration(&self, start: Instant, attrs: &[KeyValue]) {
74	95	CACHE_METRICS
75	95	.cache_operation_duration
76	95	.record(Self::duration_ms(start), attrs);
77	95	}
78
79	9	fn record_write_io(&self, bytes: Option<u64>) {
80	9	if let Some(bytes) = bytes {
81	9	CACHE_METRICS
82	9	.cache_io
83	9	.add(bytes, self.attrs.write_success());
84	9	CACHE_METRICS
85	9	.cache_entry_size
86	9	.record(bytes, self.attrs.write_success());
87	9	}0
88	9	}
89		}
90
91		impl MetricsComponent for CacheMetricsStore {
92	0	fn publish(
93	0	&self,
94	0	_kind: MetricKind,
95	0	_field_metadata: MetricFieldData,
96	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
97	0	publish!(
98	0	"cache_type",
99	0	&self.cache_type,
100	0	MetricKind::String,
101	0	"Low-cardinality cache type label emitted on cache metrics"
102		);
103	0	let _enter = group!("backend").entered();
104	0	self.backend
105	0	.publish(MetricKind::Component, MetricFieldData::default())?;
106	0	Ok(MetricPublishKnownKindData::Component)
107	0	}
108		}
109
110		#[async_trait]
111		impl StoreDriver for CacheMetricsStore {
112	12	async fn post_init(self: Arc<Self>) -> Result<(), Error> {
113		self.backend.clone().into_inner().post_init().await?;
114		Ok(())
115	12	}
116
117	22	async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
118		let start = Instant::now();
119		// Backends such as memory and filesystem cannot infer whether their
120		// copy is disposable. Require an explicit opt-in on this tier.
121		let result = if self.allow_remove {
122		self.backend.as_store_driver_pin().remove(key).await
123		} else {
124		Ok(false)
125		};
126		let attrs = match &result {
127		Ok(true) => self.attrs.delete_success(),
128		Ok(false) => self.attrs.delete_miss(),
129		Err(_) => self.attrs.delete_error(),
130		};
131		CACHE_METRICS.cache_operations.add(1, attrs);
132		self.record_duration(start, attrs);
133		result
134	22	}
135
136		/// Not counted as a delete: this tier's own copy is never dropped.
137	8	async fn remove_cached(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
138		if self.allow_remove {
139		self.backend.as_store_driver_pin().remove_cached(key).await
140		} else {
141		Ok(false)
142		}
143	8	}
144
145		async fn has_with_results(
146		self: Pin<&Self>,
147		keys: &[StoreKey<'_>],
148		results: &mut [Option<u64>],
149	36	) -> Result<(), Error> {
150		let start = Instant::now();
151		let result = self.backend.has_with_results(keys, results).await;
152		if result.is_ok() {
153	36	let hits = results.iter().filter(\|result\| result.is_some()).count();
154		let misses = results.len().saturating_sub(hits);
155		if hits > 0 {
156		CACHE_METRICS
157		.cache_operations
158		.add(hits as u64, self.attrs.read_hit());
159		}
160		if misses > 0 {
161		CACHE_METRICS
162		.cache_operations
163		.add(misses as u64, self.attrs.read_miss());
164		}
165		let duration_attrs = if hits > 0 {
166		self.attrs.read_hit()
167		} else {
168		self.attrs.read_miss()
169		};
170		self.record_duration(start, duration_attrs);
171		} else {
172		CACHE_METRICS
173		.cache_operations
174		.add(keys.len() as u64, self.attrs.read_error());
175		self.record_duration(start, self.attrs.read_error());
176		}
177		result
178	36	}
179
180		async fn list(
181		self: Pin<&Self>,
182		range: (Bound<StoreKey<'_>>, Bound<StoreKey<'_>>),
183		handler: &mut (dyn for<'a> FnMut(&'a StoreKey) -> bool + Send + Sync + '_),
184	0	) -> Result<u64, Error> {
185		self.backend.list(range, handler).await
186	0	}
187
188		async fn update(
189		self: Pin<&Self>,
190		key: StoreKey<'_>,
191		reader: DropCloserReadHalf,
192		upload_size: UploadSizeInfo,
193	10	) -> Result<u64, Error> {
194		let start = Instant::now();
195		let result = self.backend.update(key, reader, upload_size).await;
196		if let Ok(size) = &result {
197		CACHE_METRICS
198		.cache_operations
199		.add(1, self.attrs.write_success());
200		self.record_write_io(Some(*size));
201		self.record_duration(start, self.attrs.write_success());
202		} else {
203		CACHE_METRICS
204		.cache_operations
205		.add(1, self.attrs.write_error());
206		self.record_duration(start, self.attrs.write_error());
207		}
208		result
209	10	}
210
211	26	fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
212	26	self.backend.optimized_for(optimization)
213	26	}
214
215		async fn update_with_whole_file(
216		self: Pin<&Self>,
217		key: StoreKey<'_>,
218		path: OsString,
219		file: fs::FileSlot,
220		upload_size: UploadSizeInfo,
221	0	) -> Result<(u64, Option<fs::FileSlot>), Error> {
222		let start = Instant::now();
223		let bytes = Self::size_info_bytes(upload_size);
224		let result = self
225		.backend
226		.update_with_whole_file(key, path, file, upload_size)
227		.await;
228		if result.is_ok() {
229		CACHE_METRICS
230		.cache_operations
231		.add(1, self.attrs.write_success());
232		self.record_write_io(bytes);
233		self.record_duration(start, self.attrs.write_success());
234		} else {
235		CACHE_METRICS
236		.cache_operations
237		.add(1, self.attrs.write_error());
238		self.record_duration(start, self.attrs.write_error());
239		}
240		result
241	0	}
242
243	1	async fn update_oneshot(self: Pin<&Self>, key: StoreKey<'_>, data: Bytes) -> Result<(), Error> {
244		let start = Instant::now();
245		let bytes = data.len() as u64;
246		let result = self.backend.update_oneshot(key, data).await;
247		if result.is_ok() {
248		CACHE_METRICS
249		.cache_operations
250		.add(1, self.attrs.write_success());
251		self.record_write_io(Some(bytes));
252		self.record_duration(start, self.attrs.write_success());
253		} else {
254		CACHE_METRICS
255		.cache_operations
256		.add(1, self.attrs.write_error());
257		self.record_duration(start, self.attrs.write_error());
258		}
259		result
260	1	}
261
262		async fn get_part(
263		self: Pin<&Self>,
264		key: StoreKey<'_>,
265		writer: &mut DropCloserWriteHalf,
266		offset: u64,
267		length: Option<u64>,
268	26	) -> Result<(), Error> {
269		let start = Instant::now();
270		let result = self
271		.backend
272		.get_part(key, writer.borrow_mut(), offset, length)
273		.await;
274		match &result {
275		Ok(()) => {
276		CACHE_METRICS.cache_operations.add(1, self.attrs.read_hit());
277		CACHE_METRICS
278		.cache_io
279		.add(writer.get_bytes_written(), self.attrs.read_hit());
280		self.record_duration(start, self.attrs.read_hit());
281		}
282		Err(err) if err.code == Code::NotFound => {
283		CACHE_METRICS
284		.cache_operations
285		.add(1, self.attrs.read_miss());
286		self.record_duration(start, self.attrs.read_miss());
287		}
288		Err(_) => {
289		CACHE_METRICS
290		.cache_operations
291		.add(1, self.attrs.read_error());
292		self.record_duration(start, self.attrs.read_error());
293		}
294		}
295		result
296	26	}
297
298	27	fn inner_store(&self, _key: Option<StoreKey>) -> &dyn StoreDriver {
299	27	self
300	27	}
301
302	1	fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) {
303	1	self
304	1	}
305
306	0	fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
307	0	self
308	0	}
309
310	0	fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
311	0	self.backend.clone().register_health(registry);
312	0	}
313
314	0	fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
315	0	self.backend.register_remove_callback(callback)
316	0	}
317		}
318
319		#[async_trait]
320		impl HealthStatusIndicator for CacheMetricsStore {
321	0	fn get_name(&self) -> &'static str {
322	0	"CacheMetricsStore"
323	0	}
324
325		async fn check_health(
326		&self,
327		namespace: std::borrow::Cow<'static, str>,
328	0	) -> nativelink_util::health_utils::HealthStatus {
329		self.backend
330		.as_store_driver_pin()
331		.check_health(namespace)
332		.await
333	0	}
334		}