snix_castore/blobstore/combinators/
priority.rs1use std::sync::Arc;
2
3use tonic::async_trait;
4use tracing::instrument;
5
6use crate::{
7 B3Digest,
8 blobstore::{self, BlobMeta, BlobStore},
9 composition::{CompositionContext, ServiceBuilder},
10};
11
12pub struct Priority<BS> {
17 instance_name: String,
18 services: Vec<BS>,
19}
20
21impl<BS> Priority<BS> {
22 pub fn new<I: IntoIterator<Item = BS>>(instance_name: String, iter: I) -> Priority<BS> {
24 Self {
25 instance_name,
26 services: Vec::from_iter(iter),
27 }
28 }
29}
30
31#[async_trait]
32impl<BS> BlobStore for Priority<BS>
33where
34 BS: BlobStore,
35{
36 #[instrument(skip(self, digest), fields(blob.digest=%digest, instance_name=%self.instance_name))]
37 async fn get(&self, digest: &B3Digest) -> Result<Option<BlobMeta>, blobstore::Error> {
38 for (idx, service) in self.services.iter().enumerate() {
41 if let Some(chunk) = service
42 .get(digest)
43 .await
44 .map_err(|err| Error::Backend(idx, err))?
45 {
46 return Ok(Some(chunk));
47 }
48 }
49
50 Ok(None)
51 }
52
53 #[instrument(skip(self, digest), fields(blob.digest=%digest, instance_name=%self.instance_name))]
54 async fn has(&self, digest: &B3Digest) -> Result<bool, blobstore::Error> {
55 for (idx, service) in self.services.iter().enumerate() {
58 if service
59 .has(digest)
60 .await
61 .map_err(|err| Error::Backend(idx, err))?
62 {
63 return Ok(true);
64 }
65 }
66
67 Ok(false)
68 }
69
70 #[instrument(skip_all, fields(instance_name=%self.instance_name))]
71 async fn put(&self, _digest: &B3Digest, _blob_meta: BlobMeta) -> Result<(), blobstore::Error> {
72 Err(Error::Unimplemented.into())
73 }
74}
75
76#[derive(thiserror::Error, Debug)]
78pub enum Error {
79 #[error("wrong arguments: {0}")]
80 WrongConfig(&'static str),
81
82 #[error("error from service with index {0}")]
83 Backend(usize, #[source] blobstore::Error),
84
85 #[error("puts are unimplemented")]
86 Unimplemented,
87}
88
89impl From<Error> for blobstore::Error {
90 fn from(value: Error) -> Self {
91 Self(Box::new(value))
92 }
93}
94
95#[derive(serde::Deserialize, Debug)]
97#[serde(deny_unknown_fields)]
98pub struct PriorityConfig {
99 services: Vec<String>,
100}
101
102impl TryFrom<url::Url> for PriorityConfig {
103 type Error = Box<dyn std::error::Error + Send + Sync>;
104 fn try_from(url: url::Url) -> Result<Self, Self::Error> {
105 if url.has_authority() || !url.path().is_empty() {
106 return Err(Error::WrongConfig("no authority or path allowed").into());
107 }
108 Ok(serde_qs::from_str(url.query().unwrap_or_default())?)
109 }
110}
111
112#[async_trait]
113impl ServiceBuilder for PriorityConfig {
114 type Output = dyn BlobStore;
115 async fn build<'a>(
116 &'a self,
117 instance_name: &str,
118 context: &CompositionContext,
119 ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
120 let services =
121 futures::future::try_join_all(self.services.iter().map(|instance_ref| async move {
122 context.resolve::<Self::Output>(instance_ref).await
123 }))
124 .await?;
125
126 Ok(Arc::new(Priority::new(instance_name.to_string(), services)))
127 }
128}
129
130#[cfg(test)]
131mod test {
132 use std::sync::LazyLock;
133
134 use mockall::{Sequence, predicate};
135 use pretty_assertions::{assert_eq, assert_matches};
136
137 use crate::{
138 blobstore::{self, BlobMeta, BlobStore, MockBlobStore},
139 chunkstore::Chunk,
140 };
141
142 use super::{Error, Priority};
143
144 const CHUNK: Chunk = Chunk::from_static(b"foo");
145 static BLOB_META: LazyLock<BlobMeta> =
146 LazyLock::new(|| BlobMeta::from_digests_and_sizes([(CHUNK.digest(), CHUNK.len() as u64)]));
147
148 #[tokio::test]
150 async fn get_first_gets_tried_only() {
151 let mut first = MockBlobStore::new();
152 let mut last = MockBlobStore::new();
153
154 first
155 .expect_get()
156 .with(predicate::eq(CHUNK.digest()))
157 .once()
158 .returning(|_| Ok(Some(BLOB_META.clone())));
159
160 last.expect_get().never();
161
162 let uut = Priority::new("uut".to_string(), [first, last]);
163
164 assert_eq!(
165 Some(BLOB_META.clone()),
166 uut.get(&CHUNK.digest()).await.expect("to succeed")
167 )
168 }
169
170 #[tokio::test]
172 async fn get_first_then_last() {
173 let mut first = MockBlobStore::new();
174 let mut last = MockBlobStore::new();
175 let mut seq = Sequence::new();
176
177 first
178 .expect_get()
179 .with(predicate::eq(CHUNK.digest()))
180 .once()
181 .in_sequence(&mut seq)
182 .returning(|_| Ok(None));
183
184 last.expect_get()
185 .with(predicate::eq(CHUNK.digest()))
186 .once()
187 .in_sequence(&mut seq)
188 .returning(|_| Ok(Some(BLOB_META.clone())));
189
190 let uut = Priority::new("uut".to_string(), [first, last]);
191
192 assert_eq!(
193 Some(BLOB_META.clone()),
194 uut.get(&CHUNK.digest()).await.expect("to succeed")
195 )
196 }
197
198 #[tokio::test]
200 async fn get_first_then_last_not_found() {
201 let mut first = MockBlobStore::new();
202 let mut last = MockBlobStore::new();
203 let mut seq = Sequence::new();
204
205 first
206 .expect_get()
207 .with(predicate::eq(CHUNK.digest()))
208 .once()
209 .in_sequence(&mut seq)
210 .returning(|_| Ok(None));
211
212 last.expect_get()
213 .with(predicate::eq(CHUNK.digest()))
214 .once()
215 .in_sequence(&mut seq)
216 .returning(|_| Ok(None));
217
218 let uut = Priority::new("uut".to_string(), [first, last]);
219
220 assert_eq!(None, uut.get(&CHUNK.digest()).await.expect("to succeed"))
221 }
222
223 #[tokio::test]
226 async fn get_bubble_up_error_first() {
227 let mut first = MockBlobStore::new();
228 let mut last = MockBlobStore::new();
229
230 first
231 .expect_get()
232 .with(predicate::eq(CHUNK.digest()))
233 .once()
234 .returning(|_| Err(blobstore::Error("oh no".into())));
235
236 last.expect_get().never();
237
238 let uut = Priority::new("uut".to_string(), [first, last]);
239
240 let err = uut.get(&CHUNK.digest()).await.expect_err("must fail").0;
241
242 let err = err.downcast_ref::<Error>().unwrap();
243 assert_matches!(err, Error::Backend(0, _));
244 }
245
246 #[tokio::test]
248 async fn has_first_then_last() {
249 let mut first = MockBlobStore::new();
250 let mut last = MockBlobStore::new();
251 let mut seq = Sequence::new();
252
253 first
254 .expect_has()
255 .with(predicate::eq(CHUNK.digest()))
256 .once()
257 .in_sequence(&mut seq)
258 .returning(|_| Ok(false));
259
260 last.expect_has()
261 .with(predicate::eq(CHUNK.digest()))
262 .once()
263 .in_sequence(&mut seq)
264 .returning(|_| Ok(true));
265
266 let uut = Priority::new("uut".to_string(), [first, last]);
267
268 assert!(uut.has(&CHUNK.digest()).await.expect("to succeed"));
269 }
270
271 #[tokio::test]
273 async fn put_unsupported() {
274 let mut first = MockBlobStore::new();
275 first.expect_put().never();
276
277 let uut = Priority::new("uut".to_string(), [first]);
278
279 let err = uut
280 .put(&CHUNK.digest(), BLOB_META.clone())
281 .await
282 .expect_err("must fail")
283 .0;
284
285 let err = err.downcast_ref::<Error>().unwrap();
286 assert_matches!(err, Error::Unimplemented);
287 }
288}