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 csv::Terminator as CsvTerminator;
use csv_core::{ReadRecordResult, Reader as CoreReader};
use crate::{
CodecId, CodecVersion, ComponentStateEnvelope, ComponentStreamIdentity, DefaultComponentCodec,
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;
const DISCARD_OUTPUT_BYTES: usize = 4096;
const DISCARD_ENDS_LEN: usize = 1024;
const INITIAL_GROWTH_BYTES: usize = 256;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum DelimitedTerminator {
Lf,
CrLf,
}
#[derive(Clone, Copy, Debug)]
pub struct DelimitedDialect {
delimiter: u8,
quote: u8,
escape: Option<u8>,
double_quote: bool,
has_headers: bool,
flexible: bool,
terminator: DelimitedTerminator,
max_record_bytes: usize,
}
impl DelimitedDialect {
#[must_use]
pub const fn csv() -> Self {
Self {
delimiter: b',',
quote: b'"',
escape: None,
double_quote: true,
has_headers: false,
flexible: false,
terminator: DelimitedTerminator::Lf,
max_record_bytes: DEFAULT_MAX_RECORD_BYTES,
}
}
#[must_use]
pub const fn with_delimiter(mut self, delimiter: u8) -> Self {
self.delimiter = delimiter;
self
}
#[must_use]
pub const fn with_quote(mut self, quote: u8) -> Self {
self.quote = quote;
self
}
#[must_use]
pub const fn with_escape(mut self, escape: Option<u8>) -> Self {
self.escape = escape;
self
}
#[must_use]
pub const fn with_double_quote(mut self, enabled: bool) -> Self {
self.double_quote = enabled;
self
}
#[must_use]
pub const fn with_headers(mut self, enabled: bool) -> Self {
self.has_headers = enabled;
self
}
#[must_use]
pub const fn with_flexible(mut self, enabled: bool) -> Self {
self.flexible = enabled;
self
}
#[must_use]
pub const fn with_terminator(mut self, terminator: DelimitedTerminator) -> 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
}
fn core_reader(self) -> CoreReader {
let mut builder = csv_core::ReaderBuilder::new();
builder
.delimiter(self.delimiter)
.quote(self.quote)
.escape(self.escape)
.double_quote(self.double_quote);
builder.build()
}
fn writer_builder(self) -> csv::WriterBuilder {
let mut builder = csv::WriterBuilder::new();
builder
.delimiter(self.delimiter)
.quote(self.quote)
.double_quote(self.double_quote)
.flexible(self.flexible)
.terminator(match self.terminator {
DelimitedTerminator::Lf => CsvTerminator::Any(b'\n'),
DelimitedTerminator::CrLf => CsvTerminator::CRLF,
});
if let Some(escape) = self.escape {
builder.escape(escape);
}
builder
}
}
impl Default for DelimitedDialect {
fn default() -> Self {
Self::csv()
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DelimitedRecord {
fields: Vec<String>,
headers: Option<Arc<[String]>>,
}
impl DelimitedRecord {
#[must_use]
pub const fn new(fields: Vec<String>) -> Self {
Self {
fields,
headers: 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 headers = self.headers.as_ref()?;
let index = headers.iter().position(|header| header == 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 DelimitedRecord {
fn from(fields: Vec<String>) -> Self {
Self::new(fields)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct StoredPosition {
byte: u64,
line: u64,
record: u64,
}
impl StoredPosition {
const START: Self = Self {
byte: 0,
line: 1,
record: 0,
};
}
const READER_SCHEMA: &str = "oxide-batch.delimited-reader-position";
const READER_CODEC: &str = "oxide-batch.delimited-reader-position-codec";
#[derive(Clone, Copy)]
struct ReaderPositionSchema;
impl VersionedStateCodec<StoredPosition> 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: &StoredPosition) -> Result<Vec<u8>, StateCodecError> {
serde_json::to_vec(&serde_json::json!({
"byte": value.byte,
"line": value.line,
"record": value.record,
}))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode(&self, payload: &[u8]) -> Result<StoredPosition, StateCodecError> {
let value: serde_json::Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
let byte = value
.get("byte")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)?;
let line = value
.get("line")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)?;
let record = value
.get("record")
.and_then(serde_json::Value::as_u64)
.ok_or(StateCodecError::InvalidPayload)?;
Ok(StoredPosition { byte, line, record })
}
}
#[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.delimited-writer-position";
const WRITER_CODEC: &str = "oxide-batch.delimited-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 RawRecord {
Fields {
consumed: u64,
},
Oversized {
consumed: u64,
},
Eof,
}
pub struct DelimitedReader<Src> {
source: BufReader<Src>,
core: CoreReader,
output: Vec<u8>,
ends: Vec<usize>,
max_record_bytes: usize,
has_headers: bool,
flexible: bool,
first_field_count: Option<usize>,
headers: Option<Arc<[String]>>,
position: Arc<Mutex<StoredPosition>>,
seeked: bool,
}
impl<Src: Read> DelimitedReader<Src> {
fn read_raw_record(&mut self) -> io::Result<RawRecord> {
self.output.clear();
self.ends.clear();
let mut output_len = 0usize;
let mut ends_len = 0usize;
let mut consumed: u64 = 0;
let mut oversized = false;
let max_record_bytes = u64::try_from(self.max_record_bytes).unwrap_or(u64::MAX);
let mut discard_output = [0u8; DISCARD_OUTPUT_BYTES];
let mut discard_ends = [0usize; DISCARD_ENDS_LEN];
loop {
let input = self.source.fill_buf()?;
let (result, nin, nout, nend) = if oversized {
self.core
.read_record(input, &mut discard_output, &mut discard_ends)
} else {
self.core.read_record(
input,
&mut self.output[output_len..],
&mut self.ends[ends_len..],
)
};
self.source.consume(nin);
consumed += nin as u64;
if !oversized {
output_len += nout;
ends_len += nend;
if consumed > max_record_bytes {
oversized = true;
}
}
match result {
ReadRecordResult::InputEmpty => {}
ReadRecordResult::OutputFull => {
if !oversized {
let grown = if self.output.is_empty() {
INITIAL_GROWTH_BYTES
} else {
self.output.len().saturating_mul(2)
}
.min(self.max_record_bytes);
if grown <= self.output.len() {
oversized = true;
} else {
self.output.resize(grown, 0);
}
}
}
ReadRecordResult::OutputEndsFull => {
if !oversized {
let grown = self.ends.len().max(16) * 2;
self.ends.resize(grown, 0);
}
}
ReadRecordResult::Record => {
return Ok(if oversized {
RawRecord::Oversized { consumed }
} else {
self.output.truncate(output_len);
self.ends.truncate(ends_len);
RawRecord::Fields { consumed }
});
}
ReadRecordResult::End => return Ok(RawRecord::Eof),
}
}
}
fn decode_fields(&self) -> Result<Vec<String>, ReaderError> {
let mut fields = Vec::with_capacity(self.ends.len());
let mut start = 0usize;
for &end in &self.ends {
let slice = self
.output
.get(start..end)
.ok_or_else(|| ReaderError::new().with_checkpoint_advanced(true))?;
let text = std::str::from_utf8(slice)
.map_err(|_| ReaderError::new().with_checkpoint_advanced(true))?;
fields.push(text.to_owned());
start = end;
}
Ok(fields)
}
}
impl<Src: Read + Seek> DelimitedReader<Src> {
fn ensure_headers_and_seek(&mut self) -> Result<(), ReaderError> {
if self.seeked {
return Ok(());
}
self.seeked = true;
let target = *self.position.lock().unwrap_or_else(PoisonError::into_inner);
if self.has_headers {
match self.read_raw_record().map_err(|_| {
ReaderError::with_category(crate::FailureCategory::TransientInfrastructure)
})? {
RawRecord::Fields { consumed } => {
let fields = self.decode_fields()?;
if !self.flexible {
self.first_field_count = Some(fields.len());
}
self.headers = Some(fields.into_iter().collect());
if target.byte == 0 {
let mut position =
self.position.lock().unwrap_or_else(PoisonError::into_inner);
position.byte = consumed;
position.line = self.core.line();
}
}
RawRecord::Oversized { .. } => {
return Err(ReaderError::new().with_checkpoint_advanced(true));
}
RawRecord::Eof => {}
}
}
if target.byte > 0 {
self.source
.seek(SeekFrom::Start(target.byte))
.map_err(|_| ReaderError::new())?;
self.core.reset();
self.core.set_line(target.line);
}
Ok(())
}
}
impl<I, Src> crate::ItemReader<I> for DelimitedReader<Src>
where
I: From<DelimitedRecord> + 'static,
Src: Read + Seek + Send,
{
async fn read(&mut self, _context: ReadContext<'_>) -> Result<ReadOutcome<I>, ReaderError> {
self.ensure_headers_and_seek()?;
let before = *self.position.lock().unwrap_or_else(PoisonError::into_inner);
let outcome = self.read_raw_record().map_err(|_| {
ReaderError::with_category(crate::FailureCategory::TransientInfrastructure)
})?;
match outcome {
RawRecord::Eof => Ok(ReadOutcome::EndOfInput),
RawRecord::Oversized { consumed } => {
let after = StoredPosition {
byte: before.byte + consumed,
line: self.core.line(),
record: before.record + 1,
};
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = after;
Err(ReaderError::new().with_checkpoint_advanced(true))
}
RawRecord::Fields { consumed } => {
let after = StoredPosition {
byte: before.byte + consumed,
line: self.core.line(),
record: before.record + 1,
};
let field_count = self.ends.len();
if !self.flexible {
match self.first_field_count {
None => self.first_field_count = Some(field_count),
Some(expected) if expected != field_count => {
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = after;
return Err(ReaderError::new().with_checkpoint_advanced(true));
}
Some(_) => {}
}
}
let fields = match self.decode_fields() {
Ok(fields) => fields,
Err(error) => {
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = after;
return Err(error);
}
};
*self.position.lock().unwrap_or_else(PoisonError::into_inner) = after;
let decoded = DelimitedRecord {
fields,
headers: self.headers.clone(),
};
Ok(ReadOutcome::Item(decoded.into()))
}
}
}
}
pub struct DelimitedReaderStream {
position: Arc<Mutex<StoredPosition>>,
namespace: ComponentStreamIdentity,
}
impl crate::ItemStream for DelimitedReaderStream {
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::<StoredPosition>(&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) = StoredPosition::START;
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 delimited_reader<I, Src>(
source: Src,
dialect: DelimitedDialect,
identity: ComponentStreamIdentity,
) -> (
DelimitedReader<Src>,
DelimitedReaderStream,
StreamStateContract,
)
where
I: From<DelimitedRecord> + 'static,
Src: Read + Seek + Send,
{
let position = Arc::new(Mutex::new(StoredPosition::START));
let reader = DelimitedReader {
source: BufReader::new(source),
core: dialect.core_reader(),
output: Vec::new(),
ends: Vec::new(),
max_record_bytes: dialect.max_record_bytes,
has_headers: dialect.has_headers,
flexible: dialect.flexible,
first_field_count: None,
headers: None,
position: Arc::clone(&position),
seeked: false,
};
let stream = DelimitedReaderStream {
position,
namespace: identity,
};
let contract = StreamStateContract::new(reader_position_codec());
(reader, stream, contract)
}
pub fn delimited_file_reader<I>(
path: impl AsRef<Path>,
dialect: DelimitedDialect,
identity: ComponentStreamIdentity,
) -> io::Result<(
DelimitedReader<File>,
DelimitedReaderStream,
StreamStateContract,
)>
where
I: From<DelimitedRecord> + 'static,
{
let file = File::open(path)?;
Ok(delimited_reader::<I, File>(file, dialect, identity))
}
pub struct DelimitedWriter {
state: Arc<Mutex<DelimitedWriterState>>,
dialect: DelimitedDialect,
}
struct DelimitedWriterState {
file: File,
committed_bytes: u64,
}
impl<I> crate::ItemWriter<I> for DelimitedWriter
where
I: Send + Sync + Clone + Into<DelimitedRecord>,
{
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();
{
let mut csv_writer = self.dialect.writer_builder().from_writer(&mut buffer);
for item in items {
let record: DelimitedRecord = item.clone().into();
csv_writer
.write_record(record.fields())
.map_err(|_| WriterError::new())?;
}
csv_writer.flush().map_err(|_| WriterError::new())?;
}
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 DelimitedWriterStream {
state: Arc<Mutex<DelimitedWriterState>>,
namespace: ComponentStreamIdentity,
}
impl crate::ItemStream for DelimitedWriterStream {
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(
crate::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 delimited_writer(
path: impl AsRef<Path>,
dialect: DelimitedDialect,
identity: ComponentStreamIdentity,
) -> io::Result<(DelimitedWriter, DelimitedWriterStream, StreamStateContract)> {
let file = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(path)?;
let state = Arc::new(Mutex::new(DelimitedWriterState {
file,
committed_bytes: 0,
}));
let writer = DelimitedWriter {
state: Arc::clone(&state),
dialect,
};
let stream = DelimitedWriterStream {
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.delimited-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-delimited-unit-{name}-{nonce}.csv"))
}
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 (writer, stream, _contract) =
delimited_writer(&path, DelimitedDialect::csv(), 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(
&[DelimitedRecord::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();
assert_eq!(
committed_byte(&envelope),
bytes.len() as u64,
"the published checkpoint must equal the real file length, never a stale, earlier \
write's byte count"
);
let _ = std::fs::remove_file(&path);
}
}