Skip to main content

snix_castore/chunkstore/combinators/
race.rs

1use std::sync::Arc;
2
3use tonic::async_trait;
4use tracing::instrument;
5
6use crate::{
7    B3Digest,
8    chunkstore::{self, Chunk, ChunkStore},
9    combinators,
10    composition::{CompositionContext, ServiceBuilder},
11};
12
13/// Fans out requests to multiple stores in parallel, returning the first positive or erroneous answer.
14///
15/// A negative response is returned if all backends return them.
16/// Write requests are not implemented.
17pub struct Race<CS> {
18    instance_name: String,
19    services: Vec<CS>,
20}
21
22impl<CS> Race<CS> {
23    /// Construct from an iterator of services.
24    pub fn new<I: IntoIterator<Item = CS>>(instance_name: String, iter: I) -> Race<CS> {
25        Self {
26            instance_name,
27            services: Vec::from_iter(iter),
28        }
29    }
30
31    /// Add another service to the list.
32    pub fn add(&mut self, svc: CS) {
33        self.services.push(svc);
34    }
35}
36
37#[async_trait]
38impl<CS> ChunkStore for Race<CS>
39where
40    CS: ChunkStore,
41{
42    #[instrument(skip(self, digest), fields(chunk.digest=%digest, instance_name=%self.instance_name))]
43    async fn get(&self, digest: &B3Digest) -> Result<Option<Chunk>, chunkstore::Error> {
44        Ok(combinators::race::race_unary(&self.services, |svc| async {
45            // Skip over `Ok(None)` by returning None,
46            // but keep the Option<Chunk> in the returned Ok() value.
47            Some(svc.get(digest).await.transpose()?.map(Some))
48        })
49        .await
50        .map_err(Error::Racing)?)
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        Ok(combinators::race::race_unary(&self.services, |svc| async {
56            // Skip over Ok(false) by returning None.
57            match svc.has(digest).await {
58                Ok(false) => None,
59                resp => Some(resp),
60            }
61        })
62        .await
63        .map_err(Error::Racing)?)
64    }
65
66    #[instrument(skip_all, fields(instance_name=%self.instance_name))]
67    async fn put(&self, _chunk: Chunk) -> Result<B3Digest, chunkstore::Error> {
68        Err(Error::Unimplemented.into())
69    }
70}
71
72/// Error returned from [Race].
73#[derive(thiserror::Error, Debug)]
74pub enum Error {
75    #[error("wrong arguments: {0}")]
76    WrongConfig(&'static str),
77
78    #[error("error from racing")]
79    Racing(#[source] combinators::race::Error<chunkstore::Error>),
80
81    #[error("puts are unimplemented")]
82    Unimplemented,
83}
84
85impl From<Error> for chunkstore::Error {
86    fn from(value: Error) -> Self {
87        Self(Box::new(value))
88    }
89}
90
91/// Configuration for a [Race] [ChunkStore].
92#[derive(serde::Deserialize, Debug)]
93#[serde(deny_unknown_fields)]
94pub struct RaceConfig {
95    services: Vec<String>,
96}
97
98impl TryFrom<url::Url> for RaceConfig {
99    type Error = Box<dyn std::error::Error + Send + Sync>;
100    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
101        if url.has_authority() || !url.path().is_empty() {
102            return Err(Error::WrongConfig("no authority or path allowed").into());
103        }
104        Ok(serde_qs::from_str(url.query().unwrap_or_default())?)
105    }
106}
107
108#[async_trait]
109impl ServiceBuilder for RaceConfig {
110    type Output = dyn ChunkStore;
111    async fn build<'a>(
112        &'a self,
113        instance_name: &str,
114        context: &CompositionContext,
115    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
116        let services =
117            futures::future::try_join_all(self.services.iter().map(|instance_ref| async move {
118                context.resolve::<Self::Output>(instance_ref).await
119            }))
120            .await?;
121
122        Ok(Arc::new(Race::new(instance_name.to_string(), services)))
123    }
124}
125
126#[cfg(test)]
127mod test {
128    use mockall::predicate;
129    use pretty_assertions::assert_matches;
130
131    use crate::{
132        chunkstore::{self, Chunk, ChunkStore, MockChunkStore},
133        combinators,
134    };
135
136    use super::{Error, Race};
137
138    const CHUNK: Chunk = Chunk::from_static(b"foo");
139
140    /// backends are tried exhaustively if all report None.
141    #[tokio::test]
142    async fn get_tries_exhaustively_on_none() {
143        let first = {
144            let mut svc = MockChunkStore::new();
145            svc.expect_get()
146                .with(predicate::eq(CHUNK.digest()))
147                .once()
148                .returning(|_| Ok(None));
149            svc
150        };
151        let second = {
152            let mut svc = MockChunkStore::new();
153            svc.expect_get()
154                .with(predicate::eq(CHUNK.digest()))
155                .once()
156                .returning(|_| Ok(None));
157            svc
158        };
159
160        let uut = Race::new("uut".to_string(), [first, second]);
161
162        assert!(
163            uut.get(&CHUNK.digest())
164                .await
165                .expect("to succeed")
166                .is_none()
167        )
168    }
169
170    /// backends are tried exhaustively if all report false.
171    #[tokio::test]
172    async fn has_tries_exhaustively_on_false() {
173        let first = {
174            let mut svc = MockChunkStore::new();
175            svc.expect_has()
176                .with(predicate::eq(CHUNK.digest()))
177                .once()
178                .returning(|_| Ok(false));
179            svc
180        };
181        let second = {
182            let mut svc = MockChunkStore::new();
183            svc.expect_has()
184                .with(predicate::eq(CHUNK.digest()))
185                .once()
186                .returning(|_| Ok(false));
187            svc
188        };
189
190        let uut = Race::new("uut".to_string(), [first, second]);
191
192        assert!(!uut.has(&CHUNK.digest()).await.expect("to succeed"));
193    }
194
195    /// if one has it and one does not, we return the positive result.
196    #[tokio::test]
197    async fn get_returns_positive() {
198        let first = {
199            let mut svc = MockChunkStore::new();
200            svc.expect_get()
201                .with(predicate::eq(CHUNK.digest()))
202                .once()
203                .returning(|_| Ok(Some(CHUNK.clone())));
204            svc
205        };
206
207        let second = {
208            let mut svc = MockChunkStore::new();
209            svc.expect_get()
210                .with(predicate::eq(CHUNK.digest()))
211                // We cannot be certain this is called at all, so no `once()` here.
212                .returning(|_| Ok(None));
213            svc
214        };
215
216        let uut = Race::new("uut".to_string(), [first, second]);
217
218        assert_eq!(
219            Some(CHUNK.clone()),
220            uut.get(&CHUNK.digest()).await.expect("to succeed")
221        )
222    }
223
224    /// Errors are bubbled up, and the error contains the correct service index.
225    #[tokio::test]
226    async fn get_return_error() {
227        let first = {
228            let mut svc = MockChunkStore::new();
229            svc.expect_get()
230                .with(predicate::eq(CHUNK.digest()))
231                .once()
232                .returning(|_| Err(chunkstore::Error("".into())));
233            svc
234        };
235
236        // Ideally this one would be just slower than `first`.
237        let second = {
238            let mut svc = MockChunkStore::new();
239            svc.expect_get()
240                .with(predicate::eq(CHUNK.digest()))
241                // We cannot be certain this is called at all, so no `once()` here.
242                .returning(|_| Ok(None));
243            svc
244        };
245
246        let uut = Race::new("uut".to_string(), [first, second]);
247
248        let err = uut.get(&CHUNK.digest()).await.expect_err("to fail").0;
249        let err = err.downcast_ref::<Error>().unwrap();
250        assert_matches!(err, Error::Racing(combinators::race::Error::Backend(0, _)))
251    }
252
253    /// put is unsupported, and not sent to any backend.
254    #[tokio::test]
255    async fn put_unsupported() {
256        let mut first = MockChunkStore::new();
257        first.expect_put().never();
258
259        let uut = Race::new("uut".to_string(), [first]);
260
261        let err = uut.put(CHUNK.clone()).await.expect_err("must fail").0;
262
263        let err = err.downcast_ref::<Error>().unwrap();
264        assert_matches!(err, Error::Unimplemented);
265    }
266}