snix_castore/directoryservice/combinators/
priority.rs1use 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
13pub struct Priority<DS> {
18 instance_name: String,
19 services: Vec<DS>,
21}
22
23impl<DS> Priority<DS> {
24 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 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 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 return;
73 }
74 }
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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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}