use std::collections::{BTreeMap, BTreeSet};
use super::super::outbox_log::{OutboxRow, ProjectedRecord, decode_row, encode_row};
use super::super::restore_checkpoint::CheckpointError;
use super::super::state::checkpoint_body::{CheckpointBodyDecoder, CheckpointBodyEncoder, invalid};
use super::validation::{validate_batch_shape, validate_record_snapshot};
use super::{ConversationOutbox, ConversationOutboxLimits, LiveRecord, SourceBatch};
const MAGIC: &[u8; 5] = b"LCBX\x01";
fn wrapped(error: impl std::error::Error + Send + Sync + 'static) -> CheckpointError {
CheckpointError::LiveState {
source: Box::new(error),
}
}
fn write_set(
encoder: &mut CheckpointBodyEncoder,
field: &'static str,
set: &BTreeSet<u64>,
) -> Result<(), CheckpointError> {
encoder.count(field, set.len())?;
for value in set {
encoder.u64(*value);
}
Ok(())
}
fn read_key(
decoder: &mut CheckpointBodyDecoder<'_>,
field: &'static str,
previous: &mut Option<u64>,
) -> Result<u64, CheckpointError> {
let key = decoder.u64(field)?;
if previous.is_some_and(|last| last >= key) {
return Err(invalid(field));
}
*previous = Some(key);
Ok(key)
}
fn read_set(
decoder: &mut CheckpointBodyDecoder<'_>,
field: &'static str,
) -> Result<BTreeSet<u64>, CheckpointError> {
let count = decoder.count(field, 8)?;
let mut set = BTreeSet::new();
let mut previous = None;
for _ in 0..count {
set.insert(read_key(decoder, field, &mut previous)?);
}
Ok(set)
}
fn write_pairs(
encoder: &mut CheckpointBodyEncoder,
field: &'static str,
map: &BTreeMap<u64, u64>,
) -> Result<(), CheckpointError> {
encoder.count(field, map.len())?;
for (key, value) in map {
encoder.u64(*key);
encoder.u64(*value);
}
Ok(())
}
fn read_pairs(
decoder: &mut CheckpointBodyDecoder<'_>,
field: &'static str,
) -> Result<BTreeMap<u64, u64>, CheckpointError> {
let count = decoder.count(field, 16)?;
let mut map = BTreeMap::new();
let mut previous = None;
for _ in 0..count {
let key = read_key(decoder, field, &mut previous)?;
map.insert(key, decoder.u64(field)?);
}
Ok(map)
}
impl ConversationOutbox {
pub(in crate::server::participant::production) fn encode_checkpoint(
&self,
) -> Result<Vec<u8>, CheckpointError> {
let mut encoder = CheckpointBodyEncoder::new();
encoder.raw(MAGIC);
encoder.u64(self.conversation_id);
encoder.count("outbox records", self.records.len())?;
for (sequence, record) in &self.records {
encoder.u64(*sequence);
write_set(&mut encoder, "record recipients", &record.recipients)?;
encoder.u64(record.encoded_push_bytes);
}
encoder.count("outbox source batches", self.source_batches.len())?;
for (source, batch) in &self.source_batches {
encoder.u64(*source);
encoder.bytes("source batch", &batch.canonical)?;
encoder.u64(batch.first_delivery_seq);
encoder.u64(batch.last_delivery_seq);
}
encoder.count("outbox ack sources", self.ack_sources.len())?;
for (source, (participant, through)) in &self.ack_sources {
encoder.u64(*source);
encoder.u64(*participant);
encoder.u64(*through);
}
encoder.count("outbox obligations", self.all_obligations.len())?;
for (participant, sequences) in &self.all_obligations {
encoder.u64(*participant);
write_set(&mut encoder, "obligation sequences", sequences)?;
}
write_pairs(&mut encoder, "ack frontiers", &self.ack_frontiers)?;
write_pairs(&mut encoder, "marker frontiers", &self.marker_ack_frontiers)?;
write_pairs(
&mut encoder,
"next live obligations",
&self.next_live_obligations,
)?;
write_set(&mut encoder, "retired recipients", &self.retired)?;
encoder.u64(self.highest_delivery_seq);
encoder.u64(self.next_extension_sequence);
encoder.u64(self.live_recipient_obligations);
encoder.u64(self.charged_bytes);
Ok(encoder.finish())
}
pub(in crate::server::participant::production) fn decode_checkpoint(
bytes: &[u8],
conversation_id: u64,
extension_boundary: u64,
limits: ConversationOutboxLimits,
) -> Result<Self, CheckpointError> {
let mut decoder = CheckpointBodyDecoder::new(bytes);
validate_outbox_header(&mut decoder, conversation_id)?;
let mut saved_records = decode_saved_records(&mut decoder)?;
let DecodedSourceBatches {
source_batches,
projections,
last_delivery,
} = decode_source_batches(&mut decoder, conversation_id)?;
let DecodedAckSources {
ack_sources,
ack_participants,
} = decode_ack_sources(&mut decoder)?;
let indexes = decode_outbox_indexes(&mut decoder)?;
decoder.finish()?;
indexes.validate_boundaries(
extension_boundary,
limits,
last_delivery,
&ack_sources,
&ack_participants,
)?;
let DecodedOutboxIndexes {
all_obligations,
ack_frontiers,
marker_ack_frontiers,
next_live_obligations,
retired,
highest_delivery_seq,
next_extension_sequence,
live_recipient_obligations,
charged_bytes,
} = indexes;
let RestoredLiveRecords {
records,
measured_obligations,
measured_next,
measured_live,
measured_charge,
} = restore_live_records(
conversation_id,
projections,
&mut saved_records,
&retired,
&ack_frontiers,
)?;
if !saved_records.is_empty()
|| measured_live != live_recipient_obligations
|| measured_charge != charged_bytes
|| measured_next != next_live_obligations
{
return Err(invalid("outbox aggregate accounting"));
}
validate_obligation_index(
&all_obligations,
&retired,
&ack_frontiers,
highest_delivery_seq,
&measured_obligations,
)?;
let latest_source = source_batches.last_key_value().map(|(source, _)| *source);
for (source, batch) in &source_batches {
if Some(*source) != latest_source
&& records
.range(batch.first_delivery_seq..=batch.last_delivery_seq)
.next()
.is_none()
{
return Err(invalid("reclaimed source batch"));
}
}
let owner = Self {
limits,
conversation_id,
records,
source_batches,
ack_sources,
all_obligations,
ack_frontiers,
marker_ack_frontiers,
next_live_obligations,
retired,
highest_delivery_seq,
next_extension_sequence,
live_recipient_obligations,
charged_bytes,
};
if owner.encode_checkpoint()? != bytes {
return Err(invalid("canonical outbox body"));
}
Self::from_restore_parts(
conversation_id,
extension_boundary,
owner.into_restore_parts(),
limits,
)
.map_err(|source| CheckpointError::LiveState {
source: Box::new(source),
})
}
}
fn validate_outbox_header(
decoder: &mut CheckpointBodyDecoder<'_>,
conversation_id: u64,
) -> Result<(), CheckpointError> {
if decoder.take("outbox version", MAGIC.len())? != MAGIC {
return Err(invalid("outbox version"));
}
if decoder.u64("outbox conversation")? != conversation_id {
return Err(invalid("outbox conversation"));
}
Ok(())
}
type SavedRecords = BTreeMap<u64, (BTreeSet<u64>, u64)>;
struct DecodedSourceBatches {
source_batches: BTreeMap<u64, SourceBatch>,
projections: BTreeMap<u64, ProjectedRecord>,
last_delivery: Option<u64>,
}
struct DecodedAckSources {
ack_sources: BTreeMap<u64, (u64, u64)>,
ack_participants: BTreeSet<u64>,
}
struct RestoredLiveRecords {
records: BTreeMap<u64, LiveRecord>,
measured_obligations: BTreeMap<u64, BTreeSet<u64>>,
measured_next: BTreeMap<u64, u64>,
measured_live: u64,
measured_charge: u64,
}
fn decode_saved_records(
decoder: &mut CheckpointBodyDecoder<'_>,
) -> Result<SavedRecords, CheckpointError> {
let record_count = decoder.count("outbox records", 20)?;
let mut saved_records = BTreeMap::new();
let mut previous = None;
for _ in 0..record_count {
let sequence = read_key(decoder, "outbox records", &mut previous)?;
let recipients = read_set(decoder, "record recipients")?;
let charge = decoder.u64("record charge")?;
if recipients.is_empty() {
return Err(invalid("empty live record"));
}
saved_records.insert(sequence, (recipients, charge));
}
Ok(saved_records)
}
fn decode_source_batches(
decoder: &mut CheckpointBodyDecoder<'_>,
conversation_id: u64,
) -> Result<DecodedSourceBatches, CheckpointError> {
let source_count = decoder.count("outbox source batches", 28)?;
let mut source_batches = BTreeMap::new();
let mut projections = BTreeMap::new();
let mut previous = None;
let mut last_delivery = None;
for _ in 0..source_count {
let source = read_key(decoder, "source batches", &mut previous)?;
let canonical = decoder.bytes("source batch")?;
let first_delivery_seq = decoder.u64("first source delivery")?;
let last_delivery_seq = decoder.u64("last source delivery")?;
let row = decode_row(canonical).map_err(wrapped)?;
if encode_row(&row).map_err(wrapped)? != canonical {
return Err(invalid("canonical source batch"));
}
let OutboxRow::Produced(batch) = row else {
return Err(invalid("source batch kind"));
};
validate_batch_shape(&batch).map_err(wrapped)?;
if batch.source_log_sequence != source
|| batch.ordered_records.first().map(|r| r.delivery_seq) != Some(first_delivery_seq)
|| batch.ordered_records.last().map(|r| r.delivery_seq) != Some(last_delivery_seq)
{
return Err(invalid("source batch span"));
}
for record in batch.ordered_records {
validate_record_snapshot(&record, &BTreeSet::new()).map_err(wrapped)?;
if record.delivery_seq == 0
|| last_delivery.is_some_and(|last| last >= record.delivery_seq)
{
return Err(invalid("source delivery order"));
}
let measured = super::super::outbox_log::ProjectedRecord::try_new(
conversation_id,
record.delivery_seq,
record.body.clone(),
record.recipients.clone(),
record.sender,
)
.map_err(wrapped)?
.encoded_push_bytes;
if measured != record.encoded_push_bytes {
return Err(invalid("source push charge"));
}
last_delivery = Some(record.delivery_seq);
projections.insert(record.delivery_seq, record);
}
source_batches.insert(
source,
SourceBatch {
canonical: canonical.to_vec(),
first_delivery_seq,
last_delivery_seq,
},
);
}
Ok(DecodedSourceBatches {
source_batches,
projections,
last_delivery,
})
}
fn decode_ack_sources(
decoder: &mut CheckpointBodyDecoder<'_>,
) -> Result<DecodedAckSources, CheckpointError> {
let ack_count = decoder.count("outbox ack sources", 24)?;
let mut ack_sources = BTreeMap::new();
let mut previous = None;
let mut ack_participants = BTreeSet::new();
for _ in 0..ack_count {
let source = read_key(decoder, "ack source order", &mut previous)?;
let participant = decoder.u64("ack participant")?;
let through = decoder.u64("ack through")?;
if through == 0 || !ack_participants.insert(participant) {
return Err(invalid("ack source participant"));
}
ack_sources.insert(source, (participant, through));
}
Ok(DecodedAckSources {
ack_sources,
ack_participants,
})
}
fn restore_live_records(
conversation_id: u64,
projections: BTreeMap<u64, ProjectedRecord>,
saved_records: &mut SavedRecords,
retired: &BTreeSet<u64>,
ack_frontiers: &BTreeMap<u64, u64>,
) -> Result<RestoredLiveRecords, CheckpointError> {
let mut records = BTreeMap::new();
let mut measured_obligations = BTreeMap::<u64, BTreeSet<u64>>::new();
let mut measured_next = BTreeMap::new();
let mut measured_live = 0_u64;
let mut measured_charge = 0_u64;
for (sequence, projection) in projections {
let recipients: BTreeSet<_> = projection
.recipients
.iter()
.copied()
.filter(|participant| {
!retired.contains(participant)
&& sequence > ack_frontiers.get(participant).copied().unwrap_or(0)
})
.collect();
if recipients.is_empty() {
if saved_records.contains_key(&sequence) {
return Err(invalid("discharged record retained"));
}
continue;
}
let (saved_recipients, saved_charge) = saved_records
.remove(&sequence)
.ok_or_else(|| invalid("missing live record"))?;
if saved_recipients != recipients || saved_charge != projection.encoded_push_bytes {
return Err(invalid("live record snapshot"));
}
for participant in &recipients {
measured_obligations
.entry(*participant)
.or_default()
.insert(sequence);
measured_next.entry(*participant).or_insert(sequence);
}
measured_live = measured_live
.checked_add(
u64::try_from(recipients.len()).map_err(|_| invalid("live obligation count"))?,
)
.ok_or_else(|| invalid("live obligation count"))?;
measured_charge = measured_charge
.checked_add(saved_charge)
.ok_or_else(|| invalid("live record charge"))?;
records.insert(
sequence,
LiveRecord {
delivery: projection.into_delivery(conversation_id),
recipients,
encoded_push_bytes: saved_charge,
},
);
}
Ok(RestoredLiveRecords {
records,
measured_obligations,
measured_next,
measured_live,
measured_charge,
})
}
fn validate_obligation_index(
all_obligations: &BTreeMap<u64, BTreeSet<u64>>,
retired: &BTreeSet<u64>,
ack_frontiers: &BTreeMap<u64, u64>,
highest_delivery_seq: u64,
measured_obligations: &BTreeMap<u64, BTreeSet<u64>>,
) -> Result<(), CheckpointError> {
for (participant, sequences) in all_obligations {
let frontier = ack_frontiers.get(participant).copied().unwrap_or(0);
if sequences
.iter()
.any(|sequence| *sequence <= frontier || *sequence > highest_delivery_seq)
{
return Err(invalid("obligation frontier"));
}
if !retired.contains(participant)
&& measured_obligations.get(participant) != Some(sequences)
{
return Err(invalid("live obligation index"));
}
}
if measured_obligations
.keys()
.any(|participant| !all_obligations.contains_key(participant))
{
return Err(invalid("missing obligation index"));
}
Ok(())
}
struct DecodedOutboxIndexes {
all_obligations: BTreeMap<u64, BTreeSet<u64>>,
ack_frontiers: BTreeMap<u64, u64>,
marker_ack_frontiers: BTreeMap<u64, u64>,
next_live_obligations: BTreeMap<u64, u64>,
retired: BTreeSet<u64>,
highest_delivery_seq: u64,
next_extension_sequence: u64,
live_recipient_obligations: u64,
charged_bytes: u64,
}
fn decode_outbox_indexes(
decoder: &mut CheckpointBodyDecoder<'_>,
) -> Result<DecodedOutboxIndexes, CheckpointError> {
let obligation_count = decoder.count("outbox obligations", 12)?;
let mut all_obligations = BTreeMap::new();
let mut previous = None;
for _ in 0..obligation_count {
let participant = read_key(decoder, "obligation participant", &mut previous)?;
let sequences = read_set(decoder, "obligation sequences")?;
if sequences.is_empty() {
return Err(invalid("empty obligation index"));
}
all_obligations.insert(participant, sequences);
}
let ack_frontiers = read_pairs(decoder, "ack frontiers")?;
let marker_ack_frontiers = read_pairs(decoder, "marker frontiers")?;
let next_live_obligations = read_pairs(decoder, "next live obligations")?;
let retired = read_set(decoder, "retired recipients")?;
let highest_delivery_seq = decoder.u64("highest delivery")?;
let next_extension_sequence = decoder.u64("extension boundary")?;
let live_recipient_obligations = decoder.u64("live obligations")?;
let charged_bytes = decoder.u64("charged bytes")?;
Ok(DecodedOutboxIndexes {
all_obligations,
ack_frontiers,
marker_ack_frontiers,
next_live_obligations,
retired,
highest_delivery_seq,
next_extension_sequence,
live_recipient_obligations,
charged_bytes,
})
}
impl DecodedOutboxIndexes {
fn validate_boundaries(
&self,
extension_boundary: u64,
limits: ConversationOutboxLimits,
last_delivery: Option<u64>,
ack_sources: &BTreeMap<u64, (u64, u64)>,
ack_participants: &BTreeSet<u64>,
) -> Result<(), CheckpointError> {
if self.next_extension_sequence != extension_boundary
|| last_delivery.unwrap_or(0) != self.highest_delivery_seq
|| self.live_recipient_obligations > limits.max_live_recipient_obligations
{
return Err(invalid("outbox checkpoint boundary"));
}
for (participant, through) in ack_sources.values() {
if self.ack_frontiers.get(participant) != Some(through) {
return Err(invalid("ack source frontier"));
}
}
if self
.ack_frontiers
.values()
.chain(self.marker_ack_frontiers.values())
.any(|through| *through > self.highest_delivery_seq)
{
return Err(invalid("ack frontier boundary"));
}
for (participant, frontier) in &self.ack_frontiers {
if *frontier != 0 && !ack_participants.contains(participant) {
return Err(invalid("ack frontier source"));
}
}
Ok(())
}
}