use chrono::{DateTime, NaiveDate, Utc};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use job::{JobType, Jobs};
use obix::{
out::{
EventCtx, FlushOp, Handled, OutboxEventJobConfig, PersistentOutboxEvent,
SingletonSubscriber, StreamSelection, Subscription, SubscriptionStreamStatus,
},
EventSequence,
};
use cala_types::entry::EntryValues;
use crate::{
balance::{Balances, EcRollupTxn},
entry::{Entries, Entry},
ledger::error::LedgerError,
outbox::{CalaMailboxTables, ObixOutbox, OutboxEventPayload},
primitives::{EntryId, JournalId, TransactionId},
};
const EC_BALANCE_ROLLUP_JOB: JobType = JobType::new("cala.ec_balance_rollup");
const MAX_EVENTS_PER_BATCH: usize = 1_000;
pub(crate) async fn register_ec_balance_rollup(
jobs: &mut Jobs,
outbox: &ObixOutbox,
balances: &Balances,
entries: &Entries,
) -> Result<Subscription<OutboxEventPayload, CalaMailboxTables>, LedgerError> {
Ok(outbox
.register_singleton_subscriber(
jobs,
OutboxEventJobConfig::new(EC_BALANCE_ROLLUP_JOB)
.with_max_batch_size(MAX_EVENTS_PER_BATCH),
EcBalanceRollupHandler {
balances: balances.clone(),
entries: entries.clone(),
},
)
.await?)
}
struct PendingTx {
id: TransactionId,
journal_id: JournalId,
effective: NaiveDate,
created_at: DateTime<Utc>,
event: Arc<PersistentOutboxEvent<OutboxEventPayload>>,
}
impl PendingTx {
fn entry_ids(&self) -> &[EntryId] {
match &self.event.payload {
Some(OutboxEventPayload::TransactionCreated { transaction }) => &transaction.entry_ids,
_ => unreachable!(
"PendingTx is only built from a TransactionCreated event in handle_persistent"
),
}
}
}
fn entry_of(event: &PersistentOutboxEvent<OutboxEventPayload>) -> &EntryValues {
match &event.payload {
Some(OutboxEventPayload::EntryCreated { entry }) => entry,
_ => unreachable!(
"EcRollupBatch::entries only ever holds EntryCreated events, pushed in handle_persistent"
),
}
}
#[derive(Default)]
struct EcRollupBatch {
txns: Vec<PendingTx>,
entries: HashMap<TransactionId, Vec<Arc<PersistentOutboxEvent<OutboxEventPayload>>>>,
}
impl EcRollupBatch {
fn push_tx(&mut self, tx: PendingTx) {
self.txns.push(tx);
}
fn push_entry(&mut self, event: Arc<PersistentOutboxEvent<OutboxEventPayload>>) {
self.entries
.entry(entry_of(&event).transaction_id)
.or_default()
.push(event);
}
fn missing_entry_ids(&self) -> Vec<EntryId> {
self.txns
.iter()
.flat_map(|tx| {
let collected: HashSet<EntryId> = self
.entries
.get(&tx.id)
.map(|events| events.iter().map(|e| entry_of(e).id).collect())
.unwrap_or_default();
tx.entry_ids()
.iter()
.copied()
.filter(move |id| !collected.contains(id))
})
.collect()
}
fn rollup_txns<'a>(&'a self, fetched: &'a HashMap<EntryId, Entry>) -> Vec<EcRollupTxn<'a>> {
self.txns
.iter()
.map(|tx| {
let entry_ids = tx.entry_ids();
let mut entry_values: Vec<&EntryValues> = self
.entries
.get(&tx.id)
.map(|events| events.iter().map(|e| entry_of(e)).collect())
.unwrap_or_default();
if entry_values.len() != entry_ids.len() {
entry_values.extend(
entry_ids
.iter()
.filter_map(|id| fetched.get(id))
.map(Entry::values),
);
}
entry_values.sort_by_key(|e| e.sequence);
EcRollupTxn {
journal_id: tx.journal_id,
effective: tx.effective,
created_at: tx.created_at,
entries: entry_values,
}
})
.collect()
}
}
struct EcBalanceRollupHandler {
balances: Balances,
entries: Entries,
}
impl SingletonSubscriber<OutboxEventPayload> for EcBalanceRollupHandler {
const SUBSCRIPTION: StreamSelection = StreamSelection::PersistentOnly;
type Batch = EcRollupBatch;
async fn handle_persistent<'inv>(
&self,
ctx: EventCtx<'inv, Self::Batch>,
event: &Arc<PersistentOutboxEvent<OutboxEventPayload>>,
) -> Result<Handled<'inv>, Box<dyn std::error::Error + Send + Sync>> {
match &event.payload {
Some(OutboxEventPayload::TransactionCreated { transaction }) => {
let tx = PendingTx {
id: transaction.id,
journal_id: transaction.journal_id,
effective: transaction.effective,
created_at: transaction.created_at,
event: Arc::clone(event),
};
Ok(ctx.collect_with(|batch| batch.push_tx(tx)))
}
Some(OutboxEventPayload::EntryCreated { .. }) => {
let event = Arc::clone(event);
Ok(ctx.collect_with(|batch| batch.push_entry(event)))
}
_ => Ok(ctx.skip()),
}
}
#[tracing::instrument(
name = "cala_ledger.ec_rollup.flush",
skip_all,
fields(txns_count = batch.txns.len()),
err(level = "warn")
)]
async fn flush(
&self,
op: &mut FlushOp<'_>,
batch: Self::Batch,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let missing_ids = batch.missing_entry_ids();
let fetched = if missing_ids.is_empty() {
HashMap::new()
} else {
self.entries.find_all_in_op(&mut *op, &missing_ids).await?
};
let rollup_txns = batch.rollup_txns(&fetched);
self.balances.apply_ec_rollup_in_op(op, rollup_txns).await?;
Ok(())
}
}
#[cfg(feature = "fuzz")]
mod __fuzz {
use super::*;
use serde::Deserialize;
#[derive(Deserialize)]
struct FuzzTx {
id: TransactionId,
journal_id: JournalId,
effective: NaiveDate,
created_at: DateTime<Utc>,
entry_ids: Vec<EntryId>,
}
fn event(payload: OutboxEventPayload) -> Arc<PersistentOutboxEvent<OutboxEventPayload>> {
Arc::new(PersistentOutboxEvent {
id: obix::out::OutboxEventId::new(),
sequence: obix::EventSequence::from(0u64),
payload: Some(payload),
tracing_context: None,
recorded_at: Utc::now(),
commit_group: obix::CommitGroupId::from(0i64),
})
}
pub fn fuzz_batch(data: &[u8]) {
let parts: Vec<&[u8]> = data.split(|&b| b == 0xFF).collect();
if parts.len() < 2 {
return;
}
let Ok(txs) = serde_json::from_slice::<Vec<FuzzTx>>(parts[0]) else {
return;
};
let Ok(entries) = serde_json::from_slice::<Vec<EntryValues>>(parts[1]) else {
return;
};
let mut batch = EcRollupBatch::default();
for t in &txs {
batch.push_tx(PendingTx {
id: t.id,
journal_id: t.journal_id,
effective: t.effective,
created_at: t.created_at,
event: event(OutboxEventPayload::TransactionCreated {
transaction: cala_types::transaction::TransactionValues {
id: t.id,
journal_id: t.journal_id,
effective: t.effective,
created_at: t.created_at,
entry_ids: t.entry_ids.clone(),
version: 1,
modified_at: t.created_at,
tx_template_id: crate::primitives::TxTemplateId::new(),
correlation_id: String::new(),
external_id: None,
description: None,
metadata: None,
},
}),
});
}
for e in &entries {
batch.push_entry(event(OutboxEventPayload::EntryCreated { entry: e.clone() }));
}
let _missing = batch.missing_entry_ids();
let _rollup = batch.rollup_txns(&HashMap::<EntryId, Entry>::new());
}
}
#[cfg(feature = "fuzz")]
pub use __fuzz::fuzz_batch;
#[derive(Debug, Clone)]
pub struct EcRollupStatus {
pub applied: EventSequence,
pub frontier: EventSequence,
handle: Subscription<OutboxEventPayload, CalaMailboxTables>,
}
impl EcRollupStatus {
pub(crate) fn new(
status: SubscriptionStreamStatus,
handle: Subscription<OutboxEventPayload, CalaMailboxTables>,
) -> Self {
Self {
applied: status.checkpoint,
frontier: status.frontier,
handle,
}
}
#[tracing::instrument(
level = "debug",
name = "cala_ledger.ec_rollup_status.refresh",
skip_all,
fields(frontier = %self.frontier, applied, lag)
)]
pub async fn refresh(&mut self) -> Result<(), LedgerError> {
self.applied = self.handle.load().await?.checkpoint();
let span = tracing::Span::current();
span.record("applied", u64::from(self.applied));
span.record("lag", self.lag());
Ok(())
}
pub async fn await_completion(&self, timeout: std::time::Duration) -> Result<(), LedgerError> {
self.handle.await_sequence(self.frontier, timeout).await?;
Ok(())
}
pub fn lag(&self) -> u64 {
u64::from(self.frontier).saturating_sub(u64::from(self.applied))
}
pub fn is_caught_up(&self) -> bool {
self.applied >= self.frontier
}
}