Skip to main content

snix_castore/directoryservice/combinators/
priority.rs

1use std::sync::Arc;
2
3use futures::{StreamExt, TryStreamExt, stream::BoxStream};
4use tonic::async_trait;
5use tracing::instrument;
6
7use crate::{
8    B3Digest, Directory,
9    composition::{CompositionContext, ServiceBuilder},
10    directoryservice::{self, DirectoryPutter, DirectoryService, FailingPutter},
11};
12
13/// Holds references to many different directory services.
14/// Read requests try services sequentially.
15/// Any error in a service bubbles up.
16/// Write requests are not implemented.
17pub struct Priority<DS> {
18    instance_name: String,
19    // NOTE: Arc<dyn DS> implements DS too, so you can put different service types in here.
20    services: Vec<DS>,
21}
22
23impl<DS> Priority<DS> {
24    /// Construct from an iterator of services.
25    pub fn new<I: IntoIterator<Item = DS>>(instance_name: String, iter: I) -> Priority<DS> {
26        Self {
27            instance_name,
28            services: Vec::from_iter(iter),
29        }
30    }
31}
32
33#[async_trait]
34impl<DS> DirectoryService for Priority<DS>
35where
36    DS: DirectoryService,
37{
38    #[instrument(skip(self, digest), fields(directory.digest = %digest, instance_name = %self.instance_name))]
39    async fn get(&self, digest: &B3Digest) -> Result<Option<Directory>, directoryservice::Error> {
40        // traverse the list of services. If any service has it, return from there.
41        // Errors cause the combinator to bail out early.
42        for (idx, service) in self.services.iter().enumerate() {
43            if let Some(directory) = service
44                .get(digest)
45                .await
46                .map_err(|err| Error::Backend(idx, err))?
47            {
48                return Ok(Some(directory));
49            }
50        }
51
52        Ok(None)
53    }
54
55    #[instrument(skip_all, fields(directory.digest = %root_directory_digest, instance_name = %self.instance_name))]
56    fn get_recursive(
57        &self,
58        root_directory_digest: &B3Digest,
59    ) -> BoxStream<'_, Result<Directory, directoryservice::Error>> {
60        let digest = *root_directory_digest;
61        async_stream::try_stream! {
62            for (idx, service) in self.services.iter().enumerate() {
63                let mut directories_stream = service.get_recursive(&digest);
64                // Once a service said it has a closure (non-empty stream), we return everything from there, including errors.
65                if let Some(directory) = directories_stream.try_next().await.map_err(|err| { Error::Backend(idx, err)})? {
66                    yield directory;
67
68                    while let Some(directory) = directories_stream.try_next().await.map_err(|err| { Error::Backend(idx, err)})? {
69                        yield directory;
70                    }
71                    // we're done
72                    return;
73                }
74                // try the next service in the list
75            }
76        }
77        .boxed()
78    }
79
80    #[instrument(skip_all, fields(instance_name = %self.instance_name))]
81    async fn put(&self, _directory: Directory) -> Result<B3Digest, directoryservice::Error> {
82        Err(Error::Unimplemented.into())
83    }
84
85    #[instrument(skip_all)]
86    fn put_multiple_start(&self) -> Box<dyn DirectoryPutter + '_> {
87        Box::new(FailingPutter)
88    }
89}
90
91#[derive(thiserror::Error, Debug)]
92pub enum Error {
93    #[error("wrong arguments: {0}")]
94    WrongConfig(&'static str),
95
96    #[error("error from service with index {0}")]
97    Backend(usize, #[source] directoryservice::Error),
98
99    #[error("puts are unimplemented")]
100    Unimplemented,
101}
102
103impl From<Error> for directoryservice::Error {
104    fn from(value: Error) -> Self {
105        Self(Box::new(value))
106    }
107}
108
109#[derive(serde::Deserialize, Debug)]
110#[serde(deny_unknown_fields)]
111pub struct PriorityConfig {
112    services: Vec<String>,
113}
114
115impl TryFrom<url::Url> for PriorityConfig {
116    type Error = Box<dyn std::error::Error + Send + Sync>;
117    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
118        if url.has_authority() || !url.path().is_empty() {
119            return Err(Error::WrongConfig("no authority or path allowed").into());
120        }
121        Ok(serde_qs::from_str(url.query().unwrap_or_default())?)
122    }
123}
124
125#[async_trait]
126impl ServiceBuilder for PriorityConfig {
127    type Output = dyn DirectoryService;
128    async fn build<'a>(
129        &'a self,
130        instance_name: &str,
131        context: &CompositionContext,
132    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
133        let services =
134            futures::future::try_join_all(self.services.iter().map(|instance_ref| async move {
135                context.resolve::<Self::Output>(instance_ref).await
136            }))
137            .await?;
138
139        Ok(Arc::new(Priority::new(instance_name.to_string(), services)))
140    }
141}
142
143#[cfg(test)]
144mod test {
145    use mockall::{Sequence, predicate};
146    use pretty_assertions::{assert_eq, assert_matches};
147
148    use super::*;
149    use crate::{
150        directoryservice::{MockDirectoryService, failing_putter},
151        fixtures::{DIRECTORY_A, DIRECTORY_B, DIRECTORY_WITH_KEEP},
152    };
153
154    /// If first has something, last is never tried.
155    #[tokio::test]
156    async fn get_first_gets_tried_only() {
157        let mut first = MockDirectoryService::new();
158        let mut last = MockDirectoryService::new();
159
160        first
161            .expect_get()
162            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
163            .once()
164            .returning(|_| Ok(Some(DIRECTORY_WITH_KEEP.clone())));
165
166        last.expect_get().never();
167
168        let uut = Priority::new("uut".to_string(), [first, last]);
169
170        assert_eq!(
171            Some(DIRECTORY_WITH_KEEP.clone()),
172            uut.get(&DIRECTORY_WITH_KEEP.digest())
173                .await
174                .expect("to succeed")
175        )
176    }
177
178    /// If first doesn't have it, we try last.
179    #[tokio::test]
180    async fn get_first_then_last() {
181        let mut first = MockDirectoryService::new();
182        let mut last = MockDirectoryService::new();
183        let mut seq = Sequence::new();
184
185        first
186            .expect_get()
187            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
188            .once()
189            .in_sequence(&mut seq)
190            .returning(|_| Ok(None));
191
192        last.expect_get()
193            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
194            .once()
195            .in_sequence(&mut seq)
196            .returning(|_| Ok(Some(DIRECTORY_WITH_KEEP.clone())));
197
198        let uut = Priority::new("uut".to_string(), [first, last]);
199
200        assert_eq!(
201            Some(DIRECTORY_WITH_KEEP.clone()),
202            uut.get(&DIRECTORY_WITH_KEEP.digest())
203                .await
204                .expect("to succeed")
205        )
206    }
207
208    /// If none of the two have it, we return None.
209    #[tokio::test]
210    async fn get_first_then_last_not_found() {
211        let mut first = MockDirectoryService::new();
212        let mut last = MockDirectoryService::new();
213        let mut seq = Sequence::new();
214
215        first
216            .expect_get()
217            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
218            .once()
219            .in_sequence(&mut seq)
220            .returning(|_| Ok(None));
221
222        last.expect_get()
223            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
224            .once()
225            .in_sequence(&mut seq)
226            .returning(|_| Ok(None));
227
228        let uut = Priority::new("uut".to_string(), [first, last]);
229
230        assert_eq!(
231            None,
232            uut.get(&DIRECTORY_WITH_KEEP.digest())
233                .await
234                .expect("to succeed")
235        )
236    }
237
238    /// Errors are bubbled up from the first backend emitting the error,
239    /// and the error identifies the backend that emitted the error.
240    #[tokio::test]
241    async fn get_bubble_up_error_first() {
242        let mut first = MockDirectoryService::new();
243        let mut last = MockDirectoryService::new();
244
245        first
246            .expect_get()
247            .with(predicate::eq(DIRECTORY_WITH_KEEP.digest()))
248            .once()
249            .returning(|_| Err(directoryservice::Error("oh no".into())));
250
251        last.expect_get().never();
252
253        let uut = Priority::new("uut".to_string(), [first, last]);
254
255        let err = uut
256            .get(&DIRECTORY_WITH_KEEP.digest())
257            .await
258            .expect_err("must fail")
259            .0;
260
261        let err = err.downcast_ref::<Error>().unwrap();
262        assert_matches!(err, Error::Backend(0, _));
263    }
264
265    /// If the first backend responds to get_recursive, we return from there.
266    #[tokio::test]
267    async fn get_recursive_first() {
268        let mut first = MockDirectoryService::new();
269        let mut last = MockDirectoryService::new();
270
271        first
272            .expect_get_recursive()
273            .with(predicate::eq(DIRECTORY_B.digest()))
274            .once()
275            .returning(|_| {
276                futures::stream::iter([Ok(DIRECTORY_B.clone()), Ok(DIRECTORY_A.clone())]).boxed()
277            });
278        last.expect_get_recursive().never();
279
280        let uut = Priority::new("uut".to_string(), [first, last]);
281
282        let directories = uut
283            .get_recursive(&DIRECTORY_B.digest())
284            .try_collect::<Vec<_>>()
285            .await
286            .expect("to succeed");
287
288        assert_eq!(vec![DIRECTORY_B.clone(), DIRECTORY_A.clone()], directories);
289    }
290
291    /// If the first one doesn't have a directory closure, return from the next.
292    #[tokio::test]
293    async fn get_recursive_second() {
294        let mut first = MockDirectoryService::new();
295        let mut last = MockDirectoryService::new();
296        let mut seq = Sequence::new();
297
298        first
299            .expect_get_recursive()
300            .with(predicate::eq(DIRECTORY_B.digest()))
301            .once()
302            .in_sequence(&mut seq)
303            .returning(|_| futures::stream::empty().boxed());
304
305        last.expect_get_recursive()
306            .with(predicate::eq(DIRECTORY_B.digest()))
307            .once()
308            .in_sequence(&mut seq)
309            .returning(|_| {
310                futures::stream::iter([Ok(DIRECTORY_B.clone()), Ok(DIRECTORY_A.clone())]).boxed()
311            });
312
313        let uut = Priority::new("uut".to_string(), [first, last]);
314
315        let directories = uut
316            .get_recursive(&DIRECTORY_B.digest())
317            .try_collect::<Vec<_>>()
318            .await
319            .expect("to succeed");
320
321        assert_eq!(vec![DIRECTORY_B.clone(), DIRECTORY_A.clone()], directories);
322    }
323
324    /// Propagate errors from get_recursive
325    #[tokio::test]
326    async fn get_recursive_error_first() {
327        let mut first = MockDirectoryService::new();
328        let mut last = MockDirectoryService::new();
329
330        first
331            .expect_get_recursive()
332            .with(predicate::eq(DIRECTORY_B.digest()))
333            .once()
334            .returning(|_| {
335                futures::stream::iter([Err(directoryservice::Error("oh no".into()))]).boxed()
336            });
337
338        last.expect_get_recursive().never();
339
340        let uut = Priority::new("uut".to_string(), [first, last]);
341
342        let err = uut
343            .get_recursive(&DIRECTORY_B.digest())
344            .try_collect::<Vec<_>>()
345            .await
346            .expect_err("to fail")
347            .0;
348
349        let err = err.downcast_ref::<Error>().unwrap();
350        assert_matches!(err, Error::Backend(0, _));
351    }
352
353    /// put is unsupported, and not sent to the backend
354    #[tokio::test]
355    async fn put_unsupported() {
356        let mut first = MockDirectoryService::new();
357        first.expect_put().never();
358
359        let uut = Priority::new("uut".to_string(), [first]);
360
361        let err = uut
362            .put(DIRECTORY_WITH_KEEP.clone())
363            .await
364            .expect_err("must fail")
365            .0;
366
367        let err = err.downcast_ref::<Error>().unwrap();
368        assert_matches!(err, Error::Unimplemented);
369    }
370
371    /// put_recursive is unsupported, and not sent to the backend
372    #[tokio::test]
373    async fn put_recursive_unsupported() {
374        let mut first = MockDirectoryService::new();
375        first.expect_put().never();
376
377        let uut = Priority::new("uut".to_string(), [first]);
378
379        let mut handle = uut.put_multiple_start();
380        let err = handle
381            .put(DIRECTORY_WITH_KEEP.clone())
382            .await
383            .expect_err("must fail")
384            .0;
385
386        let err = err.downcast_ref::<failing_putter::Error>().unwrap();
387        assert_matches!(err, failing_putter::Error::Unimplemented);
388    }
389}