snix_castore/directoryservice/combinators/
priority.rs

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

snix_castore/directoryservice/combinators/
priority.rs
