use routers_network::Entry;
use thiserror::Error;
use crate::bus::adapter::Publisher;
use crate::event::VehicleId;
use crate::orchestrator::commit::{CommitError, Committer};
use crate::orchestrator::scheduler::CheckpointState;
use crate::protocol::ids::{Revision, SCHEMA_VERSION};
use crate::protocol::output::{CommittedOutput, ResetReason};
use crate::store::checkpoint::{CheckpointStore, StoredCheckpointState, VehicleCheckpoint};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RecoveryReport {
pub partition: u16,
pub frontier: Option<u64>,
pub prepared_found: usize,
pub prepared_finished: usize,
pub prepared_failed: Vec<VehicleId>,
}
#[derive(Debug, Error)]
pub enum RecoveryError<SE> {
#[error("the checkpoint store rejected a recovery read")]
Store(#[source] SE),
#[error("a surviving prepared commit could not be completed during restore")]
Unfinished(#[source] CommitError<SE>),
}
#[must_use]
pub fn expected_start(frontier: Option<u64>) -> Option<u64> {
frontier.map(|sequence| sequence.saturating_add(1))
}
pub async fn recover_partition<E, S, P>(
store: &S,
committer: &Committer<E, S, P>,
partition: u16,
) -> Result<RecoveryReport, RecoveryError<S::Error>>
where
E: Entry + serde::de::DeserializeOwned,
S: CheckpointStore,
P: Publisher<CommittedOutput<E>>,
{
let frontier = store
.frontier(partition)
.await
.map_err(RecoveryError::Store)?
.map(|f| f.sequence);
let prepared = store
.list_prepared(partition)
.await
.map_err(RecoveryError::Store)?;
let prepared_found = prepared.len();
let mut prepared_finished = 0;
let mut prepared_failed = Vec::new();
for (vehicle, record) in prepared {
match committer.finish_prepared(vehicle, partition, record).await {
Ok(_) => prepared_finished += 1,
Err(_) => prepared_failed.push(vehicle),
}
}
Ok(RecoveryReport {
partition,
frontier,
prepared_found,
prepared_finished,
prepared_failed,
})
}
#[derive(Clone, Debug)]
pub struct Restored<E: Entry> {
pub checkpoint: CheckpointState<E>,
pub reset: Option<ResetReason>,
pub prior: Option<Revision>,
}
pub async fn restore_vehicle<E, S, P>(
store: &S,
vehicle: VehicleId,
partition: u16,
committer: &Committer<E, S, P>,
) -> Result<Restored<E>, RecoveryError<S::Error>>
where
E: Entry + serde::de::DeserializeOwned,
S: CheckpointStore,
P: Publisher<CommittedOutput<E>>,
{
let (stored, prepared) = store.load(vehicle).await.map_err(RecoveryError::Store)?;
let stored = if let Some(prepared) = prepared {
committer
.finish_prepared(vehicle, partition, prepared)
.await
.map_err(RecoveryError::Unfinished)?;
store.load(vehicle).await.map_err(RecoveryError::Store)?.0
} else {
stored
};
let stored = match stored {
StoredCheckpointState::NeverSeen => {
return Ok(Restored {
checkpoint: CheckpointState::Absent,
reset: None,
prior: None,
});
}
StoredCheckpointState::Expired { revision } => {
return Ok(Restored {
checkpoint: CheckpointState::Absent,
reset: Some(ResetReason::StateLost),
prior: Some(revision),
});
}
StoredCheckpointState::Present(stored) => stored,
};
let prior = Some(stored.revision);
match VehicleCheckpoint::<E>::decode(&stored.bytes) {
Ok(checkpoint) if checkpoint.schema == SCHEMA_VERSION => Ok(Restored {
checkpoint: CheckpointState::Present(checkpoint),
reset: None,
prior,
}),
Ok(_) | Err(_) => Ok(Restored {
checkpoint: CheckpointState::Absent,
reset: Some(ResetReason::StateLost),
prior,
}),
}
}
#[cfg(test)]
mod tests {
use core::time::Duration;
use routers_network::mock::{MockEntryId, MockNetwork, MockNetworkBuilder};
use routers_transition::Matcher;
use routers_transition::costing::{
CostingStrategies, DefaultEmissionCost, DefaultTransitionCost,
};
use routers_transition::layer::generation::StandardGenerator;
use routers_transition::matcher::{Origin, Trip};
use routers_transition::weigh::AllCompute;
use super::*;
use crate::bus::Wire;
use crate::bus::adapter::PublishError;
use crate::bus::memory::{MemoryBus, MemoryPublisher};
use crate::event::Payload;
use crate::orchestrator::commit::{CommitConfig, Decision, Plan, plan};
use crate::orchestrator::dispatch::{DispatchConfig, build_context};
use crate::protocol::ids::{
GraphVersion, JobId, ObservationId, RegionId, Revision, SchemaVersion, SegmentId,
};
use crate::protocol::job::JobIdentity;
use crate::protocol::output::{CommittedOutput, OutputKind, TerminalReason};
use crate::store::checkpoint::{
CommitPhase, MemoryCheckpointStore, PartitionFrontier, PrepareOutcome, PreparedCommit,
};
use crate::topology::output::output_subject;
type E = MockEntryId;
type Costing = CostingStrategies<DefaultEmissionCost, DefaultTransitionCost, MockEntryId>;
const PARTITION: u16 = 7;
const SUBJECT: &str = "events.matched.v1.p.7";
const TRACE_START_US: i64 = 1_775_000_000_000_000;
fn region() -> RegionId {
RegionId::new("r1").unwrap()
}
fn graph() -> GraphVersion {
GraphVersion::new("g1").unwrap()
}
fn obs(seq: u64) -> ObservationId {
ObservationId {
partition: PARTITION,
sequence: seq,
}
}
fn committer(
store: MemoryCheckpointStore,
bus: &MemoryBus,
) -> Committer<E, MemoryCheckpointStore, MemoryPublisher<CommittedOutput<E>>> {
Committer::new(
store,
bus.publisher::<CommittedOutput<E>>(),
CommitConfig {
publish_attempts: 4,
backoff: Duration::ZERO,
},
)
}
fn terminal_plan(seq: u64) -> Plan<E> {
let id = JobIdentity {
schema: SCHEMA_VERSION,
vehicle_id: VehicleId(1),
observation: obs(seq),
base: None,
graph: graph(),
region: region(),
};
plan(
None,
Decision::Terminal {
job: id.local_decision_id(),
identity: id,
reason: TerminalReason::Unanchored,
closes_segment: false,
segment: SegmentId(seq),
},
obs(seq),
®ion(),
&graph(),
1,
)
}
async fn stage_unpublished(
store: &MemoryCheckpointStore,
bus: &MemoryBus,
vehicle: VehicleId,
seq: u64,
) {
bus.fail_next_publish(PublishError::Failed(anyhow::anyhow!("bus down")));
let c = committer(store.clone(), bus);
c.commit(vehicle, PARTITION, terminal_plan(seq), None, obs(seq))
.await
.expect_err("the armed publish failure fails the commit");
let (checkpoint, prepared) = store.load(vehicle).await.unwrap();
assert!(checkpoint.is_none(), "nothing is committed yet");
assert!(!prepared.unwrap().is_published(), "and it is unpublished");
}
async fn install_checkpoint_bytes(
store: &MemoryCheckpointStore,
bus: &MemoryBus,
vehicle: VehicleId,
seq: u64,
bytes: Vec<u8>,
) {
let output = CommittedOutput::<E>::new(
JobId(seq as u128),
vehicle,
obs(seq),
Revision(seq),
SegmentId(seq),
OutputKind::Terminal {
reason: TerminalReason::Unanchored,
closes_segment: false,
},
);
let entries: Vec<(String, String, Vec<u8>)> = vec![(
output_subject(u64::from(PARTITION)),
output.msg_id(),
output.encode().unwrap(),
)];
let prepared = PreparedCommit {
output: output.id,
output_subject: output_subject(u64::from(PARTITION)),
output_bytes: postcard::to_allocvec(&entries).unwrap(),
next_checkpoint: bytes,
next_revision: Revision(seq),
next_segment: SegmentId(seq),
expected_base: None,
phase: CommitPhase::Prepared,
raw: obs(seq),
};
assert_eq!(
store
.prepare(vehicle, PARTITION, prepared.clone())
.await
.unwrap(),
PrepareOutcome::Prepared,
);
committer(store.clone(), bus)
.finish_prepared(vehicle, PARTITION, prepared)
.await
.expect("seed commit installs the checkpoint bytes");
}
fn bent_road() -> MockNetwork {
MockNetworkBuilder::new()
.node(1, geo::point!(x: -118.15, y: 34.15))
.node(2, geo::point!(x: -118.16, y: 34.15))
.node(3, geo::point!(x: -118.17, y: 34.15))
.node(4, geo::point!(x: -118.17, y: 34.14))
.node(5, geo::point!(x: -118.18, y: 34.14))
.edge(1, 2)
.edge(2, 3)
.edge(3, 4)
.edge(4, 5)
.build()
}
fn trace_origins() -> Vec<Origin> {
[
geo::point!(x: -118.151, y: 34.1503),
geo::point!(x: -118.155, y: 34.1503),
geo::point!(x: -118.165, y: 34.1503),
geo::point!(x: -118.170, y: 34.1490),
geo::point!(x: -118.172, y: 34.1403),
geo::point!(x: -118.179, y: 34.1403),
]
.into_iter()
.enumerate()
.map(|(i, pt)| Origin::new(pt, TRACE_START_US + i as i64 * 5_000_000))
.collect()
}
fn trip_with(origins: &[Origin]) -> Trip<E> {
let net = bent_road();
let costing = Costing::default();
let generator = StandardGenerator::new(&net, &costing.emission);
let m = Matcher::new(&net, &costing, generator, AllCompute::default(), &());
let mut trip = m.begin();
for &origin in origins {
m.push(&mut trip, origin).expect("observation must anchor");
}
trip
}
fn checkpoint_at(revision: u64, segment: u64, schema: SchemaVersion) -> VehicleCheckpoint<E> {
VehicleCheckpoint {
trip: trip_with(&trace_origins()),
last_input: obs(revision),
revision: Revision(revision),
segment: SegmentId(segment),
finalized_through: None,
graph: graph(),
schema,
region: region(),
routing_version: 1,
}
}
fn continuing_head() -> Payload {
let last = *trace_origins().last().unwrap();
Payload {
vehicle_id: VehicleId(1),
timestamp: chrono::DateTime::from_timestamp_micros(last.timestamp + 5_000_000).unwrap(),
point: geo::point!(x: -118.180, y: 34.1403),
}
}
fn identity_from(
checkpoint: Option<&VehicleCheckpoint<E>>,
observation: ObservationId,
head: &Payload,
) -> JobIdentity {
let built = build_context(checkpoint, observation, head, &DispatchConfig::default())
.expect("recovery fixtures use a monotonic head");
JobIdentity {
schema: SCHEMA_VERSION,
vehicle_id: VehicleId(1),
observation,
base: built.base,
graph: graph(),
region: region(),
}
}
#[test]
fn expected_start_is_the_frontier_successor_or_the_head() {
assert_eq!(expected_start(None), None);
assert_eq!(expected_start(Some(0)), Some(1));
assert_eq!(expected_start(Some(41)), Some(42));
assert_eq!(expected_start(Some(u64::MAX)), Some(u64::MAX));
}
#[tokio::test]
async fn recovery_reports_the_stored_frontier() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
store
.set_frontier(PartitionFrontier {
partition: PARTITION,
sequence: 99,
})
.await
.unwrap();
let report = recover_partition(&store, &committer(store.clone(), &bus), PARTITION)
.await
.unwrap();
assert_eq!(report.frontier, Some(99));
assert_eq!(expected_start(report.frontier), Some(100));
assert_eq!(report.prepared_found, 0);
}
#[tokio::test]
async fn a_prepared_unpublished_record_is_published_and_promoted() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
stage_unpublished(&store, &bus, vehicle, 100).await;
let report = recover_partition(&store, &committer(store.clone(), &bus), PARTITION)
.await
.unwrap();
assert_eq!(report.prepared_found, 1);
assert_eq!(report.prepared_finished, 1);
assert!(report.prepared_failed.is_empty());
assert_eq!(bus.published(SUBJECT).len(), 1);
let (checkpoint, prepared) = store.load(vehicle).await.unwrap();
assert_eq!(checkpoint.unwrap().revision, Revision(100));
assert!(prepared.is_none(), "the prepared record is promoted away");
}
#[tokio::test]
async fn a_published_unpromoted_record_is_republished_and_promoted() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
store.fail_next(crate::store::checkpoint::Op::Promote);
committer(store.clone(), &bus)
.commit(vehicle, PARTITION, terminal_plan(100), None, obs(100))
.await
.expect_err("promotion fails");
assert_eq!(bus.published(SUBJECT).len(), 1);
assert!(store.load(vehicle).await.unwrap().1.unwrap().is_published());
let report = recover_partition(&store, &committer(store.clone(), &bus), PARTITION)
.await
.unwrap();
assert_eq!(report.prepared_finished, 1);
assert_eq!(bus.published(SUBJECT).len(), 1, "republish deduplicates");
assert_eq!(
store.load(vehicle).await.unwrap().0.unwrap().revision,
Revision(100)
);
assert!(store.load(vehicle).await.unwrap().1.is_none());
}
#[tokio::test]
async fn a_bus_failure_leaves_the_record_and_reports_the_vehicle_blocked() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
stage_unpublished(&store, &bus, vehicle, 100).await;
bus.fail_next_publish(PublishError::Failed(anyhow::anyhow!("still down")));
let report = recover_partition(&store, &committer(store.clone(), &bus), PARTITION)
.await
.unwrap();
assert_eq!(report.prepared_found, 1);
assert_eq!(report.prepared_finished, 0);
assert_eq!(report.prepared_failed, vec![vehicle]);
let (checkpoint, prepared) = store.load(vehicle).await.unwrap();
assert!(checkpoint.is_none());
assert!(prepared.is_some(), "the prepared record is retained");
}
#[tokio::test]
async fn a_read_failure_aborts_recovery() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
store.fail_next(crate::store::checkpoint::Op::Frontier);
let err = recover_partition(&store, &committer(store.clone(), &bus), PARTITION)
.await
.expect_err("a failed frontier read aborts");
assert!(matches!(err, RecoveryError::Store(_)), "got {err:?}");
}
#[tokio::test]
async fn restore_of_an_unknown_vehicle_is_absent_without_reset() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let restored: Restored<E> = restore_vehicle(
&store,
VehicleId(1),
PARTITION,
&committer(store.clone(), &bus),
)
.await
.unwrap();
assert!(matches!(restored.checkpoint, CheckpointState::Absent));
assert_eq!(restored.reset, None);
assert_eq!(restored.prior, None, "nothing stored, so no prior revision");
}
#[tokio::test]
async fn restore_of_a_clean_checkpoint_is_present_without_reset() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
let cp = checkpoint_at(100, 100, SCHEMA_VERSION);
install_checkpoint_bytes(&store, &bus, vehicle, 100, cp.encode().unwrap()).await;
let restored: Restored<E> =
restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.unwrap();
match restored.checkpoint {
CheckpointState::Present(got) => assert_eq!(got.revision, Revision(100)),
other => panic!("expected Present, got {other:?}"),
}
assert_eq!(restored.reset, None);
assert_eq!(restored.prior, Some(Revision(100)));
}
#[tokio::test]
async fn restore_of_a_corrupt_checkpoint_is_state_lost() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
install_checkpoint_bytes(&store, &bus, vehicle, 100, b"not a checkpoint".to_vec()).await;
let restored: Restored<E> =
restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.unwrap();
assert!(matches!(restored.checkpoint, CheckpointState::Absent));
assert_eq!(restored.reset, Some(ResetReason::StateLost));
assert_eq!(
restored.prior,
Some(Revision(100)),
"state was lost but the store still holds the stale revision to CAS against",
);
}
#[tokio::test]
async fn restore_of_a_stale_schema_checkpoint_is_state_lost() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
let stale = checkpoint_at(100, 100, SchemaVersion(SCHEMA_VERSION.0 + 1));
install_checkpoint_bytes(&store, &bus, vehicle, 100, stale.encode().unwrap()).await;
let restored: Restored<E> =
restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.unwrap();
assert!(matches!(restored.checkpoint, CheckpointState::Absent));
assert_eq!(restored.reset, Some(ResetReason::StateLost));
}
#[tokio::test]
async fn restore_resolves_a_surviving_prepared_commit_first() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
stage_unpublished(&store, &bus, vehicle, 100).await;
let restored: Restored<E> =
restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.unwrap();
assert_eq!(bus.published(SUBJECT).len(), 1);
match restored.checkpoint {
CheckpointState::Present(got) => assert_eq!(got.revision, Revision(100)),
other => panic!("expected Present after finishing the prepared commit, got {other:?}"),
}
assert_eq!(restored.reset, None);
assert!(store.load(vehicle).await.unwrap().1.is_none());
}
#[tokio::test]
async fn restore_blocks_when_a_prepared_commit_cannot_finish() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
stage_unpublished(&store, &bus, vehicle, 100).await;
bus.fail_next_publish(PublishError::Failed(anyhow::anyhow!("still down")));
let err = restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.expect_err("an unresolvable prepared commit blocks restore");
assert!(matches!(err, RecoveryError::Unfinished(_)), "got {err:?}");
}
#[tokio::test]
async fn a_job_rebuilt_from_the_restored_checkpoint_keeps_its_id() {
let store = MemoryCheckpointStore::new();
let bus = MemoryBus::new();
let vehicle = VehicleId(1);
let head = continuing_head();
let observation = obs(200);
let before = checkpoint_at(100, 100, SCHEMA_VERSION);
let id_before = identity_from(Some(&before), observation, &head).local_decision_id();
install_checkpoint_bytes(&store, &bus, vehicle, 100, before.encode().unwrap()).await;
let restored: Restored<E> =
restore_vehicle(&store, vehicle, PARTITION, &committer(store.clone(), &bus))
.await
.unwrap();
let CheckpointState::Present(after) = restored.checkpoint else {
panic!("the checkpoint must restore");
};
let id_after = identity_from(Some(&after), observation, &head).local_decision_id();
assert_eq!(id_before, id_after);
}
}