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

Source
// Copyright 2024 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::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use futures::try_join;
use nativelink_config::stores::SizePartitioningSpec;
use nativelink_error::{Error, ResultExt, make_input_err};
use nativelink_metric::MetricsComponent;
use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
use nativelink_util::store_trait::{
    RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
};
use tokio::join;

#[derive(Debug, MetricsComponent)]
pub struct SizePartitioningStore {
    #[metric(help = "Size to partition our data")]
    partition_size: u64,
    #[metric(group = "lower_store")]
    lower_store: Store,
    #[metric(group = "upper_store")]
    upper_store: Store,
}

impl SizePartitioningStore {
    pub fn new(spec: &SizePartitioningSpec, lower_store: Store, upper_store: Store) -> Arc<Self> {
        Arc::new(Self {
            partition_size: spec.size,
            lower_store,
            upper_store,
        })
    }
}

#[async_trait]
impl StoreDriver for SizePartitioningStore {
    async fn post_init(self: Arc<Self>) -> Result<(), Error> {
        try_join!(
            self.upper_store.clone().into_inner().post_init(),
            self.lower_store.clone().into_inner().post_init(),
        )?;
        Ok(())
    }

    async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
        let StoreKey::Digest(digest) = key.borrow() else {
            return Err(make_input_err!(
                "SizePartitioningStore only supports Digest keys, got {key:?}"
            ));
        };
        if digest.size_bytes() < self.partition_size {
            self.lower_store.as_store_driver_pin().remove(key).await
        } else {
            self.upper_store.as_store_driver_pin().remove(key).await
        }
    }

    async fn remove_cached(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
        let StoreKey::Digest(digest) = key.borrow() else {
            return Err(make_input_err!(
                "SizePartitioningStore only supports Digest keys, got {key:?}"
            ));
        };
        if digest.size_bytes() < self.partition_size {
            self.lower_store
                .as_store_driver_pin()
                .remove_cached(key)
                .await
        } else {
            self.upper_store
                .as_store_driver_pin()
                .remove_cached(key)
                .await
        }
    }

    async fn has_with_results(
        self: Pin<&Self>,
        keys: &[StoreKey<'_>],
        results: &mut [Option<u64>],
    ) -> Result<(), Error> {
        let mut non_digest_sample = None;
        let (lower_digests, upper_digests): (Vec<_>, Vec<_>) =
            keys.iter().map(StoreKey::borrow).partition(|k| {
                let StoreKey::Digest(digest) = k else {
                    non_digest_sample = Some(k.borrow().into_owned());
                    return false;
                };
                digest.size_bytes() < self.partition_size
            });
        if let Some(non_digest) = non_digest_sample {
            return Err(make_input_err!(
                "SizePartitioningStore only supports Digest keys, got {non_digest:?}"
            ));
        }
        let (lower_results, upper_results) = join!(
            self.lower_store.has_many(&lower_digests),
            self.upper_store.has_many(&upper_digests),
        );
        let mut lower_results = match lower_results {
            Ok(lower_results) => lower_results.into_iter(),
            Err(err) => match upper_results {
                Ok(_) => return Err(err),
                Err(upper_err) => return Err(err.merge(upper_err)),
            },
        };
        let mut upper_digests = upper_digests.into_iter().peekable();
        let mut upper_results = upper_results?.into_iter();
        for (digest, result) in keys.iter().zip(results.iter_mut()) {
            if Some(digest) == upper_digests.peek() {
                upper_digests.next();
                *result = upper_results
                    .next()
                    .err_tip(|| "upper_results out of sync with upper_digests")?;
            } else {
                *result = lower_results
                    .next()
                    .err_tip(|| "lower_results out of sync with lower_digests")?;
            }
        }
        Ok(())
    }

    async fn update(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        reader: DropCloserReadHalf,
        size_info: UploadSizeInfo,
    ) -> Result<u64, Error> {
        let digest = match key {
            StoreKey::Digest(digest) => digest,
            other @ StoreKey::Str(_) => {
                return Err(make_input_err!(
                    "SizePartitioningStore only supports Digest keys, got {other:?}"
                ));
            }
        };
        if digest.size_bytes() < self.partition_size {
            return self.lower_store.update(digest, reader, size_info).await;
        }
        self.upper_store.update(digest, reader, size_info).await
    }

    async fn get_part(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        writer: &mut DropCloserWriteHalf,
        offset: u64,
        length: Option<u64>,
    ) -> Result<(), Error> {
        let digest = match key {
            StoreKey::Digest(digest) => digest,
            other @ StoreKey::Str(_) => {
                return Err(make_input_err!(
                    "SizePartitioningStore only supports Digest keys, got {other:?}"
                ));
            }
        };
        if digest.size_bytes() < self.partition_size {
            return self
                .lower_store
                .get_part(digest, writer, offset, length)
                .await;
        }
        self.upper_store
            .get_part(digest, writer, offset, length)
            .await
    }

    fn inner_store(&self, key: Option<StoreKey>) -> &'_ dyn StoreDriver {
        let Some(key) = key else {
            return self;
        };
        let StoreKey::Digest(digest) = key else {
            return self;
        };
        if digest.size_bytes() < self.partition_size {
            return self.lower_store.inner_store(Some(digest));
        }
        self.upper_store.inner_store(Some(digest))
    }

    fn as_any<'a>(&'a self) -> &'a (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_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
        self.lower_store
            .register_remove_callback(callback.clone())?;
        self.upper_store.register_remove_callback(callback)?;
        Ok(())
    }
}

default_health_status_indicator!(SizePartitioningStore);

Coverage Report

Created: 2026-11-13 23:29

Line	Count	Source
1		// Copyright 2024 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::pin::Pin;
16		use std::sync::Arc;
17
18		use async_trait::async_trait;
19		use futures::try_join;
20		use nativelink_config::stores::SizePartitioningSpec;
21		use nativelink_error::{Error, ResultExt, make_input_err};
22		use nativelink_metric::MetricsComponent;
23		use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
24		use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
25		use nativelink_util::store_trait::{
26		RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
27		};
28		use tokio::join;
29
30		#[derive(Debug, MetricsComponent)]
31		pub struct SizePartitioningStore {
32		#[metric(help = "Size to partition our data")]
33		partition_size: u64,
34		#[metric(group = "lower_store")]
35		lower_store: Store,
36		#[metric(group = "upper_store")]
37		upper_store: Store,
38		}
39
40		impl SizePartitioningStore {
41	8	pub fn new(spec: &SizePartitioningSpec, lower_store: Store, upper_store: Store) -> Arc<Self> {
42	8	Arc::new(Self {
43	8	partition_size: spec.size,
44	8	lower_store,
45	8	upper_store,
46	8	})
47	8	}
48		}
49
50		#[async_trait]
51		impl StoreDriver for SizePartitioningStore {
52	4	async fn post_init(self: Arc<Self>) -> Result<(), Error> {
53		try_join!(
54		self.upper_store.clone().into_inner().post_init(),
55		self.lower_store.clone().into_inner().post_init(),
56		)?;
57		Ok(())
58	4	}
59
60	4	async fn remove(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
61		let StoreKey::Digest(digest) = key.borrow() else {
62		return Err(make_input_err!(
63		"SizePartitioningStore only supports Digest keys, got {key:?}"
64		));
65		};
66		if digest.size_bytes() < self.partition_size {
67		self.lower_store.as_store_driver_pin().remove(key).await
68		} else {
69		self.upper_store.as_store_driver_pin().remove(key).await
70		}
71	4	}
72
73	0	async fn remove_cached(self: Pin<&Self>, key: StoreKey<'_>) -> Result<bool, Error> {
74		let StoreKey::Digest(digest) = key.borrow() else {
75		return Err(make_input_err!(
76		"SizePartitioningStore only supports Digest keys, got {key:?}"
77		));
78		};
79		if digest.size_bytes() < self.partition_size {
80		self.lower_store
81		.as_store_driver_pin()
82		.remove_cached(key)
83		.await
84		} else {
85		self.upper_store
86		.as_store_driver_pin()
87		.remove_cached(key)
88		.await
89		}
90	0	}
91
92		async fn has_with_results(
93		self: Pin<&Self>,
94		keys: &[StoreKey<'_>],
95		results: &mut [Option<u64>],
96	14	) -> Result<(), Error> {
97		let mut non_digest_sample = None;
98		let (lower_digests, upper_digests): (Vec<_>, Vec<_>) =
99	14	keys.iter().map(StoreKey::borrow).partition(\|k\| {
100	14	let StoreKey::Digest(digest) = k else {
101	0	non_digest_sample = Some(k.borrow().into_owned());
102	0	return false;
103		};
104	14	digest.size_bytes() < self.partition_size
105	14	});
106		if let Some(non_digest) = non_digest_sample {
107		return Err(make_input_err!(
108		"SizePartitioningStore only supports Digest keys, got {non_digest:?}"
109		));
110		}
111		let (lower_results, upper_results) = join!(
112		self.lower_store.has_many(&lower_digests),
113		self.upper_store.has_many(&upper_digests),
114		);
115		let mut lower_results = match lower_results {
116		Ok(lower_results) => lower_results.into_iter(),
117		Err(err) => match upper_results {
118		Ok(_) => return Err(err),
119		Err(upper_err) => return Err(err.merge(upper_err)),
120		},
121		};
122		let mut upper_digests = upper_digests.into_iter().peekable();
123		let mut upper_results = upper_results?.into_iter();
124		for (digest, result) in keys.iter().zip(results.iter_mut()) {
125		if Some(digest) == upper_digests.peek() {
126		upper_digests.next();
127		*result = upper_results
128		.next()
129		.err_tip(\|\| "upper_results out of sync with upper_digests")?;
130		} else {
131		*result = lower_results
132		.next()
133		.err_tip(\|\| "lower_results out of sync with lower_digests")?;
134		}
135		}
136		Ok(())
137	14	}
138
139		async fn update(
140		self: Pin<&Self>,
141		key: StoreKey<'_>,
142		reader: DropCloserReadHalf,
143		size_info: UploadSizeInfo,
144	10	) -> Result<u64, Error> {
145		let digest = match key {
146		StoreKey::Digest(digest) => digest,
147		other @ StoreKey::Str(_) => {
148		return Err(make_input_err!(
149		"SizePartitioningStore only supports Digest keys, got {other:?}"
150		));
151		}
152		};
153		if digest.size_bytes() < self.partition_size {
154		return self.lower_store.update(digest, reader, size_info).await;
155		}
156		self.upper_store.update(digest, reader, size_info).await
157	10	}
158
159		async fn get_part(
160		self: Pin<&Self>,
161		key: StoreKey<'_>,
162		writer: &mut DropCloserWriteHalf,
163		offset: u64,
164		length: Option<u64>,
165	8	) -> Result<(), Error> {
166		let digest = match key {
167		StoreKey::Digest(digest) => digest,
168		other @ StoreKey::Str(_) => {
169		return Err(make_input_err!(
170		"SizePartitioningStore only supports Digest keys, got {other:?}"
171		));
172		}
173		};
174		if digest.size_bytes() < self.partition_size {
175		return self
176		.lower_store
177		.get_part(digest, writer, offset, length)
178		.await;
179		}
180		self.upper_store
181		.get_part(digest, writer, offset, length)
182		.await
183	8	}
184
185	6	fn inner_store(&self, key: Option<StoreKey>) -> &'_ dyn StoreDriver {
186	6	let Some(key) = key else {
187	0	return self;
188		};
189	6	let StoreKey::Digest(digest) = key else {
190	0	return self;
191		};
192	6	if digest.size_bytes() < self.partition_size {
193	6	return self.lower_store.inner_store(Some(digest));
194	0	}
195	0	self.upper_store.inner_store(Some(digest))
196	6	}
197
198	0	fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
199	0	self
200	0	}
201
202	0	fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
203	0	self
204	0	}
205
206	0	fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
207	0	self.lower_store
208	0	.register_remove_callback(callback.clone())?;
209	0	self.upper_store.register_remove_callback(callback)?;
210	0	Ok(())
211	0	}
212		}
213
214		default_health_status_indicator!(SizePartitioningStore);