use super::{
AppendBatchView, AppendSummary, AppendView, Barrier, OperationBody, OperationCodecError,
OperationKind, OperationLimits, OperationOutput, RecordPosition, decode_append_batches,
decode_append_batches_indexed, decode_append_batches_with_validated_payload,
};
use ozzy_proto::OperationId;
use smallvec::SmallVec;
use std::ops::Range;
pub(super) mod append;
mod control;
mod delay;
pub use delay::{
CARRY_FRAME_BYTES, RELEASE_FRAME_BYTES, decode_carry_batch, decode_release_batch,
decode_schedule_batch, encoded_record_bytes,
};
mod group;
mod primitives;
pub(super) use append::{
AppendTotals, TinyTable, decode_append_batch, encode_append_records, read_descriptor_head,
records_are_tiny, validate_append_descriptors, validate_tiny_descriptors,
};
use append::{decode_append, encode_append};
use control::{
decode_assign, decode_create, decode_dedup_watermark, decode_open_producer, decode_policy,
decode_producer_result_floor, decode_progress, decode_retire_producer, decode_trim,
encode_assign, encode_create, encode_dedup_watermark, encode_open_producer, encode_policy,
encode_producer_result_floor, encode_progress, encode_retire_producer, encode_trim,
};
pub(super) use primitives::{
Encoder, ReadFields, Reader, encode_atomic, enforce_limit, validate_limits,
validate_position_count,
};
pub fn encode_operation_body(
body: &OperationBody<'_>,
limits: OperationLimits,
) -> Result<Vec<u8>, OperationCodecError> {
let mut output = Vec::new();
append_operation_body(&mut output, body, limits)?;
Ok(output)
}
pub fn append_operation_body(
output: &mut Vec<u8>,
body: &OperationBody<'_>,
limits: OperationLimits,
) -> Result<Range<usize>, OperationCodecError> {
append_operation_body_to(output, body, limits)
}
pub fn append_operation_body_to(
output: &mut impl OperationOutput,
body: &OperationBody<'_>,
limits: OperationLimits,
) -> Result<Range<usize>, OperationCodecError> {
primitives::encode_atomic(output, limits, |encoder| {
match body {
OperationBody::CreatePartition(value) => encode_create(value, limits, encoder)?,
OperationBody::OpenProducer(value) => encode_open_producer(*value, encoder)?,
OperationBody::Append(value) => encode_append(value, limits, encoder)?,
OperationBody::ResolvedAppend(value) => {
super::resolved::encode_resolved_append(value, limits, encoder)?;
}
OperationBody::Progress(value) => encode_progress(*value, encoder)?,
OperationBody::Assign(value) => encode_assign(*value, encoder)?,
OperationBody::Trim(value) => encode_trim(*value, encoder)?,
OperationBody::PartitionPolicy(value) => encode_policy(*value, encoder)?,
OperationBody::Barrier(value) => encoder.id(value.operation_id.as_bytes())?,
OperationBody::ProducerResultFloor(value) => {
encode_producer_result_floor(*value, encoder)?;
}
OperationBody::DedupWatermark(value) => encode_dedup_watermark(*value, encoder)?,
OperationBody::GroupPlan(value) => group::encode(value, limits, encoder)?,
OperationBody::GroupDelete(value) => group::encode_delete(*value, encoder)?,
OperationBody::RetireProducer(value) => encode_retire_producer(*value, encoder)?,
OperationBody::Schedule(value) => delay::encode_schedule(value, limits, encoder)?,
OperationBody::Release(value) => delay::encode_release(value, limits, encoder)?,
OperationBody::Carry(value) => delay::encode_carry(value, limits, encoder)?,
}
Ok(())
})
}
pub fn validate_operation_body(
kind: OperationKind,
input: &[u8],
limits: OperationLimits,
) -> Result<(), OperationCodecError> {
decode_body::<false, true>(kind, input, limits).map(|_| ())
}
pub fn validate_operation_body_with_validated_payload(
kind: OperationKind,
input: &[u8],
limits: OperationLimits,
) -> Result<(), OperationCodecError> {
decode_body::<false, false>(kind, input, limits).map(|_| ())
}
pub fn decode_operation_body(
kind: OperationKind,
input: &[u8],
limits: OperationLimits,
) -> Result<OperationBody<'_>, OperationCodecError> {
decode_body::<true, true>(kind, input, limits)
}
pub fn decode_append_summary(
input: &[u8],
limits: OperationLimits,
) -> Result<AppendSummary, OperationCodecError> {
decode_append_summary_and_view(input, limits).map(|(summary, _)| summary)
}
pub fn decode_append_summary_and_view(
input: &[u8],
limits: OperationLimits,
) -> Result<(AppendSummary, AppendView<'_>), OperationCodecError> {
validate_limits(limits)?;
enforce_limit("operation body bytes", input.len(), limits.max_body_bytes)?;
let mut decoder = Reader::new(input);
let count = decoder.count("append batch count", limits.max_append_batches)?;
if count == 0 {
return Err(OperationCodecError::EmptyAppend);
}
let mut batches = SmallVec::with_capacity(count);
let mut totals = AppendTotals::new(limits);
for _ in 0..count {
let (_, view) = decode_append_batch::<false, true, false>(&mut decoder, &mut totals, None)?;
batches.push(view.summary);
}
decoder.finish()?;
Ok((
AppendSummary { batches },
AppendView::validated(input, limits)?,
))
}
pub fn decode_append_summary_and_batches(
input: &[u8],
limits: OperationLimits,
) -> Result<(AppendSummary, SmallVec<[AppendBatchView<'_>; 4]>), OperationCodecError> {
let views = decode_append_batches(input, limits)?;
let summary = AppendSummary {
batches: views.iter().map(|view| view.summary).collect(),
};
Ok((summary, views))
}
pub fn decode_append_summary_and_batches_indexed<'a>(
input: &'a [u8],
limits: OperationLimits,
positions: &mut Vec<RecordPosition>,
) -> Result<(AppendSummary, SmallVec<[AppendBatchView<'a>; 4]>), OperationCodecError> {
let views = decode_append_batches_indexed(input, limits, positions)?;
let summary = AppendSummary {
batches: views.iter().map(|view| view.summary).collect(),
};
Ok((summary, views))
}
pub fn decode_append_summary_and_batches_with_validated_payload(
input: &[u8],
limits: OperationLimits,
) -> Result<(AppendSummary, SmallVec<[AppendBatchView<'_>; 4]>), OperationCodecError> {
let views = decode_append_batches_with_validated_payload(input, limits)?;
let summary = AppendSummary {
batches: views.iter().map(|view| view.summary).collect(),
};
Ok((summary, views))
}
fn decode_body<const MATERIALIZE: bool, const VALIDATE_PREPARED: bool>(
kind: OperationKind,
input: &[u8],
limits: OperationLimits,
) -> Result<OperationBody<'_>, OperationCodecError> {
validate_limits(limits)?;
enforce_limit("operation body bytes", input.len(), limits.max_body_bytes)?;
let mut decoder = Reader::new(input);
let body = match kind {
OperationKind::CreatePartition => {
OperationBody::CreatePartition(decode_create(&mut decoder, limits)?)
}
OperationKind::OpenProducer => {
OperationBody::OpenProducer(decode_open_producer(&mut decoder)?)
}
OperationKind::Append => OperationBody::Append(decode_append::<
MATERIALIZE,
VALIDATE_PREPARED,
>(&mut decoder, limits)?),
OperationKind::ResolvedAppend => OperationBody::ResolvedAppend(
super::resolved::read_resolved_append::<MATERIALIZE, VALIDATE_PREPARED>(
&mut decoder,
limits,
None,
)?
.0,
),
OperationKind::Progress => OperationBody::Progress(decode_progress(&mut decoder)?),
OperationKind::Assign => OperationBody::Assign(decode_assign(&mut decoder)?),
OperationKind::Trim => OperationBody::Trim(decode_trim(&mut decoder)?),
OperationKind::PartitionPolicy => {
OperationBody::PartitionPolicy(decode_policy(&mut decoder)?)
}
OperationKind::Barrier => OperationBody::Barrier(Barrier {
operation_id: OperationId::from_bytes(decoder.id()?),
}),
OperationKind::ProducerResultFloor => {
OperationBody::ProducerResultFloor(decode_producer_result_floor(&mut decoder)?)
}
OperationKind::DedupWatermark => {
OperationBody::DedupWatermark(decode_dedup_watermark(&mut decoder)?)
}
OperationKind::GroupPlan => {
OperationBody::GroupPlan(group::decode::<MATERIALIZE>(&mut decoder, limits)?)
}
OperationKind::GroupDelete => {
OperationBody::GroupDelete(group::decode_delete(&mut decoder)?)
}
OperationKind::RetireProducer => {
OperationBody::RetireProducer(decode_retire_producer(&mut decoder)?)
}
OperationKind::Schedule | OperationKind::Release | OperationKind::Carry => {
delay::decode::<MATERIALIZE, VALIDATE_PREPARED>(kind, &mut decoder, limits)?
}
};
decoder.finish()?;
Ok(body)
}