Skip to main content

snix_castore/chunkstore/combinators/
priority.rs

1use std::sync::Arc;
2
3use tonic::async_trait;
4use tracing::instrument;
5
6use crate::{
7    B3Digest,
8    chunkstore::{self, Chunk, ChunkStore},
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<CS> {
17    instance_name: String,
18    services: Vec<CS>,
19}
20
21impl<CS> Priority<CS> {
22    /// Construct from an iterator of services.
23    pub fn new<I: IntoIterator<Item = CS>>(instance_name: String, iter: I) -> Priority<CS> {
24        Self {
25            instance_name,
26            services: Vec::from_iter(iter),
27        }
28    }
29}
30
31#[async_trait]
32impl<CS> ChunkStore for Priority<CS>
33where
34    CS: ChunkStore,
35{
36    #[instrument(skip(self, digest), fields(chunk.digest=%digest, instance_name=%self.instance_name))]
37    async fn get(&self, digest: &B3Digest) -> Result<Option<Chunk>, chunkstore::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(chunk.digest=%digest, instance_name=%self.instance_name))]
54    async fn has(&self, digest: &B3Digest) -> Result<bool, chunkstore::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, _chunk: Chunk) -> Result<B3Digest, chunkstore::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] chunkstore::Error),
84
85    #[error("puts are unimplemented")]
86    Unimplemented,
87}
88
89impl From<Error> for chunkstore::Error {
90    fn from(value: Error) -> Self {
91        Self(Box::new(value))
92    }
93}
94
95/// Configuration for a [Priority] [ChunkStore].
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 ChunkStore;
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 mockall::{Sequence, predicate};
133    use pretty_assertions::{assert_eq, assert_matches};
134
135    use crate::chunkstore::{self, Chunk, ChunkStore, MockChunkStore};
136
137    use super::{Error, Priority};
138
139    const CHUNK: Chunk = Chunk::from_static(b"foo");
140
141    /// If first has something, last is never tried.
142    #[tokio::test]
143    async fn get_first_gets_tried_only() {
144        let mut first = MockChunkStore::new();
145        let mut last = MockChunkStore::new();
146
147        first
148            .expect_get()
149            .with(predicate::eq(CHUNK.digest()))
150            .once()
151            .returning(|_| Ok(Some(CHUNK.clone())));
152
153        last.expect_get().never();
154
155        let uut = Priority::new("uut".to_string(), [first, last]);
156
157        assert_eq!(
158            Some(CHUNK.clone()),
159            uut.get(&CHUNK.digest()).await.expect("to succeed")
160        )
161    }
162
163    /// If first doesn't have it, we try last.
164    #[tokio::test]
165    async fn get_first_then_last() {
166        let mut first = MockChunkStore::new();
167        let mut last = MockChunkStore::new();
168        let mut seq = Sequence::new();
169
170        first
171            .expect_get()
172            .with(predicate::eq(CHUNK.digest()))
173            .once()
174            .in_sequence(&mut seq)
175            .returning(|_| Ok(None));
176
177        last.expect_get()
178            .with(predicate::eq(CHUNK.digest()))
179            .once()
180            .in_sequence(&mut seq)
181            .returning(|_| Ok(Some(CHUNK.clone())));
182
183        let uut = Priority::new("uut".to_string(), [first, last]);
184
185        assert_eq!(
186            Some(CHUNK.clone()),
187            uut.get(&CHUNK.digest()).await.expect("to succeed")
188        )
189    }
190
191    /// If none of the two have it, we return None.
192    #[tokio::test]
193    async fn get_first_then_last_not_found() {
194        let mut first = MockChunkStore::new();
195        let mut last = MockChunkStore::new();
196        let mut seq = Sequence::new();
197
198        first
199            .expect_get()
200            .with(predicate::eq(CHUNK.digest()))
201            .once()
202            .in_sequence(&mut seq)
203            .returning(|_| Ok(None));
204
205        last.expect_get()
206            .with(predicate::eq(CHUNK.digest()))
207            .once()
208            .in_sequence(&mut seq)
209            .returning(|_| Ok(None));
210
211        let uut = Priority::new("uut".to_string(), [first, last]);
212
213        assert_eq!(None, uut.get(&CHUNK.digest()).await.expect("to succeed"))
214    }
215
216    /// Errors are bubbled up from the first backend emitting the error,
217    /// and the error identifies the backend that emitted the error.
218    #[tokio::test]
219    async fn get_bubble_up_error_first() {
220        let mut first = MockChunkStore::new();
221        let mut last = MockChunkStore::new();
222
223        first
224            .expect_get()
225            .with(predicate::eq(CHUNK.digest()))
226            .once()
227            .returning(|_| Err(chunkstore::Error("oh no".into())));
228
229        last.expect_get().never();
230
231        let uut = Priority::new("uut".to_string(), [first, last]);
232
233        let err = uut.get(&CHUNK.digest()).await.expect_err("must fail").0;
234
235        let err = err.downcast_ref::<Error>().unwrap();
236        assert_matches!(err, Error::Backend(0, _));
237    }
238
239    /// has() returns true from the first backend that has it.
240    #[tokio::test]
241    async fn has_first_then_last() {
242        let mut first = MockChunkStore::new();
243        let mut last = MockChunkStore::new();
244        let mut seq = Sequence::new();
245
246        first
247            .expect_has()
248            .with(predicate::eq(CHUNK.digest()))
249            .once()
250            .in_sequence(&mut seq)
251            .returning(|_| Ok(false));
252
253        last.expect_has()
254            .with(predicate::eq(CHUNK.digest()))
255            .once()
256            .in_sequence(&mut seq)
257            .returning(|_| Ok(true));
258
259        let uut = Priority::new("uut".to_string(), [first, last]);
260
261        assert!(uut.has(&CHUNK.digest()).await.expect("to succeed"));
262    }
263
264    /// put is unsupported, and not sent to the backend.
265    #[tokio::test]
266    async fn put_unsupported() {
267        let mut first = MockChunkStore::new();
268        first.expect_put().never();
269
270        let uut = Priority::new("uut".to_string(), [first]);
271
272        let err = uut.put(CHUNK.clone()).await.expect_err("must fail").0;
273
274        let err = err.downcast_ref::<Error>().unwrap();
275        assert_matches!(err, Error::Unimplemented);
276    }
277}