use alloc::vec::Vec;
use bytes::Bytes;
use mnesis::Version;
use thiserror::Error;
use crate::envelope::{PendingBatch, PendingEnvelope, PersistedEnvelope};
use crate::error::AppendError;
use crate::store::{RawEventStore, Store};
use crate::stream_id::StreamKey;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Atomicity {
WholeChunk,
PerStream,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamOutcome {
Complete { version: Version },
Corrupt { reached: Option<Version> },
Mismatch {
reached: Option<Version>,
got: Version,
},
}
impl StreamOutcome {
#[must_use]
pub const fn is_complete(&self) -> bool {
matches!(self, Self::Complete { .. })
}
#[must_use]
pub const fn reached(&self) -> Option<Version> {
match self {
Self::Complete { version } => Some(*version),
Self::Corrupt { reached } | Self::Mismatch { reached, .. } => *reached,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamReport {
pub stream: StreamKey,
pub outcome: StreamOutcome,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImportReport {
streams: Vec<StreamReport>,
}
impl ImportReport {
#[must_use]
pub const fn new(streams: Vec<StreamReport>) -> Self {
Self { streams }
}
#[must_use]
pub fn streams(&self) -> &[StreamReport] {
&self.streams
}
pub fn unfinished(&self) -> impl Iterator<Item = &StreamReport> {
self.streams.iter().filter(|s| !s.outcome.is_complete())
}
#[must_use]
pub fn all_complete(&self) -> bool {
self.streams.iter().all(|s| s.outcome.is_complete())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum AbortReason {
#[error("block failed checksum")]
Corrupt,
#[error("version mismatch (expected {expected}, got {got})")]
Mismatch { expected: Version, got: Version },
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ImportError<E> {
#[error("chunk aborted at stream {stream}: {reason}")]
Aborted {
stream: StreamKey,
reason: AbortReason,
},
#[error(transparent)]
Store(E),
#[error("version overflow")]
VersionOverflow,
}
#[derive(Debug, Clone)]
pub struct StreamSection {
pub origin: Bytes,
pub blocks: Vec<ImportBlock>,
}
#[derive(Debug, Clone)]
pub enum ImportBlock {
Event(PersistedEnvelope),
Corrupt,
}
#[derive(Debug, Clone)]
pub struct PlannedAppend {
pub target: StreamKey,
pub expected_version: Option<Version>,
pub head: PendingEnvelope,
pub tail: Vec<PendingEnvelope>,
}
impl PlannedAppend {
#[must_use]
pub fn batch(&self) -> PendingBatch<'_> {
PendingBatch::from_parts(&self.head, &self.tail)
}
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum AtomicAppendError<E> {
#[error("atomic append conflict at write {index}: actual head {actual:?}")]
Conflict {
index: usize,
actual: Option<Version>,
},
#[error("atomic append store error: {0}")]
Store(#[source] E),
}
pub trait AtomicAppend: RawEventStore {
fn atomic_append_many(
&self,
writes: &[PlannedAppend],
) -> impl core::future::Future<
Output = Result<Option<Self::AllPosition>, AtomicAppendError<Self::Error>>,
> + Send;
}
impl<S: AtomicAppend> AtomicAppend for Store<S> {
async fn atomic_append_many(
&self,
writes: &[PlannedAppend],
) -> Result<Option<Self::AllPosition>, AtomicAppendError<Self::Error>> {
self.raw().atomic_append_many(writes).await
}
}
pub trait EventImporter: RawEventStore + AtomicAppend {
fn import<R>(
&self,
sections: &[StreamSection],
route: R,
atomicity: Atomicity,
) -> impl core::future::Future<Output = Result<ImportReport, ImportError<Self::Error>>> + Send
where
R: Fn(&[u8]) -> StreamKey + Send;
}
#[derive(Debug)]
enum Halt {
Complete,
Corrupt,
Gap { got: Version },
}
#[derive(Debug)]
enum SectionPlan {
Empty,
FirstCorrupt,
Run {
first: Version,
expected_version: Option<Version>,
head: PendingEnvelope,
tail: Vec<PendingEnvelope>,
last: Version,
halt: Halt,
},
}
#[derive(Debug)]
enum PlanError {
VersionOverflow,
}
fn plan_section(section: &StreamSection) -> Result<SectionPlan, PlanError> {
let mut blocks = section.blocks.iter();
let Some(first_block) = blocks.next() else {
return Ok(SectionPlan::Empty);
};
let first_event = match first_block {
ImportBlock::Corrupt => return Ok(SectionPlan::FirstCorrupt),
ImportBlock::Event(event) => event,
};
let first = first_event.version();
let expected_version = first.as_u64().checked_sub(1).and_then(Version::new);
let head = PendingEnvelope::from_persisted(first_event);
let mut tail = Vec::new();
let mut last = first;
let halt = loop {
let Some(block) = blocks.next() else {
break Halt::Complete;
};
let event = match block {
ImportBlock::Corrupt => break Halt::Corrupt,
ImportBlock::Event(event) => event,
};
let expected_next = last.next().ok_or(PlanError::VersionOverflow)?;
if event.version() != expected_next {
break Halt::Gap {
got: event.version(),
};
}
tail.push(PendingEnvelope::from_persisted(event));
last = event.version();
};
Ok(SectionPlan::Run {
first,
expected_version,
head,
tail,
last,
halt,
})
}
impl<S: RawEventStore + AtomicAppend> EventImporter for S {
async fn import<R>(
&self,
sections: &[StreamSection],
route: R,
atomicity: Atomicity,
) -> Result<ImportReport, ImportError<Self::Error>>
where
R: Fn(&[u8]) -> StreamKey + Send,
{
match atomicity {
Atomicity::PerStream => import_per_stream(self, sections, route).await,
Atomicity::WholeChunk => import_whole_chunk(self, sections, route).await,
}
}
}
pub async fn import_per_stream<S, R>(
store: &S,
sections: &[StreamSection],
route: R,
) -> Result<ImportReport, ImportError<S::Error>>
where
S: RawEventStore,
R: Fn(&[u8]) -> StreamKey + Send,
{
let mut reports = Vec::with_capacity(sections.len());
for section in sections {
let target = route(section.origin.as_ref());
let plan = match plan_section(section) {
Ok(plan) => plan,
Err(PlanError::VersionOverflow) => return Err(ImportError::VersionOverflow),
};
let outcome = match plan {
SectionPlan::Empty => continue,
SectionPlan::FirstCorrupt => StreamOutcome::Corrupt { reached: None },
SectionPlan::Run {
first,
expected_version,
head,
tail,
last,
halt,
} => match store
.append(
&target,
expected_version,
PendingBatch::from_parts(&head, &tail),
)
.await
{
Ok(_position) => match halt {
Halt::Complete => StreamOutcome::Complete { version: last },
Halt::Corrupt => StreamOutcome::Corrupt {
reached: Some(last),
},
Halt::Gap { got } => StreamOutcome::Mismatch {
reached: Some(last),
got,
},
},
Err(AppendError::Conflict { .. }) => StreamOutcome::Mismatch {
reached: None,
got: first,
},
Err(AppendError::Store(error)) => return Err(ImportError::Store(error)),
},
};
reports.push(StreamReport {
stream: target,
outcome,
});
}
Ok(ImportReport::new(reports))
}
async fn import_whole_chunk<S, R>(
store: &S,
sections: &[StreamSection],
route: R,
) -> Result<ImportReport, ImportError<S::Error>>
where
S: RawEventStore + AtomicAppend,
R: Fn(&[u8]) -> StreamKey + Send,
{
let mut writes: Vec<PlannedAppend> = Vec::with_capacity(sections.len());
let mut firsts: Vec<Version> = Vec::with_capacity(sections.len());
let mut lasts: Vec<Version> = Vec::with_capacity(sections.len());
for section in sections {
let target = route(section.origin.as_ref());
let plan = match plan_section(section) {
Ok(plan) => plan,
Err(PlanError::VersionOverflow) => return Err(ImportError::VersionOverflow),
};
let (first, expected_version, head, tail, last, halt) = match plan {
SectionPlan::Empty => continue, SectionPlan::FirstCorrupt => {
return Err(ImportError::Aborted {
stream: target,
reason: AbortReason::Corrupt,
});
}
SectionPlan::Run {
first,
expected_version,
head,
tail,
last,
halt,
} => (first, expected_version, head, tail, last, halt),
};
match halt {
Halt::Complete => {}
Halt::Corrupt => {
return Err(ImportError::Aborted {
stream: target,
reason: AbortReason::Corrupt,
});
}
Halt::Gap { got } => {
let expected = last.next().ok_or(ImportError::VersionOverflow)?;
return Err(ImportError::Aborted {
stream: target,
reason: AbortReason::Mismatch { expected, got },
});
}
}
firsts.push(first);
lasts.push(last);
writes.push(PlannedAppend {
target,
expected_version,
head,
tail,
});
}
match store.atomic_append_many(&writes).await {
Ok(_position) => {
let reports = writes
.into_iter()
.zip(lasts)
.map(|(write, last)| StreamReport {
stream: write.target,
outcome: StreamOutcome::Complete { version: last },
})
.collect();
Ok(ImportReport::new(reports))
}
Err(AtomicAppendError::Conflict { index, actual }) => {
Err(map_atomic_conflict(&firsts, &writes, index, actual))
}
Err(AtomicAppendError::Store(error)) => Err(ImportError::Store(error)),
}
}
fn map_atomic_conflict<E>(
firsts: &[Version],
writes: &[PlannedAppend],
index: usize,
actual: Option<Version>,
) -> ImportError<E> {
let (Some(&got), Some(write)) = (
firsts.get(index).or_else(|| firsts.first()),
writes.get(index).or_else(|| writes.first()),
) else {
return ImportError::VersionOverflow;
};
let expected = match actual {
Some(head) => match head.next() {
Some(next) => next,
None => return ImportError::VersionOverflow,
},
None => Version::INITIAL,
};
ImportError::Aborted {
stream: write.target.clone(),
reason: AbortReason::Mismatch { expected, got },
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
reason = "test code asserts exact values"
)]
mod plan_tests {
use super::*;
use crate::envelope::PersistedEnvelope;
use crate::value::SchemaVersion;
use bytes::Bytes;
fn v(n: u64) -> Version {
Version::new(n).expect("test version must be nonzero")
}
fn persisted(version: u64, payload: &[u8]) -> PersistedEnvelope {
let mut buf = Vec::new();
buf.extend_from_slice(b"E");
buf.extend_from_slice(payload);
let et_end = 1u32;
let pl_end = et_end + u32::try_from(payload.len()).expect("payload fits u32");
PersistedEnvelope::try_new(
v(version),
Bytes::from(buf),
SchemaVersion::INITIAL,
0..et_end,
et_end..pl_end,
None,
)
.expect("valid persisted envelope")
}
fn evt(version: u64) -> ImportBlock {
ImportBlock::Event(persisted(version, b"p"))
}
fn section(origin: &str, blocks: Vec<ImportBlock>) -> StreamSection {
StreamSection {
origin: Bytes::copy_from_slice(origin.as_bytes()),
blocks,
}
}
#[test]
fn plan_empty_section_is_empty() {
assert!(matches!(
plan_section(§ion("s", vec![])),
Ok(SectionPlan::Empty)
));
}
#[test]
fn plan_first_block_corrupt_is_first_corrupt() {
let s = section("s", vec![ImportBlock::Corrupt, evt(1)]);
assert!(matches!(plan_section(&s), Ok(SectionPlan::FirstCorrupt)));
}
#[test]
fn plan_contiguous_run_from_one_is_complete() {
let s = section("s", vec![evt(1), evt(2), evt(3)]);
let plan = plan_section(&s).expect("plans");
match plan {
SectionPlan::Run {
first,
expected_version,
last,
halt,
tail,
..
} => {
assert_eq!(first, v(1));
assert_eq!(expected_version, None); assert_eq!(last, v(3));
assert_eq!(tail.len(), 2, "head + 2 tail == the 3-event run");
assert!(matches!(halt, Halt::Complete));
}
other => panic!("expected Run, got {other:?}"),
}
}
#[test]
fn plan_run_from_midstream_sets_expected_to_first_minus_one() {
let s = section("s", vec![evt(3), evt(4)]);
match plan_section(&s).expect("plans") {
SectionPlan::Run {
first,
expected_version,
last,
halt,
..
} => {
assert_eq!(first, v(3));
assert_eq!(expected_version, Some(v(2)));
assert_eq!(last, v(4));
assert!(matches!(halt, Halt::Complete));
}
other => panic!("expected Run, got {other:?}"),
}
}
#[test]
fn plan_internal_gap_halts_with_got() {
let s = section("s", vec![evt(3), evt(4), evt(6)]);
match plan_section(&s).expect("plans") {
SectionPlan::Run {
last, halt, tail, ..
} => {
assert_eq!(last, v(4));
assert_eq!(tail.len(), 1, "head + 1 tail == the 2-event run");
assert!(matches!(halt, Halt::Gap { got } if got == v(6)));
}
other => panic!("expected Run, got {other:?}"),
}
}
#[test]
fn plan_internal_corrupt_halts_corrupt() {
let s = section("s", vec![evt(1), evt(2), ImportBlock::Corrupt, evt(3)]);
match plan_section(&s).expect("plans") {
SectionPlan::Run {
last, halt, tail, ..
} => {
assert_eq!(last, v(2));
assert_eq!(tail.len(), 1, "head + 1 tail == the 2-event run");
assert!(matches!(halt, Halt::Corrupt));
}
other => panic!("expected Run, got {other:?}"),
}
}
#[test]
fn plan_overflow_building_run_errors() {
let s = section("s", vec![evt(u64::MAX), evt(1)]);
assert!(matches!(plan_section(&s), Err(PlanError::VersionOverflow)));
}
#[test]
fn plan_run_ending_at_u64_max_with_no_successor_completes() {
let s = section("s", vec![evt(u64::MAX)]);
match plan_section(&s).expect("plans") {
SectionPlan::Run {
first,
expected_version,
tail,
last,
halt,
..
} => {
assert_eq!(first, v(u64::MAX));
assert_eq!(expected_version, Some(v(u64::MAX - 1)));
assert!(tail.is_empty(), "a one-event run is head-only");
assert_eq!(last, v(u64::MAX));
assert!(matches!(halt, Halt::Complete));
}
other => panic!("expected Run, got {other:?}"),
}
}
}