Skip to main content

snix_castore/blobstore/combinators/
priority.rs

1use 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
12/// Tries requests on multiple stores sequentially, returning the first positive or erroneous answer.
13///
14/// A negative response is returned if all backends return them.
15/// Write requests are not implemented.
16pub struct Priority<BS> {
17    instance_name: String,
18    services: Vec<BS>,
19}
20
21impl<BS> Priority<BS> {
22    /// Construct from an iterator of services.
23    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        // traverse the list of services. If any service has it, return from there.
39        // Errors cause the combinator to bail out early.
40        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        // traverse the list of services. If any service has it, return true.
56        // Errors cause the combinator to bail out early.
57        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/// Error returned from [Priority].
77#[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/// Configuration for a [Priority] [BlobStore].
96#[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    /// If first has something, last is never tried.
149    #[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    /// If first doesn't have it, we try last.
171    #[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    /// If none of the two have it, we return None.
199    #[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    /// Errors are bubbled up from the first backend emitting the error,
224    /// and the error identifies the backend that emitted the error.
225    #[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    /// has() returns true from the first backend that has it.
247    #[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    /// put is unsupported, and not sent to the backend.
272    #[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}