use chrono::{DateTime, NaiveDate, Utc};
use std::collections::{HashMap, HashSet};
use job::{JobType, Jobs};
use obix::{
out::{
EventCtx, EventSubscription, FlushOp, Handled, HandlerStreamStatus, OutboxEventHandler,
OutboxEventJobConfig, PersistentOutboxEvent, RegisteredEventHandler,
},
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<RegisteredEventHandler<OutboxEventPayload, CalaMailboxTables>, LedgerError> {
Ok(outbox
.register_event_handler(
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>,
entry_ids: Vec<EntryId>,
}
#[derive(Default)]
struct EcRollupBatch {
txns: Vec<PendingTx>,
entries: HashMap<TransactionId, Vec<EntryValues>>,
}
impl EcRollupBatch {
fn push_tx(&mut self, tx: PendingTx) {
self.txns.push(tx);
}
fn push_entry(&mut self, entry: EntryValues) {
self.entries
.entry(entry.transaction_id)
.or_default()
.push(entry);
}
fn missing_entry_ids(&self) -> Vec<EntryId> {
self.txns
.iter()
.flat_map(|tx| {
let collected: HashSet<EntryId> = self
.entries
.get(&tx.id)
.map(|entries| entries.iter().map(|e| e.id).collect())
.unwrap_or_default();
tx.entry_ids
.iter()
.copied()
.filter(move |id| !collected.contains(id))
})
.collect()
}
fn into_rollup_txns(self, mut fetched: HashMap<EntryId, Entry>) -> Vec<EcRollupTxn> {
let EcRollupBatch { txns, mut entries } = self;
txns.into_iter()
.map(|tx| {
let mut entry_values = entries.remove(&tx.id).unwrap_or_default();
if entry_values.len() != tx.entry_ids.len() {
entry_values.extend(
tx.entry_ids
.iter()
.filter_map(|id| fetched.remove(id))
.map(Entry::into_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 OutboxEventHandler<OutboxEventPayload> for EcBalanceRollupHandler {
const SUBSCRIPTION: EventSubscription = EventSubscription::PersistentOnly;
type Batch = EcRollupBatch;
async fn handle_persistent<'inv>(
&self,
ctx: EventCtx<'inv, Self::Batch>,
event: &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,
entry_ids: transaction.entry_ids.clone(),
};
Ok(ctx.collect_with(|batch| batch.push_tx(tx)))
}
Some(OutboxEventPayload::EntryCreated { entry }) => {
let entry = entry.clone();
Ok(ctx.collect_with(|batch| batch.push_entry(entry)))
}
_ => 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(op, &missing_ids).await?
};
let rollup_txns = batch.into_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>,
}
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,
entry_ids: t.entry_ids.clone(),
});
}
for e in &entries {
batch.push_entry(e.clone());
}
let _missing = batch.missing_entry_ids();
let _rollup = batch.into_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: RegisteredEventHandler<OutboxEventPayload, CalaMailboxTables>,
}
impl EcRollupStatus {
pub(crate) fn new(
status: HandlerStreamStatus,
handle: RegisteredEventHandler<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
}
}