use std::fs::File;
use std::io::{self, BufRead, BufReader, Read, Seek, SeekFrom, Write};
use std::path::Path;
use std::sync::{Arc, Mutex, PoisonError};
use crate::{
CodecId, CodecVersion, ComponentStateEnvelope, ComponentStreamIdentity, DefaultComponentCodec,
FailureCategory, ReadContext, ReadOutcome, ReaderError, RestartabilityDeclaration,
StateCodecError, StateLimits, StateSchemaId, StateSchemaVersion, StateSensitivity,
StreamCloseContext, StreamCloseError, StreamCloseOutcome, StreamOpenContext, StreamOpenError,
StreamOpenOutcome, StreamStateContract, StreamUpdateContext, StreamUpdateError,
VersionedStateCodec, WriteContext, WriteOutcome, WriterError,
};
pub const DEFAULT_MAX_RECORD_BYTES: usize = 1024 * 1024;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum FixedWidthTerminator {
Lf,
CrLf,
}
impl FixedWidthTerminator {
const fn bytes(self) -> &'static [u8] {
match self {
Self::Lf => b"\n",
Self::CrLf => b"\r\n",
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FixedWidthField {
name: Option<String>,
width: usize,
}
impl FixedWidthField {
#[must_use]
pub const fn new(width: usize) -> Self {
Self { name: None, width }
}
#[must_use]
pub fn named(name: impl Into<String>, width: usize) -> Self {
Self {
name: Some(name.into()),
width,
}
}
#[must_use]
pub const fn width(&self) -> usize {
self.width
}
#[must_use]
pub fn name(&self) -> Option<&str> {
self.name.as_deref()
}
}
#[derive(Clone, Debug)]
pub struct FixedWidthLayout {
fields: Vec<FixedWidthField>,
record_width: usize,
terminator: FixedWidthTerminator,
max_record_bytes: usize,
}
impl FixedWidthLayout {
#[must_use]
pub fn new(fields: Vec<FixedWidthField>) -> Self {
assert!(!fields.is_empty(), "a fixed-width layout needs >= 1 field");
let record_width = fields.iter().map(FixedWidthField::width).sum();
Self {
fields,
record_width,
terminator: FixedWidthTerminator::Lf,
max_record_bytes: DEFAULT_MAX_RECORD_BYTES,
}
}
#[must_use]
pub const fn with_terminator(mut self, terminator: FixedWidthTerminator) -> Self {
self.terminator = terminator;
self
}
#[must_use]
pub const fn with_max_record_bytes(mut self, max_record_bytes: usize) -> Self {
self.max_record_bytes = max_record_bytes;
self
}
#[must_use]
pub fn fields(&self) -> &[FixedWidthField] {
&self.fields
}
#[must_use]
pub const fn record_width(&self) -> usize {
self.record_width
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FixedWidthRecord {
fields: Vec<String>,
names: Option<Arc<[Option<String>]>>,
}
impl FixedWidthRecord {
#[must_use]
pub const fn new(fields: Vec<String>) -> Self {
Self {
fields,
names: None,
}
}
#[must_use]
pub fn fields(&self) -> &[String] {
&self.fields
}
#[must_use]
pub fn into_fields(self) -> Vec<String> {
self.fields
}
#[must_use]
pub fn get(&self, index: usize) -> Option<&str> {
self.fields.get(index).map(String::as_str)
}
#[must_use]
pub fn field(&self, name: &str) -> Option<&str> {
let names = self.names.as_ref()?;
let index = names
.iter()
.position(|candidate| candidate.as_deref() == Some(name))?;
self.get(index)
}
#[must_use]
pub fn len(&self) -> usize {
self.fields.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.fields.is_empty()
}
}
impl From<Vec<String>> for FixedWidthRecord {
fn from(fields: Vec<String>) -> Self {
Self::new(fields)
}
}
const READER_SCHEMA: &str = "oxide-batch.fixed-width-reader-position";
const READER_CODEC: &str = "oxide-batch.fixed-width-reader-position-codec";
#[derive(Clone, Copy)]
struct ReaderPositionSchema;
impl VersionedStateCodec<u64> for ReaderPositionSchema {
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(READER_SCHEMA).unwrap())
}
fn current_version(&self) -> StateSchemaVersion {
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema version cannot fail validation"
)]
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &u64) -> Result<Vec<u8>, StateCodecError> {
serde_json::to_vec(&serde_json::json!({ "byte": value }))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, payload: &[u8]) -> Result<u64, StateCodecError> {
let value: serde_json::Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
value
.get("byte")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)
}
}
#[allow(
clippy::unwrap_used,
reason = "fixed literal identities cannot fail validation"
)]
fn reader_position_codec() -> DefaultComponentCodec<ReaderPositionSchema> {
DefaultComponentCodec::new(
ReaderPositionSchema,
CodecId::new(READER_CODEC).unwrap(),
CodecVersion::new(1).unwrap(),
RestartabilityDeclaration::Restartable,
)
.with_sensitivity(StateSensitivity::NonSensitive)
}
const WRITER_SCHEMA: &str = "oxide-batch.fixed-width-writer-position";
const WRITER_CODEC: &str = "oxide-batch.fixed-width-writer-position-codec";
#[derive(Clone, Copy)]
struct WriterPositionSchema;
impl VersionedStateCodec<u64> for WriterPositionSchema {
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(WRITER_SCHEMA).unwrap())
}
fn current_version(&self) -> StateSchemaVersion {
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema version cannot fail validation"
)]
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &u64) -> Result<Vec<u8>, StateCodecError> {
serde_json::to_vec(&serde_json::json!({ "committed_bytes": value }))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, payload: &[u8]) -> Result<u64, StateCodecError> {
let value: serde_json::Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
value
.get("committed_bytes")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)
}
}
#[allow(
clippy::unwrap_used,
reason = "fixed literal identities cannot fail validation"
)]
fn writer_position_codec() -> DefaultComponentCodec<WriterPositionSchema> {
DefaultComponentCodec::new(
WriterPositionSchema,
CodecId::new(WRITER_CODEC).unwrap(),
CodecVersion::new(1).unwrap(),
RestartabilityDeclaration::Restartable,
)
.with_sensitivity(StateSensitivity::NonSensitive)
}
enum BoundedLine {
Line { bytes: Vec<u8>, consumed: u64 },
Eof,
TooLong { consumed: u64 },
}
fn copy_bounded(buf: &mut Vec<u8>, segment: &[u8], max_bytes: usize, too_long: &mut bool) {
if *too_long {
return;
}
let budget = max_bytes.saturating_sub(buf.len());
if segment.len() <= budget {
buf.extend_from_slice(segment);
} else {
buf.extend_from_slice(&segment[..budget]);
*too_long = true;
buf.clear();
buf.shrink_to_fit();
}
}
fn read_bounded_line<R: BufRead>(reader: &mut R, max_bytes: usize) -> io::Result<BoundedLine> {
let mut buf: Vec<u8> = Vec::new();
let mut consumed_total: u64 = 0;
let mut too_long = false;
loop {
let available = reader.fill_buf()?;
if available.is_empty() {
if consumed_total == 0 {
return Ok(BoundedLine::Eof);
}
return Ok(if too_long {
BoundedLine::TooLong {
consumed: consumed_total,
}
} else {
BoundedLine::Line {
bytes: buf,
consumed: consumed_total,
}
});
}
if let Some(newline) = available.iter().position(|&byte| byte == b'\n') {
copy_bounded(&mut buf, &available[..newline], max_bytes, &mut too_long);
let take = newline + 1;
reader.consume(take);
consumed_total += take as u64;
return Ok(if too_long {
BoundedLine::TooLong {
consumed: consumed_total,
}
} else {
if buf.last() == Some(&b'\r') {
buf.pop();
}
BoundedLine::Line {
bytes: buf,
consumed: consumed_total,
}
});
}
copy_bounded(&mut buf, available, max_bytes, &mut too_long);
let chunk_len = available.len();
reader.consume(chunk_len);
consumed_total += chunk_len as u64;
}
}
pub struct FixedWidthReader<Src> {
source: BufReader<Src>,
layout: FixedWidthLayout,
names: Option<Arc<[Option<String>]>>,
position: Arc<Mutex<u64>>,
seeked: bool,
}
impl<Src: Read + Seek> FixedWidthReader<Src> {
fn seek_if_needed(&mut self) -> Result<(), ReaderError> {
if self.seeked {
return Ok(());
}
self.seeked = true;
let target = *self.position.lock().unwrap_or_else(PoisonError::into_inner);
if target == 0 {
return Ok(());
}
self.source
.seek(SeekFrom::Start(target))
.map_err(|_| ReaderError::new())?;
Ok(())
}
}
impl<I, Src> crate::ItemReader<I> for FixedWidthReader<Src>
where
I: From<FixedWidthRecord> + 'static,
Src: Read + Seek + Send,
{
async fn read(&mut self, _context: ReadContext<'_>) -> Result<ReadOutcome<I>, ReaderError> {
self.seek_if_needed()?;
let before = *self.position.lock().unwrap_or_else(PoisonError::into_inner);
let outcome = read_bounded_line(&mut self.source, self.layout.max_record_bytes)
.map_err(|_| ReaderError::with_category(FailureCategory::TransientInfrastructure))?;
match outcome {
BoundedLine::Eof => Ok(ReadOutcome::EndOfInput),
BoundedLine::TooLong { consumed } => {
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = before + consumed;
Err(ReaderError::new().with_checkpoint_advanced(true))
}
BoundedLine::Line { bytes, consumed } => {
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = before + consumed;
if bytes.len() != self.layout.record_width() {
return Err(ReaderError::new().with_checkpoint_advanced(true));
}
let mut fields = Vec::with_capacity(self.layout.fields.len());
let mut offset = 0usize;
for field in &self.layout.fields {
let span = &bytes[offset..offset + field.width()];
let text = std::str::from_utf8(span)
.map_err(|_| ReaderError::new().with_checkpoint_advanced(true))?;
fields.push(text.to_owned());
offset += field.width();
}
let record = FixedWidthRecord {
fields,
names: self.names.clone(),
};
Ok(ReadOutcome::Item(record.into()))
}
}
}
}
pub struct FixedWidthReaderStream {
position: Arc<Mutex<u64>>,
namespace: ComponentStreamIdentity,
}
impl crate::ItemStream for FixedWidthReaderStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = reader_position_codec();
if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<u64>(&codec)
.map_err(|_| StreamOpenError::new())?;
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = restored;
Ok(StreamOpenOutcome::Restored)
} else {
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = 0;
Ok(StreamOpenOutcome::Initial)
}
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
let codec = reader_position_codec();
let current = *self.position.lock().unwrap_or_else(PoisonError::into_inner);
ComponentStateEnvelope::encode(
self.namespace.clone(),
¤t,
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
pub fn fixed_width_reader<I, Src>(
source: Src,
layout: FixedWidthLayout,
identity: ComponentStreamIdentity,
) -> (
FixedWidthReader<Src>,
FixedWidthReaderStream,
StreamStateContract,
)
where
I: From<FixedWidthRecord> + 'static,
Src: Read + Seek + Send,
{
let position = Arc::new(Mutex::new(0));
let names: Option<Arc<[Option<String>]>> =
if layout.fields.iter().any(|field| field.name.is_some()) {
Some(
layout
.fields
.iter()
.map(|field| field.name.clone())
.collect(),
)
} else {
None
};
let reader = FixedWidthReader {
source: BufReader::new(source),
layout,
names,
position: Arc::clone(&position),
seeked: false,
};
let stream = FixedWidthReaderStream {
position,
namespace: identity,
};
let contract = StreamStateContract::new(reader_position_codec());
(reader, stream, contract)
}
pub fn fixed_width_file_reader<I>(
path: impl AsRef<Path>,
layout: FixedWidthLayout,
identity: ComponentStreamIdentity,
) -> io::Result<(
FixedWidthReader<File>,
FixedWidthReaderStream,
StreamStateContract,
)>
where
I: From<FixedWidthRecord> + 'static,
{
let file = File::open(path)?;
Ok(fixed_width_reader::<I, File>(file, layout, identity))
}
pub struct FixedWidthWriter {
state: Arc<Mutex<FixedWidthWriterState>>,
layout: FixedWidthLayout,
}
struct FixedWidthWriterState {
file: File,
committed_bytes: u64,
}
impl<I> crate::ItemWriter<I> for FixedWidthWriter
where
I: Send + Sync + Clone + Into<FixedWidthRecord>,
{
async fn write<'a>(
&'a self,
items: &'a [I],
context: WriteContext<'a>,
) -> Result<WriteOutcome, WriterError> {
if context.stop_token().is_stop_requested() {
return Ok(WriteOutcome::Stopped);
}
let mut buffer: Vec<u8> = Vec::new();
for item in items {
let record: FixedWidthRecord = item.clone().into();
if record.fields.len() != self.layout.fields.len() {
return Err(WriterError::new());
}
for (field, declared) in record.fields.iter().zip(&self.layout.fields) {
if field.len() != declared.width() {
return Err(WriterError::new());
}
buffer.extend_from_slice(field.as_bytes());
}
buffer.extend_from_slice(self.layout.terminator.bytes());
}
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
state
.file
.write_all(&buffer)
.map_err(|_| WriterError::new())?;
state.file.sync_data().map_err(|_| WriterError::new())?;
let candidate_bytes = state
.file
.stream_position()
.map_err(|_| WriterError::new())?;
state.committed_bytes = candidate_bytes;
Ok(WriteOutcome::Written)
}
}
pub struct FixedWidthWriterStream {
state: Arc<Mutex<FixedWidthWriterState>>,
namespace: ComponentStreamIdentity,
}
impl crate::ItemStream for FixedWidthWriterStream {
async fn open(
&self,
context: StreamOpenContext<'_>,
) -> Result<StreamOpenOutcome, StreamOpenError> {
let codec = writer_position_codec();
let (target, outcome) = if let Some(envelope) = context.inherited_state() {
let restored = envelope
.decode::<u64>(&codec)
.map_err(|_| StreamOpenError::new())?;
(restored, StreamOpenOutcome::Restored)
} else {
(0, StreamOpenOutcome::Initial)
};
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
let actual_len = state
.file
.metadata()
.map_err(|_| StreamOpenError::new())?
.len();
if actual_len < target {
return Err(StreamOpenError::with_category(FailureCategory::Invariant));
}
if actual_len != target {
state
.file
.set_len(target)
.map_err(|_| StreamOpenError::new())?;
}
state
.file
.seek(SeekFrom::Start(target))
.map_err(|_| StreamOpenError::new())?;
state.committed_bytes = target;
Ok(outcome)
}
async fn update(
&self,
_context: StreamUpdateContext<'_>,
) -> Result<ComponentStateEnvelope, StreamUpdateError> {
let codec = writer_position_codec();
let current = self
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.committed_bytes;
ComponentStateEnvelope::encode(
self.namespace.clone(),
¤t,
&codec,
StateLimits::default(),
)
.map_err(|_| StreamUpdateError::new())
}
async fn close(
&self,
_context: StreamCloseContext<'_>,
) -> Result<StreamCloseOutcome, StreamCloseError> {
Ok(StreamCloseOutcome::Closed)
}
}
pub fn fixed_width_writer(
path: impl AsRef<Path>,
layout: FixedWidthLayout,
identity: ComponentStreamIdentity,
) -> io::Result<(
FixedWidthWriter,
FixedWidthWriterStream,
StreamStateContract,
)> {
let file = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(path)?;
let state = Arc::new(Mutex::new(FixedWidthWriterState {
file,
committed_bytes: 0,
}));
let writer = FixedWidthWriter {
state: Arc::clone(&state),
layout,
};
let stream = FixedWidthWriterStream {
state,
namespace: identity,
};
let contract = StreamStateContract::new(writer_position_codec());
Ok((writer, stream, contract))
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use std::thread;
use super::*;
use crate::{ComponentStatePayload, ItemStream, ItemWriter, StopSource};
fn identity() -> ComponentStreamIdentity {
ComponentStreamIdentity::new("oxide-batch.fixed-width-writer-unit-test")
.expect("static identity is valid")
}
fn temp_path(name: &str) -> std::path::PathBuf {
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time moves forward")
.as_nanos();
std::env::temp_dir().join(format!(
"oxide-batch-167-fixed-width-unit-{name}-{nonce}.txt"
))
}
fn committed_byte(envelope: &ComponentStateEnvelope) -> u64 {
let ComponentStatePayload::Inline(payload) = envelope.payload().unwrap() else {
panic!("writer checkpoints must use an inline payload");
};
serde_json::from_slice::<serde_json::Value>(&payload)
.unwrap()
.get("committed_bytes")
.and_then(serde_json::Value::as_u64)
.expect("writer checkpoint committed_bytes")
}
#[test]
fn concurrent_write_is_fully_serialized_by_one_lock_never_a_stale_commit() {
let path = temp_path("serialized");
let layout = FixedWidthLayout::new(vec![FixedWidthField::new(1)]);
let (writer, stream, _contract) = fixed_width_writer(&path, layout, identity()).unwrap();
let writer = Arc::new(writer);
let (_open_stop_source, open_stop_token) = StopSource::new();
futures_executor::block_on(stream.open(StreamOpenContext::new(None, &open_stop_token)))
.unwrap();
let mut guard = writer.state.lock().unwrap_or_else(PoisonError::into_inner);
let writer_for_thread = Arc::clone(&writer);
let handle = thread::spawn(move || {
let (_stop_source, stop_token) = StopSource::new();
futures_executor::block_on(writer_for_thread.write(
&[FixedWidthRecord::new(vec!["2".to_owned()])],
WriteContext::non_transactional(&stop_token),
))
.unwrap();
});
guard.file.write_all(b"1\n").unwrap();
guard.file.sync_data().unwrap();
guard.committed_bytes = guard.file.stream_position().unwrap();
drop(guard);
handle.join().unwrap();
let bytes = std::fs::read(&path).unwrap();
assert_eq!(
bytes,
b"1\n2\n".to_vec(),
"both writes must land on disk exactly once each, in the order the lock serialized \
them"
);
let envelope = futures_executor::block_on(
stream.update(crate::StreamUpdateContext::new(&open_stop_token)),
)
.unwrap();
let committed = committed_byte(&envelope);
assert_eq!(
committed,
bytes.len() as u64,
"the published checkpoint must equal the real file length, never a stale, earlier \
write's byte count"
);
assert_eq!(
committed % 2,
0,
"every committed position must land on a record_width(1) + terminator(1) boundary"
);
let _ = std::fs::remove_file(&path);
}
}