Skip to main content

snix_castore/blobstore/combinators/
race.rs

1use std::sync::Arc;
2
3use tonic::async_trait;
4use tracing::instrument;
5
6use crate::{
7    B3Digest,
8    blobstore::{self, BlobMeta, BlobStore},
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<BS> {
18    instance_name: String,
19    services: Vec<BS>,
20}
21
22impl<BS> Race<BS> {
23    /// Construct from an iterator of services.
24    pub fn new<I: IntoIterator<Item = BS>>(instance_name: String, iter: I) -> Race<BS> {
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: BS) {
33        self.services.push(svc);
34    }
35}
36
37#[async_trait]
38impl<BS> BlobStore for Race<BS>
39where
40    BS: BlobStore,
41{
42    #[instrument(skip(self, digest), fields(blob.digest=%digest, instance_name=%self.instance_name))]
43    async fn get(&self, digest: &B3Digest) -> Result<Option<BlobMeta>, blobstore::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(blob.digest=%digest, instance_name=%self.instance_name))]
54    async fn has(&self, digest: &B3Digest) -> Result<bool, blobstore::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, _digest: &B3Digest, _blob_meta: BlobMeta) -> Result<(), blobstore::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<blobstore::Error>),
80
81    #[error("puts are unimplemented")]
82    Unimplemented,
83}
84
85impl From<Error> for blobstore::Error {
86    fn from(value: Error) -> Self {
87        Self(Box::new(value))
88    }
89}
90
91/// Configuration for a [Race] [BlobStore].
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 BlobStore;
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 std::sync::LazyLock;
129
130    use mockall::predicate;
131    use pretty_assertions::assert_matches;
132
133    use crate::{
134        blobstore::{self, BlobMeta, BlobStore, MockBlobStore},
135        chunkstore::Chunk,
136        combinators,
137    };
138
139    use super::{Error, Race};
140
141    const CHUNK: Chunk = Chunk::from_static(b"foo");
142    static BLOB_META: LazyLock<BlobMeta> =
143        LazyLock::new(|| BlobMeta::from_digests_and_sizes([(CHUNK.digest(), CHUNK.len() as u64)]));
144
145    /// backends are tried exhaustively if all report None.
146    #[tokio::test]
147    async fn get_tries_exhaustively_on_none() {
148        let first = {
149            let mut svc = MockBlobStore::new();
150            svc.expect_get()
151                .with(predicate::eq(CHUNK.digest()))
152                .once()
153                .returning(|_| Ok(None));
154            svc
155        };
156        let second = {
157            let mut svc = MockBlobStore::new();
158            svc.expect_get()
159                .with(predicate::eq(CHUNK.digest()))
160                .once()
161                .returning(|_| Ok(None));
162            svc
163        };
164
165        let uut = Race::new("uut".to_string(), [first, second]);
166
167        assert!(
168            uut.get(&CHUNK.digest())
169                .await
170                .expect("to succeed")
171                .is_none()
172        )
173    }
174
175    /// backends are tried exhaustively if all report false.
176    #[tokio::test]
177    async fn has_tries_exhaustively_on_false() {
178        let first = {
179            let mut svc = MockBlobStore::new();
180            svc.expect_has()
181                .with(predicate::eq(CHUNK.digest()))
182                .once()
183                .returning(|_| Ok(false));
184            svc
185        };
186        let second = {
187            let mut svc = MockBlobStore::new();
188            svc.expect_has()
189                .with(predicate::eq(CHUNK.digest()))
190                .once()
191                .returning(|_| Ok(false));
192            svc
193        };
194
195        let uut = Race::new("uut".to_string(), [first, second]);
196
197        assert!(!uut.has(&CHUNK.digest()).await.expect("to succeed"));
198    }
199
200    /// if one has it and one does not, we return the positive result.
201    #[tokio::test]
202    async fn get_returns_positive() {
203        let first = {
204            let mut svc = MockBlobStore::new();
205            svc.expect_get()
206                .with(predicate::eq(CHUNK.digest()))
207                .once()
208                .returning(|_| Ok(Some(BLOB_META.clone())));
209            svc
210        };
211
212        let second = {
213            let mut svc = MockBlobStore::new();
214            svc.expect_get()
215                .with(predicate::eq(CHUNK.digest()))
216                // We cannot be certain this is called at all, so no `once()` here.
217                .returning(|_| Ok(None));
218            svc
219        };
220
221        let uut = Race::new("uut".to_string(), [first, second]);
222
223        assert_eq!(
224            Some(BLOB_META.clone()),
225            uut.get(&CHUNK.digest()).await.expect("to succeed")
226        )
227    }
228
229    /// Errors are bubbled up, and the error contains the correct service index.
230    #[tokio::test]
231    async fn get_return_error() {
232        let first = {
233            let mut svc = MockBlobStore::new();
234            svc.expect_get()
235                .with(predicate::eq(CHUNK.digest()))
236                .once()
237                .returning(|_| Err(blobstore::Error("".into())));
238            svc
239        };
240
241        // Ideally this one would be just slower than `first`.
242        let second = {
243            let mut svc = MockBlobStore::new();
244            svc.expect_get()
245                .with(predicate::eq(CHUNK.digest()))
246                // We cannot be certain this is called at all, so no `once()` here.
247                .returning(|_| Ok(None));
248            svc
249        };
250
251        let uut = Race::new("uut".to_string(), [first, second]);
252
253        let err = uut.get(&CHUNK.digest()).await.expect_err("to fail").0;
254        let err = err.downcast_ref::<Error>().unwrap();
255        assert_matches!(err, Error::Racing(combinators::race::Error::Backend(0, _)))
256    }
257
258    /// put is unsupported, and not sent to any backend.
259    #[tokio::test]
260    async fn put_unsupported() {
261        let mut first = MockBlobStore::new();
262        first.expect_put().never();
263
264        let uut = Race::new("uut".to_string(), [first]);
265
266        let err = uut
267            .put(&CHUNK.digest(), BLOB_META.clone())
268            .await
269            .expect_err("must fail")
270            .0;
271
272        let err = err.downcast_ref::<Error>().unwrap();
273        assert_matches!(err, Error::Unimplemented);
274    }
275}