use std::collections::HashMap;
use std::sync::Arc;
use alkblobs::{Batch, Pin};
use alkgit::{GitRefs, RefLine, RefOutcome, RefUpdate, RegistryError, RepoRecord, StorageError};
use async_trait::async_trait;
use crate::reactor::{Reactor, RegistryEvent};
use crate::registry::InMemoryRegistry;
pub(crate) enum PinCustody {
Pins(Vec<Pin>),
Batch(Batch),
}
struct PendingEntry {
updates: Vec<RefUpdate>,
custody: PinCustody,
}
pub(crate) struct PendingPinRegistry {
repos: parking_lot::Mutex<HashMap<String, Vec<PendingEntry>>>,
}
impl Default for PendingPinRegistry {
fn default() -> Self {
Self::new()
}
}
impl PendingPinRegistry {
pub(crate) fn new() -> Self {
PendingPinRegistry {
repos: parking_lot::Mutex::new(HashMap::new()),
}
}
pub(crate) fn retain(&self, repo_id: &str, updates: Vec<RefUpdate>, custody: PinCustody) {
self.repos
.lock()
.entry(repo_id.to_owned())
.or_default()
.push(PendingEntry { updates, custody });
}
pub(crate) fn drain(&self, repo_id: &str, updates: &[RefUpdate]) -> Option<PinCustody> {
let mut repos = self.repos.lock();
let drained = {
let entries = repos.get_mut(repo_id)?;
let index = entries.iter().position(|entry| entry.updates == updates)?;
Some(entries.remove(index))
};
if drained.is_some() {
self.remove_empty_locked(&mut repos, repo_id);
}
drained.map(|entry| entry.custody)
}
pub(crate) fn prune_stale(&self, repo_id: &str, updates: &[RefUpdate]) -> usize {
let mut repos = self.repos.lock();
let Some(entries) = repos.get_mut(repo_id) else {
return 0;
};
let before = entries.len();
entries.retain(|entry| entry.updates == updates);
let pruned = before - entries.len();
self.remove_empty_locked(&mut repos, repo_id);
pruned
}
#[cfg(test)]
pub(crate) fn pending_count(&self, repo_id: &str) -> usize {
self.repos.lock().get(repo_id).map(Vec::len).unwrap_or(0)
}
fn remove_empty_locked(&self, repos: &mut HashMap<String, Vec<PendingEntry>>, repo_id: &str) {
if repos.get(repo_id).is_some_and(Vec::is_empty) {
repos.remove(repo_id);
}
}
}
#[async_trait]
pub(crate) trait RegistryMutations: Send + Sync {
async fn record_mutation(&self, repo_id: &str) -> Result<i64, RegistryError>;
}
pub(crate) struct MemMutationLegs {
durable: Arc<InMemoryRegistry>,
reactor: Arc<Reactor>,
}
impl MemMutationLegs {
pub(crate) fn new(durable: Arc<InMemoryRegistry>, reactor: Arc<Reactor>) -> Self {
MemMutationLegs { durable, reactor }
}
}
#[async_trait]
impl RegistryMutations for MemMutationLegs {
async fn record_mutation(&self, repo_id: &str) -> Result<i64, RegistryError> {
let seq = self.durable.record_mutation(repo_id).await?;
self.reactor
.announce(&RegistryEvent::Mutation {
seq,
repo_id: repo_id.to_owned(),
})
.await;
Ok(seq)
}
}
pub(crate) struct CommittedRefs {
inner: Arc<dyn GitRefs>,
pins: Arc<PendingPinRegistry>,
mutations: Arc<dyn RegistryMutations>,
}
impl CommittedRefs {
pub(crate) fn new(
inner: Arc<dyn GitRefs>,
pins: Arc<PendingPinRegistry>,
mutations: Arc<dyn RegistryMutations>,
) -> Self {
CommittedRefs {
inner,
pins,
mutations,
}
}
fn release_drained(&self, repo_id: &str, updates: &[RefUpdate]) {
match self.pins.drain(repo_id, updates) {
Some(PinCustody::Pins(pins)) => drop(pins),
Some(PinCustody::Batch(batch)) => drop(batch),
None => {}
}
}
}
#[async_trait]
impl GitRefs for CommittedRefs {
async fn list_refs(&self, repo: &RepoRecord) -> Result<Vec<RefLine>, StorageError> {
self.inner.list_refs(repo).await
}
async fn apply_updates(
&self,
repo: &RepoRecord,
updates: Vec<RefUpdate>,
atomic: bool,
) -> Result<Vec<RefOutcome>, StorageError> {
let list = updates.clone();
if !list.is_empty() {
self.pins.prune_stale(&repo.repo_id, &list);
}
match self.inner.apply_updates(repo, updates, atomic).await {
Err(e) => {
self.release_drained(&repo.repo_id, &list);
Err(e)
}
Ok(outcomes) => {
let aborted = atomic && outcomes.iter().any(|outcome| outcome.result.is_err());
self.release_drained(&repo.repo_id, &list);
if !aborted && !list.is_empty() {
let _ = self.mutations.record_mutation(&repo.repo_id).await;
}
Ok(outcomes)
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use alkstore::Store as _;
use gix_hash::ObjectId;
use gix_object::Kind;
use alkgit::ATOMIC_PUSH_FAILURE_REASON;
use super::*;
use crate::odb::{key_to_oid, ObjectFront};
use crate::refs::{InMemoryRefTable, CAS_STALE_REASON};
use crate::registry::{CachedRegistryProvider, RegistryWrite};
use crate::testutil::record as repo;
fn up(name: &str, expected: Option<u8>, new: Option<u8>) -> RefUpdate {
RefUpdate {
name: name.to_owned(),
expected: expected.map(|b| ObjectId::from([b; 20])),
new: new.map(|b| ObjectId::from([b; 20])),
}
}
fn up_oid(name: &str, expected: Option<ObjectId>, new: Option<ObjectId>) -> RefUpdate {
RefUpdate {
name: name.to_owned(),
expected,
new,
}
}
fn mem_front() -> ObjectFront {
ObjectFront::new(Arc::new(alkblobs::Store::mem_only()))
}
fn closure_of(
front: &ObjectFront,
rows: &Arc<InMemoryRefTable>,
) -> Arc<crate::liveness::RefClosure> {
crate::liveness::RefClosure::new(
front.clone(),
rows.clone() as Arc<dyn crate::liveness::RefRows>,
Arc::new(AtomicU64::new(0)),
)
}
fn pending_count(pins: &PendingPinRegistry, repo_id: &str) -> usize {
pins.repos.lock().get(repo_id).map(Vec::len).unwrap_or(0)
}
fn entry_present(pins: &PendingPinRegistry, repo_id: &str, updates: &[RefUpdate]) -> bool {
pins.repos
.lock()
.get(repo_id)
.is_some_and(|entries| entries.iter().any(|entry| entry.updates == updates))
}
struct LegProbe {
called: AtomicU64,
fail: bool,
seqs: parking_lot::Mutex<Vec<i64>>,
}
impl LegProbe {
fn failing() -> Arc<Self> {
Arc::new(LegProbe {
called: AtomicU64::new(0),
fail: true,
seqs: parking_lot::Mutex::new(Vec::new()),
})
}
fn recording() -> Arc<Self> {
Arc::new(LegProbe {
called: AtomicU64::new(0),
fail: false,
seqs: parking_lot::Mutex::new(Vec::new()),
})
}
fn calls(&self) -> u64 {
self.called.load(Ordering::Relaxed)
}
fn seqs(&self) -> Vec<i64> {
self.seqs.lock().clone()
}
}
#[async_trait]
impl RegistryMutations for LegProbe {
async fn record_mutation(&self, _repo_id: &str) -> Result<i64, RegistryError> {
let call = self.called.fetch_add(1, Ordering::Relaxed) as i64 + 1;
if self.fail {
return Err(RegistryError::Io(
"marker leg suppressed (probe)".to_owned(),
));
}
self.seqs.lock().push(call);
Ok(call)
}
}
struct ProbedRefs {
probe: Box<dyn Fn() + Send + Sync>,
outcome: Result<(), StorageError>,
}
#[async_trait]
impl GitRefs for ProbedRefs {
async fn list_refs(&self, _repo: &RepoRecord) -> Result<Vec<RefLine>, StorageError> {
Ok(Vec::new())
}
async fn apply_updates(
&self,
_repo: &RepoRecord,
_updates: Vec<RefUpdate>,
_atomic: bool,
) -> Result<Vec<RefOutcome>, StorageError> {
(self.probe)();
match &self.outcome {
Err(e) => Err(e.clone()),
Ok(()) => Ok(Vec::new()),
}
}
}
#[tokio::test(flavor = "multi_thread")]
async fn drain_correlates_by_content_one_entry_per_apply() {
let pins = Arc::new(PendingPinRegistry::new());
pins.retain(
"r",
vec![up("refs/heads/a", None, Some(1))],
PinCustody::Pins(Vec::new()),
);
pins.retain(
"r",
vec![up("refs/heads/b", None, Some(2))],
PinCustody::Pins(Vec::new()),
);
pins.retain(
"other",
vec![up("refs/heads/a", None, Some(1))],
PinCustody::Pins(Vec::new()),
);
assert_eq!(pending_count(&pins, "r"), 2);
assert!(
pins.drain("r", &[up("refs/heads/c", None, Some(9))])
.is_none(),
"a non-matching list drains nothing"
);
assert!(
pins.drain("r", &[up("refs/heads/a", None, Some(1))])
.is_some(),
"the repo's own matching entry drains"
);
assert!(
pins.drain("r", &[up("refs/heads/a", None, Some(1))])
.is_none(),
"the other repo's twin is not this repo's entry"
);
assert!(
pending_count(&pins, "r") > 0,
"the other list's entry stays"
);
assert!(
pins.drain("r", &[up("refs/heads/b", None, Some(2))])
.is_some(),
"the other list drains by its content"
);
assert_eq!(
pending_count(&pins, "r"),
0,
"a fully drained repo tier is gone"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn twin_entries_drain_one_per_apply() {
let pins = Arc::new(PendingPinRegistry::new());
let list = vec![up("refs/heads/main", None, Some(1))];
pins.retain("repo", list.clone(), PinCustody::Pins(Vec::new()));
pins.retain("repo", list.clone(), PinCustody::Pins(Vec::new()));
assert!(pins.drain("repo", &list).is_some());
assert!(pins.drain("repo", &list).is_some());
assert!(pins.drain("repo", &list).is_none(), "exactly two drains");
}
#[tokio::test(flavor = "multi_thread")]
async fn prune_stale_releases_the_non_matching_class() {
let front = mem_front();
let store = front.store().clone();
let rows = Arc::new(InMemoryRefTable::new());
let source = closure_of(&front, &rows);
source.register(&store).expect("the closure registers");
let (_, pin_a) = front.put(Kind::Blob, b"bridged a\n").await.unwrap();
let (_, pin_b) = front.put(Kind::Blob, b"bridged b\n").await.unwrap();
let key_a = pin_a.key().clone();
let key_b = pin_b.key().clone();
let pins = Arc::new(PendingPinRegistry::new());
pins.retain(
"repo",
vec![up("refs/heads/gone", None, Some(1))],
PinCustody::Pins(vec![pin_a]),
);
pins.retain(
"repo",
vec![up("refs/heads/gone", None, Some(1))],
PinCustody::Pins(vec![pin_b]),
);
pins.retain(
"repo",
vec![up("refs/heads/keep", None, Some(2))],
PinCustody::Pins(Vec::new()),
);
let pruned = pins.prune_stale("repo", &[up("refs/heads/keep", None, Some(2))]);
assert_eq!(pruned, 2, "both non-matching entries pruned");
assert!(
pins.drain("repo", &[up("refs/heads/keep", None, Some(2))])
.is_some(),
"the matching entry stayed for the drain"
);
assert_eq!(pending_count(&pins, "repo"), 0);
let report = store.sweep().await.expect("the sweep runs");
assert!(report.aborted.is_none());
assert_eq!(
report.deleted, 2,
"the pruned custody's objects are sweep-eligible"
);
assert!(front
.stat(&key_to_oid(&key_a).unwrap())
.await
.unwrap()
.is_none());
assert!(front
.stat(&key_to_oid(&key_b).unwrap())
.await
.unwrap()
.is_none());
}
#[tokio::test(flavor = "multi_thread")]
async fn a_pending_entry_holds_across_the_sweep_until_drained() {
let front = mem_front();
let store = front.store().clone();
let rows = Arc::new(InMemoryRefTable::new());
let source = closure_of(&front, &rows);
source.register(&store).expect("the closure registers");
let (_, pin) = front.put(Kind::Blob, b"bridged\n").await.unwrap();
let oid = key_to_oid(pin.key()).unwrap();
let pins = Arc::new(PendingPinRegistry::new());
let list = vec![up("refs/heads/main", None, Some(1))];
pins.retain("repo", list.clone(), PinCustody::Pins(vec![pin]));
let report = store.sweep().await.expect("the sweep runs");
assert!(report.aborted.is_none());
assert_eq!(
report.deleted, 0,
"the bridged push's objects are not reclaimable"
);
assert!(
front.stat(&oid).await.unwrap().is_some(),
"held across the window"
);
let custody = pins.drain("repo", &list).expect("the entry drains");
drop(custody);
assert!(
front.stat(&oid).await.unwrap().is_some(),
"release alone reclaims nothing"
);
let report = store.sweep().await.expect("the sweep runs");
assert_eq!(
report.deleted, 1,
"the released objects recover the crash way"
);
assert!(front.stat(&oid).await.unwrap().is_none());
}
#[tokio::test(flavor = "multi_thread")]
async fn the_drain_lands_only_after_the_cas_returns() {
let front = mem_front();
let pins = Arc::new(PendingPinRegistry::new());
let list = vec![up("refs/heads/main", None, Some(1))];
let (_, pin) = front.put(Kind::Blob, b"held\n").await.unwrap();
pins.retain("repo", list.clone(), PinCustody::Pins(vec![pin]));
let saw_held = Arc::new(AtomicBool::new(false));
let during_cas = pins.clone();
let key_at_cas = list.clone();
let saw = saw_held.clone();
let probe = move || {
saw.store(
entry_present(&during_cas, "repo", &key_at_cas),
Ordering::Relaxed,
);
};
let inner: Arc<dyn GitRefs> = Arc::new(ProbedRefs {
probe: Box::new(probe),
outcome: Ok(()),
});
let wrapped = CommittedRefs::new(inner, pins.clone(), LegProbe::recording());
let outcomes = wrapped
.apply_updates(&repo("repo"), list.clone(), false)
.await
.expect("the committed CAS");
assert!(outcomes.is_empty());
assert!(
saw_held.load(Ordering::Relaxed),
"the matching entry's guards were registry-held across the CAS"
);
assert!(
!entry_present(&pins, "repo", &list),
"the drain landed only after the CAS returned"
);
assert_eq!(pending_count(&pins, "repo"), 0);
}
#[tokio::test(flavor = "multi_thread")]
async fn the_committed_cas_drains_and_rides_the_legs() {
let registry = Arc::new(InMemoryRegistry::new());
registry.put(&repo("repo")).await.unwrap();
let provider = Arc::new(CachedRegistryProvider::new());
let mem_store = Arc::new(alkstore_mem::MemStore::new());
let reactor = Arc::new(Reactor::new(
mem_store.clone(),
registry.clone(),
provider.clone(),
));
let legs = Arc::new(MemMutationLegs::new(registry.clone(), reactor.clone()));
let front = mem_front();
let (_, pin) = front.put(Kind::Blob, b"landed\n").await.unwrap();
let key = pin.key().clone();
let pins = Arc::new(PendingPinRegistry::new());
let inner: Arc<dyn GitRefs> = Arc::new(InMemoryRefTable::new());
let wrapped = CommittedRefs::new(inner, pins.clone(), legs);
let list = vec![up("refs/heads/main", None, Some(1))];
pins.retain("repo", list.clone(), PinCustody::Pins(vec![pin]));
let outcomes = wrapped
.apply_updates(&repo("repo"), list, false)
.await
.expect("the committed CAS");
assert!(outcomes.iter().all(|outcome| outcome.result.is_ok()));
assert_eq!(pending_count(&pins, "repo"), 0, "the entry drained");
let seq = registry.watermark().await.unwrap();
assert_eq!(
seq, 2,
"the marker write anchored the watermark (1 = the put)"
);
let stream = mem_store
.stream(crate::reactor::EVENTS_STREAM)
.await
.unwrap();
let mut rx = stream.subscribe("gate").await.unwrap();
let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("the announce lands inside the deadline")
.expect("the feed delivered")
.expect("the event is clean");
match RegistryEvent::decode_event(&event) {
Some(RegistryEvent::Mutation { seq: 2, repo_id }) => {
assert_eq!(repo_id, "repo");
}
other => panic!("the Mutation announce rides the marker write: {other:?}"),
}
assert!(front
.stat(&key_to_oid(&key).unwrap())
.await
.unwrap()
.is_some());
}
#[tokio::test(flavor = "multi_thread")]
async fn the_next_apply_prunes_the_stale_entry() {
let front = mem_front();
let store = front.store().clone();
let rows = Arc::new(InMemoryRefTable::new());
let source = closure_of(&front, &rows);
source.register(&store).expect("the closure registers");
let legs = LegProbe::recording();
let pins = Arc::new(PendingPinRegistry::new());
let inner: Arc<dyn GitRefs> = Arc::new(InMemoryRefTable::new());
let wrapped = CommittedRefs::new(inner, pins.clone(), legs);
let (_, stale_pin) = front.put(Kind::Blob, b"never applied\n").await.unwrap();
let stale_oid = key_to_oid(stale_pin.key()).unwrap();
pins.retain(
"repo",
vec![up("refs/heads/stale", None, Some(1))],
PinCustody::Pins(vec![stale_pin]),
);
let (_, live_pin) = front
.put(Kind::Blob, b"the live push's object\n")
.await
.unwrap();
let live_oid = key_to_oid(live_pin.key()).unwrap();
let live_list = vec![up_oid("refs/heads/live", None, Some(live_oid))];
let outcomes = wrapped
.apply_updates(&repo("repo"), live_list, false)
.await
.expect("the next apply runs");
assert!(outcomes.iter().all(|outcome| outcome.result.is_ok()));
assert_eq!(
pending_count(&pins, "repo"),
0,
"the stale entry pruned alongside the drain"
);
let report = store.sweep().await.expect("the sweep runs");
assert_eq!(report.deleted, 1, "the pruned guards recovered GC-decreed");
assert!(front.stat(&stale_oid).await.unwrap().is_none());
assert!(
front.stat(&live_oid).await.unwrap().is_some(),
"the landed ref's object covered"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn the_aborted_cas_drops_the_guards_and_skips_the_legs() {
let front = mem_front();
let store = front.store().clone();
let rows = Arc::new(InMemoryRefTable::new());
let source = closure_of(&front, &rows);
source.register(&store).expect("the closure registers");
let legs = LegProbe::recording();
let pins = Arc::new(PendingPinRegistry::new());
let inner: Arc<dyn GitRefs> = Arc::new(InMemoryRefTable::new());
let wrapped = CommittedRefs::new(inner, pins.clone(), legs.clone());
let (_, base_pin) = front.put(Kind::Blob, b"the seed's object\n").await.unwrap();
let base_oid = key_to_oid(base_pin.key()).unwrap();
let seed = wrapped
.apply_updates(
&repo("repo"),
vec![up_oid("refs/heads/base", None, Some(base_oid))],
false,
)
.await
.unwrap();
assert!(seed.iter().all(|outcome| outcome.result.is_ok()));
assert_eq!(legs.calls(), 1);
let (_, pin) = front.put(Kind::Blob, b"aborted\n").await.unwrap();
let aborted_oid = key_to_oid(pin.key()).unwrap();
let list = vec![
up_oid(
"refs/heads/aborted-sibling",
None,
Some(ObjectId::from([9; 20])),
),
up_oid(
"refs/heads/base",
Some(ObjectId::from([8; 20])),
Some(ObjectId::from([9; 20])),
),
];
pins.retain("repo", list.clone(), PinCustody::Pins(vec![pin]));
let outcomes = wrapped
.apply_updates(&repo("repo"), list, true)
.await
.expect("the aborted CAS is a business outcome");
assert_eq!(
outcomes[0].result,
Err(ATOMIC_PUSH_FAILURE_REASON.to_owned())
);
assert_eq!(pending_count(&pins, "repo"), 0, "the drained entry is gone");
assert_eq!(legs.calls(), 1, "the aborted CAS carries no legs");
let report = store.sweep().await.expect("the sweep runs");
assert_eq!(report.deleted, 1, "the dropped guards' objects are orphans");
assert!(front.stat(&aborted_oid).await.unwrap().is_none());
assert!(
front.stat(&base_oid).await.unwrap().is_some(),
"the seeded ref's object covered"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_substrate_err_drops_the_guards_and_skips_the_legs() {
let legs = LegProbe::recording();
let pins = Arc::new(PendingPinRegistry::new());
let list = vec![up("refs/heads/main", None, Some(1))];
pins.retain("repo", list.clone(), PinCustody::Pins(Vec::new()));
let inner: Arc<dyn GitRefs> = Arc::new(ProbedRefs {
probe: Box::new(|| ()),
outcome: Err(StorageError::Io("engine failure".to_owned())),
});
let wrapped = CommittedRefs::new(inner, pins.clone(), legs.clone());
let err = wrapped
.apply_updates(&repo("repo"), list, false)
.await
.expect_err("the substrate failure propagates");
assert!(err.to_string().contains("engine failure"));
assert_eq!(pending_count(&pins, "repo"), 0, "the guards dropped");
assert_eq!(legs.calls(), 0, "the Err ending carries no legs");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_failed_marker_write_never_fails_the_push() {
let legs = LegProbe::failing();
let pins = Arc::new(PendingPinRegistry::new());
let inner: Arc<dyn GitRefs> = Arc::new(InMemoryRefTable::new());
let wrapped = CommittedRefs::new(inner, pins.clone(), legs.clone());
let outcomes = wrapped
.apply_updates(
&repo("repo"),
vec![up("refs/heads/main", None, Some(1))],
false,
)
.await
.expect("the committed CAS is the push outcome");
assert!(outcomes.iter().all(|outcome| outcome.result.is_ok()));
assert_eq!(legs.calls(), 1);
assert_eq!(pending_count(&pins, "repo"), 0);
}
#[tokio::test(flavor = "multi_thread")]
async fn the_leg_firing_rule_pins_the_committed_transactions() {
let legs = LegProbe::recording();
let pins = Arc::new(PendingPinRegistry::new());
let inner: Arc<dyn GitRefs> = Arc::new(InMemoryRefTable::new());
let wrapped = CommittedRefs::new(inner, pins.clone(), legs.clone());
let outcomes = wrapped
.apply_updates(&repo("repo"), Vec::new(), false)
.await
.expect("the empty call is a committed nothing");
assert!(outcomes.is_empty());
assert_eq!(legs.calls(), 0, "the empty-update call carries no legs");
let denied = wrapped
.apply_updates(
&repo("repo"),
vec![up("refs/heads/base", Some(3), Some(4))],
false,
)
.await
.expect("denials are outcomes");
assert_eq!(denied[0].result, Err(CAS_STALE_REASON.to_owned()));
assert_eq!(legs.calls(), 1);
assert_eq!(legs.seqs(), vec![1], "one marker per committed push");
}
#[tokio::test(flavor = "multi_thread")]
async fn list_refs_delegates_verbatim() {
let rows = Arc::new(InMemoryRefTable::new());
let list = vec![up("refs/heads/main", None, Some(7))];
rows.apply_updates(&repo("repo"), list, false)
.await
.unwrap();
let wrapped = CommittedRefs::new(
rows.clone() as Arc<dyn GitRefs>,
Arc::new(PendingPinRegistry::new()),
LegProbe::recording(),
);
let lines = wrapped.list_refs(&repo("repo")).await.unwrap();
assert_eq!(lines.len(), 1);
assert_eq!(lines[0].name, "refs/heads/main");
assert_eq!(lines[0].oid, Some(ObjectId::from([7; 20])));
}
}