Skip to main content

mkit_server/store/
read.rs

1//! Typed readers over any [`NamespaceStore`]: the `store::keys` layouts and
2//! `store::codec` values in one place, so no backend reimplements them.
3
4use mkit_core::hash::Hash;
5use std::collections::BTreeMap;
6
7use super::codec;
8use super::error::StoreError;
9use super::keys::{self, ParsedKey};
10use super::kv::{Cursor, Key, NamespaceStore};
11use super::partition::Partition;
12use crate::pipeline::ShardMap;
13use crate::quota::{QuotaScope, QuotaState};
14use crate::replay::{ReplayKey, ReplayRecord};
15use crate::repo::{RepoId, RepoName};
16
17/// One page of [`list_refs`].
18#[derive(Debug, Clone, PartialEq, Eq, Default)]
19pub struct RefPage {
20    /// Full ref names and ids, in name order.
21    pub refs: Vec<(String, Hash)>,
22    /// Resume point, if more refs may follow.
23    pub next: Option<Cursor>,
24}
25
26/// A ref's id, if it exists.
27pub async fn read_ref<S: NamespaceStore>(
28    store: &S,
29    p: &Partition,
30    repo: &RepoName,
31    name: &str,
32) -> Result<Option<Hash>, StoreError> {
33    let value = store.get(p, &keys::ref_key(repo, name)).await?;
34    value.as_ref().map(codec::decode_ref_id).transpose()
35}
36
37/// Up to `limit` refs of `repo` whose names start with `prefix` as raw
38/// bytes, after `after`. Names are returned in full. `ListRefs` passes
39/// [`crate::refs::list_scan_prefix`] so it matches at a component boundary
40/// (SPEC-REFS §4), then strips it.
41pub async fn list_refs<S: NamespaceStore>(
42    store: &S,
43    p: &Partition,
44    repo: &RepoName,
45    prefix: &str,
46    after: Option<&Cursor>,
47    limit: u32,
48) -> Result<RefPage, StoreError> {
49    let (start, end) = keys::ref_prefix_range(repo, prefix);
50    let page = store.scan(p, &start, &end, after, limit).await?;
51    let refs = page
52        .entries
53        .iter()
54        .map(|(key, value)| match keys::parse(key) {
55            Some(ParsedKey::Ref { name, .. }) => Ok((name, codec::decode_ref_id(value)?)),
56            _ => Err(StoreError::Corrupt("malformed ref key".into())),
57        })
58        .collect::<Result<_, _>>()?;
59    Ok(RefPage {
60        refs,
61        next: page.next,
62    })
63}
64
65/// Repository-scoped pack membership, with an optional read-your-writes ref.
66/// Unknown or malformed hints are ignored (STC §7.9). Membership is checked
67/// in this repository's index first, then in its strongly consistent ref shard.
68pub async fn is_member<S: NamespaceStore>(
69    store: &S,
70    shards: &dyn ShardMap,
71    repo: &RepoId,
72    pack: &Hash,
73    hint: Option<&str>,
74) -> Result<bool, StoreError> {
75    let key = keys::membership(&repo.name, pack);
76    let index = shards.membership(repo, &crate::store::BlobKey::pack(*pack));
77    if store.get(&index, &key).await?.is_some() {
78        return Ok(true);
79    }
80    if let Some(name) = hint
81        && crate::refs::validate_ref_name(name)
82        && crate::refs::is_served_ref_name(name)
83    {
84        return Ok(store
85            .get(&shards.ref_shard(repo, name), &key)
86            .await?
87            .is_some());
88    }
89    Ok(false)
90}
91
92/// Batch membership for one repository: check the consuming ref shard's
93/// local rows first, then unresolved packs in their membership partitions.
94/// This is the hook point for §12.2's deleted-repository generation rule.
95pub async fn members_many<S: NamespaceStore>(
96    store: &S,
97    shards: &dyn ShardMap,
98    repo: &RepoId,
99    ref_shard: &Partition,
100    packs: &[Hash],
101) -> Result<Vec<bool>, StoreError> {
102    // One MKPL check must fit the indexed lookup cap, including on Workers.
103    if packs.len() > super::index::MAX_LOOKUP_IDS {
104        return Err(StoreError::Invalid("object index limit exceeded".into()));
105    }
106    let keys: Vec<_> = packs
107        .iter()
108        .map(|pack| keys::membership(&repo.name, pack))
109        .collect();
110    let local = store.get_many(ref_shard, &keys).await?;
111    if local.len() != packs.len() {
112        return Err(StoreError::Corrupt("short local membership read".into()));
113    }
114    let mut found: Vec<_> = local.into_iter().map(|value| value.is_some()).collect();
115    let mut grouped: BTreeMap<Partition, Vec<usize>> = BTreeMap::new();
116    for (i, pack) in packs.iter().enumerate() {
117        if !found[i] {
118            grouped
119                .entry(shards.membership(repo, &crate::store::BlobKey::pack(*pack)))
120                .or_default()
121                .push(i);
122        }
123    }
124    for (partition, indices) in grouped {
125        let wanted: Vec<_> = indices.iter().map(|i| keys[*i].clone()).collect();
126        let values = store.get_many(&partition, &wanted).await?;
127        if values.len() != indices.len() {
128            return Err(StoreError::Corrupt("short membership index read".into()));
129        }
130        for (i, value) in indices.into_iter().zip(values) {
131            found[i] = value.is_some();
132        }
133    }
134    Ok(found)
135}
136
137/// The replay record for `scope`, if any (PRD §5.4 stage 0).
138pub async fn replay_lookup<S: NamespaceStore>(
139    store: &S,
140    p: &Partition,
141    scope: &ReplayKey,
142) -> Result<Option<ReplayRecord>, StoreError> {
143    let value = store.get(p, &keys::replay(&scope.0)).await?;
144    value.as_ref().map(codec::decode_replay_record).transpose()
145}
146
147/// The current quota window of `scope`, if any.
148pub async fn quota_state<S: NamespaceStore>(
149    store: &S,
150    p: &Partition,
151    scope: &QuotaScope,
152) -> Result<Option<QuotaState>, StoreError> {
153    let value = store.get(p, &keys::quota(scope)).await?;
154    value.as_ref().map(codec::decode_quota_state).transpose()
155}
156
157/// The grant epoch; an absent key is epoch 0.
158pub async fn grant_epoch<S: NamespaceStore>(store: &S, p: &Partition) -> Result<u64, StoreError> {
159    let value = store.get(p, &keys::grant_epoch()).await?;
160    value.as_ref().map_or(Ok(0), codec::decode_u64)
161}
162
163/// How long past its envelope's expiry a replay record is kept. An envelope
164/// verifies while the verifier's clock is at most its expiry, so a record
165/// pruned on a clock that runs ahead could let a replay through on a clock
166/// that runs behind. The grace exceeds any tolerated skew: the commit-
167/// deadline margin (SPEC-WRITE-GRANTS §5.5, 5 s) and the auth v2 clock
168/// lead (`mkit_core::write_auth::MAX_CLOCK_LEAD_MS`, 30 s).
169pub const REPLAY_PRUNE_GRACE_MS: u64 = 60_000;
170
171// The grace covers the largest skew the auth v2 verifier tolerates.
172const _: () =
173    assert!(REPLAY_PRUNE_GRACE_MS >= mkit_core::write_auth::MAX_CLOCK_LEAD_MS.unsigned_abs());
174
175/// Up to `limit` replay records whose envelopes expired before
176/// `now_ms - REPLAY_PRUNE_GRACE_MS`, as `(index key, record key)` pairs to
177/// delete. A record is written once per scope and never revived after its
178/// envelope expired, so the deletes need no precondition.
179pub async fn expired_replay_keys<S: NamespaceStore>(
180    store: &S,
181    p: &Partition,
182    now_ms: u64,
183    limit: u32,
184) -> Result<Vec<(Key, Key)>, StoreError> {
185    let (start, end) = keys::replay_expiry_before(now_ms.saturating_sub(REPLAY_PRUNE_GRACE_MS));
186    let page = store.scan(p, &start, &end, None, limit).await?;
187    page.entries
188        .into_iter()
189        .map(|(index, _)| match keys::parse(&index) {
190            Some(ParsedKey::ReplayExpiry { scope, .. }) => Ok((index, keys::replay(&scope))),
191            _ => Err(StoreError::Corrupt("malformed replay expiry key".into())),
192        })
193        .collect()
194}
195
196/// Up to `limit` quota windows of length `window_ms` that ended by
197/// `now_ms`, as `(index key, quota key)` pairs. A quota key is rewritten
198/// when its scope opens a new window, so the caller guards each quota-key
199/// delete with `Equals` on the value it read.
200pub async fn stale_quota_keys<S: NamespaceStore>(
201    store: &S,
202    p: &Partition,
203    now_ms: u64,
204    window_ms: u64,
205    limit: u32,
206) -> Result<Vec<(Key, Key)>, StoreError> {
207    let Some(last_stale_start) = now_ms.checked_sub(window_ms) else {
208        return Ok(Vec::new());
209    };
210    let (start, end) = keys::quota_window_before(last_stale_start.saturating_add(1));
211    let page = store.scan(p, &start, &end, None, limit).await?;
212    page.entries
213        .into_iter()
214        .map(|(index, _)| {
215            let quota = keys::quota_for_window(&index)
216                .ok_or_else(|| StoreError::Corrupt("malformed quota window key".into()))?;
217            Ok((index, quota))
218        })
219        .collect()
220}
221
222#[cfg(test)]
223mod tests {
224    use futures_executor::block_on;
225
226    use super::*;
227    use crate::memory::MemoryKv;
228    use crate::repo::NamespaceKey;
229    use crate::store::{Batch, Value};
230
231    fn ns() -> Partition {
232        Partition::Namespace(NamespaceKey::deployment_default())
233    }
234
235    #[test]
236    fn typed_readers_decode_their_layouts() {
237        let kv = MemoryKv::default();
238        let repo = RepoName::new("r").unwrap();
239        let scope = QuotaScope::for_signer(&NamespaceKey::deployment_default(), &[1; 32]);
240        let old = ReplayKey([1; 32]);
241        let batch = Batch::new()
242            .put(
243                keys::ref_key(&repo, "refs/heads/a"),
244                codec::encode_ref_id(&[1; 32]),
245            )
246            .put(
247                keys::ref_key(&repo, "refs/heads/b"),
248                codec::encode_ref_id(&[2; 32]),
249            )
250            .put(
251                keys::ref_key(&repo, "refs/tags/t"),
252                codec::encode_ref_id(&[3; 32]),
253            )
254            .put(keys::grant_epoch(), codec::encode_u64(4))
255            .put(keys::replay_expiry(100, &old.0), Value::default())
256            .put(keys::replay_expiry(200, &[2; 32]), Value::default())
257            .put(keys::quota_window(0, &scope), Value::default());
258        block_on(async {
259            assert_eq!(grant_epoch(&kv, &ns()).await.unwrap(), 0);
260            kv.apply(&ns(), batch).await.unwrap();
261            assert_eq!(
262                read_ref(&kv, &ns(), &repo, "refs/heads/b").await.unwrap(),
263                Some([2; 32])
264            );
265            assert_eq!(
266                read_ref(&kv, &ns(), &repo, "refs/heads/c").await.unwrap(),
267                None
268            );
269            let page = list_refs(&kv, &ns(), &repo, "refs/heads/", None, 1)
270                .await
271                .unwrap();
272            assert_eq!(page.refs, vec![("refs/heads/a".into(), [1; 32])]);
273            let rest = list_refs(&kv, &ns(), &repo, "refs/heads/", page.next.as_ref(), 10)
274                .await
275                .unwrap();
276            assert_eq!((rest.refs.len(), rest.next), (1, None));
277            assert_eq!(grant_epoch(&kv, &ns()).await.unwrap(), 4);
278            assert_eq!(replay_lookup(&kv, &ns(), &old).await.unwrap(), None);
279            assert_eq!(quota_state(&kv, &ns(), &scope).await.unwrap(), None);
280            assert_eq!(
281                expired_replay_keys(&kv, &ns(), 200 + REPLAY_PRUNE_GRACE_MS, 10)
282                    .await
283                    .unwrap(),
284                vec![(keys::replay_expiry(100, &old.0), keys::replay(&old.0))]
285            );
286            // Within the grace window nothing is pruned, even long past expiry.
287            let within = expired_replay_keys(&kv, &ns(), 100 + REPLAY_PRUNE_GRACE_MS, 10).await;
288            assert!(within.unwrap().is_empty());
289            assert!(
290                stale_quota_keys(&kv, &ns(), 99, 100, 10)
291                    .await
292                    .unwrap()
293                    .is_empty()
294            );
295            assert_eq!(
296                stale_quota_keys(&kv, &ns(), 100, 100, 10).await.unwrap(),
297                vec![(keys::quota_window(0, &scope), keys::quota(&scope))]
298            );
299        });
300    }
301}
302
303#[cfg(test)]
304mod membership_tests {
305    use super::*;
306    use crate::pipeline::D34Shards;
307    use crate::store::{Batch, BatchOutcome, PartitionStats, ScanPage, StoreCapabilities, Value};
308    use crate::{MemoryKv, NamespaceKey};
309    use futures_executor::block_on;
310    use std::sync::Mutex;
311
312    #[derive(Default)]
313    struct Observed {
314        inner: MemoryKv,
315        gets: Mutex<Vec<(Partition, Key)>>,
316    }
317    impl NamespaceStore for Observed {
318        fn capabilities(&self) -> StoreCapabilities {
319            self.inner.capabilities()
320        }
321        async fn get(&self, p: &Partition, k: &Key) -> Result<Option<Value>, StoreError> {
322            self.gets.lock().unwrap().push((p.clone(), k.clone()));
323            self.inner.get(p, k).await
324        }
325        async fn scan(
326            &self,
327            p: &Partition,
328            start: &Key,
329            end: &Key,
330            after: Option<&Cursor>,
331            limit: u32,
332        ) -> Result<ScanPage, StoreError> {
333            self.inner.scan(p, start, end, after, limit).await
334        }
335        async fn apply(&self, p: &Partition, batch: Batch) -> Result<BatchOutcome, StoreError> {
336            self.inner.apply(p, batch).await
337        }
338        async fn stats(&self, p: &Partition) -> Result<PartitionStats, StoreError> {
339            self.inner.stats(p).await
340        }
341        async fn probe(&self) -> Result<(), StoreError> {
342            self.inner.probe().await
343        }
344    }
345    fn repo(name: &str) -> RepoId {
346        RepoId {
347            namespace: NamespaceKey::deployment_default(),
348            name: RepoName::new(name).unwrap(),
349        }
350    }
351    const PACK: Hash = [7; 32];
352    const REF: &str = "refs/heads/main";
353    fn plant(kv: &Observed, p: &Partition, repo: &RepoId) {
354        assert_eq!(
355            block_on(kv.inner.apply(
356                p,
357                Batch::new().put(keys::membership(&repo.name, &PACK), Value::default())
358            ))
359            .unwrap(),
360            BatchOutcome::Committed
361        );
362    }
363    #[test]
364    fn index_hit_needs_only_one_get() {
365        let kv = Observed::default();
366        let r = repo("a");
367        let shards = D34Shards;
368        plant(
369            &kv,
370            &shards.membership(&r, &crate::store::BlobKey::pack(PACK)),
371            &r,
372        );
373        assert!(block_on(is_member(&kv, &shards, &r, &PACK, Some(REF))).unwrap());
374        assert_eq!(kv.gets.lock().unwrap().len(), 1);
375    }
376    #[test]
377    fn ref_hint_sees_unrelayed_membership() {
378        let kv = Observed::default();
379        let r = repo("a");
380        let shards = D34Shards;
381        plant(&kv, &shards.ref_shard(&r, REF), &r);
382        assert!(!block_on(is_member(&kv, &shards, &r, &PACK, None)).unwrap());
383        kv.gets.lock().unwrap().clear();
384        assert!(block_on(is_member(&kv, &shards, &r, &PACK, Some(REF))).unwrap());
385        assert_eq!(kv.gets.lock().unwrap()[1].0, shards.ref_shard(&r, REF));
386    }
387    #[test]
388    fn missing_membership_and_unknown_hint_are_false() {
389        let kv = Observed::default();
390        let r = repo("a");
391        assert!(!block_on(is_member(&kv, &D34Shards, &r, &PACK, None)).unwrap());
392        assert!(!block_on(is_member(&kv, &D34Shards, &r, &PACK, Some(REF))).unwrap());
393    }
394    #[test]
395    fn invalid_and_unserved_hints_do_not_read_a_ref_shard() {
396        for hint in [
397            "bad ref",
398            "refs/heads/../main",
399            "main",
400            "refsx/heads/main",
401            "",
402        ] {
403            let kv = Observed::default();
404            let r = repo("a");
405            assert!(!block_on(is_member(&kv, &D34Shards, &r, &PACK, Some(hint))).unwrap());
406            // The mandatory index lookup still happens; the hint adds no call.
407            assert_eq!(kv.gets.lock().unwrap().len(), 1, "{hint}");
408        }
409    }
410    #[test]
411    fn index_and_ref_hint_never_cross_repository_or_namespace() {
412        let kv = Observed::default();
413        let a = repo("a");
414        let b = repo("b");
415        let shards = D34Shards;
416        plant(
417            &kv,
418            &shards.membership(&a, &crate::store::BlobKey::pack(PACK)),
419            &a,
420        );
421        plant(&kv, &shards.ref_shard(&a, REF), &a);
422        assert!(!block_on(is_member(&kv, &shards, &b, &PACK, Some(REF))).unwrap());
423        let other = RepoId {
424            namespace: NamespaceKey::from_stored(format!("ed25519-{}", "b".repeat(64))),
425            name: a.name.clone(),
426        };
427        assert!(!block_on(is_member(&kv, &shards, &other, &PACK, Some(REF))).unwrap());
428    }
429}