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()
101 || self
102 .repos
103 .values()
104 .any(|r| r.dir.is_some() && r.bytes.is_none())
105 }
106}
107
108#[derive(Default)]
109struct State {
110 catalog: Option<Option<gunnar_core::BlobId>>,
112 ns: BTreeMap<String, Ns>,
113 by_dir: HashMap<PathBuf, (String, String)>,
115 open_last: HashSet<PathBuf>,
116 reconcile_after: String,
118}
119
120#[derive(Default)]
122struct Work {
123 repos: Vec<(String, String, PathBuf)>,
124 tombstones: BTreeSet<String>,
125}
126
127pub type Page = (Vec<(String, Held)>, bool, bool);
130
131pub type Exact = (Held, Vec<(String, RepoBytes, Option<PathBuf>)>);
134
135#[derive(Debug, Clone, Default, PartialEq, Eq)]
137pub struct Held {
138 pub bytes: RepoBytes,
139 pub repos: u64,
140 pub tombstoned_bytes: u64,
141 pub objects: u64,
142}
143
144pub struct UsageLedger {
146 epoch: String,
147 seq: AtomicU64,
148 marks: Mutex<Marks>,
149 state: Mutex<State>,
150 settling: Mutex<()>,
152 drift: AtomicU64,
153}
154
155impl std::fmt::Debug for UsageLedger {
156 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157 f.debug_struct("UsageLedger")
158 .field("epoch", &self.epoch)
159 .field("seq", &self.seq.load(Ordering::Relaxed))
160 .finish_non_exhaustive()
161 }
162}
163
164fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
165 m.lock().unwrap_or_else(|p| p.into_inner())
168}
169
170pub fn tombstoned_bytes(stores: &ServerStores, namespace: &str) -> u64 {
175 let dir = stores.repos_root().join(namespace);
176 let Ok(entries) = std::fs::read_dir(&dir) else {
177 return 0;
178 };
179 let mut total = 0;
180 for entry in entries.flatten() {
181 let name = entry.file_name().to_string_lossy().into_owned();
182 let Some(store) = name.strip_suffix(".tombstone") else {
183 continue;
184 };
185 if !store.ends_with(".git") {
186 continue;
187 }
188 let Some(record) = ServerStores::read_tombstone(&entry.path()) else {
189 continue;
190 };
191 if record.reason != TombstoneReason::Deleted {
192 continue;
193 }
194 total += dir_bytes(&dir.join(store));
195 }
196 total
197}
198
199fn parse_cursor(cursor: &str) -> Option<(&str, i64, u64)> {
201 let mut parts = cursor.split('.');
202 let (epoch, day, seq) = (parts.next()?, parts.next()?, parts.next()?);
203 if parts.next().is_some() || epoch.is_empty() {
204 return None;
205 }
206 Some((epoch, day.parse().ok()?, seq.parse().ok()?))
207}
208
209impl Default for UsageLedger {
210 fn default() -> Self {
211 Self::new()
212 }
213}
214
215impl UsageLedger {
216 pub fn new() -> Self {
217 static LEDGERS: AtomicU64 = AtomicU64::new(0);
218 let n = LEDGERS.fetch_add(1, Ordering::Relaxed);
219 let epoch = format!("{:x}{:x}{n:x}", gunnar_core::now_ms(), std::process::id());
220 UsageLedger {
221 epoch,
222 seq: AtomicU64::new(0),
223 marks: Mutex::new(Marks::default()),
224 state: Mutex::new(State::default()),
225 settling: Mutex::new(()),
226 drift: AtomicU64::new(0),
227 }
228 }
229
230 pub fn epoch(&self) -> &str {
232 &self.epoch
233 }
234
235 pub fn drift_corrections(&self) -> u64 {
239 self.drift.load(Ordering::Relaxed)
240 }
241
242 pub fn mark_repo(&self, namespace: &str, name: &str) {
246 lock(&self.marks)
247 .repos
248 .insert((namespace.to_owned(), name.to_owned()));
249 }
250
251 pub fn mark_namespace(&self, namespace: &str) {
254 lock(&self.marks).namespaces.insert(namespace.to_owned());
255 }
256
257 pub fn mark_flow(&self, namespace: &str) {
259 lock(&self.marks).flow.insert(namespace.to_owned());
260 }
261
262 pub fn mark_all(&self) {
264 lock(&self.marks).all = true;
265 }
266
267 fn bump(&self, ns: &mut Ns) {
268 ns.changed_seq = self.seq.fetch_add(1, Ordering::SeqCst) + 1;
269 }
270
271 pub fn cursor_now(&self, now_unix_ms: i64) -> String {
273 format!(
274 "{}.{}.{}",
275 self.epoch,
276 gunnar_core::day_of(now_unix_ms),
277 self.seq.load(Ordering::SeqCst)
278 )
279 }
280
281 fn honour(&self, cursor: &str, now_unix_ms: i64) -> Option<u64> {
283 let (epoch, day, seq) = parse_cursor(cursor)?;
284 (epoch == self.epoch && day == gunnar_core::day_of(now_unix_ms)).then_some(seq)
285 }
286
287 fn refresh_index(
291 &self,
292 stores: &ServerStores,
293 st: &mut State,
294 marks: &mut Marks,
295 ) -> gunnar_core::Result<()> {
296 let catalog = stores.catalog();
297 let version = catalog.version()?;
298 if st.catalog == Some(version) {
299 return Ok(());
300 }
301 let initial = st.catalog.is_none();
302 let mut want: BTreeMap<String, BTreeSet<String>> = catalog
303 .namespaces()?
304 .into_iter()
305 .map(|n| (n.name, BTreeSet::new()))
306 .collect();
307 for repo in catalog.list(None)? {
308 want.entry(repo.namespace.clone())
309 .or_default()
310 .insert(repo.name.clone());
311 }
312 let gone: Vec<String> = st
313 .ns
314 .keys()
315 .filter(|n| !want.contains_key(*n))
316 .cloned()
317 .collect();
318 for name in gone {
319 if let Some(ns) = st.ns.remove(&name) {
320 for row in ns.repos.values() {
321 if let Some(dir) = &row.dir {
322 st.by_dir.remove(dir);
323 }
324 }
325 }
326 }
327 for (name, set) in want {
328 let fresh = !st.ns.contains_key(&name);
329 let mut ns = st.ns.remove(&name).unwrap_or_default();
330 let old: BTreeSet<String> = ns.repos.keys().cloned().collect();
331 if old != set {
332 for removed in old.difference(&set) {
333 if let Some(row) = ns.repos.remove(removed) {
334 if let Some(dir) = row.dir {
335 st.by_dir.remove(&dir);
336 }
337 }
338 }
339 for added in set.difference(&old) {
340 let dir = RepoRef::new(&name, added)
341 .ok()
342 .and_then(|path| stores.repo_dir(&path).ok());
343 if let Some(dir) = &dir {
344 st.by_dir.insert(dir.clone(), (name.clone(), added.clone()));
345 }
346 ns.repos.insert(added.clone(), RepoRow { dir, bytes: None });
347 }
348 marks.namespaces.insert(name.clone());
351 if !initial {
352 self.bump(&mut ns);
353 }
354 } else if fresh && !initial {
355 self.bump(&mut ns);
356 }
357 st.ns.insert(name, ns);
358 }
359 st.catalog = Some(version);
360 Ok(())
361 }
362
363 fn collect(&self, stores: &ServerStores, st: &mut State, marks: Marks, sweep: bool) -> Work {
366 let mut work = Work::default();
367 let mut repos: BTreeSet<(String, String)> = BTreeSet::new();
368 let mut whole: BTreeSet<String> = BTreeSet::new();
369 if marks.all {
370 whole.extend(st.ns.keys().cloned());
371 }
372 whole.extend(marks.namespaces);
373 for name in marks.flow {
374 if let Some(ns) = st.ns.get_mut(&name) {
375 self.bump(ns);
376 whole.insert(name);
377 }
378 }
379 repos.extend(marks.repos);
380 let now_open: HashMap<PathBuf, u64> =
388 stores.engine().open_object_counts().into_iter().collect();
389 let full = sweep;
390 let mut objects: HashMap<String, u64> = HashMap::new();
391 for (dir, count) in &now_open {
392 if let Some((ns, repo)) = st.by_dir.get(dir) {
393 *objects.entry(ns.clone()).or_default() += count;
394 if full || !st.open_last.contains(dir) {
395 repos.insert((ns.clone(), repo.clone()));
396 }
397 }
398 }
399 for dir in st.open_last.iter().filter(|d| !now_open.contains_key(*d)) {
400 if let Some((ns, repo)) = st.by_dir.get(dir) {
401 repos.insert((ns.clone(), repo.clone()));
402 }
403 }
404 let names: Vec<String> = st
405 .ns
406 .iter()
407 .filter(|(_, ns)| ns.objects > 0)
408 .map(|(n, _)| n.clone())
409 .chain(objects.keys().cloned())
410 .collect();
411 for name in names {
412 let now = objects.get(&name).copied().unwrap_or(0);
413 if let Some(ns) = st.ns.get_mut(&name) {
414 if ns.objects != now {
415 ns.objects = now;
416 self.bump(ns);
417 }
418 }
419 }
420 st.open_last = now_open.into_keys().collect();
421 for name in &whole {
422 if let Some(ns) = st.ns.get(name) {
423 for (repo, row) in &ns.repos {
424 if let Some(dir) = &row.dir {
425 work.repos.push((name.clone(), repo.clone(), dir.clone()));
426 }
427 }
428 work.tombstones.insert(name.clone());
429 }
430 }
431 for (name, repo) in repos {
432 if whole.contains(&name) {
433 continue;
434 }
435 if let Some(dir) = st
436 .ns
437 .get(&name)
438 .and_then(|ns| ns.repos.get(&repo))
439 .and_then(|row| row.dir.clone())
440 {
441 work.repos.push((name, repo, dir));
442 }
443 }
444 work
445 }
446
447 fn carry_out(&self, stores: &ServerStores, work: Work) -> BTreeSet<String> {
451 let repos: Vec<(String, String, RepoBytes)> = work
452 .repos
453 .into_iter()
454 .map(|(ns, repo, dir)| (ns, repo, repo_bytes(&dir)))
455 .collect();
456 let tombstones: Vec<(String, u64)> = work
457 .tombstones
458 .into_iter()
459 .map(|ns| {
460 let t = tombstoned_bytes(stores, &ns);
461 (ns, t)
462 })
463 .collect();
464 let mut changed = BTreeSet::new();
465 let mut st = lock(&self.state);
466 for (name, repo, bytes) in repos {
467 let Some(ns) = st.ns.get_mut(&name) else {
468 continue;
469 };
470 let Some(row) = ns.repos.get_mut(&repo) else {
471 continue;
472 };
473 if row.bytes != Some(bytes) {
474 row.bytes = Some(bytes);
475 self.bump(ns);
476 changed.insert(name);
477 }
478 }
479 for (name, bytes) in tombstones {
480 let Some(ns) = st.ns.get_mut(&name) else {
481 continue;
482 };
483 if ns.tombstoned != Some(bytes) {
484 ns.tombstoned = Some(bytes);
485 self.bump(ns);
486 changed.insert(name);
487 }
488 }
489 changed
490 }
491
492 pub fn settle(&self, stores: &ServerStores, sweep: bool) -> gunnar_core::Result<()> {
496 let _one = lock(&self.settling);
497 let marks = std::mem::take(&mut *lock(&self.marks));
498 let work = {
499 let mut st = lock(&self.state);
500 let mut marks = marks;
501 if let Err(err) = self.refresh_index(stores, &mut st, &mut marks) {
502 let mut back = lock(&self.marks);
504 back.repos.extend(marks.repos);
505 back.namespaces.extend(marks.namespaces);
506 back.flow.extend(marks.flow);
507 back.all |= marks.all;
508 return Err(err);
509 }
510 self.collect(stores, &mut st, marks, sweep)
511 };
512 self.carry_out(stores, work);
513 Ok(())
514 }
515
516 fn measure_unmeasured(&self, stores: &ServerStores, names: &[String]) {
518 let work = {
519 let st = lock(&self.state);
520 let mut work = Work::default();
521 for name in names {
522 let Some(ns) = st.ns.get(name) else { continue };
523 if !ns.unmeasured() {
524 continue;
525 }
526 for (repo, row) in &ns.repos {
527 if let (Some(dir), None) = (&row.dir, &row.bytes) {
528 work.repos.push((name.clone(), repo.clone(), dir.clone()));
529 }
530 }
531 if ns.tombstoned.is_none() {
532 work.tombstones.insert(name.clone());
533 }
534 }
535 work
536 };
537 self.carry_out(stores, work);
538 }
539
540 fn held_of(ns: &Ns) -> Held {
541 let mut held = Held {
542 tombstoned_bytes: ns.tombstoned.unwrap_or(0),
543 objects: ns.objects,
544 ..Held::default()
545 };
546 for row in ns.repos.values() {
547 held.repos += 1;
548 if let Some(b) = row.bytes {
549 held.bytes.pack += b.pack;
550 held.bytes.cache += b.cache;
551 held.bytes.lfs += b.lfs;
552 }
553 }
554 held
555 }
556
557 pub fn page(
564 &self,
565 stores: &ServerStores,
566 after: &str,
567 page_size: usize,
568 changed_since: &str,
569 now_unix_ms: i64,
570 ) -> gunnar_core::Result<Page> {
571 self.settle(stores, after.is_empty())?;
572 let since = self.honour(changed_since, now_unix_ms);
573 let page_size = page_size.max(1);
574 let (picked, more) = {
575 let st = lock(&self.state);
576 let range = if after.is_empty() {
577 st.ns.range::<str, _>((Bound::Unbounded, Bound::Unbounded))
578 } else {
579 st.ns
580 .range::<str, _>((Bound::Excluded(after), Bound::Unbounded))
581 };
582 let mut picked = Vec::new();
583 let mut more = false;
584 for (name, ns) in range {
585 if since.is_some_and(|s| ns.changed_seq <= s) {
586 continue;
587 }
588 if picked.len() == page_size {
589 more = true;
590 break;
591 }
592 picked.push(name.clone());
593 }
594 (picked, more)
595 };
596 self.measure_unmeasured(stores, &picked);
597 let st = lock(&self.state);
598 let rows = picked
599 .into_iter()
600 .filter_map(|name| st.ns.get(&name).map(|ns| (name.clone(), Self::held_of(ns))))
601 .collect();
602 Ok((rows, more, since.is_some()))
603 }
604
605 pub fn exact(
610 &self,
611 stores: &ServerStores,
612 namespace: &str,
613 ) -> gunnar_core::Result<Option<Exact>> {
614 {
617 let _one = lock(&self.settling);
618 let mut st = lock(&self.state);
619 let mut marks = Marks::default();
620 self.refresh_index(stores, &mut st, &mut marks)?;
621 if !marks.namespaces.is_empty() {
624 lock(&self.marks).namespaces.extend(marks.namespaces);
625 }
626 }
627 let work = {
628 let st = lock(&self.state);
629 let Some(ns) = st.ns.get(namespace) else {
630 return Ok(None);
631 };
632 let mut work = Work::default();
633 for (repo, row) in &ns.repos {
634 if let Some(dir) = &row.dir {
635 work.repos
636 .push((namespace.to_owned(), repo.clone(), dir.clone()));
637 }
638 }
639 work.tombstones.insert(namespace.to_owned());
640 work
641 };
642 self.carry_out(stores, work);
643 let st = lock(&self.state);
644 let Some(ns) = st.ns.get(namespace) else {
645 return Ok(None);
646 };
647 let rows = ns
648 .repos
649 .iter()
650 .map(|(name, row)| (name.clone(), row.bytes.unwrap_or_default(), row.dir.clone()))
651 .collect();
652 Ok(Some((Self::held_of(ns), rows)))
653 }
654
655 pub fn reconcile_step(
666 &self,
667 stores: &ServerStores,
668 tick: Duration,
669 ) -> gunnar_core::Result<(usize, usize)> {
670 self.settle(stores, false)?;
671 let slice = {
672 let mut st = lock(&self.state);
673 let total = st.ns.len();
674 if total == 0 {
675 return Ok((0, 0));
676 }
677 let ticks = (RECONCILE_CYCLE.as_secs() / tick.as_secs().max(1)).max(1) as usize;
678 let want = total.div_ceil(ticks).max(1);
679 let after = st.reconcile_after.clone();
680 let mut slice: Vec<String> = st
681 .ns
682 .range::<str, _>((
683 if after.is_empty() {
684 Bound::Unbounded
685 } else {
686 Bound::Excluded(after.as_str())
687 },
688 Bound::Unbounded,
689 ))
690 .take(want)
691 .map(|(n, _)| n.clone())
692 .collect();
693 if slice.len() < want {
694 let have: BTreeSet<String> = slice.iter().cloned().collect();
696 slice.extend(
697 st.ns
698 .keys()
699 .filter(|n| !have.contains(*n))
700 .take(want - slice.len())
701 .cloned(),
702 );
703 }
704 st.reconcile_after = slice.last().cloned().unwrap_or_default();
705 slice
706 };
707 let work = {
708 let st = lock(&self.state);
709 let mut work = Work::default();
710 for name in &slice {
711 let Some(ns) = st.ns.get(name) else { continue };
712 if ns.unmeasured() {
713 continue;
715 }
716 for (repo, row) in &ns.repos {
717 if let Some(dir) = &row.dir {
718 work.repos.push((name.clone(), repo.clone(), dir.clone()));
719 }
720 }
721 work.tombstones.insert(name.clone());
722 }
723 work
724 };
725 let looked = slice.len();
726 let corrected = self.carry_out(stores, work);
727 if !corrected.is_empty() {
728 self.drift
729 .fetch_add(corrected.len() as u64, Ordering::Relaxed);
730 tracing::info!(
731 target: gunnar_core::observe::TARGET_STORE,
732 event = "usage.reconciled",
733 corrected = corrected.len(),
734 first = %corrected.iter().next().cloned().unwrap_or_default(),
735 "the usage ledger's reconcile pass found {} namespace(s) whose bytes had moved without a mark; corrected",
736 corrected.len()
737 );
738 }
739 Ok((looked, corrected.len()))
740 }
741}
742
743#[cfg(test)]
744mod tests {
745 use super::*;
746
747 #[test]
748 fn a_cursor_is_honoured_only_by_its_own_ledger_on_its_own_day() {
749 let ledger = UsageLedger::new();
750 let other = UsageLedger::new();
751 let noon = gunnar_core::midnight_of(20_000) + 12 * 3_600_000;
752 let cursor = ledger.cursor_now(noon);
753 assert_eq!(ledger.honour(&cursor, noon), Some(0));
754 assert_eq!(
755 other.honour(&cursor, noon),
756 None,
757 "another ledger (a restart) answers everything"
758 );
759 assert_eq!(
760 ledger.honour(&cursor, noon + 86_400_000),
761 None,
762 "the day turned: the *_today meters reset"
763 );
764 assert_eq!(ledger.honour("", noon), None);
765 assert_eq!(ledger.honour("x.y.z", noon), None);
766 assert_ne!(ledger.epoch(), other.epoch());
767 }
768}