1use 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#[derive(Debug, Clone, PartialEq, Eq, Default)]
19pub struct RefPage {
20 pub refs: Vec<(String, Hash)>,
22 pub next: Option<Cursor>,
24}
25
26pub 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
37pub 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
65pub 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
92pub 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 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
137pub 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
147pub 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
157pub 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
163pub const REPLAY_PRUNE_GRACE_MS: u64 = 60_000;
170
171const _: () =
173 assert!(REPLAY_PRUNE_GRACE_MS >= mkit_core::write_auth::MAX_CLOCK_LEAD_MS.unsigned_abs());
174
175pub 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
196pub 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 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 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}