use super::{AppendResolution, DurableChannel, MessageEnvelope, UnresolvedAppend};
use crate::durability::{DurabilityError, HEAD_CHECKPOINT_EVERY, advance_head};
#[derive(Clone, Debug)]
pub(super) struct PendingBatch {
rows: Vec<UnresolvedAppend>,
}
impl DurableChannel {
pub(crate) fn pending_batch_evidence(&self) -> Option<UnresolvedAppend> {
self.pending_batch.as_ref()?.rows.first().cloned()
}
pub(crate) async fn publish_batch(
&mut self,
envelopes: &[MessageEnvelope],
) -> Result<u64, DurabilityError> {
self.resolve_pending_batch().await?;
self.resolve_pending().await?;
let first = envelopes
.first()
.ok_or_else(|| DurabilityError::ConfigError("durable batch is empty".to_owned()))?;
let partition = self.partition_for(first);
if envelopes
.iter()
.any(|envelope| self.partition_for(envelope) != partition)
{
return Err(DurabilityError::ConfigError(
"durable batch spans partitions".to_owned(),
));
}
let start = self.sequence_for_append(partition)?;
let count = u64::try_from(envelopes.len()).map_err(|error| {
DurabilityError::ConfigError(format!("durable batch size overflow: {error}"))
})?;
let next = start.checked_add(count).ok_or_else(|| {
DurabilityError::ConfigError("durable batch sequence overflow".to_owned())
})?;
let rows = envelopes
.iter()
.enumerate()
.map(|(index, envelope)| {
let offset = u64::try_from(index).map_err(|error| {
DurabilityError::ConfigError(format!("durable batch index overflow: {error}"))
})?;
Ok(UnresolvedAppend::new(
self.channel_id.clone(),
partition,
start + offset,
envelope.serialize()?,
envelope.idempotency_key.clone(),
))
})
.collect::<Result<Vec<_>, DurabilityError>>()?;
let payloads = rows.iter().map(|row| row.payload().to_vec()).collect();
self.pending_batch = Some(PendingBatch { rows });
let stream_key = self.stream_key_for(partition);
match self.store.append_all(&stream_key, payloads, start).await {
Ok(assigned) if assigned == next => self.finish_batch(partition, start, next).await?,
Err(conflict @ DurabilityError::SequenceConflict { actual, .. }) => {
self.pending_batch = None;
self.refresh_sequence(partition, &stream_key, start, actual)
.await?;
return Err(conflict);
}
outcome => {
return Err(self.batch_unresolved(format!("batch append outcome: {outcome:?}")));
}
}
Ok(start)
}
fn batch_unresolved(&mut self, cause: String) -> DurabilityError {
if let Some(batch) = self.pending_batch.as_mut() {
for row in &mut batch.rows {
row.set_cause(cause.clone());
}
}
let Some(row) = self.pending_batch_evidence() else {
return DurabilityError::ConfigError(
"durable batch lost its append evidence".to_owned(),
);
};
DurabilityError::AppendUnresolved {
stream_key: self.stream_key_for(row.partition()),
sequence: row.sequence(),
cause,
}
}
pub(crate) async fn resolve_pending_batch(
&mut self,
) -> Result<Option<AppendResolution>, DurabilityError> {
let Some(batch) = self.pending_batch.as_ref() else {
return Ok(None);
};
let first = batch.rows.first().ok_or_else(|| {
DurabilityError::ConfigError("pending durable batch has no rows".to_owned())
})?;
let partition = first.partition();
let start = first.sequence();
let stream_key = self.stream_key_for(partition);
let mut matches = 0;
let mut observed_head = start;
for expected in &batch.rows {
let row = self
.store
.read_at(&stream_key, expected.sequence())
.await
.map_err(|error| DurabilityError::PendingAppendUnreadable {
stream_key: stream_key.clone(),
sequence: expected.sequence(),
append_error: expected.cause().to_owned(),
read_error: error.to_string(),
})?;
if let Some(row) = row {
observed_head =
observed_head.max(row.sequence.checked_add(1).ok_or_else(|| {
DurabilityError::ConfigError(
"durable batch recovery sequence overflow".to_owned(),
)
})?);
if row.sequence == expected.sequence() && row.payload == expected.payload() {
matches += 1;
}
}
}
if matches == 0 {
self.refresh_sequence(partition, &stream_key, start, observed_head)
.await?;
self.pending_batch = None;
return Ok(Some(AppendResolution::Absent));
}
if matches != batch.rows.len() {
return Err(DurabilityError::ConfigError(format!(
"partial durable batch: {matches}/{} reserved rows match",
batch.rows.len()
)));
}
let next = start
.checked_add(u64::try_from(batch.rows.len()).map_err(|error| {
DurabilityError::ConfigError(format!(
"durable batch recovery size overflow: {error}"
))
})?)
.ok_or_else(|| {
DurabilityError::ConfigError("durable batch recovery head overflow".to_owned())
})?;
self.finish_batch(partition, start, next).await?;
Ok(Some(AppendResolution::Committed { position: start }))
}
async fn finish_batch(
&mut self,
partition: usize,
start: u64,
next: u64,
) -> Result<(), DurabilityError> {
let stream_key = self.stream_key_for(partition);
self.advance_partition(partition, &stream_key, next).await?;
if next % HEAD_CHECKPOINT_EVERY != 0
&& start / HEAD_CHECKPOINT_EVERY < next / HEAD_CHECKPOINT_EVERY
{
if let Err(error) = advance_head(self.store.as_ref(), &stream_key, next).await {
tracing::warn!(stream = %stream_key, head = next, %error, "durable batch head checkpoint failed");
}
}
if let Err(error) = self.store.flush().await {
return Err(self.batch_unresolved(format!("durable batch flush failed: {error}")));
}
self.pending_batch = None;
Ok(())
}
}