use std::collections::{BTreeMap, HashMap};
use std::fmt::Write as _;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use liminal::durability::bridge::block_on;
use liminal::durability::settle::apply_routed_batch_settled;
use liminal::durability::{DurabilityError, DurableStore, RoutedBatchItem};
use super::log::OperationLog;
use super::operation_batch::OperationBatch;
use super::outbox::ConversationOutboxLimits;
use super::outbox_log::OutboxLog;
use super::restore_checkpoint_history::{
CheckpointHistoryNode, CheckpointHistoryRoot, HistoryNodeKind, ROOT_ENCODED_BYTES,
};
use super::restore_checkpoint_locator::CheckpointLocator;
use super::restore_from_checkpoint::{CheckpointDecision, load_checkpoint};
use super::state::ConversationAuthority;
use super::state::checkpoint_owner::ConversationAuthorityRestoreContext;
pub(super) const CHECKPOINT_BATCH_BYTES: usize = 16 * 1024 * 1024;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PublicationBound {
pub pinned: usize,
pub store: Option<usize>,
pub bound: usize,
}
impl PublicationBound {
#[must_use]
pub fn item_limit(&self) -> usize {
self.store.unwrap_or(usize::MAX)
}
}
#[cfg(test)]
pub(super) static TEST_BATCH_BYTES: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
#[cfg(test)]
fn pinned_batch_bytes() -> usize {
match TEST_BATCH_BYTES.load(Ordering::SeqCst) {
0 => CHECKPOINT_BATCH_BYTES,
set => set,
}
}
#[cfg(not(test))]
const fn pinned_batch_bytes() -> usize {
CHECKPOINT_BATCH_BYTES
}
pub fn publication_bound(store: &dyn DurableStore) -> PublicationBound {
let store_bound = store.max_request_bytes();
let pinned = pinned_batch_bytes();
PublicationBound {
pinned,
store: store_bound,
bound: store_bound.map_or(pinned, |admitted| admitted.min(pinned)),
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub(super) enum CheckpointPublishError {
#[error("checkpoint_item_exceeds_store_bound: key={key} bytes={bytes} bound={bound}")]
ItemExceedsStoreBound {
key: String,
bytes: usize,
bound: usize,
},
}
fn key_name(key: &[u8]) -> String {
match std::str::from_utf8(key) {
Ok(text) if text.chars().all(|c| !c.is_control()) => text.to_owned(),
_ => key.iter().fold(String::new(), |mut hex, byte| {
if write!(hex, "{byte:02x}").is_err() {
hex.push_str("??");
}
hex
}),
}
}
pub(super) fn split_history(
nodes: BTreeMap<Vec<u8>, Vec<u8>>,
bound: usize,
item_limit: usize,
) -> Result<Vec<Vec<RoutedBatchItem>>, CheckpointPublishError> {
let mut batches: Vec<Vec<RoutedBatchItem>> = Vec::new();
let mut current: Vec<RoutedBatchItem> = Vec::new();
let mut current_bytes = 0_usize;
for (key, value) in nodes {
let bytes = value.len();
if bytes > item_limit {
return Err(CheckpointPublishError::ItemExceedsStoreBound {
key: key_name(&key),
bytes,
bound: item_limit,
});
}
if current_bytes.saturating_add(bytes) > bound && !current.is_empty() {
batches.push(std::mem::take(&mut current));
current_bytes = 0;
}
current_bytes = current_bytes.saturating_add(bytes);
current.push(RoutedBatchItem {
key,
expected: None,
value,
ttl: None,
});
}
if !current.is_empty() {
batches.push(current);
}
Ok(batches)
}
pub(super) fn batch_bytes(items: &[RoutedBatchItem]) -> usize {
items.iter().map(|item| item.value.len()).sum()
}
pub(super) fn commit_count(
node_bytes: impl IntoIterator<Item = usize>,
tail: usize,
bound: usize,
) -> usize {
let mut batches = 0_usize;
let mut current_bytes = 0_usize;
let mut open = false;
for bytes in node_bytes {
if open && current_bytes.saturating_add(bytes) > bound {
batches += 1;
current_bytes = 0;
}
open = true;
current_bytes = current_bytes.saturating_add(bytes);
}
if open {
batches += 1;
}
if tail == 0 {
return batches;
}
if open && current_bytes.saturating_add(tail) <= bound {
batches
} else {
batches + 1
}
}
pub(super) struct CheckpointPublication {
pub(super) conversation_id: u64,
pub(super) route: String,
pub(super) history: Vec<Vec<RoutedBatchItem>>,
pub(super) final_batch: OperationBatch,
pub(super) bound: PublicationBound,
pub(super) bytes: usize,
}
impl CheckpointPublication {
pub(super) fn new(
conversation_id: u64,
route: &str,
mut nodes: BTreeMap<Vec<u8>, Vec<u8>>,
mut final_batch: OperationBatch,
bound: PublicationBound,
) -> Result<Self, PublicationLayoutError> {
nodes.extend(
final_batch
.take_body_parts()
.into_iter()
.map(|part| (part.key, part.value)),
);
let tail = final_batch.value_bytes();
if tail > bound.item_limit() {
return Err(CheckpointPublishError::ItemExceedsStoreBound {
key: "checkpoint body".to_owned(),
bytes: tail,
bound: bound.item_limit(),
}
.into());
}
let node_bytes: usize = nodes.values().map(Vec::len).sum();
let mut history = split_history(nodes, bound.bound, bound.item_limit())?;
if let Some(last) = history.last() {
if batch_bytes(last).saturating_add(tail) <= bound.bound {
if let Some(last) = history.pop() {
for item in last {
final_batch.stage(item)?;
}
}
}
}
Ok(Self {
conversation_id,
route: route.to_owned(),
history,
final_batch,
bound,
bytes: node_bytes.saturating_add(tail),
})
}
pub(super) fn commits(&self) -> usize {
self.history.len().saturating_add(1)
}
}
#[derive(Debug, thiserror::Error)]
pub(super) enum PublicationLayoutError {
#[error(transparent)]
Refused(#[from] CheckpointPublishError),
#[error(transparent)]
Store(#[from] DurabilityError),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublicationRefusal {
pub batch_index: usize,
pub bytes: usize,
pub bound: usize,
pub refusal: String,
}
impl std::fmt::Display for PublicationRefusal {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"checkpoint publication refused at batch {}: {} bytes against a bound of {}: {}",
self.batch_index, self.bytes, self.bound, self.refusal
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PublishedCheckpoint {
pub commits: usize,
pub bytes: usize,
}
pub(super) async fn commit_publication(
store: &dyn DurableStore,
publication: CheckpointPublication,
) -> Result<PublishedCheckpoint, PublicationRefusal> {
let CheckpointPublication {
conversation_id,
route,
history,
mut final_batch,
bound,
bytes,
} = publication;
let total = history.len().saturating_add(1);
for (index, items) in history.into_iter().enumerate() {
let batch_bytes = batch_bytes(&items);
if let Err(error) = apply_routed_batch_settled(store, &route, items).await {
return Err(PublicationRefusal {
batch_index: index.saturating_add(1),
bytes: batch_bytes,
bound: bound.bound,
refusal: error.to_string(),
});
}
}
let final_bytes = final_batch.value_bytes();
let prune = final_batch.take_prune();
if let Err(error) = final_batch.commit(store, || ()).await {
return Err(PublicationRefusal {
batch_index: total,
bytes: final_bytes,
bound: bound.bound,
refusal: error.to_string(),
});
}
if let Some(plan) = prune {
super::prune::report_cut(store, conversation_id, &plan).await;
}
super::restore_checkpoint_history::retire_v1_history(store, conversation_id).await;
Ok(PublishedCheckpoint {
commits: total,
bytes,
})
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PublicationState {
InFlight,
Owed(PublicationRefusal),
Published(PublishedCheckpoint),
}
#[derive(Debug, Clone, Default)]
pub struct PublicationLedger {
states: Arc<Mutex<HashMap<u64, PublicationState>>>,
bound_logged: Arc<AtomicBool>,
}
impl PublicationLedger {
fn states(&self) -> Option<std::sync::MutexGuard<'_, HashMap<u64, PublicationState>>> {
self.states.lock().map_or_else(
|_| {
tracing::warn!(
"checkpoint publication ledger lock is poisoned; no publication is handed off and the owner keeps serving from its replay"
);
None
},
Some,
)
}
pub(super) fn begin(&self, conversation_id: u64) -> bool {
let Some(mut states) = self.states() else {
return false;
};
if matches!(
states.get(&conversation_id),
Some(PublicationState::InFlight)
) {
return false;
}
states.insert(conversation_id, PublicationState::InFlight);
true
}
pub(super) fn in_flight(&self, conversation_id: u64) -> bool {
self.states().is_none_or(|states| {
matches!(
states.get(&conversation_id),
Some(PublicationState::InFlight)
)
})
}
pub(super) fn finish(&self, conversation_id: u64, state: PublicationState) {
if let Some(mut states) = self.states() {
states.insert(conversation_id, state);
}
}
#[cfg(test)]
pub fn checkpoint_owed(&self, conversation_id: u64) -> bool {
self.states().is_some_and(|states| {
matches!(
states.get(&conversation_id),
Some(PublicationState::Owed(_))
)
})
}
#[cfg(test)]
pub fn state(&self, conversation_id: u64) -> Option<PublicationState> {
self.states()
.and_then(|states| states.get(&conversation_id).cloned())
}
pub(super) fn log_bound_once(&self, bound: PublicationBound) {
if bound.bound >= bound.pinned {
return;
}
if self
.bound_logged
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok()
{
tracing::info!(
pinned_batch_bytes = bound.pinned,
store_max_request_bytes = bound.bound,
"checkpoint batch bound: the store admits less than the pinned batch bound, so the store's bound is used"
);
}
}
}
#[derive(Debug, Default)]
pub struct CheckpointPublisher {
ledger: PublicationLedger,
threads: Mutex<Vec<std::thread::JoinHandle<()>>>,
crossing_observer: std::sync::OnceLock<Arc<dyn crate::server::connection::ConnectionNotifier>>,
}
impl CheckpointPublisher {
pub fn install_crossing_observer(
&self,
observer: Arc<dyn crate::server::connection::ConnectionNotifier>,
) {
if self.crossing_observer.set(observer).is_err() {
tracing::warn!("checkpoint crossing observer was already installed; the first stays");
}
}
pub fn crossing_committed(&self, crossing: &crate::server::connection::CheckpointCrossing) {
if let Some(observer) = self.crossing_observer.get() {
observer.on_checkpoint_crossing(crossing);
}
}
pub const fn ledger(&self) -> &PublicationLedger {
&self.ledger
}
pub(super) fn spawn(
&self,
store: Arc<dyn DurableStore>,
publication: CheckpointPublication,
) -> Result<(), std::io::Error> {
let conversation_id = publication.conversation_id;
let ledger = self.ledger.clone();
let commits = publication.commits();
let bytes = publication.bytes;
let bound = publication.bound;
tracing::info!(
conversation_id,
batches = commits,
bytes,
bound = bound.bound,
"checkpoint publication handed off the owner's write path"
);
let spawned = std::thread::Builder::new()
.name(format!("checkpoint-publish-{conversation_id}"))
.spawn(move || {
let outcome = block_on(commit_publication(store.as_ref(), publication));
let state = match outcome {
Ok(Ok(published)) => {
tracing::info!(
conversation_id,
commits = published.commits,
bytes = published.bytes,
"checkpoint published in batches"
);
PublicationState::Published(published)
}
Ok(Err(refusal)) => {
log_refusal(conversation_id, &refusal);
PublicationState::Owed(refusal)
}
Err(bridge) => {
let refusal = PublicationRefusal {
batch_index: 0,
bytes,
bound: bound.bound,
refusal: bridge.to_string(),
};
log_refusal(conversation_id, &refusal);
PublicationState::Owed(refusal)
}
};
ledger.finish(conversation_id, state);
});
match spawned {
Ok(handle) => {
if let Ok(mut threads) = self.threads.lock() {
threads.push(handle);
} else {
tracing::warn!(
conversation_id,
"checkpoint publisher handle list lock is poisoned; the publication runs unjoined"
);
}
Ok(())
}
Err(error) => {
let refusal = PublicationRefusal {
batch_index: 0,
bytes,
bound: bound.bound,
refusal: format!("publisher thread could not start: {error}"),
};
log_refusal(conversation_id, &refusal);
self.ledger
.finish(conversation_id, PublicationState::Owed(refusal));
Err(error)
}
}
}
pub(super) fn refuse_before_submission(
&self,
conversation_id: u64,
refusal: PublicationRefusal,
) {
log_refusal(conversation_id, &refusal);
self.ledger
.finish(conversation_id, PublicationState::Owed(refusal));
}
#[cfg(test)]
pub fn join_all(&self) -> Result<(), String> {
let handles: Vec<_> = std::mem::take(
&mut *self
.threads
.lock()
.map_err(|_| "the publisher handle list lock is poisoned".to_owned())?,
);
for handle in handles {
handle
.join()
.map_err(|_| "a checkpoint publisher thread panicked".to_owned())?;
}
Ok(())
}
}
fn log_refusal(conversation_id: u64, refusal: &PublicationRefusal) {
tracing::warn!(
conversation_id,
batch_index = refusal.batch_index,
bytes = refusal.bytes,
bound = refusal.bound,
refusal = %refusal.refusal,
"checkpoint publication refused by the store; the owner keeps serving from its replay, the checkpoint is owed, and the next crossing retries"
);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CheckpointCensus {
pub conversation_id: u64,
pub nodes: BTreeMap<&'static str, (usize, usize)>,
pub history_nodes: usize,
pub history_bytes: usize,
pub tokens: usize,
pub committed_admissions: usize,
pub slots: usize,
pub attach_provenance: usize,
pub retired: usize,
pub locator: Option<(u64, u64)>,
pub log_rows_behind: u64,
pub log_rows_ahead: u64,
pub outbox_rows_behind: u64,
pub outbox_rows_ahead: u64,
pub body: Result<usize, String>,
pub bound: PublicationBound,
pub batches: usize,
}
impl CheckpointCensus {
pub fn lines(&self) -> Vec<String> {
let mut lines = vec![format!("conversation: {}", self.conversation_id)];
for (kind, (count, bytes)) in &self.nodes {
lines.push(format!("{kind}: {count} nodes, {bytes} bytes"));
}
lines.push(format!(
"history to publish: {} nodes, {} bytes",
self.history_nodes, self.history_bytes
));
lines.push(format!("tokens: {}", self.tokens));
lines.push(format!(
"committed admissions: {}",
self.committed_admissions
));
lines.push(format!("slots: {}", self.slots));
lines.push(format!(
"attach provenance entries: {}",
self.attach_provenance
));
lines.push(format!("retired identities: {}", self.retired));
match self.locator {
Some((covered, extension)) => lines.push(format!(
"latest locator: covered log sequence {covered}, extension boundary {extension}"
)),
None => lines.push("latest locator: absent".to_owned()),
}
lines.push(format!(
"operation log rows: {} behind the locator, {} ahead",
self.log_rows_behind, self.log_rows_ahead
));
lines.push(format!(
"outbox rows: {} behind the locator, {} ahead",
self.outbox_rows_behind, self.outbox_rows_ahead
));
match &self.body {
Ok(bytes) => lines.push(format!("checkpoint body and locator: {bytes} bytes")),
Err(refusal) => lines.push(format!("checkpoint body: not capturable: {refusal}")),
}
lines.push(format!(
"batch bound: {} bytes (pinned {}, store {})",
self.bound.bound,
self.bound.pinned,
self.bound
.store
.map_or_else(|| "unbounded".to_owned(), |bound| bound.to_string())
));
lines.push(format!("batches a publication commits: {}", self.batches));
lines
}
}
#[derive(Debug, thiserror::Error)]
pub enum CensusError {
#[error("the conversation could not be replayed: {0}")]
Replay(String),
#[error("the checkpoint locator could not be read: {0}")]
Locator(String),
#[error("the history could not be captured: {0}")]
Capture(String),
#[error("a captured node could not be decoded: {0}")]
Node(String),
#[error("the outbox limits could not be built: {0}")]
Limits(String),
}
pub fn checkpoint_census(
store: &Arc<dyn DurableStore>,
config: &crate::config::types::ParticipantConfig,
conversation_id: u64,
) -> Result<CheckpointCensus, CensusError> {
let limits =
ConversationOutboxLimits::try_new(config.max_retained_record_rows, config.identity_slots)
.map_err(|error| CensusError::Limits(error.to_string()))?;
let log = OperationLog::new(Arc::clone(store), conversation_id);
let outbox_log = OutboxLog::new(Arc::clone(store), conversation_id);
let (mut owner, audit, schema, prefix) = block_on(ConversationAuthority::replay_for_rebuild(
conversation_id,
&log,
&outbox_log,
config,
limits,
))
.map_err(|error| CensusError::Replay(error.to_string()))?
.map_err(|error| CensusError::Replay(error.to_string()))?;
let (root, locator) = latest_root_and_locator(store, config, conversation_id)?;
let history = block_on(
super::restore_checkpoint_history_capture::capture_identity_history(
store.as_ref(),
&owner,
root,
),
)
.map_err(|error| CensusError::Capture(error.to_string()))?
.map_err(|error| CensusError::Capture(error.to_string()))?;
let nodes = count_nodes_by_kind(&history)?;
let bound = publication_bound(store.as_ref());
let body = owner
.retain_checkpoint_metadata(history.root, &audit)
.map_err(|error| error.to_string())
.and_then(|()| checkpoint_body_bytes(store.as_ref(), &owner, schema, history.root, prefix));
let batches = commit_count(
history.nodes().values().map(Vec::len),
body.clone().unwrap_or(0),
bound.bound,
);
let log_head = owner.next_log_sequence;
let outbox_head = owner.outbox.as_ref().map_or(
0,
super::outbox::ConversationOutbox::next_extension_sequence,
);
let (log_rows_behind, outbox_rows_behind) = locator.map_or((0, 0), |(covered, extension)| {
(
covered.saturating_add(1).min(log_head),
extension.min(outbox_head),
)
});
Ok(CheckpointCensus {
conversation_id,
nodes,
history_nodes: history.node_count(),
history_bytes: history.byte_total(),
tokens: owner.tokens.len(),
committed_admissions: owner.committed_admissions.len(),
slots: owner.slots.len(),
attach_provenance: owner
.slots
.values()
.map(|slot| slot.attach_provenance.len())
.sum(),
retired: owner.retired.len(),
locator,
log_rows_behind,
log_rows_ahead: log_head.saturating_sub(log_rows_behind),
outbox_rows_behind,
outbox_rows_ahead: outbox_head.saturating_sub(outbox_rows_behind),
body,
bound,
batches,
})
}
fn latest_root_and_locator(
store: &Arc<dyn DurableStore>,
config: &crate::config::types::ParticipantConfig,
conversation_id: u64,
) -> Result<(CheckpointHistoryRoot, Option<(u64, u64)>), CensusError> {
let loaded = block_on(
load_checkpoint::<ConversationAuthority, ROOT_ENCODED_BYTES>(
store.as_ref(),
conversation_id,
|positions, root| ConversationAuthorityRestoreContext::new(*positions, root, config),
|locator, root| {
CheckpointHistoryRoot::decode_canonical(root, locator.conversation_id())
.map(|_| ())
.map_err(
|source| super::restore_checkpoint::CheckpointError::LiveState {
source: Box::new(source),
},
)
},
),
)
.map_err(|error| CensusError::Locator(error.to_string()))?
.map_err(|error| CensusError::Locator(error.to_string()))?;
Ok(match loaded.decision {
CheckpointDecision::Restore(checkpoint) => {
let root =
CheckpointHistoryRoot::decode_canonical(checkpoint.history_root(), conversation_id)
.map_err(|error| CensusError::Locator(error.to_string()))?;
let locator = loaded
.observed_locator
.as_deref()
.map(|bytes| CheckpointLocator::decode_canonical(bytes, conversation_id))
.transpose()
.map_err(|error| CensusError::Locator(error.to_string()))?
.map(|locator| (locator.covered_sequence(), locator.extension_boundary()));
(root, locator)
}
CheckpointDecision::Rebuild(_) => (CheckpointHistoryRoot::empty(conversation_id), None),
})
}
fn count_nodes_by_kind(
history: &super::restore_checkpoint_history_capture::PreparedCheckpointHistory,
) -> Result<BTreeMap<&'static str, (usize, usize)>, CensusError> {
let mut nodes: BTreeMap<&'static str, (usize, usize)> = BTreeMap::new();
for kind in [
HistoryNodeKind::RetainedValue,
HistoryNodeKind::RetiredHistory,
HistoryNodeKind::IndexPage,
] {
nodes.insert(kind.name(), (0, 0));
}
for (key, body) in history.nodes() {
let node_root = CheckpointHistoryRoot::from_routed_node_key(key).ok_or_else(|| {
CensusError::Node(format!("not a history node key: {}", key_name(key)))
})?;
let node = CheckpointHistoryNode::decode_canonical(body, node_root)
.map_err(|error| CensusError::Node(error.to_string()))?;
let entry = nodes.entry(node.kind().name()).or_insert((0, 0));
entry.0 = entry.0.saturating_add(1);
entry.1 = entry.1.saturating_add(body.len());
}
Ok(nodes)
}
fn checkpoint_body_bytes(
store: &dyn DurableStore,
owner: &ConversationAuthority,
schema: super::restore_checkpoint::CheckpointSchema,
root: CheckpointHistoryRoot,
prefix: [u8; 32],
) -> Result<usize, String> {
let extension = owner.outbox.as_ref().map_or(
0,
super::outbox::ConversationOutbox::next_extension_sequence,
);
let positions = owner
.checkpoint_positions(schema, owner.next_log_sequence, extension)
.map_err(|error| error.to_string())?;
let route =
super::restore_checkpoint_locator::checkpoint_conversation_route(owner.conversation_id);
let mut scratch = OperationBatch::new(&route);
block_on(scratch.stage_borrowed_checkpoint(
store,
positions,
&root.encode_canonical(),
owner,
prefix,
))
.map_err(|error| error.to_string())?
.map_err(|error| error.to_string())?;
Ok(scratch.value_bytes())
}