use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::ops::Bound;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, MutexGuard};
use std::time::Duration;
use gunnar_core::RepoRef;
use crate::control_plane::{repo_bytes, RepoBytes};
use crate::stores::{dir_bytes, ServerStores, TombstoneReason};
pub const RECONCILE_CYCLE: Duration = Duration::from_secs(30 * 60);
#[derive(Default)]
struct Marks {
repos: HashSet<(String, String)>,
namespaces: HashSet<String>,
flow: HashSet<String>,
all: bool,
}
#[derive(Debug, Clone, Default)]
struct RepoRow {
dir: Option<PathBuf>,
bytes: Option<RepoBytes>,
}
#[derive(Debug, Default)]
struct Ns {
repos: BTreeMap<String, RepoRow>,
tombstoned: Option<u64>,
objects: u64,
changed_seq: u64,
}
impl Ns {
fn unmeasured(&self) -> bool {
self.tombstoned.is_none()
|| self
.repos
.values()
.any(|r| r.dir.is_some() && r.bytes.is_none())
}
}
#[derive(Default)]
struct State {
catalog: Option<Option<gunnar_core::BlobId>>,
ns: BTreeMap<String, Ns>,
by_dir: HashMap<PathBuf, (String, String)>,
open_last: HashSet<PathBuf>,
reconcile_after: String,
}
#[derive(Default)]
struct Work {
repos: Vec<(String, String, PathBuf)>,
tombstones: BTreeSet<String>,
}
pub type Page = (Vec<(String, Held)>, bool, bool);
pub type Exact = (Held, Vec<(String, RepoBytes, Option<PathBuf>)>);
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Held {
pub bytes: RepoBytes,
pub repos: u64,
pub tombstoned_bytes: u64,
pub objects: u64,
}
pub struct UsageLedger {
epoch: String,
seq: AtomicU64,
marks: Mutex<Marks>,
state: Mutex<State>,
settling: Mutex<()>,
drift: AtomicU64,
}
impl std::fmt::Debug for UsageLedger {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UsageLedger")
.field("epoch", &self.epoch)
.field("seq", &self.seq.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(|p| p.into_inner())
}
pub fn tombstoned_bytes(stores: &ServerStores, namespace: &str) -> u64 {
let dir = stores.repos_root().join(namespace);
let Ok(entries) = std::fs::read_dir(&dir) else {
return 0;
};
let mut total = 0;
for entry in entries.flatten() {
let name = entry.file_name().to_string_lossy().into_owned();
let Some(store) = name.strip_suffix(".tombstone") else {
continue;
};
if !store.ends_with(".git") {
continue;
}
let Some(record) = ServerStores::read_tombstone(&entry.path()) else {
continue;
};
if record.reason != TombstoneReason::Deleted {
continue;
}
total += dir_bytes(&dir.join(store));
}
total
}
fn parse_cursor(cursor: &str) -> Option<(&str, i64, u64)> {
let mut parts = cursor.split('.');
let (epoch, day, seq) = (parts.next()?, parts.next()?, parts.next()?);
if parts.next().is_some() || epoch.is_empty() {
return None;
}
Some((epoch, day.parse().ok()?, seq.parse().ok()?))
}
impl Default for UsageLedger {
fn default() -> Self {
Self::new()
}
}
impl UsageLedger {
pub fn new() -> Self {
static LEDGERS: AtomicU64 = AtomicU64::new(0);
let n = LEDGERS.fetch_add(1, Ordering::Relaxed);
let epoch = format!("{:x}{:x}{n:x}", gunnar_core::now_ms(), std::process::id());
UsageLedger {
epoch,
seq: AtomicU64::new(0),
marks: Mutex::new(Marks::default()),
state: Mutex::new(State::default()),
settling: Mutex::new(()),
drift: AtomicU64::new(0),
}
}
pub fn epoch(&self) -> &str {
&self.epoch
}
pub fn drift_corrections(&self) -> u64 {
self.drift.load(Ordering::Relaxed)
}
pub fn mark_repo(&self, namespace: &str, name: &str) {
lock(&self.marks)
.repos
.insert((namespace.to_owned(), name.to_owned()));
}
pub fn mark_namespace(&self, namespace: &str) {
lock(&self.marks).namespaces.insert(namespace.to_owned());
}
pub fn mark_flow(&self, namespace: &str) {
lock(&self.marks).flow.insert(namespace.to_owned());
}
pub fn mark_all(&self) {
lock(&self.marks).all = true;
}
fn bump(&self, ns: &mut Ns) {
ns.changed_seq = self.seq.fetch_add(1, Ordering::SeqCst) + 1;
}
pub fn cursor_now(&self, now_unix_ms: i64) -> String {
format!(
"{}.{}.{}",
self.epoch,
gunnar_core::day_of(now_unix_ms),
self.seq.load(Ordering::SeqCst)
)
}
fn honour(&self, cursor: &str, now_unix_ms: i64) -> Option<u64> {
let (epoch, day, seq) = parse_cursor(cursor)?;
(epoch == self.epoch && day == gunnar_core::day_of(now_unix_ms)).then_some(seq)
}
fn refresh_index(
&self,
stores: &ServerStores,
st: &mut State,
marks: &mut Marks,
) -> gunnar_core::Result<()> {
let catalog = stores.catalog();
let version = catalog.version()?;
if st.catalog == Some(version) {
return Ok(());
}
let initial = st.catalog.is_none();
let mut want: BTreeMap<String, BTreeSet<String>> = catalog
.namespaces()?
.into_iter()
.map(|n| (n.name, BTreeSet::new()))
.collect();
for repo in catalog.list(None)? {
want.entry(repo.namespace.clone())
.or_default()
.insert(repo.name.clone());
}
let gone: Vec<String> = st
.ns
.keys()
.filter(|n| !want.contains_key(*n))
.cloned()
.collect();
for name in gone {
if let Some(ns) = st.ns.remove(&name) {
for row in ns.repos.values() {
if let Some(dir) = &row.dir {
st.by_dir.remove(dir);
}
}
}
}
for (name, set) in want {
let fresh = !st.ns.contains_key(&name);
let mut ns = st.ns.remove(&name).unwrap_or_default();
let old: BTreeSet<String> = ns.repos.keys().cloned().collect();
if old != set {
for removed in old.difference(&set) {
if let Some(row) = ns.repos.remove(removed) {
if let Some(dir) = row.dir {
st.by_dir.remove(&dir);
}
}
}
for added in set.difference(&old) {
let dir = RepoRef::new(&name, added)
.ok()
.and_then(|path| stores.repo_dir(&path).ok());
if let Some(dir) = &dir {
st.by_dir.insert(dir.clone(), (name.clone(), added.clone()));
}
ns.repos.insert(added.clone(), RepoRow { dir, bytes: None });
}
marks.namespaces.insert(name.clone());
if !initial {
self.bump(&mut ns);
}
} else if fresh && !initial {
self.bump(&mut ns);
}
st.ns.insert(name, ns);
}
st.catalog = Some(version);
Ok(())
}
fn collect(&self, stores: &ServerStores, st: &mut State, marks: Marks, sweep: bool) -> Work {
let mut work = Work::default();
let mut repos: BTreeSet<(String, String)> = BTreeSet::new();
let mut whole: BTreeSet<String> = BTreeSet::new();
if marks.all {
whole.extend(st.ns.keys().cloned());
}
whole.extend(marks.namespaces);
for name in marks.flow {
if let Some(ns) = st.ns.get_mut(&name) {
self.bump(ns);
whole.insert(name);
}
}
repos.extend(marks.repos);
let now_open: HashMap<PathBuf, u64> =
stores.engine().open_object_counts().into_iter().collect();
let full = sweep;
let mut objects: HashMap<String, u64> = HashMap::new();
for (dir, count) in &now_open {
if let Some((ns, repo)) = st.by_dir.get(dir) {
*objects.entry(ns.clone()).or_default() += count;
if full || !st.open_last.contains(dir) {
repos.insert((ns.clone(), repo.clone()));
}
}
}
for dir in st.open_last.iter().filter(|d| !now_open.contains_key(*d)) {
if let Some((ns, repo)) = st.by_dir.get(dir) {
repos.insert((ns.clone(), repo.clone()));
}
}
let names: Vec<String> = st
.ns
.iter()
.filter(|(_, ns)| ns.objects > 0)
.map(|(n, _)| n.clone())
.chain(objects.keys().cloned())
.collect();
for name in names {
let now = objects.get(&name).copied().unwrap_or(0);
if let Some(ns) = st.ns.get_mut(&name) {
if ns.objects != now {
ns.objects = now;
self.bump(ns);
}
}
}
st.open_last = now_open.into_keys().collect();
for name in &whole {
if let Some(ns) = st.ns.get(name) {
for (repo, row) in &ns.repos {
if let Some(dir) = &row.dir {
work.repos.push((name.clone(), repo.clone(), dir.clone()));
}
}
work.tombstones.insert(name.clone());
}
}
for (name, repo) in repos {
if whole.contains(&name) {
continue;
}
if let Some(dir) = st
.ns
.get(&name)
.and_then(|ns| ns.repos.get(&repo))
.and_then(|row| row.dir.clone())
{
work.repos.push((name, repo, dir));
}
}
work
}
fn carry_out(&self, stores: &ServerStores, work: Work) -> BTreeSet<String> {
let repos: Vec<(String, String, RepoBytes)> = work
.repos
.into_iter()
.map(|(ns, repo, dir)| (ns, repo, repo_bytes(&dir)))
.collect();
let tombstones: Vec<(String, u64)> = work
.tombstones
.into_iter()
.map(|ns| {
let t = tombstoned_bytes(stores, &ns);
(ns, t)
})
.collect();
let mut changed = BTreeSet::new();
let mut st = lock(&self.state);
for (name, repo, bytes) in repos {
let Some(ns) = st.ns.get_mut(&name) else {
continue;
};
let Some(row) = ns.repos.get_mut(&repo) else {
continue;
};
if row.bytes != Some(bytes) {
row.bytes = Some(bytes);
self.bump(ns);
changed.insert(name);
}
}
for (name, bytes) in tombstones {
let Some(ns) = st.ns.get_mut(&name) else {
continue;
};
if ns.tombstoned != Some(bytes) {
ns.tombstoned = Some(bytes);
self.bump(ns);
changed.insert(name);
}
}
changed
}
pub fn settle(&self, stores: &ServerStores, sweep: bool) -> gunnar_core::Result<()> {
let _one = lock(&self.settling);
let marks = std::mem::take(&mut *lock(&self.marks));
let work = {
let mut st = lock(&self.state);
let mut marks = marks;
if let Err(err) = self.refresh_index(stores, &mut st, &mut marks) {
let mut back = lock(&self.marks);
back.repos.extend(marks.repos);
back.namespaces.extend(marks.namespaces);
back.flow.extend(marks.flow);
back.all |= marks.all;
return Err(err);
}
self.collect(stores, &mut st, marks, sweep)
};
self.carry_out(stores, work);
Ok(())
}
fn measure_unmeasured(&self, stores: &ServerStores, names: &[String]) {
let work = {
let st = lock(&self.state);
let mut work = Work::default();
for name in names {
let Some(ns) = st.ns.get(name) else { continue };
if !ns.unmeasured() {
continue;
}
for (repo, row) in &ns.repos {
if let (Some(dir), None) = (&row.dir, &row.bytes) {
work.repos.push((name.clone(), repo.clone(), dir.clone()));
}
}
if ns.tombstoned.is_none() {
work.tombstones.insert(name.clone());
}
}
work
};
self.carry_out(stores, work);
}
fn held_of(ns: &Ns) -> Held {
let mut held = Held {
tombstoned_bytes: ns.tombstoned.unwrap_or(0),
objects: ns.objects,
..Held::default()
};
for row in ns.repos.values() {
held.repos += 1;
if let Some(b) = row.bytes {
held.bytes.pack += b.pack;
held.bytes.cache += b.cache;
held.bytes.lfs += b.lfs;
}
}
held
}
pub fn page(
&self,
stores: &ServerStores,
after: &str,
page_size: usize,
changed_since: &str,
now_unix_ms: i64,
) -> gunnar_core::Result<Page> {
self.settle(stores, after.is_empty())?;
let since = self.honour(changed_since, now_unix_ms);
let page_size = page_size.max(1);
let (picked, more) = {
let st = lock(&self.state);
let range = if after.is_empty() {
st.ns.range::<str, _>((Bound::Unbounded, Bound::Unbounded))
} else {
st.ns
.range::<str, _>((Bound::Excluded(after), Bound::Unbounded))
};
let mut picked = Vec::new();
let mut more = false;
for (name, ns) in range {
if since.is_some_and(|s| ns.changed_seq <= s) {
continue;
}
if picked.len() == page_size {
more = true;
break;
}
picked.push(name.clone());
}
(picked, more)
};
self.measure_unmeasured(stores, &picked);
let st = lock(&self.state);
let rows = picked
.into_iter()
.filter_map(|name| st.ns.get(&name).map(|ns| (name.clone(), Self::held_of(ns))))
.collect();
Ok((rows, more, since.is_some()))
}
pub fn exact(
&self,
stores: &ServerStores,
namespace: &str,
) -> gunnar_core::Result<Option<Exact>> {
{
let _one = lock(&self.settling);
let mut st = lock(&self.state);
let mut marks = Marks::default();
self.refresh_index(stores, &mut st, &mut marks)?;
if !marks.namespaces.is_empty() {
lock(&self.marks).namespaces.extend(marks.namespaces);
}
}
let work = {
let st = lock(&self.state);
let Some(ns) = st.ns.get(namespace) else {
return Ok(None);
};
let mut work = Work::default();
for (repo, row) in &ns.repos {
if let Some(dir) = &row.dir {
work.repos
.push((namespace.to_owned(), repo.clone(), dir.clone()));
}
}
work.tombstones.insert(namespace.to_owned());
work
};
self.carry_out(stores, work);
let st = lock(&self.state);
let Some(ns) = st.ns.get(namespace) else {
return Ok(None);
};
let rows = ns
.repos
.iter()
.map(|(name, row)| (name.clone(), row.bytes.unwrap_or_default(), row.dir.clone()))
.collect();
Ok(Some((Self::held_of(ns), rows)))
}
pub fn reconcile_step(
&self,
stores: &ServerStores,
tick: Duration,
) -> gunnar_core::Result<(usize, usize)> {
self.settle(stores, false)?;
let slice = {
let mut st = lock(&self.state);
let total = st.ns.len();
if total == 0 {
return Ok((0, 0));
}
let ticks = (RECONCILE_CYCLE.as_secs() / tick.as_secs().max(1)).max(1) as usize;
let want = total.div_ceil(ticks).max(1);
let after = st.reconcile_after.clone();
let mut slice: Vec<String> = st
.ns
.range::<str, _>((
if after.is_empty() {
Bound::Unbounded
} else {
Bound::Excluded(after.as_str())
},
Bound::Unbounded,
))
.take(want)
.map(|(n, _)| n.clone())
.collect();
if slice.len() < want {
let have: BTreeSet<String> = slice.iter().cloned().collect();
slice.extend(
st.ns
.keys()
.filter(|n| !have.contains(*n))
.take(want - slice.len())
.cloned(),
);
}
st.reconcile_after = slice.last().cloned().unwrap_or_default();
slice
};
let work = {
let st = lock(&self.state);
let mut work = Work::default();
for name in &slice {
let Some(ns) = st.ns.get(name) else { continue };
if ns.unmeasured() {
continue;
}
for (repo, row) in &ns.repos {
if let Some(dir) = &row.dir {
work.repos.push((name.clone(), repo.clone(), dir.clone()));
}
}
work.tombstones.insert(name.clone());
}
work
};
let looked = slice.len();
let corrected = self.carry_out(stores, work);
if !corrected.is_empty() {
self.drift
.fetch_add(corrected.len() as u64, Ordering::Relaxed);
tracing::info!(
target: gunnar_core::observe::TARGET_STORE,
event = "usage.reconciled",
corrected = corrected.len(),
first = %corrected.iter().next().cloned().unwrap_or_default(),
"the usage ledger's reconcile pass found {} namespace(s) whose bytes had moved without a mark; corrected",
corrected.len()
);
}
Ok((looked, corrected.len()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_cursor_is_honoured_only_by_its_own_ledger_on_its_own_day() {
let ledger = UsageLedger::new();
let other = UsageLedger::new();
let noon = gunnar_core::midnight_of(20_000) + 12 * 3_600_000;
let cursor = ledger.cursor_now(noon);
assert_eq!(ledger.honour(&cursor, noon), Some(0));
assert_eq!(
other.honour(&cursor, noon),
None,
"another ledger (a restart) answers everything"
);
assert_eq!(
ledger.honour(&cursor, noon + 86_400_000),
None,
"the day turned: the *_today meters reset"
);
assert_eq!(ledger.honour("", noon), None);
assert_eq!(ledger.honour("x.y.z", noon), None);
assert_ne!(ledger.epoch(), other.epoch());
}
}