snix_castore/chunkstore/combinators/
priority.rs1use 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
12pub struct Priority<CS> {
17 instance_name: String,
18 services: Vec<CS>,
19}
20
21impl<CS> Priority<CS> {
22 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 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 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#[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#[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 #[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 #[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 #[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 #[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 #[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 #[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}