use std::{
collections::{HashMap, HashSet, VecDeque},
sync::{Arc, RwLock},
};
use tracing::instrument;
use crate::primitives::{AccountId, AccountSetId, JournalId};
use super::{
error::AccountSetError,
repo::{AccountSetRepo, SetGraphNode},
};
const TIMER_REFRESH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
const COLD_EPOCH: i64 = -1;
#[derive(Debug, Clone, Copy)]
struct SetMeta {
journal_id: JournalId,
eventually_consistent: bool,
}
#[derive(Debug)]
struct GraphSnapshot {
epoch: i64,
parents: HashMap<AccountSetId, Vec<AccountSetId>>,
meta: HashMap<AccountSetId, SetMeta>,
}
impl GraphSnapshot {
fn cold() -> Self {
Self {
epoch: COLD_EPOCH,
parents: HashMap::new(),
meta: HashMap::new(),
}
}
}
struct Overlay {
parents: HashMap<AccountSetId, Vec<AccountSetId>>,
meta: HashMap<AccountSetId, SetMeta>,
}
fn index_nodes(
nodes: Vec<SetGraphNode>,
) -> (
HashMap<AccountSetId, Vec<AccountSetId>>,
HashMap<AccountSetId, SetMeta>,
) {
let mut parents: HashMap<AccountSetId, Vec<AccountSetId>> = HashMap::new();
let mut meta = HashMap::new();
for node in nodes {
meta.insert(
node.id,
SetMeta {
journal_id: node.journal_id,
eventually_consistent: node.eventually_consistent,
},
);
if let Some(parent_id) = node.parent_id {
parents.entry(node.id).or_default().push(parent_id);
}
}
(parents, meta)
}
#[derive(Clone)]
pub(super) struct SetGraphCache {
inner: Arc<SetGraphCacheInner>,
}
struct SetGraphCacheInner {
repo: AccountSetRepo,
snapshot: RwLock<Arc<GraphSnapshot>>,
refresh_lock: tokio::sync::Mutex<()>,
}
impl std::fmt::Debug for SetGraphCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SetGraphCache").finish_non_exhaustive()
}
}
impl SetGraphCache {
pub(super) fn new(repo: AccountSetRepo) -> Self {
let inner = Arc::new(SetGraphCacheInner {
repo,
snapshot: RwLock::new(Arc::new(GraphSnapshot::cold())),
refresh_lock: tokio::sync::Mutex::new(()),
});
let weak = Arc::downgrade(&inner);
tokio::spawn(async move {
let mut interval = tokio::time::interval(TIMER_REFRESH_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
interval.tick().await;
let Some(inner) = weak.upgrade() else { break };
if let Err(error) = Self::refresh(&inner).await {
tracing::warn!(%error, "set_graph_cache timer refresh failed");
}
}
});
Self { inner }
}
#[instrument(
level = "debug",
name = "account_set.fetch_mappings_in_op",
skip(self, op, entry_pairs),
fields(accounts = entry_pairs.0.len(), path = tracing::field::Empty),
err(level = "warn")
)]
pub(super) async fn fetch_mappings_in_op(
&self,
op: &mut impl es_entity::AtomicOperation,
journal_id: JournalId,
entry_pairs: &(Vec<AccountId>, Vec<&str>),
) -> Result<HashMap<AccountId, Vec<AccountSetId>>, AccountSetError> {
let span = tracing::Span::current();
let account_ids: Vec<AccountId> = {
let mut ids = entry_pairs.0.clone();
ids.sort_unstable();
ids.dedup();
ids
};
let probe = self
.inner
.repo
.probe_direct_memberships_in_op(op, &account_ids)
.await?;
if probe.seeds.is_empty() {
span.record("path", "no_memberships");
return Ok(HashMap::new());
}
let snapshot = self.load();
if snapshot.epoch != probe.epoch {
span.record(
"path",
if snapshot.epoch == COLD_EPOCH {
"fallback_cold"
} else {
"fallback_epoch"
},
);
self.spawn_refresh();
return self
.inner
.repo
.walk_mappings_and_lock_in_op(&mut *op, journal_id, entry_pairs)
.await;
}
let missing: Vec<AccountSetId> = {
let mut missing: Vec<_> = probe
.seeds
.iter()
.map(|(_, set_id)| *set_id)
.filter(|set_id| !snapshot.meta.contains_key(set_id))
.collect();
missing.sort_unstable();
missing.dedup();
missing
};
let overlay = if missing.is_empty() {
None
} else {
let nodes = self
.inner
.repo
.fetch_set_graph_nodes_in_op(op, &missing)
.await?;
let (parents, meta) = index_nodes(nodes);
Some(Overlay { parents, meta })
};
match Self::expand(
&snapshot,
overlay.as_ref(),
journal_id,
&probe.seeds,
entry_pairs,
) {
Some((mappings, lock_pairs)) => {
span.record(
"path",
if overlay.is_some() {
"supplement"
} else {
"memory"
},
);
self.inner
.repo
.lock_resolved_ancestors_in_op(op, journal_id, &lock_pairs)
.await?;
Ok(mappings)
}
None => {
span.record("path", "fallback_unknown");
self.inner
.repo
.walk_mappings_and_lock_in_op(&mut *op, journal_id, entry_pairs)
.await
}
}
}
#[instrument(
level = "debug",
name = "account_set.assert_no_double_membership",
skip(self, op, account_set_ids, account_ids),
fields(pairs = account_ids.len(), path = tracing::field::Empty),
err(level = "warn")
)]
pub(super) async fn assert_no_double_membership_in_op(
&self,
op: &mut impl es_entity::AtomicOperation,
account_set_ids: &[AccountSetId],
account_ids: &[AccountId],
) -> Result<(), AccountSetError> {
let span = tracing::Span::current();
let distinct_account_ids: Vec<AccountId> = {
let mut ids = account_ids.to_vec();
ids.sort_unstable();
ids.dedup();
ids
};
let probe = self
.inner
.repo
.probe_direct_memberships_in_op(op, &distinct_account_ids)
.await?;
let snapshot = self.load();
if snapshot.epoch != probe.epoch {
span.record(
"path",
if snapshot.epoch == COLD_EPOCH {
"fallback_cold"
} else {
"fallback_epoch"
},
);
self.spawn_refresh();
return self
.inner
.repo
.assert_no_double_membership(op, account_set_ids, account_ids)
.await;
}
let missing: Vec<AccountSetId> = {
let mut missing: Vec<_> = account_set_ids
.iter()
.chain(probe.seeds.iter().map(|(_, set_id)| set_id))
.copied()
.filter(|set_id| !snapshot.meta.contains_key(set_id))
.collect();
missing.sort_unstable();
missing.dedup();
missing
};
let overlay = if missing.is_empty() {
None
} else {
let nodes = self
.inner
.repo
.fetch_set_graph_nodes_in_op(op, &missing)
.await?;
let (parents, meta) = index_nodes(nodes);
Some(Overlay { parents, meta })
};
match Self::count_membership_paths(
&snapshot,
overlay.as_ref(),
account_set_ids,
account_ids,
&probe.seeds,
) {
Some(false) => {
span.record(
"path",
if overlay.is_some() {
"supplement"
} else {
"memory"
},
);
Ok(())
}
Some(true) => Err(AccountSetError::MemberAlreadyAdded),
None => {
span.record("path", "fallback_unknown");
self.inner
.repo
.assert_no_double_membership(op, account_set_ids, account_ids)
.await
}
}
}
fn count_membership_paths(
snapshot: &GraphSnapshot,
overlay: Option<&Overlay>,
new_set_ids: &[AccountSetId],
new_account_ids: &[AccountId],
existing_seeds: &[(AccountId, AccountSetId)],
) -> Option<bool> {
let known = |set_id: &AccountSetId| -> bool {
snapshot.meta.contains_key(set_id)
|| overlay.is_some_and(|o| o.meta.contains_key(set_id))
};
let parents_of = |set_id: &AccountSetId| -> &[AccountSetId] {
snapshot
.parents
.get(set_id)
.or_else(|| overlay.and_then(|o| o.parents.get(set_id)))
.map(Vec::as_slice)
.unwrap_or(&[])
};
let mut per_account: HashMap<AccountId, Vec<AccountSetId>> = HashMap::new();
for (set_id, account_id) in new_set_ids.iter().zip(new_account_ids) {
per_account.entry(*account_id).or_default().push(*set_id);
}
for (account_id, set_id) in existing_seeds {
per_account.entry(*account_id).or_default().push(*set_id);
}
for seeds in per_account.into_values() {
let mut path_counts: HashMap<AccountSetId, u32> = HashMap::new();
let mut queue: VecDeque<AccountSetId> = seeds.into();
while let Some(set_id) = queue.pop_front() {
if !known(&set_id) {
return None;
}
let count = path_counts.entry(set_id).or_default();
*count += 1;
if *count > 1 {
return Some(true);
}
queue.extend(parents_of(&set_id));
}
}
Some(false)
}
fn load(&self) -> Arc<GraphSnapshot> {
self.inner
.snapshot
.read()
.expect("set_graph_cache snapshot lock poisoned")
.clone()
}
#[allow(clippy::type_complexity)]
fn expand<'c>(
snapshot: &GraphSnapshot,
overlay: Option<&Overlay>,
journal_id: JournalId,
seeds: &[(AccountId, AccountSetId)],
(entry_account_ids, entry_currencies): &(Vec<AccountId>, Vec<&'c str>),
) -> Option<(
HashMap<AccountId, Vec<AccountSetId>>,
(Vec<AccountSetId>, Vec<&'c str>),
)> {
let meta_of = |set_id: &AccountSetId| -> Option<SetMeta> {
snapshot
.meta
.get(set_id)
.or_else(|| overlay.and_then(|o| o.meta.get(set_id)))
.copied()
};
let parents_of = |set_id: &AccountSetId| -> &[AccountSetId] {
snapshot
.parents
.get(set_id)
.or_else(|| overlay.and_then(|o| o.parents.get(set_id)))
.map(Vec::as_slice)
.unwrap_or(&[])
};
let mut per_account: HashMap<AccountId, Vec<AccountSetId>> = HashMap::new();
for (account_id, set_id) in seeds {
per_account.entry(*account_id).or_default().push(*set_id);
}
let mut mappings = HashMap::new();
let mut non_ec: HashMap<AccountId, Vec<AccountSetId>> = HashMap::new();
for (account_id, seed_sets) in per_account {
let mut visited: HashSet<AccountSetId> = HashSet::new();
let mut queue: VecDeque<AccountSetId> = seed_sets.into();
let mut ancestors = Vec::new();
let mut non_ec_ancestors = Vec::new();
while let Some(set_id) = queue.pop_front() {
if !visited.insert(set_id) {
continue;
}
let meta = meta_of(&set_id)?;
if meta.journal_id == journal_id {
ancestors.push(set_id);
if !meta.eventually_consistent {
non_ec_ancestors.push(set_id);
}
}
queue.extend(parents_of(&set_id));
}
if !non_ec_ancestors.is_empty() {
non_ec.insert(account_id, non_ec_ancestors);
}
if !ancestors.is_empty() {
mappings.insert(account_id, ancestors);
}
}
let mut lock_pairs: Vec<(AccountSetId, &str)> = entry_account_ids
.iter()
.zip(entry_currencies.iter())
.flat_map(|(account_id, currency)| {
non_ec
.get(account_id)
.into_iter()
.flatten()
.map(move |set_id| (*set_id, *currency))
})
.collect();
lock_pairs.sort_unstable();
lock_pairs.dedup();
Some((mappings, lock_pairs.into_iter().unzip()))
}
fn spawn_refresh(&self) {
let inner = Arc::clone(&self.inner);
tokio::spawn(async move {
if let Err(error) = Self::refresh(&inner).await {
tracing::warn!(%error, "set_graph_cache refresh failed");
}
});
}
#[instrument(
level = "debug",
name = "cala_ledger.set_graph_cache.refresh",
skip_all,
fields(epoch = tracing::field::Empty, sets = tracing::field::Empty),
err(level = "warn")
)]
async fn refresh(inner: &SetGraphCacheInner) -> Result<(), AccountSetError> {
let Ok(_guard) = inner.refresh_lock.try_lock() else {
return Ok(());
};
let data = inner.repo.fetch_set_graph().await?;
let (parents, meta) = index_nodes(data.nodes);
let new = GraphSnapshot {
epoch: data.epoch,
parents,
meta,
};
tracing::Span::current().record("epoch", new.epoch);
tracing::Span::current().record("sets", new.meta.len());
let mut current = inner
.snapshot
.write()
.expect("set_graph_cache snapshot lock poisoned");
if new.epoch >= current.epoch {
*current = Arc::new(new);
}
Ok(())
}
}