use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::sync::Arc;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex as AsyncMutex;
use crate::{
ComponentStateEnvelope, ComponentStatePayload, ComponentStreamIdentity, ContentIdentity,
DefaultComponentCodec, ExternalStateReference, FailureCategory, ItemReader, ItemStream,
ItemWriter, ReadContext, ReadOutcome, ReaderError, RestartabilityDeclaration, StateCodecError,
StateLimits, StateSchemaId, StateSchemaVersion, StateSensitivity, StreamCloseContext,
StreamCloseError, StreamCloseOutcome, StreamOpenContext, StreamOpenError, StreamOpenOutcome,
StreamRuntimeOutcome, StreamStateContract, StreamUpdateContext, StreamUpdateError,
VersionedStateCodec, WriteContext, WriteOutcome, WriterError,
};
const MAX_RESOURCE_IDENTITY_BYTES: usize = 1024;
#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ResourceIdentity(String);
impl ResourceIdentity {
pub fn new(value: impl Into<String>) -> Result<Self, MultiResourceConfigError> {
let value = value.into();
if value.is_empty() {
return Err(MultiResourceConfigError::EmptyResourceIdentity);
}
if value.len() > MAX_RESOURCE_IDENTITY_BYTES {
return Err(MultiResourceConfigError::ResourceIdentityTooLong {
max_bytes: MAX_RESOURCE_IDENTITY_BYTES,
});
}
if value.chars().any(char::is_control) {
return Err(MultiResourceConfigError::MalformedResourceIdentity);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Debug for ResourceIdentity {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_tuple("ResourceIdentity")
.field(&self.0)
.finish()
}
}
impl fmt::Display for ResourceIdentity {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub struct ResourceSetRevision([u8; 32]);
impl ResourceSetRevision {
#[must_use]
pub fn of<'a>(resources: impl IntoIterator<Item = &'a ResourceIdentity>) -> Self {
let mut hasher = Sha256::new();
hasher.update(b"oxide-batch.resource-set\0");
for resource in resources {
hasher.update(resource.as_str().as_bytes());
hasher.update([0u8]);
}
Self(hasher.finalize().into())
}
fn to_hex(self) -> String {
hex_encode(&self.0)
}
fn from_hex(hex: &str) -> Option<Self> {
let bytes = hex_decode(hex)?;
let array: [u8; 32] = bytes.try_into().ok()?;
Some(Self(array))
}
}
impl fmt::Debug for ResourceSetRevision {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "ResourceSetRevision({})", self.to_hex())
}
}
#[derive(Clone)]
pub struct ResourceSet {
resources: Arc<[ResourceIdentity]>,
revision: ResourceSetRevision,
}
impl ResourceSet {
#[must_use]
pub fn new(resources: Vec<ResourceIdentity>) -> Self {
let revision = ResourceSetRevision::of(resources.iter());
Self {
resources: Arc::from(resources),
revision,
}
}
#[must_use]
pub fn resources(&self) -> &[ResourceIdentity] {
&self.resources
}
#[must_use]
pub const fn revision(&self) -> ResourceSetRevision {
self.revision
}
#[must_use]
pub fn len(&self) -> usize {
self.resources.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.resources.is_empty()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum MultiResourceConfigError {
EmptyResourceIdentity,
ResourceIdentityTooLong {
max_bytes: usize,
},
MalformedResourceIdentity,
}
impl fmt::Display for MultiResourceConfigError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::EmptyResourceIdentity => {
formatter.write_str("resource identity must not be empty")
}
Self::ResourceIdentityTooLong { max_bytes } => {
write!(
formatter,
"resource identity exceeds {max_bytes} UTF-8 bytes"
)
}
Self::MalformedResourceIdentity => {
formatter.write_str("resource identity contains a control character")
}
}
}
}
impl std::error::Error for MultiResourceConfigError {}
fn hex_encode(bytes: &[u8]) -> String {
use fmt::Write as _;
let mut out = String::with_capacity(bytes.len() * 2);
for byte in bytes {
let _ = write!(out, "{byte:02x}");
}
out
}
fn hex_decode(text: &str) -> Option<Vec<u8>> {
if !text.len().is_multiple_of(2) {
return None;
}
let mut out = Vec::with_capacity(text.len() / 2);
let bytes = text.as_bytes();
let mut i = 0;
while i < bytes.len() {
let hi = (bytes[i] as char).to_digit(16)?;
let lo = (bytes[i + 1] as char).to_digit(16)?;
out.push(u8::try_from((hi << 4) | lo).ok()?);
i += 2;
}
Some(out)
}
#[derive(Clone)]
struct DelegateEnvelopeColumns {
namespace: String,
schema_id: String,
schema_version: u32,
codec_id: String,
codec_version: u32,
checksum_algorithm: u16,
checksum_algorithm_version: u16,
checksum: [u8; 32],
is_external: bool,
payload_inline_hex: Option<String>,
payload_external_content_id_hex: Option<String>,
payload_external_len: Option<u64>,
}
impl DelegateEnvelopeColumns {
fn from_envelope(envelope: &ComponentStateEnvelope) -> Result<Self, StateCodecError> {
let payload = envelope
.payload()
.map_err(|_| StateCodecError::InvalidPayload)?;
let (payload_inline_hex, payload_external_content_id_hex, payload_external_len) =
match payload {
ComponentStatePayload::Inline(bytes) => (Some(hex_encode(&bytes)), None, None),
ComponentStatePayload::External(reference) => (
None,
Some(hex_encode(reference.content_id().as_bytes())),
Some(reference.encoded_len()),
),
};
Ok(Self {
namespace: envelope.namespace().as_str().to_owned(),
schema_id: envelope.schema_id().as_str().to_owned(),
schema_version: envelope.schema_version().get(),
codec_id: envelope.codec_id().as_str().to_owned(),
codec_version: envelope.codec_version().get(),
checksum_algorithm: envelope.checksum_algorithm(),
checksum_algorithm_version: envelope.checksum_algorithm_version(),
checksum: envelope.checksum(),
is_external: envelope.is_external(),
payload_inline_hex,
payload_external_content_id_hex,
payload_external_len,
})
}
fn namespace_matches(&self, expected: &ComponentStreamIdentity) -> bool {
self.namespace == expected.as_str()
}
fn to_envelope(
&self,
namespace: ComponentStreamIdentity,
limits: StateLimits,
) -> Option<ComponentStateEnvelope> {
let payload = if self.is_external {
let content_id_hex = self.payload_external_content_id_hex.as_deref()?;
let bytes = hex_decode(content_id_hex)?;
let array: [u8; 32] = bytes.try_into().ok()?;
let content_id = ContentIdentity::from_bytes(array);
let len = self.payload_external_len?;
ComponentStatePayload::External(ExternalStateReference::new(content_id, len))
} else {
let hex = self.payload_inline_hex.as_deref()?;
ComponentStatePayload::Inline(hex_decode(hex)?)
};
ComponentStateEnvelope::from_durable(
namespace,
&self.schema_id,
self.schema_version,
&self.codec_id,
self.codec_version,
self.checksum_algorithm,
self.checksum_algorithm_version,
self.checksum,
payload,
limits,
)
.ok()
}
}
#[derive(Clone)]
struct MultiResourceState {
resource_set_revision: ResourceSetRevision,
resource_index: u32,
delegate: Option<DelegateEnvelopeColumns>,
resource_batches_written: u64,
}
fn delegate_columns_to_json(delegate: &DelegateEnvelopeColumns) -> serde_json::Value {
serde_json::json!({
"namespace": delegate.namespace,
"schema_id": delegate.schema_id,
"schema_version": delegate.schema_version,
"codec_id": delegate.codec_id,
"codec_version": delegate.codec_version,
"checksum_algorithm": delegate.checksum_algorithm,
"checksum_algorithm_version": delegate.checksum_algorithm_version,
"checksum": hex_encode(&delegate.checksum),
"is_external": delegate.is_external,
"payload_inline": delegate.payload_inline_hex,
"payload_external_content_id": delegate.payload_external_content_id_hex,
"payload_external_len": delegate.payload_external_len,
})
}
fn delegate_columns_from_json(value: &serde_json::Value) -> Option<DelegateEnvelopeColumns> {
let checksum_hex = value.get("checksum")?.as_str()?;
let checksum_bytes = hex_decode(checksum_hex)?;
let checksum: [u8; 32] = checksum_bytes.try_into().ok()?;
Some(DelegateEnvelopeColumns {
namespace: value.get("namespace")?.as_str()?.to_owned(),
schema_id: value.get("schema_id")?.as_str()?.to_owned(),
schema_version: u32::try_from(value.get("schema_version")?.as_u64()?).ok()?,
codec_id: value.get("codec_id")?.as_str()?.to_owned(),
codec_version: u32::try_from(value.get("codec_version")?.as_u64()?).ok()?,
checksum_algorithm: u16::try_from(value.get("checksum_algorithm")?.as_u64()?).ok()?,
checksum_algorithm_version: u16::try_from(
value.get("checksum_algorithm_version")?.as_u64()?,
)
.ok()?,
checksum,
is_external: value.get("is_external")?.as_bool()?,
payload_inline_hex: value
.get("payload_inline")
.and_then(serde_json::Value::as_str)
.map(str::to_owned),
payload_external_content_id_hex: value
.get("payload_external_content_id")
.and_then(serde_json::Value::as_str)
.map(str::to_owned),
payload_external_len: value
.get("payload_external_len")
.and_then(serde_json::Value::as_u64),
})
}
const MULTI_RESOURCE_SCHEMA: &str = "oxide-batch.multi-resource-position";
const MULTI_RESOURCE_CODEC: &str = "oxide-batch.multi-resource-position-codec";
#[derive(Clone, Copy)]
struct MultiResourcePositionSchema;
impl VersionedStateCodec<MultiResourceState> for MultiResourcePositionSchema {
fn schema_id(&self) -> &StateSchemaId {
static SCHEMA: std::sync::OnceLock<StateSchemaId> = std::sync::OnceLock::new();
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema identity cannot fail validation"
)]
SCHEMA.get_or_init(|| StateSchemaId::new(MULTI_RESOURCE_SCHEMA).unwrap())
}
fn current_version(&self) -> StateSchemaVersion {
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema version cannot fail validation"
)]
StateSchemaVersion::new(2).unwrap()
}
fn encode(&self, value: &MultiResourceState) -> Result<Vec<u8>, StateCodecError> {
let delegate = value.delegate.as_ref().map(delegate_columns_to_json);
serde_json::to_vec(&serde_json::json!({
"resource_set_revision": value.resource_set_revision.to_hex(),
"resource_index": value.resource_index,
"delegate": delegate,
"resource_batches_written": value.resource_batches_written,
}))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, payload: &[u8]) -> Result<MultiResourceState, StateCodecError> {
let value: serde_json::Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
let revision_hex = value
.get("resource_set_revision")
.and_then(serde_json::Value::as_str)
.ok_or(StateCodecError::InvalidPayload)?;
let resource_set_revision =
ResourceSetRevision::from_hex(revision_hex).ok_or(StateCodecError::InvalidPayload)?;
let resource_index = u32::try_from(
value
.get("resource_index")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)?,
)
.map_err(|_| StateCodecError::InvalidPayload)?;
let delegate = match value.get("delegate") {
None | Some(serde_json::Value::Null) => None,
Some(delegate_value) => Some(
delegate_columns_from_json(delegate_value)
.ok_or(StateCodecError::InvalidPayload)?,
),
};
let resource_batches_written = value
.get("resource_batches_written")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0);
Ok(MultiResourceState {
resource_set_revision,
resource_index,
delegate,
resource_batches_written,
})
}
}
fn multi_resource_position_codec(
restartability: RestartabilityDeclaration,
) -> DefaultComponentCodec<MultiResourcePositionSchema> {
#[allow(
clippy::unwrap_used,
reason = "fixed literal identities cannot fail validation"
)]
DefaultComponentCodec::new(
MultiResourcePositionSchema,
crate::CodecId::new(MULTI_RESOURCE_CODEC).unwrap(),
crate::CodecVersion::new(1).unwrap(),
restartability,
)
.with_sensitivity(StateSensitivity::NonSensitive)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MultiResourceOpenError {
resource_ordinal: u32,
category: FailureCategory,
}
impl MultiResourceOpenError {
#[must_use]
pub const fn new(resource_ordinal: u32) -> Self {
Self {
resource_ordinal,
category: FailureCategory::UserComponent,
}
}
#[must_use]
pub const fn with_category(resource_ordinal: u32, category: FailureCategory) -> Self {
Self {
resource_ordinal,
category,
}
}
#[must_use]
pub const fn resource_ordinal(self) -> u32 {
self.resource_ordinal
}
#[must_use]
pub const fn category(self) -> FailureCategory {
self.category
}
}
impl fmt::Display for MultiResourceOpenError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
formatter,
"multi-resource open failed at resource ordinal {}",
self.resource_ordinal
)
}
}
impl std::error::Error for MultiResourceOpenError {}
pub trait MultiResourceReaderOpener<I>: Send + Sync {
type Reader: ItemReader<I>;
type Stream: ItemStream;
fn open<'a>(
&'a self,
resource: &'a ResourceIdentity,
resource_ordinal: u32,
delegate_identity: &'a ComponentStreamIdentity,
) -> impl Future<
Output = Result<(Self::Reader, Self::Stream, StreamStateContract), MultiResourceOpenError>,
> + Send
+ 'a;
}
struct StreamHandle<S> {
index: u32,
stream: S,
retired: bool,
}
struct ReaderShared<I, O: MultiResourceReaderOpener<I>> {
opener: O,
resources: ResourceSet,
identity: ComponentStreamIdentity,
handle: AsyncMutex<Option<StreamHandle<O::Stream>>>,
pending_reader: std::sync::Mutex<Option<O::Reader>>,
poisoned: AsyncMutex<Option<FailureCategory>>,
limits: StateLimits,
_marker: PhantomData<fn() -> I>,
}
pub struct MultiResourceReader<I, O: MultiResourceReaderOpener<I>> {
shared: Arc<ReaderShared<I, O>>,
reader: Option<O::Reader>,
started: bool,
transitioning: bool,
poisoned: Option<FailureCategory>,
}
impl<I, O> ItemReader<I> for MultiResourceReader<I, O>
where
I: Send + 'static,
O: MultiResourceReaderOpener<I> + 'static,
{
async fn read(&mut self, context: ReadContext<'_>) -> Result<ReadOutcome<I>, ReaderError> {
if let Some(category) = self.poisoned {
return Err(ReaderError::with_category(category));
}
if !self.started {
self.started = true;
let taken = self
.shared
.pending_reader
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
match taken {
Some(reader) => self.reader = Some(reader),
None => {
return Ok(ReadOutcome::EndOfInput);
}
}
}
loop {
if !self.transitioning {
let Some(reader) = self.reader.as_mut() else {
return Ok(ReadOutcome::EndOfInput);
};
match reader.read(context).await? {
ReadOutcome::Item(item) => return Ok(ReadOutcome::Item(item)),
ReadOutcome::Stopped => return Ok(ReadOutcome::Stopped),
ReadOutcome::EndOfInput => {
let mut guard = self.shared.handle.lock().await;
let handle = guard.as_mut().ok_or_else(|| {
ReaderError::with_category(FailureCategory::Invariant)
})?;
let closed = handle
.stream
.close(StreamCloseContext::new(
context.stop_token(),
StreamRuntimeOutcome::ResourceBoundary,
))
.await;
handle.retired = true;
drop(guard);
self.reader = None;
match closed {
Ok(_) => self.transitioning = true,
Err(error) => {
let category = error.category();
self.poisoned = Some(category);
*self.shared.poisoned.lock().await = Some(category);
return Err(ReaderError::with_category(category));
}
}
}
}
}
let current_index = {
let guard = self.shared.handle.lock().await;
guard
.as_ref()
.ok_or_else(|| ReaderError::with_category(FailureCategory::Invariant))?
.index
};
let next_index = current_index + 1;
if next_index as usize >= self.shared.resources.len() {
*self.shared.handle.lock().await = None;
self.transitioning = false;
return Ok(ReadOutcome::EndOfInput);
}
let next_resource = &self.shared.resources.resources()[next_index as usize];
let (new_reader, stream, _contract) = self
.shared
.opener
.open(next_resource, next_index, &self.shared.identity)
.await
.map_err(|error| ReaderError::with_category(error.category()))?;
stream
.open(StreamOpenContext::new(None, context.stop_token()))
.await
.map_err(|error| ReaderError::with_category(error.category()))?;
let mut guard = self.shared.handle.lock().await;
self.reader = Some(new_reader);
*guard = Some(StreamHandle {
index: next_index,
stream,
retired: false,
});
drop(guard);
self.transitioning = false;
}
}
}
pub struct MultiResourceReaderStream<I, O: MultiResourceReaderOpener<I>> {
shared: Arc<ReaderShared<I, O>>,
}
impl<I, O> ItemStream for MultiResourceReaderStream<I, O>
where
I: Send + 'static,
O: MultiResourceReaderOpener<I> + 'static,
{
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let restored = context.inherited_state().is_some();
let (resource_index, inherited_delegate) = match context.inherited_state() {
Some(envelope) => {
let state = envelope
.decode::<MultiResourceState>(&codec)
.map_err(|_| StreamOpenError::new())?;
if state.resource_set_revision != self.shared.resources.revision() {
return Err(StreamOpenError::with_category(
FailureCategory::UnsupportedCapability,
));
}
(state.resource_index, state.delegate)
}
None => (0, None),
};
let resources_len = u32::try_from(self.shared.resources.len()).unwrap_or(u32::MAX);
if resource_index > resources_len {
return Err(StreamOpenError::with_category(FailureCategory::Invariant));
}
if resource_index == resources_len {
*self.shared.handle.lock().await = None;
*self
.shared
.pending_reader
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
return Ok(if restored {
StreamOpenOutcome::Restored
} else {
StreamOpenOutcome::Initial
});
}
let resource = &self.shared.resources.resources()[resource_index as usize];
let (reader, stream, contract) = self
.shared
.opener
.open(resource, resource_index, &self.shared.identity)
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
match inherited_delegate {
Some(columns) => {
if !columns.namespace_matches(&self.shared.identity) {
return Err(StreamOpenError::with_category(FailureCategory::Invariant));
}
let inner = columns
.to_envelope(self.shared.identity.clone(), self.shared.limits)
.ok_or_else(StreamOpenError::new)?;
let validated = contract
.validate_for_open(&inner)
.map_err(|_| StreamOpenError::new())?;
stream
.open(StreamOpenContext::new(
Some(&validated),
context.stop_token(),
))
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
}
None => {
stream
.open(StreamOpenContext::new(None, context.stop_token()))
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
}
}
*self
.shared
.pending_reader
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(reader);
*self.shared.handle.lock().await = Some(StreamHandle {
index: resource_index,
stream,
retired: false,
});
Ok(if restored {
StreamOpenOutcome::Restored
} else {
StreamOpenOutcome::Initial
})
}
async fn update(
&self,
context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
if let Some(category) = *self.shared.poisoned.lock().await {
return Err(StreamUpdateError::with_category(category));
}
let guard = self.shared.handle.lock().await;
let (resource_index, delegate) = match guard.as_ref() {
Some(handle) => {
let envelope = handle
.stream
.update(StreamUpdateContext::new(context.stop_token()))
.await
.map_err(|error| StreamUpdateError::with_category(error.category()))?;
if envelope.namespace() != &self.shared.identity {
return Err(StreamUpdateError::with_category(FailureCategory::Invariant));
}
let columns = DelegateEnvelopeColumns::from_envelope(&envelope)
.map_err(|_| StreamUpdateError::new())?;
(handle.index, Some(columns))
}
None => (
u32::try_from(self.shared.resources.len()).unwrap_or(u32::MAX),
None,
),
};
drop(guard);
let state = MultiResourceState {
resource_set_revision: self.shared.resources.revision(),
resource_index,
delegate,
resource_batches_written: 0,
};
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
ComponentStateEnvelope::encode(
self.shared.identity.clone(),
&state,
&codec,
self.shared.limits,
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
let guard = self.shared.handle.lock().await;
if let Some(handle) = guard.as_ref()
&& !handle.retired
{
handle
.stream
.close(StreamCloseContext::new(
context.stop_token(),
context.outcome(),
))
.await
.map_err(|error| StreamCloseError::with_category(error.category()))?;
}
Ok(StreamCloseOutcome::Closed)
}
}
pub fn multi_resource_reader<I, O>(
resources: ResourceSet,
opener: O,
identity: ComponentStreamIdentity,
restartability: RestartabilityDeclaration,
) -> (
MultiResourceReader<I, O>,
MultiResourceReaderStream<I, O>,
StreamStateContract,
)
where
I: Send + 'static,
O: MultiResourceReaderOpener<I> + 'static,
{
let shared = Arc::new(ReaderShared {
opener,
resources,
identity,
handle: AsyncMutex::new(None),
pending_reader: std::sync::Mutex::new(None),
poisoned: AsyncMutex::new(None),
limits: StateLimits::default(),
_marker: PhantomData,
});
let reader = MultiResourceReader {
shared: Arc::clone(&shared),
reader: None,
started: false,
transitioning: false,
poisoned: None,
};
let stream = MultiResourceReaderStream { shared };
let contract = StreamStateContract::new(multi_resource_position_codec(restartability));
(reader, stream, contract)
}
pub trait MultiResourceWriterOpener<O>: Send + Sync {
type Writer: ItemWriter<O>;
type Stream: ItemStream;
fn open<'a>(
&'a self,
resource: &'a ResourceIdentity,
resource_ordinal: u32,
delegate_identity: &'a ComponentStreamIdentity,
) -> impl Future<
Output = Result<(Self::Writer, Self::Stream, StreamStateContract), MultiResourceOpenError>,
> + Send
+ 'a;
}
pub trait RolloverPolicy: Send + Sync {
fn should_roll_over(&self, resource_batches_written: u64) -> bool;
}
#[derive(Clone, Copy, Debug)]
pub struct BatchCountRollover {
max_batches_per_resource: u64,
}
impl BatchCountRollover {
#[must_use]
pub const fn new(max_batches_per_resource: u64) -> Self {
Self {
max_batches_per_resource,
}
}
}
impl RolloverPolicy for BatchCountRollover {
fn should_roll_over(&self, resource_batches_written: u64) -> bool {
resource_batches_written >= self.max_batches_per_resource
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct NoRollover;
impl RolloverPolicy for NoRollover {
fn should_roll_over(&self, _resource_batches_written: u64) -> bool {
false
}
}
struct ActiveWriteResource<W, S> {
index: u32,
writer: W,
stream: S,
batches_written: u64,
retired: bool,
}
enum WriterState<W, S> {
Active(ActiveWriteResource<W, S>),
Poisoned(FailureCategory),
}
struct WriterShared<O, Opener: MultiResourceWriterOpener<O>, P> {
opener: Opener,
resources: ResourceSet,
identity: ComponentStreamIdentity,
rollover: P,
active: AsyncMutex<Option<WriterState<Opener::Writer, Opener::Stream>>>,
limits: StateLimits,
_marker: PhantomData<fn() -> O>,
}
pub struct MultiResourceWriter<O, Opener: MultiResourceWriterOpener<O>, P> {
shared: Arc<WriterShared<O, Opener, P>>,
}
impl<O, Opener, P> ItemWriter<O> for MultiResourceWriter<O, Opener, P>
where
O: Sync + 'static,
Opener: MultiResourceWriterOpener<O> + 'static,
P: RolloverPolicy + 'static,
{
async fn write<'a>(
&'a self,
items: &'a [O],
mut context: WriteContext<'a>,
) -> Result<WriteOutcome, WriterError> {
let stop = context.stop_token();
if stop.is_stop_requested() {
return Ok(WriteOutcome::Stopped);
}
if self.shared.resources.is_empty() {
return Err(WriterError::with_category(FailureCategory::Invariant));
}
let mut guard = self.shared.active.lock().await;
let (should_roll_over, current_index) = match guard.as_ref() {
Some(WriterState::Poisoned(category)) => {
return Err(WriterError::with_category(*category));
}
Some(WriterState::Active(active)) => (
self.shared
.rollover
.should_roll_over(active.batches_written),
active.index,
),
None => return Err(WriterError::with_category(FailureCategory::Invariant)),
};
if should_roll_over {
let next_index = current_index + 1;
if (next_index as usize) < self.shared.resources.len() {
let mut boundary_close_failed = None;
if let Some(WriterState::Active(active)) = guard.as_mut()
&& !active.retired
{
let closed = active
.stream
.close(StreamCloseContext::new(
stop,
StreamRuntimeOutcome::ResourceBoundary,
))
.await;
active.retired = true;
if let Err(error) = closed {
boundary_close_failed = Some(error.category());
}
}
if let Some(category) = boundary_close_failed {
*guard = Some(WriterState::Poisoned(category));
return Err(WriterError::with_category(category));
}
let next_resource = &self.shared.resources.resources()[next_index as usize];
let (writer, stream, _contract) = self
.shared
.opener
.open(next_resource, next_index, &self.shared.identity)
.await
.map_err(|error| WriterError::with_category(error.category()))?;
stream
.open(StreamOpenContext::new(None, stop))
.await
.map_err(|error| WriterError::with_category(error.category()))?;
*guard = Some(WriterState::Active(ActiveWriteResource {
index: next_index,
writer,
stream,
batches_written: 0,
retired: false,
}));
}
}
let Some(WriterState::Active(active)) = guard.as_mut() else {
return Err(WriterError::with_category(FailureCategory::Invariant));
};
let enlisted = context.is_enlisted();
let delegate_context = if enlisted {
let transaction: &mut dyn crate::BusinessTransaction =
context.transaction().ok_or_else(WriterError::new)?;
WriteContext::enlisted(stop, transaction)
} else {
WriteContext::non_transactional(stop)
};
let outcome = active.writer.write(items, delegate_context).await?;
if matches!(outcome, WriteOutcome::Written) {
active.batches_written += 1;
}
Ok(outcome)
}
}
pub struct MultiResourceWriterStream<O, Opener: MultiResourceWriterOpener<O>, P> {
shared: Arc<WriterShared<O, Opener, P>>,
}
impl<O, Opener, P> ItemStream for MultiResourceWriterStream<O, Opener, P>
where
O: Sync + 'static,
Opener: MultiResourceWriterOpener<O> + 'static,
P: RolloverPolicy + 'static,
{
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let restored = context.inherited_state().is_some();
let (resource_index, inherited_delegate, resource_batches_written) =
match context.inherited_state() {
Some(envelope) => {
let state = envelope
.decode::<MultiResourceState>(&codec)
.map_err(|_| StreamOpenError::new())?;
if state.resource_set_revision != self.shared.resources.revision() {
return Err(StreamOpenError::with_category(
FailureCategory::UnsupportedCapability,
));
}
(
state.resource_index,
state.delegate,
state.resource_batches_written,
)
}
None => (0, None, 0),
};
let resources_len = u32::try_from(self.shared.resources.len()).unwrap_or(u32::MAX);
if resource_index >= resources_len {
return Err(StreamOpenError::with_category(FailureCategory::Invariant));
}
let resource = &self.shared.resources.resources()[resource_index as usize];
let (writer, stream, contract) = self
.shared
.opener
.open(resource, resource_index, &self.shared.identity)
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
match inherited_delegate {
Some(columns) => {
if !columns.namespace_matches(&self.shared.identity) {
return Err(StreamOpenError::with_category(FailureCategory::Invariant));
}
let inner = columns
.to_envelope(self.shared.identity.clone(), self.shared.limits)
.ok_or_else(StreamOpenError::new)?;
let validated = contract
.validate_for_open(&inner)
.map_err(|_| StreamOpenError::new())?;
stream
.open(StreamOpenContext::new(
Some(&validated),
context.stop_token(),
))
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
}
None => {
stream
.open(StreamOpenContext::new(None, context.stop_token()))
.await
.map_err(|error| StreamOpenError::with_category(error.category()))?;
}
}
*self.shared.active.lock().await = Some(WriterState::Active(ActiveWriteResource {
index: resource_index,
writer,
stream,
batches_written: resource_batches_written,
retired: false,
}));
Ok(if restored {
StreamOpenOutcome::Restored
} else {
StreamOpenOutcome::Initial
})
}
async fn update(
&self,
context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
let guard = self.shared.active.lock().await;
let active = match guard.as_ref() {
Some(WriterState::Active(active)) => active,
Some(WriterState::Poisoned(category)) => {
return Err(StreamUpdateError::with_category(*category));
}
None => return Err(StreamUpdateError::with_category(FailureCategory::Invariant)),
};
let envelope = active
.stream
.update(StreamUpdateContext::new(context.stop_token()))
.await
.map_err(|error| StreamUpdateError::with_category(error.category()))?;
if envelope.namespace() != &self.shared.identity {
return Err(StreamUpdateError::with_category(FailureCategory::Invariant));
}
let columns = DelegateEnvelopeColumns::from_envelope(&envelope)
.map_err(|_| StreamUpdateError::new())?;
let state = MultiResourceState {
resource_set_revision: self.shared.resources.revision(),
resource_index: active.index,
delegate: Some(columns),
resource_batches_written: active.batches_written,
};
drop(guard);
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
ComponentStateEnvelope::encode(
self.shared.identity.clone(),
&state,
&codec,
self.shared.limits,
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
let guard = self.shared.active.lock().await;
if let Some(WriterState::Active(active)) = guard.as_ref()
&& !active.retired
{
active
.stream
.close(StreamCloseContext::new(
context.stop_token(),
context.outcome(),
))
.await
.map_err(|error| StreamCloseError::with_category(error.category()))?;
}
Ok(StreamCloseOutcome::Closed)
}
}
pub type MultiResourceWriterTriple<O, Opener, P> = (
MultiResourceWriter<O, Opener, P>,
MultiResourceWriterStream<O, Opener, P>,
StreamStateContract,
);
pub fn multi_resource_writer<O, Opener, P>(
resources: ResourceSet,
opener: Opener,
identity: ComponentStreamIdentity,
rollover: P,
restartability: RestartabilityDeclaration,
) -> Result<MultiResourceWriterTriple<O, Opener, P>, WriterConfigError>
where
O: Sync + 'static,
Opener: MultiResourceWriterOpener<O> + 'static,
P: RolloverPolicy + 'static,
{
if resources.is_empty() {
return Err(WriterConfigError::EmptyResourceSet);
}
let shared = Arc::new(WriterShared {
opener,
resources,
identity,
rollover,
active: AsyncMutex::new(None),
limits: StateLimits::default(),
_marker: PhantomData,
});
let writer = MultiResourceWriter {
shared: Arc::clone(&shared),
};
let stream = MultiResourceWriterStream { shared };
let contract = StreamStateContract::new(multi_resource_position_codec(restartability));
Ok((writer, stream, contract))
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum WriterConfigError {
EmptyResourceSet,
}
impl fmt::Display for WriterConfigError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::EmptyResourceSet => {
formatter.write_str("multi-resource writer requires a nonempty resource set")
}
}
}
}
impl std::error::Error for WriterConfigError {}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use std::collections::{HashMap, HashSet};
use super::*;
use crate::StopSource;
fn identity(name: &str) -> ComponentStreamIdentity {
ComponentStreamIdentity::new(format!("oxide-batch.multi-resource-unit-test.{name}"))
.expect("static identity is valid")
}
fn resource_set(names: &[&str]) -> ResourceSet {
ResourceSet::new(
names
.iter()
.map(|name| ResourceIdentity::new((*name).to_owned()).unwrap())
.collect(),
)
}
#[test]
fn resource_identity_rejects_one_byte_past_its_ceiling() {
assert!(ResourceIdentity::new("x".repeat(MAX_RESOURCE_IDENTITY_BYTES)).is_ok());
let error = ResourceIdentity::new("x".repeat(MAX_RESOURCE_IDENTITY_BYTES + 1))
.expect_err("one byte past the ceiling must be refused, not silently truncated");
assert_eq!(
error,
MultiResourceConfigError::ResourceIdentityTooLong {
max_bytes: MAX_RESOURCE_IDENTITY_BYTES
}
);
}
#[test]
fn resource_identity_rejects_empty_and_control_characters() {
assert_eq!(
ResourceIdentity::new(""),
Err(MultiResourceConfigError::EmptyResourceIdentity)
);
assert_eq!(
ResourceIdentity::new("a\nb"),
Err(MultiResourceConfigError::MalformedResourceIdentity)
);
}
#[derive(Clone, Copy, Eq, PartialEq)]
struct VecPosition(u64);
const VEC_READ_SCHEMA: &str = "oxide-batch.multi-resource-unit-test.vec-read-position";
const VEC_READ_CODEC: &str = "oxide-batch.multi-resource-unit-test.vec-read-position-codec";
#[derive(Clone, Copy)]
struct VecReadSchema;
impl VersionedStateCodec<VecPosition> for VecReadSchema {
fn schema_id(&self) -> &StateSchemaId {
static SCHEMA: std::sync::OnceLock<StateSchemaId> = std::sync::OnceLock::new();
SCHEMA.get_or_init(|| StateSchemaId::new(VEC_READ_SCHEMA).unwrap())
}
fn current_version(&self) -> StateSchemaVersion {
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &VecPosition) -> Result<Vec<u8>, StateCodecError> {
serde_json::to_vec(&serde_json::json!({ "ordinal": value.0 }))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, payload: &[u8]) -> Result<VecPosition, StateCodecError> {
let value: serde_json::Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
let ordinal = value
.get("ordinal")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)?;
Ok(VecPosition(ordinal))
}
}
fn vec_read_codec() -> DefaultComponentCodec<VecReadSchema> {
DefaultComponentCodec::new(
VecReadSchema,
crate::CodecId::new(VEC_READ_CODEC).unwrap(),
crate::CodecVersion::new(1).unwrap(),
RestartabilityDeclaration::Restartable,
)
.with_sensitivity(StateSensitivity::NonSensitive)
}
struct VecItemReader {
items: Vec<u64>,
ordinal: Arc<AsyncMutex<u64>>,
}
impl ItemReader<u64> for VecItemReader {
async fn read(
&mut self,
_context: ReadContext<'_>,
) -> Result<ReadOutcome<u64>, ReaderError> {
let mut ordinal = self.ordinal.lock().await;
let index = usize::try_from(*ordinal).unwrap_or(usize::MAX);
match self.items.get(index) {
Some(item) => {
let item = *item;
*ordinal += 1;
Ok(ReadOutcome::Item(item))
}
None => Ok(ReadOutcome::EndOfInput),
}
}
}
struct VecItemReaderStream {
ordinal: Arc<AsyncMutex<u64>>,
namespace: ComponentStreamIdentity,
}
impl ItemStream for VecItemReaderStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = vec_read_codec();
if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<VecPosition>(&codec)
.map_err(|_| StreamOpenError::new())?;
*self.ordinal.lock().await = restored.0;
Ok(StreamOpenOutcome::Restored)
} else {
*self.ordinal.lock().await = 0;
Ok(StreamOpenOutcome::Initial)
}
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
let codec = vec_read_codec();
let ordinal = *self.ordinal.lock().await;
ComponentStateEnvelope::encode(
self.namespace.clone(),
&VecPosition(ordinal),
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
struct TestReaderOpener {
data: HashMap<String, Vec<u64>>,
fail: HashSet<String>,
}
impl MultiResourceReaderOpener<u64> for TestReaderOpener {
type Reader = VecItemReader;
type Stream = VecItemReaderStream;
async fn open(
&self,
resource: &ResourceIdentity,
resource_ordinal: u32,
delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Reader, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
if self.fail.contains(resource.as_str()) {
return Err(MultiResourceOpenError::new(resource_ordinal));
}
let items = self
.data
.get(resource.as_str())
.cloned()
.unwrap_or_default();
let ordinal = Arc::new(AsyncMutex::new(0));
let reader = VecItemReader {
items,
ordinal: Arc::clone(&ordinal),
};
let stream = VecItemReaderStream {
ordinal,
namespace: delegate_identity.clone(),
};
Ok((reader, stream, StreamStateContract::new(vec_read_codec())))
}
}
fn stop() -> (StopSource, crate::StopToken) {
StopSource::new()
}
fn read_context(stop: &crate::StopToken) -> ReadContext<'_> {
ReadContext::new(stop)
}
#[test]
fn ordered_traversal_across_resources_reads_in_declared_order() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1, 2]);
data.insert("b".to_owned(), vec![3]);
data.insert("c".to_owned(), vec![4, 5]);
let opener = TestReaderOpener {
data,
fail: HashSet::new(),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b", "c"]),
opener,
identity("ordered"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
let mut collected = Vec::new();
loop {
match reader.read(read_context(&token)).await.unwrap() {
ReadOutcome::Item(item) => collected.push(item),
ReadOutcome::EndOfInput => break,
ReadOutcome::Stopped => panic!("unexpected stop"),
}
}
assert_eq!(collected, vec![1, 2, 3, 4, 5]);
});
}
#[test]
fn empty_resource_set_reader_returns_end_of_input_immediately() {
let opener = TestReaderOpener {
data: HashMap::new(),
fail: HashSet::new(),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&[]),
opener,
identity("empty"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::EndOfInput
);
});
}
#[test]
fn restart_mid_resource_resumes_at_committed_position_not_resource_start() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1, 2]);
data.insert("b".to_owned(), vec![10, 20, 30]);
let opener = || TestReaderOpener {
data: data.clone(),
fail: HashSet::new(),
};
let resources = resource_set(&["a", "b"]);
let (_source, token) = stop();
let committed_envelope = futures_executor::block_on(async {
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resources.clone(),
opener(),
identity("restart"),
RestartabilityDeclaration::Restartable,
);
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(1)
);
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(2)
);
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(10)
);
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
futures_executor::block_on(async {
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resources,
opener(),
identity("restart"),
RestartabilityDeclaration::Restartable,
);
stream
.open(StreamOpenContext::new(Some(&committed_envelope), &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(20)
);
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(30)
);
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::EndOfInput
);
});
}
#[test]
fn resource_set_revision_mismatch_on_restart_is_rejected() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
data.insert("b".to_owned(), vec![2]);
let (_source, token) = stop();
let committed_envelope = futures_executor::block_on(async {
let opener = TestReaderOpener {
data: data.clone(),
fail: HashSet::new(),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b"]),
opener,
identity("revision"),
RestartabilityDeclaration::Restartable,
);
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
reader.read(read_context(&token)).await.unwrap();
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
futures_executor::block_on(async {
let opener = TestReaderOpener {
data,
fail: HashSet::new(),
};
let (_reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["z", "a", "b"]),
opener,
identity("revision"),
RestartabilityDeclaration::Restartable,
);
let result = stream
.open(StreamOpenContext::new(Some(&committed_envelope), &token))
.await;
assert!(
result.is_err(),
"changed resource set must be rejected, not silently resumed"
);
});
}
#[test]
fn resource_open_failure_does_not_advance_past_the_failure_point() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
data.insert("c".to_owned(), vec![3]);
let mut fail = HashSet::new();
fail.insert("b".to_owned());
let opener = TestReaderOpener { data, fail };
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b", "c"]),
opener,
identity("open-failure"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(1)
);
let result = reader.read(read_context(&token)).await;
assert!(
result.is_err(),
"resource b's open failure must surface, not be swallowed"
);
let retried = reader.read(read_context(&token)).await;
assert!(
retried.is_err(),
"a retry must fail at the same resource again, never skip to c"
);
});
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum LifecycleEvent {
Open(String),
Update(String),
Close(String, RecordedOutcome),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RecordedOutcome {
Committed,
Failed,
Stopped,
Unknown,
ResourceBoundary,
}
impl From<StreamRuntimeOutcome> for RecordedOutcome {
fn from(outcome: StreamRuntimeOutcome) -> Self {
match outcome {
StreamRuntimeOutcome::Committed => Self::Committed,
StreamRuntimeOutcome::Failed => Self::Failed,
StreamRuntimeOutcome::Stopped => Self::Stopped,
StreamRuntimeOutcome::Unknown => Self::Unknown,
StreamRuntimeOutcome::ResourceBoundary => Self::ResourceBoundary,
}
}
}
type Log = Arc<std::sync::Mutex<Vec<LifecycleEvent>>>;
fn log_event(log: &Log, event: LifecycleEvent) {
log.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(event);
}
struct RecordingReaderStream {
ordinal: Arc<AsyncMutex<u64>>,
namespace: ComponentStreamIdentity,
resource: String,
log: Log,
fail_close: Arc<std::sync::atomic::AtomicBool>,
}
impl ItemStream for RecordingReaderStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
log_event(&self.log, LifecycleEvent::Open(self.resource.clone()));
let codec = vec_read_codec();
if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<VecPosition>(&codec)
.map_err(|_| StreamOpenError::new())?;
*self.ordinal.lock().await = restored.0;
Ok(StreamOpenOutcome::Restored)
} else {
*self.ordinal.lock().await = 0;
Ok(StreamOpenOutcome::Initial)
}
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
log_event(&self.log, LifecycleEvent::Update(self.resource.clone()));
let codec = vec_read_codec();
let ordinal = *self.ordinal.lock().await;
ComponentStateEnvelope::encode(
self.namespace.clone(),
&VecPosition(ordinal),
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
log_event(
&self.log,
LifecycleEvent::Close(self.resource.clone(), context.outcome().into()),
);
if self.fail_close.load(std::sync::atomic::Ordering::SeqCst) {
return Err(StreamCloseError::new());
}
Ok(StreamCloseOutcome::Closed)
}
}
struct RecordingReaderOpener {
data: HashMap<String, Vec<u64>>,
log: Log,
fail_open: HashSet<String>,
fail_close: HashMap<String, Arc<std::sync::atomic::AtomicBool>>,
}
impl MultiResourceReaderOpener<u64> for RecordingReaderOpener {
type Reader = VecItemReader;
type Stream = RecordingReaderStream;
async fn open(
&self,
resource: &ResourceIdentity,
resource_ordinal: u32,
delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Reader, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
if self.fail_open.contains(resource.as_str()) {
return Err(MultiResourceOpenError::new(resource_ordinal));
}
let items = self
.data
.get(resource.as_str())
.cloned()
.unwrap_or_default();
let ordinal = Arc::new(AsyncMutex::new(0));
let reader = VecItemReader {
items,
ordinal: Arc::clone(&ordinal),
};
let fail_close = self
.fail_close
.get(resource.as_str())
.cloned()
.unwrap_or_else(|| Arc::new(std::sync::atomic::AtomicBool::new(false)));
let stream = RecordingReaderStream {
ordinal,
namespace: delegate_identity.clone(),
resource: resource.as_str().to_owned(),
log: Arc::clone(&self.log),
fail_close,
};
Ok((reader, stream, StreamStateContract::new(vec_read_codec())))
}
}
#[test]
fn reader_closes_each_delegate_exactly_once_with_resource_boundary_outcome_in_order() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1, 2]);
data.insert("b".to_owned(), vec![3]);
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let opener = RecordingReaderOpener {
data,
log: Arc::clone(&log),
fail_open: HashSet::new(),
fail_close: HashMap::new(),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b"]),
opener,
identity("lifecycle-order"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
loop {
match reader.read(read_context(&token)).await.unwrap() {
ReadOutcome::Item(_) => {}
ReadOutcome::EndOfInput => break,
ReadOutcome::Stopped => panic!("unexpected stop"),
}
}
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
LifecycleEvent::Open("b".to_owned()),
LifecycleEvent::Close("b".to_owned(), RecordedOutcome::ResourceBoundary),
],
"the intermediate delegate (\"a\") and the final delegate (\"b\") must \
each open exactly once and close exactly once, in order, with \
`ResourceBoundary` -- never `Committed`, which would falsely \
claim the enclosing step attempt itself had committed"
);
}
#[test]
fn reader_resource_boundary_close_failure_poisons_the_reader_without_advancing_or_retrying() {
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
data.insert("b".to_owned(), vec![2]);
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let mut fail_close = HashMap::new();
fail_close.insert(
"a".to_owned(),
Arc::new(std::sync::atomic::AtomicBool::new(true)),
);
let opener = RecordingReaderOpener {
data,
log: Arc::clone(&log),
fail_open: HashSet::new(),
fail_close,
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b"]),
opener,
identity("boundary-close-failure"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(1)
);
let envelope = stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap();
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let state = envelope.decode::<MultiResourceState>(&codec).unwrap();
assert_eq!(
state.resource_index, 0,
"checkpoint advancement must not silently occur after a failed boundary close"
);
let result = reader.read(read_context(&token)).await;
assert!(result.is_err(), "a boundary close failure must surface");
for _ in 0..3 {
let retried = reader.read(read_context(&token)).await;
assert!(
retried.is_err(),
"a poisoned instance must keep failing, never silently recover"
);
}
let post_poison_update = stream.update(StreamUpdateContext::new(&token)).await;
assert!(
post_poison_update.is_err(),
"update must fail closed once poisoned, not expose a stale checkpoint"
);
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Update("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
],
"\"a\" must be closed exactly once, ever -- never retried after a \
failure, and \"b\" must never be opened"
);
}
#[test]
fn outer_terminal_close_does_not_double_close_a_reader_delegate_already_retired_at_a_boundary()
{
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
let mut fail_open = HashSet::new();
fail_open.insert("b".to_owned());
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let opener = RecordingReaderOpener {
data,
log: Arc::clone(&log),
fail_open,
fail_close: HashMap::new(),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b"]),
opener,
identity("no-double-close"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
assert_eq!(
reader.read(read_context(&token)).await.unwrap(),
ReadOutcome::Item(1)
);
let result = reader.read(read_context(&token)).await;
assert!(result.is_err());
stream
.close(StreamCloseContext::new(
&token,
StreamRuntimeOutcome::Failed,
))
.await
.unwrap();
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
],
"\"a\" must be closed exactly once total -- the outer terminal \
close must recognize it is already retired and skip it, not \
close it a second time with the step's real outcome"
);
}
#[test]
fn reader_update_fails_closed_when_delegate_reports_the_wrong_namespace() {
struct WrongNamespaceStream {
wrong_namespace: ComponentStreamIdentity,
}
impl ItemStream for WrongNamespaceStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
Ok(if context.inherited_state().is_some() {
StreamOpenOutcome::Restored
} else {
StreamOpenOutcome::Initial
})
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
ComponentStateEnvelope::encode(
self.wrong_namespace.clone(),
&VecPosition(0),
&vec_read_codec(),
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
struct WrongNamespaceOpener {
wrong_namespace: ComponentStreamIdentity,
}
impl MultiResourceReaderOpener<u64> for WrongNamespaceOpener {
type Reader = VecItemReader;
type Stream = WrongNamespaceStream;
async fn open(
&self,
_resource: &ResourceIdentity,
_resource_ordinal: u32,
_delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Reader, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
Ok((
VecItemReader {
items: vec![1],
ordinal: Arc::new(AsyncMutex::new(0)),
},
WrongNamespaceStream {
wrong_namespace: self.wrong_namespace.clone(),
},
StreamStateContract::new(vec_read_codec()),
))
}
}
let opener = WrongNamespaceOpener {
wrong_namespace: identity("attacker-namespace"),
};
let (mut reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a"]),
opener,
identity("expected-namespace"),
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
reader.read(read_context(&token)).await.unwrap();
let result = stream.update(StreamUpdateContext::new(&token)).await;
assert!(
result.is_err(),
"a nested delegate reporting the wrong namespace must fail \
the outer update closed, not be embedded into an outer \
candidate that could reach durable checkpoint state"
);
});
}
#[test]
fn reader_open_fails_closed_on_a_hand_crafted_mismatched_namespace_record() {
let expected_identity = identity("restore-namespace-check");
let wrong_identity = identity("restore-namespace-check-wrong");
let delegate_envelope = ComponentStateEnvelope::encode(
wrong_identity,
&VecPosition(0),
&vec_read_codec(),
StateLimits::default(),
)
.unwrap();
let columns = DelegateEnvelopeColumns::from_envelope(&delegate_envelope).unwrap();
let resources = resource_set(&["a"]);
let state = MultiResourceState {
resource_set_revision: resources.revision(),
resource_index: 0,
delegate: Some(columns),
resource_batches_written: 0,
};
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let outer_envelope = ComponentStateEnvelope::encode(
expected_identity.clone(),
&state,
&codec,
StateLimits::default(),
)
.unwrap();
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
let opener = TestReaderOpener {
data,
fail: HashSet::new(),
};
let (_reader, stream, _contract) = multi_resource_reader::<u64, _>(
resources,
opener,
expected_identity,
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
let result = futures_executor::block_on(
stream.open(StreamOpenContext::new(Some(&outer_envelope), &token)),
);
assert!(
result.is_err(),
"a stored delegate envelope under the wrong namespace must fail \
closed on restore, never be silently reconstructed under the \
expected identity"
);
}
#[derive(Clone, Copy)]
struct LegacyV1MultiResourceSchema;
impl VersionedStateCodec<MultiResourceState> for LegacyV1MultiResourceSchema {
fn schema_id(&self) -> &StateSchemaId {
static SCHEMA: std::sync::OnceLock<StateSchemaId> = std::sync::OnceLock::new();
SCHEMA.get_or_init(|| StateSchemaId::new(MULTI_RESOURCE_SCHEMA).unwrap())
}
fn current_version(&self) -> StateSchemaVersion {
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &MultiResourceState) -> Result<Vec<u8>, StateCodecError> {
let delegate = value.delegate.as_ref().map(|delegate| {
serde_json::json!({
"schema_id": delegate.schema_id,
"schema_version": delegate.schema_version,
"codec_id": delegate.codec_id,
"codec_version": delegate.codec_version,
"checksum_algorithm": delegate.checksum_algorithm,
"checksum_algorithm_version": delegate.checksum_algorithm_version,
"checksum": hex_encode(&delegate.checksum),
"is_external": delegate.is_external,
"payload_inline": delegate.payload_inline_hex,
"payload_external_content_id": delegate.payload_external_content_id_hex,
"payload_external_len": delegate.payload_external_len,
})
});
serde_json::to_vec(&serde_json::json!({
"resource_set_revision": value.resource_set_revision.to_hex(),
"resource_index": value.resource_index,
"delegate": delegate,
"resource_batches_written": value.resource_batches_written,
}))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, _payload: &[u8]) -> Result<MultiResourceState, StateCodecError> {
unreachable!(
"only used to construct a stale v1 fixture; the current (v2) \
codec has no declared upgrade edge from v1, so the shared \
upgrade-chain walk rejects a recorded v1 envelope before \
any codec's `decode` runs on it"
)
}
}
#[test]
fn stale_v1_delegate_record_without_a_namespace_fails_closed_on_restore() {
let identity = identity("legacy-v1-fail-closed");
let inner_envelope = ComponentStateEnvelope::encode(
identity.clone(),
&VecPosition(0),
&vec_read_codec(),
StateLimits::default(),
)
.unwrap();
let columns = DelegateEnvelopeColumns::from_envelope(&inner_envelope).unwrap();
let resources = resource_set(&["a"]);
let state = MultiResourceState {
resource_set_revision: resources.revision(),
resource_index: 0,
delegate: Some(columns),
resource_batches_written: 0,
};
#[allow(clippy::unwrap_used)]
let legacy_codec = DefaultComponentCodec::new(
LegacyV1MultiResourceSchema,
crate::CodecId::new(MULTI_RESOURCE_CODEC).unwrap(),
crate::CodecVersion::new(1).unwrap(),
RestartabilityDeclaration::Restartable,
);
let legacy_envelope = ComponentStateEnvelope::encode(
identity.clone(),
&state,
&legacy_codec,
StateLimits::default(),
)
.unwrap();
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1]);
let opener = TestReaderOpener {
data,
fail: HashSet::new(),
};
let (_reader, stream, _contract) = multi_resource_reader::<u64, _>(
resources,
opener,
identity,
RestartabilityDeclaration::Restartable,
);
let (_source, token) = stop();
let result = futures_executor::block_on(
stream.open(StreamOpenContext::new(Some(&legacy_envelope), &token)),
);
assert!(
result.is_err(),
"a version-1 durable record (pre-#177, no delegate namespace) \
must fail closed on restore -- never silently accepted, and \
never migrated by inventing a namespace it never carried"
);
}
struct VecItemWriter {
resource: String,
sink: Arc<std::sync::Mutex<HashMap<String, Vec<u64>>>>,
committed_len: Arc<AsyncMutex<u64>>,
}
impl ItemWriter<u64> for VecItemWriter {
async fn write<'a>(
&'a self,
items: &'a [u64],
context: WriteContext<'a>,
) -> Result<WriteOutcome, WriterError> {
if context.stop_token().is_stop_requested() {
return Ok(WriteOutcome::Stopped);
}
let mut committed = self.committed_len.lock().await;
self.sink
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.entry(self.resource.clone())
.or_default()
.extend_from_slice(items);
*committed += u64::try_from(items.len()).unwrap();
Ok(WriteOutcome::Written)
}
}
struct VecItemWriterStream {
committed_len: Arc<AsyncMutex<u64>>,
namespace: ComponentStreamIdentity,
}
impl ItemStream for VecItemWriterStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = vec_read_codec();
if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<VecPosition>(&codec)
.map_err(|_| StreamOpenError::new())?;
*self.committed_len.lock().await = restored.0;
Ok(StreamOpenOutcome::Restored)
} else {
*self.committed_len.lock().await = 0;
Ok(StreamOpenOutcome::Initial)
}
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
let codec = vec_read_codec();
let committed = *self.committed_len.lock().await;
ComponentStateEnvelope::encode(
self.namespace.clone(),
&VecPosition(committed),
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
struct TestWriterOpener {
sink: Arc<std::sync::Mutex<HashMap<String, Vec<u64>>>>,
}
impl MultiResourceWriterOpener<u64> for TestWriterOpener {
type Writer = VecItemWriter;
type Stream = VecItemWriterStream;
async fn open(
&self,
resource: &ResourceIdentity,
_resource_ordinal: u32,
delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Writer, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
let committed_len = Arc::new(AsyncMutex::new(0));
let writer = VecItemWriter {
resource: resource.as_str().to_owned(),
sink: Arc::clone(&self.sink),
committed_len: Arc::clone(&committed_len),
};
let stream = VecItemWriterStream {
committed_len,
namespace: delegate_identity.clone(),
};
Ok((writer, stream, StreamStateContract::new(vec_read_codec())))
}
}
fn write_context(stop: &crate::StopToken) -> WriteContext<'_> {
WriteContext::non_transactional(stop)
}
#[test]
fn rollover_writes_batches_to_successive_resources_in_order() {
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b", "c"]),
opener,
identity("rollover"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1, 2], write_context(&token)).await.unwrap();
writer.write(&[3], write_context(&token)).await.unwrap();
writer.write(&[4, 5], write_context(&token)).await.unwrap();
});
let sink = sink.lock().unwrap();
assert_eq!(sink.get("a"), Some(&vec![1, 2]));
assert_eq!(sink.get("b"), Some(&vec![3]));
assert_eq!(sink.get("c"), Some(&vec![4, 5]));
}
#[test]
fn writer_restart_mid_resource_resumes_committed_position() {
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let (_source, token) = stop();
let committed_envelope = futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-restart"),
NoRollover,
RestartabilityDeclaration::Restartable,
)
.unwrap();
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1, 2], write_context(&token)).await.unwrap();
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
let second_attempt_envelope = futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-restart"),
NoRollover,
RestartabilityDeclaration::Restartable,
)
.unwrap();
stream
.open(StreamOpenContext::new(Some(&committed_envelope), &token))
.await
.unwrap();
writer.write(&[3], write_context(&token)).await.unwrap();
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let state = second_attempt_envelope
.decode::<MultiResourceState>(&codec)
.unwrap();
assert_eq!(
state.resource_index, 0,
"must still be resource \"a\", index 0"
);
let delegate = state
.delegate
.expect("a written resource always has delegate state");
let inner = delegate
.to_envelope(identity("writer-restart"), StateLimits::default())
.expect("embedded delegate envelope must decode");
let inner_position = inner.decode::<VecPosition>(&vec_read_codec()).unwrap();
assert_eq!(
inner_position.0, 3,
"committed length must resume from 2 and advance by this attempt's batch, not restart from 0"
);
let sink = sink.lock().unwrap();
assert_eq!(sink.get("a"), Some(&vec![1, 2, 3]));
}
#[test]
fn rollover_counter_survives_restart_and_still_caps_batches_per_resource() {
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let (_source, token) = stop();
let committed_envelope = futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("rollover-restart"),
BatchCountRollover::new(2),
RestartabilityDeclaration::Restartable,
)
.unwrap();
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("rollover-restart"),
BatchCountRollover::new(2),
RestartabilityDeclaration::Restartable,
)
.unwrap();
stream
.open(StreamOpenContext::new(Some(&committed_envelope), &token))
.await
.unwrap();
writer.write(&[2], write_context(&token)).await.unwrap();
writer.write(&[3], write_context(&token)).await.unwrap();
});
let sink = sink.lock().unwrap();
assert_eq!(
sink.get("a"),
Some(&vec![1, 2]),
"resource \"a\" must never exceed its 2-batch cap across the restart"
);
assert_eq!(
sink.get("b"),
Some(&vec![3]),
"the third batch must have rolled over into resource \"b\""
);
}
#[test]
fn stale_resource_set_revision_is_rejected_on_writer_restart() {
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let (_source, token) = stop();
let committed_envelope = futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-revision"),
NoRollover,
RestartabilityDeclaration::Restartable,
)
.unwrap();
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap()
});
futures_executor::block_on(async {
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (_writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b", "c"]),
opener,
identity("writer-revision"),
NoRollover,
RestartabilityDeclaration::Restartable,
)
.unwrap();
let result = stream
.open(StreamOpenContext::new(Some(&committed_envelope), &token))
.await;
assert!(
result.is_err(),
"changed resource set must be rejected on writer restart too"
);
});
}
struct RecordingWriterStream {
committed_len: Arc<AsyncMutex<u64>>,
namespace: ComponentStreamIdentity,
resource: String,
log: Log,
fail_close: Arc<std::sync::atomic::AtomicBool>,
}
impl ItemStream for RecordingWriterStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
log_event(&self.log, LifecycleEvent::Open(self.resource.clone()));
let codec = vec_read_codec();
if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<VecPosition>(&codec)
.map_err(|_| StreamOpenError::new())?;
*self.committed_len.lock().await = restored.0;
Ok(StreamOpenOutcome::Restored)
} else {
*self.committed_len.lock().await = 0;
Ok(StreamOpenOutcome::Initial)
}
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
log_event(&self.log, LifecycleEvent::Update(self.resource.clone()));
let codec = vec_read_codec();
let committed = *self.committed_len.lock().await;
ComponentStateEnvelope::encode(
self.namespace.clone(),
&VecPosition(committed),
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
log_event(
&self.log,
LifecycleEvent::Close(self.resource.clone(), context.outcome().into()),
);
if self.fail_close.load(std::sync::atomic::Ordering::SeqCst) {
return Err(StreamCloseError::new());
}
Ok(StreamCloseOutcome::Closed)
}
}
struct RecordingWriterOpener {
sink: Arc<std::sync::Mutex<HashMap<String, Vec<u64>>>>,
log: Log,
fail_open: HashSet<String>,
fail_close: HashMap<String, Arc<std::sync::atomic::AtomicBool>>,
}
impl MultiResourceWriterOpener<u64> for RecordingWriterOpener {
type Writer = VecItemWriter;
type Stream = RecordingWriterStream;
async fn open(
&self,
resource: &ResourceIdentity,
resource_ordinal: u32,
delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Writer, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
if self.fail_open.contains(resource.as_str()) {
return Err(MultiResourceOpenError::new(resource_ordinal));
}
let committed_len = Arc::new(AsyncMutex::new(0));
let writer = VecItemWriter {
resource: resource.as_str().to_owned(),
sink: Arc::clone(&self.sink),
committed_len: Arc::clone(&committed_len),
};
let fail_close = self
.fail_close
.get(resource.as_str())
.cloned()
.unwrap_or_else(|| Arc::new(std::sync::atomic::AtomicBool::new(false)));
let stream = RecordingWriterStream {
committed_len,
namespace: delegate_identity.clone(),
resource: resource.as_str().to_owned(),
log: Arc::clone(&self.log),
fail_close,
};
Ok((writer, stream, StreamStateContract::new(vec_read_codec())))
}
}
#[test]
fn writer_rollover_closes_outgoing_delegate_exactly_once_with_resource_boundary_outcome() {
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let opener = RecordingWriterOpener {
sink: Arc::clone(&sink),
log: Arc::clone(&log),
fail_open: HashSet::new(),
fail_close: HashMap::new(),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-lifecycle-order"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
writer.write(&[2], write_context(&token)).await.unwrap();
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
LifecycleEvent::Open("b".to_owned()),
],
"rollover must close the outgoing delegate (\"a\") exactly once, \
with `ResourceBoundary` -- never `Committed` -- before opening \
the next resource (\"b\")"
);
}
#[test]
fn writer_resource_boundary_close_failure_poisons_the_writer_without_rolling_over_or_retrying()
{
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let mut fail_close = HashMap::new();
fail_close.insert(
"a".to_owned(),
Arc::new(std::sync::atomic::AtomicBool::new(true)),
);
let opener = RecordingWriterOpener {
sink: Arc::clone(&sink),
log: Arc::clone(&log),
fail_open: HashSet::new(),
fail_close,
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-boundary-close-failure"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
let envelope = stream
.update(StreamUpdateContext::new(&token))
.await
.unwrap();
let codec = multi_resource_position_codec(RestartabilityDeclaration::Restartable);
let state = envelope.decode::<MultiResourceState>(&codec).unwrap();
assert_eq!(
state.resource_index, 0,
"checkpoint advancement must not silently occur after a failed boundary close"
);
let result = writer.write(&[2], write_context(&token)).await;
assert!(result.is_err(), "a boundary close failure must surface");
for _ in 0..3 {
let retried = writer.write(&[2], write_context(&token)).await;
assert!(
retried.is_err(),
"a poisoned instance must keep failing, never silently recover"
);
}
let post_poison_update = stream.update(StreamUpdateContext::new(&token)).await;
assert!(
post_poison_update.is_err(),
"update must fail closed once poisoned, not expose a stale checkpoint"
);
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Update("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
],
"\"a\" must be closed exactly once, ever -- never retried after a \
failure, and \"b\" must never be opened"
);
let sink = sink.lock().unwrap();
assert_eq!(
sink.get("a"),
Some(&vec![1]),
"the item from the failed rollover attempt must never have \
reached \"b\""
);
assert_eq!(
sink.get("b"),
None,
"\"b\" must never receive any items once the writer is poisoned"
);
}
#[test]
fn outer_terminal_close_does_not_double_close_a_writer_delegate_already_retired_at_a_boundary()
{
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let log: Log = Arc::new(std::sync::Mutex::new(Vec::new()));
let mut fail_open = HashSet::new();
fail_open.insert("b".to_owned());
let opener = RecordingWriterOpener {
sink: Arc::clone(&sink),
log: Arc::clone(&log),
fail_open,
fail_close: HashMap::new(),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("writer-no-double-close"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
let result = writer.write(&[2], write_context(&token)).await;
assert!(result.is_err());
stream
.close(StreamCloseContext::new(
&token,
StreamRuntimeOutcome::Failed,
))
.await
.unwrap();
});
let events = log.lock().unwrap().clone();
assert_eq!(
events,
vec![
LifecycleEvent::Open("a".to_owned()),
LifecycleEvent::Close("a".to_owned(), RecordedOutcome::ResourceBoundary),
],
"\"a\" must be closed exactly once total -- the outer terminal \
close must recognize it is already retired and skip it, not \
close it a second time with the step's real outcome"
);
}
#[test]
fn concurrent_writes_never_race_the_poisoning_transition() {
struct BlockingCloseStream {
namespace: ComponentStreamIdentity,
proceed: Arc<tokio::sync::Notify>,
}
impl ItemStream for BlockingCloseStream {
async fn open(
&self,
_context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
Ok(StreamOpenOutcome::Initial)
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
ComponentStateEnvelope::encode(
self.namespace.clone(),
&VecPosition(0),
&vec_read_codec(),
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
self.proceed.notified().await;
Err(StreamCloseError::new())
}
}
struct BlockingOpener {
sink: Arc<std::sync::Mutex<HashMap<String, Vec<u64>>>>,
proceed: Arc<tokio::sync::Notify>,
}
impl MultiResourceWriterOpener<u64> for BlockingOpener {
type Writer = VecItemWriter;
type Stream = BlockingCloseStream;
async fn open(
&self,
resource: &ResourceIdentity,
_resource_ordinal: u32,
delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Writer, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
let writer = VecItemWriter {
resource: resource.as_str().to_owned(),
sink: Arc::clone(&self.sink),
committed_len: Arc::new(AsyncMutex::new(0)),
};
let stream = BlockingCloseStream {
namespace: delegate_identity.clone(),
proceed: Arc::clone(&self.proceed),
};
Ok((writer, stream, StreamStateContract::new(vec_read_codec())))
}
}
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let proceed = Arc::new(tokio::sync::Notify::new());
let opener = BlockingOpener {
sink: Arc::clone(&sink),
proceed: Arc::clone(&proceed),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("concurrent-writer-race"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
});
let waker = std::task::Waker::noop();
let mut cx = std::task::Context::from_waker(waker);
let mut a_fut = std::pin::pin!(writer.write(&[2], write_context(&token)));
assert!(matches!(
a_fut.as_mut().poll(&mut cx),
std::task::Poll::Pending
));
let mut b_fut = std::pin::pin!(writer.write(&[3], write_context(&token)));
assert!(matches!(
b_fut.as_mut().poll(&mut cx),
std::task::Poll::Pending
));
proceed.notify_one();
let a_result = loop {
if let std::task::Poll::Ready(result) = a_fut.as_mut().poll(&mut cx) {
break result;
}
};
assert!(
a_result.is_err(),
"the closing write must surface the boundary close failure"
);
let b_result = loop {
if let std::task::Poll::Ready(result) = b_fut.as_mut().poll(&mut cx) {
break result;
}
};
assert!(
b_result.is_err(),
"a concurrent write must never observe a stale pre-poisoning \
state and must never roll over into \"b\" once the writer is \
poisoned"
);
let sink = sink.lock().unwrap();
assert_eq!(
sink.get("b"),
None,
"\"b\" must never receive any items once the writer is \
poisoned, even under a concurrent write"
);
}
#[test]
fn writer_update_fails_closed_when_delegate_reports_the_wrong_namespace() {
struct WrongNamespaceWriterStream {
wrong_namespace: ComponentStreamIdentity,
}
impl ItemStream for WrongNamespaceWriterStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
Ok(if context.inherited_state().is_some() {
StreamOpenOutcome::Restored
} else {
StreamOpenOutcome::Initial
})
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
ComponentStateEnvelope::encode(
self.wrong_namespace.clone(),
&VecPosition(0),
&vec_read_codec(),
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
struct WrongNamespaceWriterOpener {
sink: Arc<std::sync::Mutex<HashMap<String, Vec<u64>>>>,
wrong_namespace: ComponentStreamIdentity,
}
impl MultiResourceWriterOpener<u64> for WrongNamespaceWriterOpener {
type Writer = VecItemWriter;
type Stream = WrongNamespaceWriterStream;
async fn open(
&self,
resource: &ResourceIdentity,
_resource_ordinal: u32,
_delegate_identity: &ComponentStreamIdentity,
) -> Result<(Self::Writer, Self::Stream, StreamStateContract), MultiResourceOpenError>
{
Ok((
VecItemWriter {
resource: resource.as_str().to_owned(),
sink: Arc::clone(&self.sink),
committed_len: Arc::new(AsyncMutex::new(0)),
},
WrongNamespaceWriterStream {
wrong_namespace: self.wrong_namespace.clone(),
},
StreamStateContract::new(vec_read_codec()),
))
}
}
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let opener = WrongNamespaceWriterOpener {
sink,
wrong_namespace: identity("attacker-namespace-writer"),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a"]),
opener,
identity("expected-namespace-writer"),
NoRollover,
RestartabilityDeclaration::Restartable,
)
.unwrap();
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
writer.write(&[1], write_context(&token)).await.unwrap();
let result = stream.update(StreamUpdateContext::new(&token)).await;
assert!(
result.is_err(),
"a nested delegate reporting the wrong namespace must fail \
the outer update closed, not be embedded into an outer \
candidate that could reach durable checkpoint state"
);
});
}
#[test]
fn peek_reader_over_multi_resource_reader_crosses_resource_boundary_without_corrupting_order() {
use crate::item_components::PeekReader;
let mut data = HashMap::new();
data.insert("a".to_owned(), vec![1, 2]);
data.insert("b".to_owned(), vec![3, 4]);
let opener = TestReaderOpener {
data,
fail: HashSet::new(),
};
let (reader, stream, _contract) = multi_resource_reader::<u64, _>(
resource_set(&["a", "b"]),
opener,
identity("peek-composition"),
RestartabilityDeclaration::Restartable,
);
let mut peeking = PeekReader::new(reader);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
let mut collected = Vec::new();
loop {
match peeking.peek(ReadContext::new(&token)).await.unwrap() {
crate::item_components::PeekOutcome::Item(_) => {}
crate::item_components::PeekOutcome::EndOfInput => break,
crate::item_components::PeekOutcome::Stopped => panic!("unexpected stop"),
}
match peeking.read(ReadContext::new(&token)).await.unwrap() {
ReadOutcome::Item(item) => collected.push(item),
ReadOutcome::EndOfInput => break,
ReadOutcome::Stopped => panic!("unexpected stop"),
}
}
assert_eq!(collected, vec![1, 2, 3, 4]);
});
}
#[test]
fn synchronized_writer_over_multi_resource_writer_preserves_rollover_order() {
use crate::item_components::SynchronizedWriter;
let sink = Arc::new(std::sync::Mutex::new(HashMap::new()));
let opener = TestWriterOpener {
sink: Arc::clone(&sink),
};
let (writer, stream, _contract) = multi_resource_writer::<u64, _, _>(
resource_set(&["a", "b"]),
opener,
identity("sync-composition"),
BatchCountRollover::new(1),
RestartabilityDeclaration::Restartable,
)
.unwrap();
let synchronized = SynchronizedWriter::new(writer);
let (_source, token) = stop();
futures_executor::block_on(async {
stream
.open(StreamOpenContext::new(None, &token))
.await
.unwrap();
synchronized
.write(&[1], write_context(&token))
.await
.unwrap();
synchronized
.write(&[2], write_context(&token))
.await
.unwrap();
});
let sink = sink.lock().unwrap();
assert_eq!(sink.get("a"), Some(&vec![1]));
assert_eq!(sink.get("b"), Some(&vec![2]));
}
}