/build/source/nativelink-store/src/cache_metrics_store.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.

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,
    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(),
            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 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-09-18 20:40

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