use super::{JobIdentity, shot_seed};
use crate::{Error, Result};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::collections::{BTreeMap, BTreeSet};
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EventLimits {
pub max_events: usize,
pub max_pending: usize,
pub max_checkpoint_bytes: usize,
}
impl EventLimits {
fn validate(self) -> Result<()> {
if self.max_events == 0
|| self.max_events > 1_048_576
|| self.max_pending == 0
|| self.max_pending > self.max_events
|| self.max_checkpoint_bytes == 0
{
return Err(Error::InvalidValue("invalid event budgets".into()));
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Event<E> {
pub id: u64,
pub tick: u64,
pub cause: Option<u64>,
pub operation: E,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Snapshot<S, E> {
version: u32,
identity: JobIdentity,
shot: u64,
limits: EventLimits,
watermark: u64,
next_id: u64,
state: S,
pending: Vec<Event<E>>,
completed: Vec<Event<E>>,
}
pub struct EventMachine<S, E> {
snapshot: Snapshot<S, E>,
queue: BTreeMap<(u64, u64), Event<E>>,
}
impl<S: Clone + Serialize + DeserializeOwned, E: Clone + Serialize + DeserializeOwned>
EventMachine<S, E>
{
pub fn new(
identity: JobIdentity,
shot: u64,
state: S,
initial: Vec<(u64, E)>,
limits: EventLimits,
) -> Result<Self> {
limits.validate()?;
identity.validate()?;
if shot >= identity.attempts
|| initial.len() > limits.max_pending
|| initial.len() > limits.max_events
{
return Err(Error::InvalidValue(
"event initial shot/queue exceeds budget".into(),
));
}
let queue = initial
.into_iter()
.enumerate()
.map(|(id, (tick, operation))| {
(
(tick, id as u64),
Event {
id: id as u64,
tick,
cause: None,
operation,
},
)
})
.collect::<BTreeMap<_, _>>();
Ok(Self {
snapshot: Snapshot {
version: 1,
identity,
shot,
limits,
watermark: 0,
next_id: queue.len() as u64,
state,
pending: vec![],
completed: vec![],
},
queue,
})
}
pub fn state(&self) -> &S {
&self.snapshot.state
}
pub fn completed(&self) -> &[Event<E>] {
&self.snapshot.completed
}
pub fn pending(&self) -> impl Iterator<Item = &Event<E>> {
self.queue.values()
}
pub fn watermark(&self) -> u64 {
self.snapshot.watermark
}
pub fn advance<F>(&mut self, through: u64, max_steps: usize, mut apply: F) -> Result<usize>
where
F: FnMut(&mut S, &Event<E>, u64) -> Result<Vec<(u64, E)>>,
{
if through < self.snapshot.watermark || max_steps == 0 {
return Err(Error::InvalidValue(
"invalid causal window/step budget".into(),
));
}
let mut steps = 0;
while steps < max_steps {
let Some((&key, event)) = self.queue.first_key_value() else {
break;
};
if event.tick > through {
break;
}
let mut candidate = self.snapshot.state.clone();
let domain = shot_seed(
self.snapshot.identity.seed,
self.snapshot.shot,
0x4556_454e_5453_0001,
);
let seed = shot_seed(domain, event.id, 0);
let emitted = apply(&mut candidate, event, seed)?;
let total = self
.snapshot
.next_id
.checked_add(emitted.len() as u64)
.ok_or(Error::Overflow)?;
let pending = self.queue.len() - 1 + emitted.len();
if total > self.snapshot.limits.max_events as u64
|| pending > self.snapshot.limits.max_pending
{
return Err(Error::InvalidValue(
"event generation or pending queue budget exhausted".into(),
));
}
if emitted.iter().any(|(tick, _)| *tick < event.tick) {
return Err(Error::InvalidValue(
"event emitted before its causal parent".into(),
));
}
let event = self.queue.remove(&key).unwrap();
for (tick, operation) in emitted {
let id = self.snapshot.next_id;
self.snapshot.next_id += 1;
self.queue.insert(
(tick, id),
Event {
id,
tick,
cause: Some(event.id),
operation,
},
);
}
self.snapshot.state = candidate;
self.snapshot.watermark = event.tick;
self.snapshot.completed.push(event);
steps += 1;
}
if self
.queue
.first_key_value()
.is_none_or(|(key, _)| key.0 > through)
{
self.snapshot.watermark = through;
}
Ok(steps)
}
pub fn checkpoint(&self) -> Result<Vec<u8>> {
struct Bounded {
bytes: Vec<u8>,
limit: usize,
}
impl std::io::Write for Bounded {
fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
if b.len() > self.limit.saturating_sub(self.bytes.len()) {
return Err(std::io::Error::other("event checkpoint budget"));
}
self.bytes.extend_from_slice(b);
Ok(b.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let mut snapshot = self.snapshot.clone();
snapshot.pending = self.queue.values().cloned().collect();
let mut out = Bounded {
bytes: vec![],
limit: snapshot.limits.max_checkpoint_bytes,
};
serde_json::to_writer(&mut out, &snapshot)
.map_err(|e| Error::InvalidValue(e.to_string()))?;
Ok(out.bytes)
}
pub fn restore<F>(
bytes: &[u8],
expected: &JobIdentity,
max_bytes: usize,
validate_state: F,
) -> Result<Self>
where
F: FnOnce(&S) -> Result<()>,
{
if bytes.len() > max_bytes {
return Err(Error::InvalidValue(
"event checkpoint admission budget".into(),
));
}
let mut s: Snapshot<S, E> =
serde_json::from_slice(bytes).map_err(|e| Error::InvalidValue(e.to_string()))?;
s.limits.validate()?;
s.identity.validate()?;
if s.version != 1
|| &s.identity != expected
|| s.shot >= s.identity.attempts
|| bytes.len() > s.limits.max_checkpoint_bytes
|| s.next_id > s.limits.max_events as u64
|| s.pending.len() > s.limits.max_pending
|| s.pending.len() + s.completed.len() != s.next_id as usize
{
return Err(Error::InvalidValue(
"event checkpoint identity/count mismatch".into(),
));
}
let mut known = BTreeMap::new();
let mut ids = BTreeSet::new();
let mut previous = None;
for event in &s.completed {
let key = (event.tick, event.id);
if previous.is_some_and(|p| p >= key) || event.tick > s.watermark {
return Err(Error::InvalidValue("invalid committed event order".into()));
}
Self::validate_event(event, &known, &mut ids, s.next_id)?;
known.insert(event.id, event.tick);
previous = Some(key);
}
let mut queue = BTreeMap::new();
for event in s.pending.drain(..) {
if event.tick < s.watermark || previous.is_some_and(|key| key >= (event.tick, event.id))
{
return Err(Error::InvalidValue(
"pending event precedes the committed frontier".into(),
));
}
Self::validate_event(&event, &known, &mut ids, s.next_id)?;
queue.insert((event.tick, event.id), event);
}
validate_state(&s.state)?;
Ok(Self { snapshot: s, queue })
}
fn validate_event(
event: &Event<E>,
known: &BTreeMap<u64, u64>,
ids: &mut BTreeSet<u64>,
next_id: u64,
) -> Result<()> {
if event.id >= next_id
|| !ids.insert(event.id)
|| event.cause.is_some_and(|id| {
id >= event.id || known.get(&id).is_none_or(|&tick| tick > event.tick)
})
{
return Err(Error::InvalidValue(
"invalid event identity/causal parent".into(),
));
}
Ok(())
}
}