aspero_engine/
lookup.rs

1//! Lookups by identifier: routing across providers, the cache, and
2//! sharing calls in flight.
3
4use aspero_core::{Error, ErrorCategory, FieldSet, ProviderId, SourceRecord, Usage, WorkId};
5
6use crate::inflight::Answer;
7use crate::{Engine, Entry, Served, add_usage};
8
9/// The outcome of one provider's attempt at one ID.
10#[derive(Debug, Clone)]
11#[non_exhaustive]
12pub enum Outcome {
13    Found,
14    NotFound,
15    Failed(Error),
16}
17
18#[derive(Debug, Clone)]
19pub struct Attempt {
20    pub provider: ProviderId,
21    pub outcome: Outcome,
22    pub served: Served,
23}
24
25#[derive(Debug, Clone)]
26pub struct LookupResult {
27    /// As requested.
28    pub id: WorkId,
29    pub record: Option<SourceRecord>,
30    /// Every provider tried, in order. Empty: no enabled provider accepts
31    /// this kind of identifier.
32    pub attempts: Vec<Attempt>,
33}
34
35impl LookupResult {
36    /// No provider found it, and at least one could answer authoritatively
37    /// (rather than every attempt failing).
38    pub fn is_not_found(&self) -> bool {
39        self.record.is_none()
40            && self
41                .attempts
42                .iter()
43                .any(|a| matches!(a.outcome, Outcome::NotFound))
44    }
45}
46
47#[derive(Debug, Clone)]
48pub struct LookupReport {
49    /// One per requested ID, in request order.
50    pub results: Vec<LookupResult>,
51    /// Combined usage per provider that reported any.
52    pub usage: Vec<(ProviderId, Usage)>,
53}
54
55impl Engine {
56    /// Looks up every ID, falling through providers by priority. Never fails
57    /// as a whole: per-ID outcomes are in the report.
58    pub async fn lookup(&self, ids: &[WorkId], fields: FieldSet) -> LookupReport {
59        let mut results: Vec<LookupResult> = ids
60            .iter()
61            .map(|id| LookupResult {
62                id: id.clone(),
63                record: None,
64                attempts: Vec::new(),
65            })
66            .collect();
67        let mut usage: Vec<(ProviderId, Usage)> = Vec::new();
68
69        for entry in &self.entries {
70            let caps = entry.provider.capabilities();
71            let eligible: Vec<usize> = (0..results.len())
72                .filter(|&i| results[i].record.is_none())
73                .filter(|&i| caps.lookup_kinds.contains(&results[i].id.kind()))
74                .collect();
75            let ids: Vec<WorkId> = eligible.iter().map(|&i| results[i].id.clone()).collect();
76            let answers = self.lookup_at(entry, &ids, fields, &mut usage).await;
77            for (&i, (attempt, record)) in eligible.iter().zip(answers) {
78                results[i].record = record;
79                results[i].attempts.push(attempt);
80            }
81        }
82        LookupReport { results, usage }
83    }
84
85    /// Asks one provider about `ids` (all of kinds it accepts): fresh cached
86    /// answers first, the rest in batches with retries, and stale cached
87    /// records for any that fail. One answer per ID, in order.
88    pub(crate) async fn lookup_at(
89        &self,
90        entry: &Entry,
91        ids: &[WorkId],
92        fields: FieldSet,
93        usage: &mut Vec<(ProviderId, Usage)>,
94    ) -> Vec<(Attempt, Option<SourceRecord>)> {
95        let provider = entry.provider.id();
96        let caps = entry.provider.capabilities();
97        let attempt = |outcome, served| Attempt {
98            provider: provider.clone(),
99            outcome,
100            served,
101        };
102        if self.options.offline {
103            return self.offline_lookups(&provider, ids, fields);
104        }
105        let mut answers: Vec<Option<Answer>> = vec![None; ids.len()];
106
107        // Fresh cached answers first (one cache read for the whole set).
108        let mut pending = Vec::new();
109        let mut from_cache = 0;
110        for (i, cached) in self
111            .cached_lookups(&provider, ids, fields)
112            .into_iter()
113            .enumerate()
114        {
115            answers[i] = match cached {
116                Some(Some(record)) => {
117                    from_cache += 1;
118                    Some((attempt(Outcome::Found, Served::Cache), Some(record)))
119                }
120                Some(None) => {
121                    from_cache += 1;
122                    Some((attempt(Outcome::NotFound, Served::Cache), None))
123                }
124                None => {
125                    pending.push(i);
126                    None
127                }
128            };
129        }
130
131        if from_cache > 0 {
132            tracing::info!(target: "aspero::cache", provider = %provider, ids = from_cache, "lookups served from cache");
133        }
134        // Ask about the rest, unless another caller is already asking.
135        let wanted: Vec<(usize, WorkId)> = pending.iter().map(|&i| (i, ids[i].clone())).collect();
136        let (mut leads, joined) = self.inflight.claim(&provider, &wanted, fields);
137        let led = leads.indices();
138        let mut chunks = led.chunks(caps.max_batch.max(1));
139        while let Some(chunk) = chunks.next() {
140            leads.publish(&answers);
141            let batch: Vec<WorkId> = chunk.iter().map(|&i| ids[i].clone()).collect();
142            let answer = self
143                .call(entry, || entry.provider.lookup(&batch, fields))
144                .await;
145
146            let error = match answer {
147                Ok(resp) => {
148                    if let Some(u) = resp.usage {
149                        add_usage(usage, &provider, u);
150                    }
151                    // The contract is one answer per ID, in order; anything
152                    // else is a provider bug, not a miss.
153                    let in_order = resp.value.len() == batch.len()
154                        && resp.value.iter().zip(&batch).all(|(l, id)| &l.id == id);
155                    if in_order {
156                        let stored: Vec<_> = resp
157                            .value
158                            .iter()
159                            .map(|l| (&l.id, l.record.as_ref()))
160                            .collect();
161                        self.remember_lookups(&provider, &stored);
162                        for (&i, lookup) in chunk.iter().zip(resp.value) {
163                            let outcome = match lookup.record {
164                                Some(_) => Outcome::Found,
165                                None => Outcome::NotFound,
166                            };
167                            answers[i] = Some((attempt(outcome, Served::Live), lookup.record));
168                        }
169                        continue;
170                    }
171                    Error::decode(ContractViolation)
172                }
173                Err(error) => error,
174            };
175
176            // Out of budget or credentials: don't spend more calls on this
177            // provider in this lookup.
178            let give_up = matches!(
179                error.category(),
180                ErrorCategory::Throttled | ErrorCategory::Auth
181            );
182            let affected: Vec<usize> = if give_up {
183                chunk
184                    .iter()
185                    .chain(chunks.by_ref().flatten())
186                    .copied()
187                    .collect()
188            } else {
189                chunk.to_vec()
190            };
191            // Stale-on-error: an old answer beats none.
192            let affected_ids: Vec<WorkId> = affected.iter().map(|&i| ids[i].clone()).collect();
193            for (&i, stale) in affected
194                .iter()
195                .zip(self.stale_lookups(&provider, &affected_ids))
196            {
197                answers[i] = Some(match stale {
198                    Some(record) => (attempt(Outcome::Found, Served::Stale), Some(record)),
199                    None => (attempt(Outcome::Failed(error.clone()), Served::Live), None),
200                });
201            }
202        }
203        leads.publish(&answers);
204        drop(leads);
205
206        // Answers others fetched; ask again for any they gave up on.
207        let mut abandoned = Vec::new();
208        for (i, call) in joined {
209            match call.answer().await {
210                Some(answer) => answers[i] = Some(answer),
211                None => abandoned.push(i),
212            }
213        }
214        if !abandoned.is_empty() {
215            let retry: Vec<WorkId> = abandoned.iter().map(|&i| ids[i].clone()).collect();
216            let again = Box::pin(self.lookup_at(entry, &retry, fields, usage)).await;
217            for (i, answer) in abandoned.into_iter().zip(again) {
218                answers[i] = Some(answer);
219            }
220        }
221        answers
222            .into_iter()
223            .map(|a| a.expect("every ID answered from cache, provider or failure"))
224            .collect()
225    }
226
227    // Cache access. Cache failures are never fatal: a broken cache behaves
228    // like an empty one, and failed writes are skipped.
229
230    /// For each ID, a fresh cached answer covering `fields`:
231    /// `Some(Some(record))`, `Some(None)` for a fresh "not found", or `None`
232    /// to ask the provider.
233    fn cached_lookups(
234        &self,
235        provider: &ProviderId,
236        ids: &[WorkId],
237        fields: FieldSet,
238    ) -> Vec<Option<Option<SourceRecord>>> {
239        let now = (self.clock)();
240        let policy = self.cache_policy;
241        self.cache_entries(provider, ids)
242            .into_iter()
243            .map(|entry| {
244                let entry = entry?;
245                match entry.value {
246                    Some(ref record)
247                        if entry.is_fresh(now, policy.records)
248                            && fields.is_subset(record.fields) =>
249                    {
250                        Some(entry.value)
251                    }
252                    None if entry.is_fresh(now, policy.misses) => Some(None),
253                    _ => None,
254                }
255            })
256            .collect()
257    }
258
259    /// Every answer from the cache, however old: offline.
260    fn offline_lookups(
261        &self,
262        provider: &ProviderId,
263        ids: &[WorkId],
264        fields: FieldSet,
265    ) -> Vec<Answer> {
266        let fresh = self.cached_lookups(provider, ids, fields);
267        let any = self.stale_lookups(provider, ids);
268        let attempt = |outcome, served| Attempt {
269            provider: provider.clone(),
270            outcome,
271            served,
272        };
273        fresh
274            .into_iter()
275            .zip(any)
276            .map(|answer| match answer {
277                (Some(Some(record)), _) => (attempt(Outcome::Found, Served::Cache), Some(record)),
278                (Some(None), _) => (attempt(Outcome::NotFound, Served::Cache), None),
279                (None, Some(record)) => (attempt(Outcome::Found, Served::Stale), Some(record)),
280                (None, None) => (attempt(Outcome::Failed(Error::Offline), Served::Live), None),
281            })
282            .collect()
283    }
284
285    /// For each ID, any cached record, however old.
286    fn stale_lookups(&self, provider: &ProviderId, ids: &[WorkId]) -> Vec<Option<SourceRecord>> {
287        self.cache_entries(provider, ids)
288            .into_iter()
289            .map(|entry| entry.and_then(|e| e.value))
290            .collect()
291    }
292
293    fn cache_entries(
294        &self,
295        provider: &ProviderId,
296        ids: &[WorkId],
297    ) -> Vec<Option<aspero_store::CachedLookup>> {
298        self.cache
299            .as_ref()
300            .and_then(|c| c.get_lookups(provider, ids).ok())
301            .unwrap_or_else(|| vec![None; ids.len()])
302    }
303
304    fn remember_lookups(
305        &self,
306        provider: &ProviderId,
307        answers: &[(&WorkId, Option<&SourceRecord>)],
308    ) {
309        if let Some(cache) = &self.cache {
310            let _ = cache.put_lookups(provider, answers, (self.clock)());
311        }
312    }
313}
314
315#[derive(Debug, thiserror::Error)]
316#[error("provider answered a batch lookup out of order or incompletely")]
317struct ContractViolation;
aspero_engine/lookup.rs

aspero_engine/
lookup.rs
