1use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
55use std::ops::Bound;
56use std::path::PathBuf;
57use std::sync::atomic::{AtomicU64, Ordering};
58use std::sync::{Mutex, MutexGuard};
59use std::time::Duration;
60
61use gunnar_core::RepoRef;
62
63use crate::control_plane::{repo_bytes, RepoBytes};
64use crate::stores::{dir_bytes, ServerStores, TombstoneReason};
65
66pub const RECONCILE_CYCLE: Duration = Duration::from_secs(30 * 60);
68
69#[derive(Default)]
71struct Marks {
72 repos: HashSet<(String, String)>,
73 namespaces: HashSet<String>,
74 flow: HashSet<String>,
75 all: bool,
76}
77
78#[derive(Debug, Clone, Default)]
79struct RepoRow {
80 dir: Option<PathBuf>,
83 bytes: Option<RepoBytes>,
85}
86
87#[derive(Debug, Default)]
88struct Ns {
89 repos: BTreeMap<String, RepoRow>,
90 tombstoned: Option<u64>,
92 objects: u64,
94 changed_seq: u64,
96}
97
98impl Ns {
99 fn unmeasured(&self) -> bool {
100 self.tombstoned.is_none() || self.repos.values().any(|r| r.dir.is_some() && r.bytes.is_none())
101 }
102}
103
104#[derive(Default)]
105struct State {
106 catalog: Option<Option<gunnar_core::BlobId>>,
108 ns: BTreeMap<String, Ns>,
109 by_dir: HashMap<PathBuf, (String, String)>,
111 open_last: HashSet<PathBuf>,
112 reconcile_after: String,
114}
115
116#[derive(Default)]
118struct Work {
119 repos: Vec<(String, String, PathBuf)>,
120 tombstones: BTreeSet<String>,
121}
122
123pub type Page = (Vec<(String, Held)>, bool, bool);
126
127pub type Exact = (Held, Vec<(String, RepoBytes, Option<PathBuf>)>);
130
131#[derive(Debug, Clone, Default, PartialEq, Eq)]
133pub struct Held {
134 pub bytes: RepoBytes,
135 pub repos: u64,
136 pub tombstoned_bytes: u64,
137 pub objects: u64,
138}
139
140pub struct UsageLedger {
142 epoch: String,
143 seq: AtomicU64,
144 marks: Mutex<Marks>,
145 state: Mutex<State>,
146 settling: Mutex<()>,
148 drift: AtomicU64,
149}
150
151impl std::fmt::Debug for UsageLedger {
152 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153 f.debug_struct("UsageLedger")
154 .field("epoch", &self.epoch)
155 .field("seq", &self.seq.load(Ordering::Relaxed))
156 .finish_non_exhaustive()
157 }
158}
159
160fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
161 m.lock().unwrap_or_else(|p| p.into_inner())
164}
165
166pub fn tombstoned_bytes(stores: &ServerStores, namespace: &str) -> u64 {
171 let dir = stores.repos_root().join(namespace);
172 let Ok(entries) = std::fs::read_dir(&dir) else {
173 return 0;
174 };
175 let mut total = 0;
176 for entry in entries.flatten() {
177 let name = entry.file_name().to_string_lossy().into_owned();
178 let Some(store) = name.strip_suffix(".tombstone") else {
179 continue;
180 };
181 if !store.ends_with(".git") {
182 continue;
183 }
184 let Some(record) = ServerStores::read_tombstone(&entry.path()) else {
185 continue;
186 };
187 if record.reason != TombstoneReason::Deleted {
188 continue;
189 }
190 total += dir_bytes(&dir.join(store));
191 }
192 total
193}
194
195fn parse_cursor(cursor: &str) -> Option<(&str, i64, u64)> {
197 let mut parts = cursor.split('.');
198 let (epoch, day, seq) = (parts.next()?, parts.next()?, parts.next()?);
199 if parts.next().is_some() || epoch.is_empty() {
200 return None;
201 }
202 Some((epoch, day.parse().ok()?, seq.parse().ok()?))
203}
204
205impl Default for UsageLedger {
206 fn default() -> Self {
207 Self::new()
208 }
209}
210
211impl UsageLedger {
212 pub fn new() -> Self {
213 static LEDGERS: AtomicU64 = AtomicU64::new(0);
214 let n = LEDGERS.fetch_add(1, Ordering::Relaxed);
215 let epoch = format!("{:x}{:x}{n:x}", gunnar_core::now_ms(), std::process::id());
216 UsageLedger {
217 epoch,
218 seq: AtomicU64::new(0),
219 marks: Mutex::new(Marks::default()),
220 state: Mutex::new(State::default()),
221 settling: Mutex::new(()),
222 drift: AtomicU64::new(0),
223 }
224 }
225
226 pub fn epoch(&self) -> &str {
228 &self.epoch
229 }
230
231 pub fn drift_corrections(&self) -> u64 {
235 self.drift.load(Ordering::Relaxed)
236 }
237
238 pub fn mark_repo(&self, namespace: &str, name: &str) {
242 lock(&self.marks).repos.insert((namespace.to_owned(), name.to_owned()));
243 }
244
245 pub fn mark_namespace(&self, namespace: &str) {
248 lock(&self.marks).namespaces.insert(namespace.to_owned());
249 }
250
251 pub fn mark_flow(&self, namespace: &str) {
253 lock(&self.marks).flow.insert(namespace.to_owned());
254 }
255
256 pub fn mark_all(&self) {
258 lock(&self.marks).all = true;
259 }
260
261 fn bump(&self, ns: &mut Ns) {
262 ns.changed_seq = self.seq.fetch_add(1, Ordering::SeqCst) + 1;
263 }
264
265 pub fn cursor_now(&self, now_unix_ms: i64) -> String {
267 format!("{}.{}.{}", self.epoch, gunnar_core::day_of(now_unix_ms), self.seq.load(Ordering::SeqCst))
268 }
269
270 fn honour(&self, cursor: &str, now_unix_ms: i64) -> Option<u64> {
272 let (epoch, day, seq) = parse_cursor(cursor)?;
273 (epoch == self.epoch && day == gunnar_core::day_of(now_unix_ms)).then_some(seq)
274 }
275
276 fn refresh_index(&self, stores: &ServerStores, st: &mut State, marks: &mut Marks) -> gunnar_core::Result<()> {
280 let catalog = stores.catalog();
281 let version = catalog.version()?;
282 if st.catalog == Some(version) {
283 return Ok(());
284 }
285 let initial = st.catalog.is_none();
286 let mut want: BTreeMap<String, BTreeSet<String>> =
287 catalog.namespaces()?.into_iter().map(|n| (n.name, BTreeSet::new())).collect();
288 for repo in catalog.list(None)? {
289 want.entry(repo.namespace.clone()).or_default().insert(repo.name.clone());
290 }
291 let gone: Vec<String> = st.ns.keys().filter(|n| !want.contains_key(*n)).cloned().collect();
292 for name in gone {
293 if let Some(ns) = st.ns.remove(&name) {
294 for row in ns.repos.values() {
295 if let Some(dir) = &row.dir {
296 st.by_dir.remove(dir);
297 }
298 }
299 }
300 }
301 for (name, set) in want {
302 let fresh = !st.ns.contains_key(&name);
303 let mut ns = st.ns.remove(&name).unwrap_or_default();
304 let old: BTreeSet<String> = ns.repos.keys().cloned().collect();
305 if old != set {
306 for removed in old.difference(&set) {
307 if let Some(row) = ns.repos.remove(removed) {
308 if let Some(dir) = row.dir {
309 st.by_dir.remove(&dir);
310 }
311 }
312 }
313 for added in set.difference(&old) {
314 let dir = RepoRef::new(&name, added).ok().and_then(|path| stores.repo_dir(&path).ok());
315 if let Some(dir) = &dir {
316 st.by_dir.insert(dir.clone(), (name.clone(), added.clone()));
317 }
318 ns.repos.insert(added.clone(), RepoRow { dir, bytes: None });
319 }
320 marks.namespaces.insert(name.clone());
323 if !initial {
324 self.bump(&mut ns);
325 }
326 } else if fresh && !initial {
327 self.bump(&mut ns);
328 }
329 st.ns.insert(name, ns);
330 }
331 st.catalog = Some(version);
332 Ok(())
333 }
334
335 fn collect(&self, stores: &ServerStores, st: &mut State, marks: Marks, sweep: bool) -> Work {
338 let mut work = Work::default();
339 let mut repos: BTreeSet<(String, String)> = BTreeSet::new();
340 let mut whole: BTreeSet<String> = BTreeSet::new();
341 if marks.all {
342 whole.extend(st.ns.keys().cloned());
343 }
344 whole.extend(marks.namespaces);
345 for name in marks.flow {
346 if let Some(ns) = st.ns.get_mut(&name) {
347 self.bump(ns);
348 whole.insert(name);
349 }
350 }
351 repos.extend(marks.repos);
352 let now_open: HashMap<PathBuf, u64> = stores.engine().open_object_counts().into_iter().collect();
360 let full = sweep;
361 let mut objects: HashMap<String, u64> = HashMap::new();
362 for (dir, count) in &now_open {
363 if let Some((ns, repo)) = st.by_dir.get(dir) {
364 *objects.entry(ns.clone()).or_default() += count;
365 if full || !st.open_last.contains(dir) {
366 repos.insert((ns.clone(), repo.clone()));
367 }
368 }
369 }
370 for dir in st.open_last.iter().filter(|d| !now_open.contains_key(*d)) {
371 if let Some((ns, repo)) = st.by_dir.get(dir) {
372 repos.insert((ns.clone(), repo.clone()));
373 }
374 }
375 let names: Vec<String> = st.ns.iter().filter(|(_, ns)| ns.objects > 0).map(|(n, _)| n.clone()).chain(objects.keys().cloned()).collect();
376 for name in names {
377 let now = objects.get(&name).copied().unwrap_or(0);
378 if let Some(ns) = st.ns.get_mut(&name) {
379 if ns.objects != now {
380 ns.objects = now;
381 self.bump(ns);
382 }
383 }
384 }
385 st.open_last = now_open.into_keys().collect();
386 for name in &whole {
387 if let Some(ns) = st.ns.get(name) {
388 for (repo, row) in &ns.repos {
389 if let Some(dir) = &row.dir {
390 work.repos.push((name.clone(), repo.clone(), dir.clone()));
391 }
392 }
393 work.tombstones.insert(name.clone());
394 }
395 }
396 for (name, repo) in repos {
397 if whole.contains(&name) {
398 continue;
399 }
400 if let Some(dir) = st.ns.get(&name).and_then(|ns| ns.repos.get(&repo)).and_then(|row| row.dir.clone()) {
401 work.repos.push((name, repo, dir));
402 }
403 }
404 work
405 }
406
407 fn carry_out(&self, stores: &ServerStores, work: Work) -> BTreeSet<String> {
411 let repos: Vec<(String, String, RepoBytes)> =
412 work.repos.into_iter().map(|(ns, repo, dir)| (ns, repo, repo_bytes(&dir))).collect();
413 let tombstones: Vec<(String, u64)> =
414 work.tombstones.into_iter().map(|ns| { let t = tombstoned_bytes(stores, &ns); (ns, t) }).collect();
415 let mut changed = BTreeSet::new();
416 let mut st = lock(&self.state);
417 for (name, repo, bytes) in repos {
418 let Some(ns) = st.ns.get_mut(&name) else { continue };
419 let Some(row) = ns.repos.get_mut(&repo) else { continue };
420 if row.bytes != Some(bytes) {
421 row.bytes = Some(bytes);
422 self.bump(ns);
423 changed.insert(name);
424 }
425 }
426 for (name, bytes) in tombstones {
427 let Some(ns) = st.ns.get_mut(&name) else { continue };
428 if ns.tombstoned != Some(bytes) {
429 ns.tombstoned = Some(bytes);
430 self.bump(ns);
431 changed.insert(name);
432 }
433 }
434 changed
435 }
436
437 pub fn settle(&self, stores: &ServerStores, sweep: bool) -> gunnar_core::Result<()> {
441 let _one = lock(&self.settling);
442 let marks = std::mem::take(&mut *lock(&self.marks));
443 let work = {
444 let mut st = lock(&self.state);
445 let mut marks = marks;
446 if let Err(err) = self.refresh_index(stores, &mut st, &mut marks) {
447 let mut back = lock(&self.marks);
449 back.repos.extend(marks.repos);
450 back.namespaces.extend(marks.namespaces);
451 back.flow.extend(marks.flow);
452 back.all |= marks.all;
453 return Err(err);
454 }
455 self.collect(stores, &mut st, marks, sweep)
456 };
457 self.carry_out(stores, work);
458 Ok(())
459 }
460
461 fn measure_unmeasured(&self, stores: &ServerStores, names: &[String]) {
463 let work = {
464 let st = lock(&self.state);
465 let mut work = Work::default();
466 for name in names {
467 let Some(ns) = st.ns.get(name) else { continue };
468 if !ns.unmeasured() {
469 continue;
470 }
471 for (repo, row) in &ns.repos {
472 if let (Some(dir), None) = (&row.dir, &row.bytes) {
473 work.repos.push((name.clone(), repo.clone(), dir.clone()));
474 }
475 }
476 if ns.tombstoned.is_none() {
477 work.tombstones.insert(name.clone());
478 }
479 }
480 work
481 };
482 self.carry_out(stores, work);
483 }
484
485 fn held_of(ns: &Ns) -> Held {
486 let mut held = Held {
487 tombstoned_bytes: ns.tombstoned.unwrap_or(0),
488 objects: ns.objects,
489 ..Held::default()
490 };
491 for row in ns.repos.values() {
492 held.repos += 1;
493 if let Some(b) = row.bytes {
494 held.bytes.pack += b.pack;
495 held.bytes.cache += b.cache;
496 held.bytes.lfs += b.lfs;
497 }
498 }
499 held
500 }
501
502 pub fn page(
509 &self,
510 stores: &ServerStores,
511 after: &str,
512 page_size: usize,
513 changed_since: &str,
514 now_unix_ms: i64,
515 ) -> gunnar_core::Result<Page> {
516 self.settle(stores, after.is_empty())?;
517 let since = self.honour(changed_since, now_unix_ms);
518 let page_size = page_size.max(1);
519 let (picked, more) = {
520 let st = lock(&self.state);
521 let range = if after.is_empty() {
522 st.ns.range::<str, _>((Bound::Unbounded, Bound::Unbounded))
523 } else {
524 st.ns.range::<str, _>((Bound::Excluded(after), Bound::Unbounded))
525 };
526 let mut picked = Vec::new();
527 let mut more = false;
528 for (name, ns) in range {
529 if since.is_some_and(|s| ns.changed_seq <= s) {
530 continue;
531 }
532 if picked.len() == page_size {
533 more = true;
534 break;
535 }
536 picked.push(name.clone());
537 }
538 (picked, more)
539 };
540 self.measure_unmeasured(stores, &picked);
541 let st = lock(&self.state);
542 let rows = picked
543 .into_iter()
544 .filter_map(|name| st.ns.get(&name).map(|ns| (name.clone(), Self::held_of(ns))))
545 .collect();
546 Ok((rows, more, since.is_some()))
547 }
548
549 pub fn exact(&self, stores: &ServerStores, namespace: &str) -> gunnar_core::Result<Option<Exact>> {
554 {
557 let _one = lock(&self.settling);
558 let mut st = lock(&self.state);
559 let mut marks = Marks::default();
560 self.refresh_index(stores, &mut st, &mut marks)?;
561 if !marks.namespaces.is_empty() {
564 lock(&self.marks).namespaces.extend(marks.namespaces);
565 }
566 }
567 let work = {
568 let st = lock(&self.state);
569 let Some(ns) = st.ns.get(namespace) else {
570 return Ok(None);
571 };
572 let mut work = Work::default();
573 for (repo, row) in &ns.repos {
574 if let Some(dir) = &row.dir {
575 work.repos.push((namespace.to_owned(), repo.clone(), dir.clone()));
576 }
577 }
578 work.tombstones.insert(namespace.to_owned());
579 work
580 };
581 self.carry_out(stores, work);
582 let st = lock(&self.state);
583 let Some(ns) = st.ns.get(namespace) else {
584 return Ok(None);
585 };
586 let rows = ns.repos.iter().map(|(name, row)| (name.clone(), row.bytes.unwrap_or_default(), row.dir.clone())).collect();
587 Ok(Some((Self::held_of(ns), rows)))
588 }
589
590 pub fn reconcile_step(&self, stores: &ServerStores, tick: Duration) -> gunnar_core::Result<(usize, usize)> {
601 self.settle(stores, false)?;
602 let slice = {
603 let mut st = lock(&self.state);
604 let total = st.ns.len();
605 if total == 0 {
606 return Ok((0, 0));
607 }
608 let ticks = (RECONCILE_CYCLE.as_secs() / tick.as_secs().max(1)).max(1) as usize;
609 let want = total.div_ceil(ticks).max(1);
610 let after = st.reconcile_after.clone();
611 let mut slice: Vec<String> = st
612 .ns
613 .range::<str, _>((if after.is_empty() { Bound::Unbounded } else { Bound::Excluded(after.as_str()) }, Bound::Unbounded))
614 .take(want)
615 .map(|(n, _)| n.clone())
616 .collect();
617 if slice.len() < want {
618 let have: BTreeSet<String> = slice.iter().cloned().collect();
620 slice.extend(st.ns.keys().filter(|n| !have.contains(*n)).take(want - slice.len()).cloned());
621 }
622 st.reconcile_after = slice.last().cloned().unwrap_or_default();
623 slice
624 };
625 let work = {
626 let st = lock(&self.state);
627 let mut work = Work::default();
628 for name in &slice {
629 let Some(ns) = st.ns.get(name) else { continue };
630 if ns.unmeasured() {
631 continue;
633 }
634 for (repo, row) in &ns.repos {
635 if let Some(dir) = &row.dir {
636 work.repos.push((name.clone(), repo.clone(), dir.clone()));
637 }
638 }
639 work.tombstones.insert(name.clone());
640 }
641 work
642 };
643 let looked = slice.len();
644 let corrected = self.carry_out(stores, work);
645 if !corrected.is_empty() {
646 self.drift.fetch_add(corrected.len() as u64, Ordering::Relaxed);
647 tracing::info!(
648 target: gunnar_core::observe::TARGET_STORE,
649 event = "usage.reconciled",
650 corrected = corrected.len(),
651 first = %corrected.iter().next().cloned().unwrap_or_default(),
652 "the usage ledger's reconcile pass found {} namespace(s) whose bytes had moved without a mark; corrected",
653 corrected.len()
654 );
655 }
656 Ok((looked, corrected.len()))
657 }
658}
659
660#[cfg(test)]
661mod tests {
662 use super::*;
663
664 #[test]
665 fn a_cursor_is_honoured_only_by_its_own_ledger_on_its_own_day() {
666 let ledger = UsageLedger::new();
667 let other = UsageLedger::new();
668 let noon = gunnar_core::midnight_of(20_000) + 12 * 3_600_000;
669 let cursor = ledger.cursor_now(noon);
670 assert_eq!(ledger.honour(&cursor, noon), Some(0));
671 assert_eq!(other.honour(&cursor, noon), None, "another ledger (a restart) answers everything");
672 assert_eq!(ledger.honour(&cursor, noon + 86_400_000), None, "the day turned: the *_today meters reset");
673 assert_eq!(ledger.honour("", noon), None);
674 assert_eq!(ledger.honour("x.y.z", noon), None);
675 assert_ne!(ledger.epoch(), other.epoch());
676 }
677}