use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use encoding_rs::{Encoding, UTF_8};
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::sync::Notify;
use crate::buffer::{LineTerminator, OutputBufferPolicy, OverflowMode};
pub(crate) type LineHandler = Arc<dyn Fn(&str) + Send + Sync>;
pub(crate) struct SharedLines {
inner: Mutex<Inner>,
notify: Notify,
count: AtomicUsize,
dropped: AtomicUsize,
read_error: Mutex<Option<std::io::Error>>,
}
struct Inner {
lines: VecDeque<String>,
max_lines: Option<usize>,
max_bytes: Option<usize>,
bytes: usize,
seen_bytes: usize,
mode: OverflowMode,
closed: bool,
overflowed: bool,
discarding: bool,
}
impl Inner {
fn over_backlog(&self) -> bool {
self.max_lines.is_some_and(|n| self.lines.len() > n)
|| self
.max_bytes
.is_some_and(|b| self.bytes > b || self.lines.len() > b)
}
fn would_fit(&self, len: usize) -> bool {
self.max_lines.is_none_or(|n| self.lines.len() < n)
&& self
.max_bytes
.is_none_or(|b| self.bytes + len <= b && self.lines.len() < b)
}
}
pub(crate) enum Popped {
Line(String),
Empty,
Closed,
}
impl SharedLines {
pub(crate) fn new(policy: &OutputBufferPolicy) -> Arc<Self> {
Arc::new(Self {
inner: Mutex::new(Inner {
lines: VecDeque::new(),
max_lines: policy.max_lines,
max_bytes: policy.max_bytes,
bytes: 0,
seen_bytes: 0,
mode: policy.overflow,
closed: false,
overflowed: false,
discarding: false,
}),
notify: Notify::new(),
count: AtomicUsize::new(0),
dropped: AtomicUsize::new(0),
read_error: Mutex::new(None),
})
}
pub(crate) fn push(&self, line: String) {
let total_lines = self.count.fetch_add(1, Ordering::Relaxed) + 1;
let mut policy_dropped = false;
{
let mut inner = self.inner.lock().expect("SharedLines poisoned");
inner.seen_bytes = inner.seen_bytes.saturating_add(line.len());
if inner.discarding {
} else {
match inner.mode {
OverflowMode::Error => {
let over = match (inner.max_lines, inner.max_bytes) {
(None, None) => true,
(lines_cap, bytes_cap) => {
lines_cap.is_some_and(|n| total_lines > n)
|| bytes_cap.is_some_and(|b| {
inner.seen_bytes > b || total_lines > b
})
}
};
if over {
inner.overflowed = true;
policy_dropped = true;
} else {
inner.bytes += line.len();
inner.lines.push_back(line);
}
}
OverflowMode::DropOldest => {
inner.bytes += line.len();
inner.lines.push_back(line);
while inner.over_backlog() {
match inner.lines.pop_front() {
Some(old) => {
inner.bytes = inner.bytes.saturating_sub(old.len());
policy_dropped = true;
}
None => break,
}
}
}
OverflowMode::DropNewest => {
if inner.would_fit(line.len()) {
inner.bytes += line.len();
inner.lines.push_back(line);
} else {
policy_dropped = true;
}
}
}
}
}
if policy_dropped {
self.dropped.fetch_add(1, Ordering::Relaxed);
}
self.notify.notify_one();
}
pub(crate) fn byte_cap(&self) -> Option<usize> {
self.inner.lock().expect("SharedLines poisoned").max_bytes
}
pub(crate) fn record_oversized_line(&self, byte_len: usize) {
self.count.fetch_add(1, Ordering::Relaxed);
{
let mut inner = self.inner.lock().expect("SharedLines poisoned");
inner.seen_bytes = inner.seen_bytes.saturating_add(byte_len);
if matches!(inner.mode, OverflowMode::Error) {
inner.overflowed = true;
}
}
self.dropped.fetch_add(1, Ordering::Relaxed);
self.notify.notify_one();
}
fn close(&self) {
self.inner
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.closed = true;
self.notify.notify_one();
}
pub(crate) fn close_now(&self) {
self.close();
}
pub(crate) fn start_discarding(&self) {
let mut inner = self.inner.lock().unwrap_or_else(|p| p.into_inner());
inner.discarding = true;
inner.lines.clear();
inner.bytes = 0;
}
pub(crate) fn count(&self) -> usize {
self.count.load(Ordering::Relaxed)
}
pub(crate) fn seen_bytes(&self) -> usize {
self.inner
.lock()
.unwrap_or_else(|p| p.into_inner())
.seen_bytes
}
pub(crate) fn dropped(&self) -> usize {
self.dropped.load(Ordering::Relaxed)
}
pub(crate) fn set_read_error(&self, err: std::io::Error) {
let mut slot = self.read_error.lock().unwrap_or_else(|p| p.into_inner());
if slot.is_none() {
*slot = Some(err);
}
}
pub(crate) fn take_read_error(&self) -> Option<std::io::Error> {
self.read_error
.lock()
.unwrap_or_else(|p| p.into_inner())
.take()
}
pub(crate) fn overflowed(&self) -> bool {
self.inner
.lock()
.unwrap_or_else(|p| p.into_inner())
.overflowed
}
pub(crate) fn drain(&self) -> Vec<String> {
let mut inner = self.inner.lock().expect("SharedLines poisoned");
inner.bytes = 0;
inner.lines.drain(..).collect()
}
pub(crate) fn try_pop(&self) -> Popped {
let mut inner = self.inner.lock().expect("SharedLines poisoned");
if let Some(line) = inner.lines.pop_front() {
inner.bytes = inner.bytes.saturating_sub(line.len());
Popped::Line(line)
} else if inner.closed {
Popped::Closed
} else {
Popped::Empty
}
}
pub(crate) async fn changed(self: Arc<Self>) {
self.notify.notified().await;
}
}
pub(crate) type TeeSink = Arc<tokio::sync::Mutex<Box<dyn tokio::io::AsyncWrite + Send + Unpin>>>;
#[derive(Clone)]
pub(crate) struct StreamConfig {
pub encoding: &'static Encoding,
pub handler: Option<LineHandler>,
pub tee: Option<TeeSink>,
pub terminator: LineTerminator,
}
impl StreamConfig {
pub(crate) fn new() -> Self {
Self {
encoding: UTF_8,
handler: None,
tee: None,
terminator: LineTerminator::Newline,
}
}
pub(crate) fn with_encoding(mut self, encoding: &'static Encoding) -> Self {
self.encoding = encoding;
self
}
}
#[cfg(test)]
pub(crate) async fn pump_lines<R>(
reader: R,
encoding: &'static Encoding,
handler: Option<LineHandler>,
sink: Arc<SharedLines>,
) where
R: AsyncRead + Unpin,
{
pump_lines_core(
reader,
StreamConfig {
encoding,
handler,
tee: None,
terminator: LineTerminator::Newline,
},
sink,
)
.await
}
#[cfg(test)]
pub(crate) async fn pump_lines_term<R>(
reader: R,
encoding: &'static Encoding,
terminator: LineTerminator,
sink: Arc<SharedLines>,
) where
R: AsyncRead + Unpin,
{
pump_lines_core(
reader,
StreamConfig {
encoding,
handler: None,
tee: None,
terminator,
},
sink,
)
.await
}
pub(crate) async fn pump_lines_core<R>(mut reader: R, config: StreamConfig, sink: Arc<SharedLines>)
where
R: AsyncRead + Unpin,
{
let StreamConfig {
encoding,
handler,
tee,
terminator,
} = config;
struct CloseOnDrop(Arc<SharedLines>);
impl Drop for CloseOnDrop {
fn drop(&mut self) {
self.0.close();
}
}
let sink = CloseOnDrop(sink);
let mut handler = handler;
let mut tee = tee;
async fn emit(
handler: &mut Option<LineHandler>,
tee: &mut Option<TeeSink>,
sink: &SharedLines,
line: String,
) {
if let Some(h) = handler {
let invoked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| h(&line)));
if invoked.is_err() {
*handler = None;
#[cfg(feature = "tracing")]
tracing::warn!(
target: "processkit",
"line handler panicked; disabled for the rest of the run"
);
}
}
if let Some(t) = tee {
use tokio::io::AsyncWriteExt;
let mut w = t.lock().await;
let wrote = async {
w.write_all(line.as_bytes()).await?;
w.write_all(b"\n").await
}
.await;
drop(w);
if wrote.is_err() {
*tee = None;
#[cfg(feature = "tracing")]
tracing::warn!(
target: "processkit",
"tee writer errored; disabled for the rest of the run"
);
}
}
sink.push(line);
}
struct Term {
content_len: usize,
resume: usize,
}
fn next_terminator(pending: &str, terminator: LineTerminator, eof: bool) -> Option<Term> {
let bytes = pending.as_bytes();
match terminator {
LineTerminator::Newline => {
let nl = pending.find('\n')?;
let content_len = if nl > 0 && bytes[nl - 1] == b'\r' {
nl - 1
} else {
nl
};
Some(Term {
content_len,
resume: nl + 1,
})
}
LineTerminator::CarriageReturn => {
let pos = bytes.iter().position(|&b| b == b'\n' || b == b'\r')?;
if bytes[pos] == b'\n' {
Some(Term {
content_len: pos,
resume: pos + 1,
})
} else {
match bytes.get(pos + 1) {
Some(b'\n') => Some(Term {
content_len: pos,
resume: pos + 2,
}),
Some(_) => Some(Term {
content_len: pos,
resume: pos + 1,
}),
None => eof.then_some(Term {
content_len: pos,
resume: pos + 1,
}),
}
}
}
}
}
fn skip_over_cap_len(sub: &str) -> usize {
if sub.ends_with('\r') {
sub.len() - 1 } else {
sub.len()
}
}
const CHUNK: usize = 8192;
let cap = sink.0.byte_cap();
let mut decoder = encoding.new_decoder_with_bom_removal();
let mut pending = String::new(); let mut chunk = [0u8; CHUNK];
let mut oversized: Option<usize> = None;
loop {
let (n, eof, read_err) = match reader.read(&mut chunk).await {
Ok(0) => (0, true, None),
Ok(n) => (n, false, None),
Err(e) if crate::running::is_broken_pipe(&e) => (0, true, None),
Err(e) => (0, true, Some(e)),
};
let errored = read_err.is_some();
let last = eof && !errored;
if let Some(need) = decoder.max_utf8_buffer_length(n) {
pending.reserve(need);
}
let _ = decoder.decode_to_string(&chunk[..n], &mut pending, last);
let mut start = 0usize;
loop {
let sub = &pending[start..];
if let Some(skipped) = oversized {
match next_terminator(sub, terminator, eof) {
Some(term) => {
let line_len = skipped.saturating_add(term.content_len);
start += term.resume;
oversized = None;
sink.0.record_oversized_line(line_len);
}
None => {
let advance = skip_over_cap_len(sub);
start += advance;
oversized = Some(skipped.saturating_add(advance));
break;
}
}
} else {
match next_terminator(sub, terminator, eof) {
Some(term) => {
let len = term.content_len;
if cap.is_none_or(|c| len <= c) {
let line = sub[..len].to_owned(); start += term.resume;
emit(&mut handler, &mut tee, &sink.0, line).await;
} else {
start += term.resume;
sink.0.record_oversized_line(len);
}
}
None if cap
.is_some_and(|c| sub.len() - usize::from(sub.ends_with('\r')) > c) =>
{
let advance = skip_over_cap_len(sub);
start += advance;
oversized = Some(advance);
break;
}
None => break,
}
}
}
if start > 0 {
pending.drain(..start);
}
if eof {
if let Some(skipped) = oversized.take() {
let line_len = skipped.saturating_add(pending.len());
pending.clear();
sink.0.record_oversized_line(line_len);
} else if !pending.is_empty() {
let line = std::mem::take(&mut pending);
if cap.is_some_and(|c| line.len() > c) {
sink.0.record_oversized_line(line.len());
} else {
emit(&mut handler, &mut tee, &sink.0, line).await;
}
}
if let Some(t) = &tee {
use tokio::io::AsyncWriteExt;
let _ = t.lock().await.flush().await;
}
if let Some(e) = read_err {
sink.0.set_read_error(e);
}
break;
}
}
}
#[cfg(any(test, fuzzing))]
struct ChunkedReader {
chunks: VecDeque<Vec<u8>>,
err_at_end: Option<std::io::Error>,
}
#[cfg(any(test, fuzzing))]
impl ChunkedReader {
fn new(chunks: impl IntoIterator<Item = Vec<u8>>) -> Self {
Self {
chunks: chunks.into_iter().collect(),
err_at_end: None,
}
}
#[allow(dead_code, reason = "only exercised by the hand-written unit tests")]
fn erroring(chunks: impl IntoIterator<Item = Vec<u8>>) -> Self {
Self::erroring_with(chunks, std::io::Error::other("boom"))
}
#[allow(dead_code, reason = "only exercised by the hand-written unit tests")]
fn erroring_with(chunks: impl IntoIterator<Item = Vec<u8>>, err: std::io::Error) -> Self {
Self {
chunks: chunks.into_iter().collect(),
err_at_end: Some(err),
}
}
}
#[cfg(any(test, fuzzing))]
impl AsyncRead for ChunkedReader {
fn poll_read(
mut self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
if let Some(chunk) = self.chunks.pop_front() {
let n = chunk.len().min(buf.remaining());
buf.put_slice(&chunk[..n]);
if n < chunk.len() {
self.chunks.push_front(chunk[n..].to_vec());
}
std::task::Poll::Ready(Ok(()))
} else if let Some(err) = self.err_at_end.take() {
std::task::Poll::Ready(Err(err))
} else {
std::task::Poll::Ready(Ok(())) }
}
}
#[cfg(any(test, fuzzing))]
fn to_chunks(bytes: &[u8], sizes: &[usize]) -> Vec<Vec<u8>> {
let mut out = Vec::new();
let mut pos = 0;
let mut i = 0;
while pos < bytes.len() {
let size = sizes[i % sizes.len()].max(1);
let end = (pos + size).min(bytes.len());
out.push(bytes[pos..end].to_vec());
pos = end;
i += 1;
}
out
}
#[cfg(fuzzing)]
pub fn fuzz_decode_pump_lines(raw: &[u8], chunk_sizes: &[u8], encoding_idx: u8) {
const ENCODINGS: [&Encoding; 4] = [
encoding_rs::UTF_8,
encoding_rs::SHIFT_JIS,
encoding_rs::UTF_16LE,
encoding_rs::WINDOWS_1252,
];
let encoding = ENCODINGS[encoding_idx as usize % ENCODINGS.len()];
let sizes: Vec<usize> = if chunk_sizes.is_empty() {
vec![raw.len().max(1)]
} else {
chunk_sizes.iter().map(|&b| b as usize + 1).collect()
};
let chunks = to_chunks(raw, &sizes);
let reader = ChunkedReader::new(chunks);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.expect("current-thread runtime");
rt.block_on(pump_lines_core(
reader,
StreamConfig {
encoding,
handler: None,
tee: None,
terminator: LineTerminator::Newline,
},
sink.clone(),
));
let lines = sink.drain();
assert!(lines.len() <= sink.count(), "backlog exceeds lines seen");
}
#[cfg(test)]
mod tests {
use super::*;
use crate::buffer::OutputBufferPolicy;
#[tokio::test]
async fn pumps_utf8_lines_and_counts() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(
&b"one\ntwo\nthree\n"[..],
encoding_rs::UTF_8,
None,
sink.clone(),
)
.await;
assert_eq!(sink.count(), 3);
assert_eq!(sink.drain(), vec!["one", "two", "three"]);
}
#[tokio::test]
async fn decodes_shift_jis() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(
&[0x82, 0xA0, b'\n'][..],
encoding_rs::SHIFT_JIS,
None,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["\u{3042}"]);
}
#[tokio::test]
async fn drop_oldest_keeps_tail_but_counts_all() {
let sink = SharedLines::new(&OutputBufferPolicy::bounded(2));
pump_lines(&b"a\nb\nc\nd\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 4, "every line is counted");
assert_eq!(sink.drain(), vec!["c", "d"], "only the newest two retained");
}
#[tokio::test]
async fn drop_newest_keeps_head() {
let policy = OutputBufferPolicy::bounded(2).with_overflow(OverflowMode::DropNewest);
let sink = SharedLines::new(&policy);
pump_lines(&b"a\nb\nc\nd\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["a", "b"]);
}
#[tokio::test]
async fn fail_loud_sets_overflow_once_full_but_retains_the_cap() {
let sink = SharedLines::new(&OutputBufferPolicy::fail_loud(2));
pump_lines(&b"a\nb\nc\nd\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(sink.overflowed(), "third line must trip the fail-loud flag");
assert_eq!(sink.count(), 4, "every line is still counted");
assert_eq!(sink.drain(), vec!["a", "b"], "retains up to the cap");
}
#[tokio::test]
async fn fail_loud_under_the_cap_does_not_overflow() {
let sink = SharedLines::new(&OutputBufferPolicy::fail_loud(5));
pump_lines(&b"a\nb\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(!sink.overflowed(), "two lines under a 5-line cap is fine");
}
#[tokio::test]
async fn fail_loud_zero_errors_on_the_first_line() {
let sink = SharedLines::new(&OutputBufferPolicy::fail_loud(0));
pump_lines(&b"oops\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(sink.overflowed(), "any line is over a 0-line ceiling");
assert!(sink.drain().is_empty(), "still retains nothing");
}
#[tokio::test]
async fn unbounded_with_error_mode_is_zero_tolerance_not_inert() {
let sink =
SharedLines::new(&OutputBufferPolicy::unbounded().with_overflow(OverflowMode::Error));
pump_lines(&b"anything\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(
sink.overflowed(),
"unbounded + Error must fail loud on any output, not be inert"
);
assert!(sink.drain().is_empty(), "zero-tolerance retains nothing");
}
#[tokio::test]
async fn unbounded_without_error_mode_retains_everything() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"a\nb\nc\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(!sink.overflowed());
assert_eq!(sink.drain(), ["a", "b", "c"]);
}
#[tokio::test]
async fn dropped_counts_policy_drops_not_consumer_pops() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"a\nb\nc\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 3);
assert_eq!(sink.dropped(), 0, "unbounded policy discards nothing");
assert!(matches!(sink.try_pop(), Popped::Line(_)));
assert!(matches!(sink.try_pop(), Popped::Line(_)));
assert_eq!(
sink.dropped(),
0,
"a streaming consumer's pops are not truncation"
);
let bounded = SharedLines::new(&OutputBufferPolicy::bounded(2));
pump_lines(
&b"a\nb\nc\nd\n"[..],
encoding_rs::UTF_8,
None,
bounded.clone(),
)
.await;
assert_eq!(
bounded.dropped(),
2,
"DropOldest discarded the two oldest lines"
);
let newest = SharedLines::new(
&OutputBufferPolicy::bounded(2).with_overflow(OverflowMode::DropNewest),
);
pump_lines(
&b"a\nb\nc\nd\n"[..],
encoding_rs::UTF_8,
None,
newest.clone(),
)
.await;
assert_eq!(
newest.dropped(),
2,
"DropNewest discarded the two newest lines"
);
}
#[tokio::test]
async fn start_discarding_drops_the_backlog_and_retains_nothing_after() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
sink.push("a".into());
sink.push("b".into());
sink.start_discarding();
assert!(sink.drain().is_empty(), "the buffered backlog is dropped");
sink.push("c".into());
assert!(
sink.drain().is_empty(),
"lines pushed after discarding are not retained"
);
assert_eq!(sink.count(), 3, "every line is still counted");
}
#[tokio::test]
async fn bounded_zero_without_error_mode_never_overflows() {
let sink = SharedLines::new(&OutputBufferPolicy::bounded(0));
pump_lines(&b"a\nb\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(!sink.overflowed());
}
#[tokio::test]
async fn handler_sees_every_line_even_when_nothing_retained() {
let seen = Arc::new(Mutex::new(Vec::new()));
let captured = seen.clone();
let handler: LineHandler =
Arc::new(move |line: &str| captured.lock().unwrap().push(line.to_owned()));
let sink = SharedLines::new(&OutputBufferPolicy::bounded(0));
pump_lines(
&b"x\ny\n"[..],
encoding_rs::UTF_8,
Some(handler),
sink.clone(),
)
.await;
assert_eq!(sink.count(), 2);
assert!(
sink.drain().is_empty(),
"retain-nothing policy keeps no lines"
);
assert_eq!(*seen.lock().unwrap(), vec!["x", "y"]);
}
#[tokio::test]
async fn crlf_only_line_is_one_empty_line() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"\r\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 1);
assert_eq!(sink.drain(), vec![""]);
}
#[tokio::test]
async fn final_line_without_a_trailing_newline_is_emitted() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"alpha\nomega"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 2, "the un-terminated tail still counts");
assert_eq!(sink.drain(), vec!["alpha", "omega"]);
}
#[tokio::test]
async fn empty_reader_closes_with_no_lines() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b""[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 0);
assert!(sink.drain().is_empty());
assert!(
matches!(sink.try_pop(), Popped::Closed),
"the sink must close on EOF so a streaming consumer ends"
);
}
#[tokio::test]
async fn invalid_multibyte_decodes_lossily_not_fatally() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(
&[0x82, b'\n'][..],
encoding_rs::SHIFT_JIS,
None,
sink.clone(),
)
.await;
let lines = sink.drain();
assert_eq!(lines.len(), 1);
assert!(
lines[0].contains('\u{FFFD}'),
"invalid bytes decode to the replacement char: {lines:?}"
);
}
#[tokio::test]
async fn panicking_handler_is_isolated_and_capture_completes() {
use std::sync::atomic::{AtomicUsize, Ordering};
let calls = Arc::new(AtomicUsize::new(0));
let handler: LineHandler = {
let calls = calls.clone();
Arc::new(move |_: &str| {
if calls.fetch_add(1, Ordering::SeqCst) == 1 {
panic!("boom on the second line");
}
})
};
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let task = tokio::spawn(pump_lines(
&b"1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n"[..],
encoding_rs::UTF_8,
Some(handler),
sink.clone(),
));
task.await
.expect("the pump task must survive a handler panic");
assert_eq!(sink.count(), 10, "every line captured despite the panic");
assert_eq!(
sink.drain(),
(1..=10).map(|n| n.to_string()).collect::<Vec<_>>()
);
assert_eq!(
calls.load(Ordering::SeqCst),
2,
"the handler is disabled after its panic (called for lines 1 and 2 only)"
);
assert!(
matches!(sink.try_pop(), Popped::Closed),
"sink closes normally after the drain"
);
}
#[tokio::test]
async fn utf16le_lines_decode_and_split_correctly() {
let bytes = [
0x41, 0x00, 0x42, 0x00, 0x0A, 0x00, 0x43, 0x00, 0x44, 0x00, 0x0A, 0x00, ];
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&bytes[..], encoding_rs::UTF_16LE, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["AB", "CD"]);
}
#[tokio::test]
async fn utf16le_code_unit_split_across_reads_is_reassembled() {
let reader = ChunkedReader::new([vec![0x41, 0x00, 0x42], vec![0x00, 0x0A, 0x00]]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(reader, encoding_rs::UTF_16LE, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["AB"]);
}
#[tokio::test]
async fn utf16le_leading_bom_is_removed_once() {
let bytes = [0xFF, 0xFE, 0x41, 0x00, 0x0A, 0x00];
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&bytes[..], encoding_rs::UTF_16LE, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["A"]);
}
#[tokio::test]
async fn utf8_leading_bom_is_removed_once_not_per_line() {
let bytes = [0xEF, 0xBB, 0xBF, b'h', b'i', b'\n', b'b', b'y', b'e', b'\n'];
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&bytes[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["hi", "bye"]);
}
#[tokio::test]
async fn strips_exactly_one_trailing_cr_not_all() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"data\r\r\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["data\r"]);
}
#[tokio::test]
async fn lone_trailing_cr_at_eof_is_kept_as_content() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"tail\r"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["tail\r"]);
}
#[tokio::test]
async fn mid_stream_read_error_flushes_the_partial_tail() {
let reader = ChunkedReader::erroring([b"done\npart".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 2, "the partial tail still counts");
assert_eq!(sink.drain(), vec!["done", "part"]);
assert!(
sink.take_read_error().is_some(),
"the OS read error is recorded on the sink for a consuming finisher"
);
}
#[tokio::test]
async fn legacy_line_starting_with_bom_bytes_is_not_resniffed() {
let bytes = [0xFF, 0xFE, b'x', b'\n'];
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&bytes[..], encoding_rs::WINDOWS_1252, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["\u{00FF}\u{00FE}x"]);
}
#[tokio::test]
async fn fail_loud_trips_on_total_even_when_streamed_dry() {
let sink = SharedLines::new(&OutputBufferPolicy::fail_loud(2));
sink.push("a".into());
assert!(matches!(sink.try_pop(), Popped::Line(_))); sink.push("b".into());
assert!(matches!(sink.try_pop(), Popped::Line(_))); assert!(!sink.overflowed(), "two lines is within the cap");
sink.push("c".into()); assert!(
sink.overflowed(),
"the 3rd line trips the ceiling even though the backlog was drained dry"
);
}
#[tokio::test]
async fn max_bytes_drop_oldest_evicts_to_fit_the_byte_cap() {
let policy = OutputBufferPolicy::unbounded().with_max_bytes(5);
let sink = SharedLines::new(&policy);
pump_lines(&b"aa\nbb\ncc\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["bb", "cc"]);
assert_eq!(sink.count(), 3, "every line is still counted");
}
#[tokio::test]
async fn max_bytes_drops_a_single_oversized_line_whole() {
let policy = OutputBufferPolicy::unbounded().with_max_bytes(3);
let sink = SharedLines::new(&policy);
pump_lines(
&b"toolong\nok\n"[..],
encoding_rs::UTF_8,
None,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["ok"], "the oversized line was dropped");
assert_eq!(sink.count(), 2);
assert!(sink.dropped() >= 1);
}
#[tokio::test]
async fn max_bytes_fail_loud_trips_on_byte_total() {
let policy = OutputBufferPolicy::unbounded()
.with_overflow(OverflowMode::Error)
.with_max_bytes(4);
let sink = SharedLines::new(&policy);
pump_lines(&b"ab\ncd\nef\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(
sink.overflowed(),
"6 cumulative bytes over a 4-byte ceiling must trip it"
);
}
#[tokio::test]
async fn max_bytes_under_the_cap_does_not_trip_or_drop() {
let policy = OutputBufferPolicy::fail_loud(10).with_max_bytes(100);
let sink = SharedLines::new(&policy);
pump_lines(&b"ab\ncd\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert!(!sink.overflowed());
assert_eq!(sink.dropped(), 0);
assert_eq!(sink.drain(), vec!["ab", "cd"]);
}
#[tokio::test]
async fn max_bytes_drop_newest_keeps_head_within_byte_cap() {
let policy = OutputBufferPolicy::unbounded()
.with_overflow(OverflowMode::DropNewest)
.with_max_bytes(4);
let sink = SharedLines::new(&policy);
pump_lines(&b"ab\ncd\nef\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["ab", "cd"]);
}
#[tokio::test]
async fn max_bytes_bounds_a_flood_of_empty_lines_drop_oldest() {
let policy = OutputBufferPolicy::unbounded().with_max_bytes(100);
let sink = SharedLines::new(&policy);
let flood = "\n".repeat(10_000);
pump_lines(flood.as_bytes(), encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 10_000, "every empty line is still counted");
let retained = sink.drain();
assert!(
retained.len() <= 100,
"the backlog must stay bounded by the byte cap even for all-empty lines, got {}",
retained.len()
);
assert!(
sink.dropped() > 0,
"the flood must be evicted, not retained without bound"
);
}
#[tokio::test]
async fn max_bytes_bounds_a_flood_of_empty_lines_drop_newest() {
let policy = OutputBufferPolicy::unbounded()
.with_overflow(OverflowMode::DropNewest)
.with_max_bytes(100);
let sink = SharedLines::new(&policy);
let flood = "\n".repeat(10_000);
pump_lines(flood.as_bytes(), encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.count(), 10_000, "every empty line is still counted");
let retained = sink.drain();
assert!(
retained.len() <= 100,
"the retained head must stay bounded by the byte cap even for all-empty lines, got {}",
retained.len()
);
assert!(
sink.dropped() > 0,
"the excess of the flood must be dropped, not retained without bound"
);
}
#[tokio::test]
async fn max_bytes_error_mode_trips_on_a_flood_of_empty_lines() {
let policy = OutputBufferPolicy::unbounded()
.with_overflow(OverflowMode::Error)
.with_max_bytes(100);
let sink = SharedLines::new(&policy);
let flood = "\n".repeat(10_000);
pump_lines(flood.as_bytes(), encoding_rs::UTF_8, None, sink.clone()).await;
assert!(
sink.overflowed(),
"a flood of empty lines under a byte cap must trip OverflowMode::Error"
);
assert!(
sink.dropped() > 0,
"the excess lines must be flagged dropped"
);
}
#[tokio::test]
async fn max_bytes_skips_an_over_cap_line_streamed_across_reads_without_buffering_it() {
let reader = ChunkedReader::new([vec![b'X'; 50_000], b"\n".to_vec(), b"tail\n".to_vec()]);
let policy = OutputBufferPolicy::unbounded().with_max_bytes(8);
let sink = SharedLines::new(&policy);
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(
sink.drain(),
vec!["tail"],
"the over-cap flood is dropped; the small line is kept"
);
assert_eq!(sink.count(), 2, "both lines are counted");
assert!(sink.dropped() >= 1, "the over-cap line is a truncation");
}
#[tokio::test]
async fn over_cap_crlf_line_byte_count_is_stable_across_a_read_boundary() {
let content = vec![b'X'; 10];
let mut one = content.clone();
one.extend_from_slice(b"\r\ntail\n");
let single = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(4));
pump_lines(&one[..], encoding_rs::UTF_8, None, single.clone()).await;
let mut first = content.clone();
first.push(b'\r');
let reader = ChunkedReader::new([first, b"\ntail\n".to_vec()]);
let split = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(4));
pump_lines(reader, encoding_rs::UTF_8, None, split.clone()).await;
assert_eq!(
split.seen_bytes(),
single.seen_bytes(),
"the CRLF terminator must not be counted only when it lands at a chunk end"
);
assert_eq!(
single.seen_bytes(),
14,
"over-cap content (10) excluding the CRLF, plus the retained 'tail' (4)"
);
assert_eq!(split.drain(), vec!["tail"], "the over-cap line is dropped");
assert_eq!(single.drain(), vec!["tail"]);
}
#[tokio::test]
async fn over_cap_skip_keeps_a_lone_cr_as_content_across_reads() {
let mut first = vec![b'X'; 10];
first.push(b'\r');
let reader = ChunkedReader::new([first, b"YYYYY\n".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(4));
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(
sink.seen_bytes(),
16,
"a lone CR before non-newline content is counted, not dropped"
);
assert!(
sink.drain().is_empty(),
"the over-cap line is dropped whole"
);
}
#[tokio::test]
async fn crlf_line_at_exactly_the_cap_is_retained_regardless_of_read_boundary() {
let single = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(&b"ab\r\n"[..], encoding_rs::UTF_8, None, single.clone()).await;
let reader = ChunkedReader::new([b"ab\r".to_vec(), b"\n".to_vec()]);
let split = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(reader, encoding_rs::UTF_8, None, split.clone()).await;
assert_eq!(
single.drain(),
vec!["ab"],
"one-chunk: an at-cap CRLF line is retained"
);
assert_eq!(
split.drain(),
vec!["ab"],
"split CRLF must retain the at-cap line identically — not drop it"
);
}
#[tokio::test]
async fn over_cap_unterminated_tail_at_eof_is_dropped_not_delivered() {
let seen = Arc::new(Mutex::new(Vec::new()));
let captured = seen.clone();
let handler: LineHandler =
Arc::new(move |line: &str| captured.lock().unwrap().push(line.to_owned()));
let sink = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(
&b"ab\r"[..],
encoding_rs::UTF_8,
Some(handler),
sink.clone(),
)
.await;
assert!(
sink.drain().is_empty(),
"an over-cap unterminated tail is not retained"
);
assert!(
seen.lock().unwrap().is_empty(),
"an over-cap line is never delivered to the handler"
);
assert!(
sink.dropped() >= 1,
"the over-cap tail counts as a truncation"
);
}
#[tokio::test]
async fn error_mode_byte_cap_drains_a_post_trip_flood_without_retaining() {
let policy = OutputBufferPolicy::unbounded()
.with_overflow(OverflowMode::Error)
.with_max_bytes(3);
let sink = SharedLines::new(&policy);
let reader =
ChunkedReader::new([b"abcd\n".to_vec(), vec![b'Z'; 20_000], b"\nmore\n".to_vec()]);
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert!(
sink.overflowed(),
"the over-cap first line trips the ceiling"
);
assert!(
sink.drain().is_empty(),
"nothing is retained under the fail-loud ceiling"
);
}
#[tokio::test]
async fn byte_cap_judges_a_crlf_line_like_its_lf_twin_at_the_boundary() {
let lf = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(&b"ab\n"[..], encoding_rs::UTF_8, None, lf.clone()).await;
assert_eq!(
lf.drain(),
vec!["ab"],
"LF: 2-byte content fits a 2-byte cap"
);
let crlf = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(&b"ab\r\n"[..], encoding_rs::UTF_8, None, crlf.clone()).await;
assert_eq!(
crlf.drain(),
vec!["ab"],
"CRLF: the same 2-byte content must also fit (the '\\r' is a terminator)"
);
assert_eq!(crlf.dropped(), 0, "nothing was over-cap");
let over = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines(&b"abc\r\n"[..], encoding_rs::UTF_8, None, over.clone()).await;
assert!(
over.drain().is_empty(),
"3-byte content exceeds the 2-byte cap"
);
assert!(over.dropped() >= 1);
}
#[tokio::test]
async fn read_error_after_incomplete_multibyte_does_not_fabricate_a_phantom_char() {
let reader = ChunkedReader::erroring([b"ok\n\xC3".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(
sink.drain(),
vec!["ok"],
"the truncated lead byte produces no phantom line"
);
assert_eq!(sink.count(), 1);
assert!(
sink.take_read_error().is_some(),
"the read error is recorded even though the truncated multibyte tail is dropped"
);
}
#[tokio::test]
async fn clean_eof_records_no_read_error() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(&b"a\nb\n"[..], encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["a", "b"]);
assert!(
sink.take_read_error().is_none(),
"a clean EOF is a complete capture, not an incomplete one"
);
}
#[tokio::test]
async fn read_error_on_a_line_boundary_keeps_the_line_and_records_the_error() {
let reader = ChunkedReader::erroring([b"done\n".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["done"], "the completed line is retained");
assert!(
sink.take_read_error().is_some(),
"even a boundary-aligned read error is recorded"
);
}
#[tokio::test]
async fn broken_pipe_read_is_treated_as_clean_eof_not_an_incomplete_capture() {
let reader = ChunkedReader::erroring_with(
[b"tail\n".to_vec()],
std::io::Error::from(std::io::ErrorKind::BrokenPipe),
);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()).await;
assert_eq!(sink.drain(), vec!["tail"]);
assert!(
sink.take_read_error().is_none(),
"a broken-pipe read is the normal end of a stream, not an incomplete capture"
);
}
#[tokio::test]
async fn concurrent_read_errors_on_both_streams_are_each_recorded() {
let out_sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let err_sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let out = tokio::spawn(pump_lines(
ChunkedReader::erroring([b"o1\no2".to_vec()]),
encoding_rs::UTF_8,
None,
out_sink.clone(),
));
let err = tokio::spawn(pump_lines(
ChunkedReader::erroring([b"e1\ne2".to_vec()]),
encoding_rs::UTF_8,
None,
err_sink.clone(),
));
out.await.expect("stdout pump");
err.await.expect("stderr pump");
assert_eq!(out_sink.drain(), vec!["o1", "o2"]);
assert_eq!(err_sink.drain(), vec!["e1", "e2"]);
assert!(
out_sink.take_read_error().is_some(),
"stdout error recorded"
);
assert!(
err_sink.take_read_error().is_some(),
"stderr error recorded"
);
assert!(
matches!(out_sink.try_pop(), Popped::Closed),
"each sink still closes so a streaming consumer ends"
);
assert!(matches!(err_sink.try_pop(), Popped::Closed));
}
#[tokio::test]
async fn cr_mode_splits_progress_frames_live() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"Progress: 0%\rProgress: 50%\rProgress: 100%\n"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(sink.count(), 3, "three frames, not one accumulated line");
assert_eq!(
sink.drain(),
vec!["Progress: 0%", "Progress: 50%", "Progress: 100%"]
);
}
#[tokio::test]
async fn cr_mode_leading_cr_and_unterminated_tail() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"\rA\rB"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["", "A", "B"]);
}
#[tokio::test]
async fn cr_mode_crlf_is_a_single_terminator_no_empty_lines() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"a\r\nb\r\n"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(
sink.drain(),
vec!["a", "b"],
"no empty line between CR and LF"
);
}
#[tokio::test]
async fn cr_mode_mixed_terminators() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"a\rb\nc\r\nd"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["a", "b", "c", "d"]);
}
#[tokio::test]
async fn cr_mode_crlf_split_across_reads_stays_one_terminator() {
let reader = ChunkedReader::new([b"a\r".to_vec(), b"\nb\r".to_vec(), b"\n".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
reader,
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["a", "b"], "split CRLF is one terminator");
}
#[tokio::test]
async fn cr_mode_lone_cr_at_read_boundary_is_a_frame_terminator() {
let reader = ChunkedReader::new([b"a\r".to_vec(), b"b\n".to_vec()]);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
reader,
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["a", "b"]);
}
#[tokio::test]
async fn cr_mode_trailing_cr_at_eof_terminates_the_frame() {
let cr = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"tail\r"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
cr.clone(),
)
.await;
assert_eq!(
cr.drain(),
vec!["tail"],
"CR mode: trailing `\\r` terminates"
);
let nl = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"tail\r"[..],
encoding_rs::UTF_8,
LineTerminator::Newline,
nl.clone(),
)
.await;
assert_eq!(
nl.drain(),
vec!["tail\r"],
"Newline mode: trailing `\\r` is content"
);
}
#[tokio::test]
async fn cr_mode_default_newline_is_unchanged_lone_cr_is_content() {
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_term(
&b"a\rb\n"[..],
encoding_rs::UTF_8,
LineTerminator::Newline,
sink.clone(),
)
.await;
assert_eq!(
sink.drain(),
vec!["a\rb"],
"Newline mode keeps the inner `\\r`"
);
}
#[tokio::test]
async fn cr_mode_byte_cap_skips_an_over_cap_frame_but_keeps_small_ones() {
let reader = ChunkedReader::new([vec![b'X'; 50_000], b"\rtail\n".to_vec()]);
let policy = OutputBufferPolicy::unbounded().with_max_bytes(8);
let sink = SharedLines::new(&policy);
pump_lines_term(
reader,
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
sink.clone(),
)
.await;
assert_eq!(
sink.drain(),
vec!["tail"],
"the over-cap frame is dropped; the small frame is kept"
);
assert_eq!(sink.count(), 2, "both frames are counted");
assert!(sink.dropped() >= 1, "the over-cap frame is a truncation");
}
#[tokio::test]
async fn cr_mode_at_cap_frame_retained_regardless_of_read_boundary() {
let single = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines_term(
&b"ab\r"[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
single.clone(),
)
.await;
let reader = ChunkedReader::new([b"ab".to_vec(), b"\r".to_vec()]);
let split = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(2));
pump_lines_term(
reader,
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
split.clone(),
)
.await;
assert_eq!(
single.drain(),
vec!["ab"],
"one-chunk: at-cap CR frame retained"
);
assert_eq!(
split.drain(),
vec!["ab"],
"split CR must retain the at-cap frame identically — not drop it"
);
assert_eq!(split.dropped(), 0, "nothing was over-cap");
}
#[tokio::test]
async fn cr_mode_over_cap_frame_byte_count_is_stable_across_a_read_boundary() {
let content = vec![b'X'; 10];
let mut one = content.clone();
one.extend_from_slice(b"\rtail\n");
let single = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(4));
pump_lines_term(
&one[..],
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
single.clone(),
)
.await;
let mut first = content.clone();
first.push(b'\r');
let reader = ChunkedReader::new([first, b"tail\n".to_vec()]);
let split = SharedLines::new(&OutputBufferPolicy::unbounded().with_max_bytes(4));
pump_lines_term(
reader,
encoding_rs::UTF_8,
LineTerminator::CarriageReturn,
split.clone(),
)
.await;
assert_eq!(
split.seen_bytes(),
single.seen_bytes(),
"the CR terminator must not be counted only when it lands at a chunk end"
);
assert_eq!(
single.seen_bytes(),
14,
"over-cap content (10) excluding the CR, plus the retained 'tail' (4)"
);
assert_eq!(split.drain(), vec!["tail"]);
}
#[tokio::test]
async fn cr_mode_handler_and_tee_see_each_frame() {
let buf = Arc::new(Mutex::new(Vec::new()));
let seen = Arc::new(Mutex::new(Vec::new()));
let captured = seen.clone();
let handler: LineHandler =
Arc::new(move |line: &str| captured.lock().unwrap().push(line.to_owned()));
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
&b"50%\r100%\n"[..],
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: Some(handler),
tee: Some(tee_of(VecSink(buf.clone()))),
terminator: LineTerminator::CarriageReturn,
},
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["50%", "100%"], "buffer sees frames");
assert_eq!(
*seen.lock().unwrap(),
vec!["50%", "100%"],
"handler sees frames"
);
assert_eq!(
String::from_utf8(buf.lock().unwrap().clone()).unwrap(),
"50%\n100%\n",
"the tee writes each frame followed by a newline"
);
}
#[derive(Clone)]
struct VecSink(Arc<Mutex<Vec<u8>>>);
impl tokio::io::AsyncWrite for VecSink {
fn poll_write(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
self.0.lock().unwrap().extend_from_slice(buf);
std::task::Poll::Ready(Ok(buf.len()))
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
}
fn tee_of(sink: impl tokio::io::AsyncWrite + Send + Unpin + 'static) -> TeeSink {
Arc::new(tokio::sync::Mutex::new(Box::new(sink)))
}
async fn assert_buffered_tee_flushes_after_read_error(stream: &str) {
let buf = Arc::new(Mutex::new(Vec::new()));
let buffered = tokio::io::BufWriter::with_capacity(1024, VecSink(buf.clone()));
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
ChunkedReader::erroring([b"complete\npartial".to_vec()]),
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: None,
tee: Some(tee_of(buffered)),
terminator: LineTerminator::Newline,
},
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["complete", "partial"]);
assert_eq!(
&*buf.lock().unwrap(),
b"complete\npartial\n",
"{stream} tee must flush through its buffering writer after a read error"
);
}
#[tokio::test]
async fn tee_writes_each_decoded_line_plus_newline_to_the_async_sink() {
let buf = Arc::new(Mutex::new(Vec::new()));
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
&b"one\ntwo\n"[..],
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: None,
tee: Some(tee_of(VecSink(buf.clone()))),
terminator: LineTerminator::Newline,
},
sink.clone(),
)
.await;
assert_eq!(sink.drain(), vec!["one", "two"], "capture is unaffected");
let teed = String::from_utf8(buf.lock().unwrap().clone()).unwrap();
assert_eq!(teed, "one\ntwo\n", "the tee got each line + a newline");
}
#[tokio::test]
async fn tee_write_error_is_isolated_and_capture_continues() {
struct ErrSink;
impl tokio::io::AsyncWrite for ErrSink {
fn poll_write(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
_buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
std::task::Poll::Ready(Err(std::io::Error::other("nope")))
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
}
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
&b"a\nb\nc\n"[..],
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: None,
tee: Some(tee_of(ErrSink)),
terminator: LineTerminator::Newline,
},
sink.clone(),
)
.await;
assert_eq!(
sink.drain(),
vec!["a", "b", "c"],
"capture survives a tee write error"
);
}
#[tokio::test]
async fn stdout_buffered_tee_flushes_after_read_error() {
assert_buffered_tee_flushes_after_read_error("stdout").await;
}
#[tokio::test]
async fn stderr_buffered_tee_flushes_after_read_error() {
assert_buffered_tee_flushes_after_read_error("stderr").await;
}
#[tokio::test]
async fn tee_flush_error_is_isolated_and_capture_completes() {
struct FlushErrSink;
impl tokio::io::AsyncWrite for FlushErrSink {
fn poll_write(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
std::task::Poll::Ready(Ok(buf.len()))
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Err(std::io::Error::other("nope")))
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
}
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
&b"a\nb\n"[..],
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: None,
tee: Some(tee_of(FlushErrSink)),
terminator: LineTerminator::Newline,
},
sink.clone(),
)
.await;
assert_eq!(
sink.drain(),
vec!["a", "b"],
"capture survives a tee flush error"
);
}
#[tokio::test]
async fn tee_and_line_handler_both_fire_independently() {
let buf = Arc::new(Mutex::new(Vec::new()));
let seen = Arc::new(Mutex::new(Vec::new()));
let captured = seen.clone();
let handler: LineHandler =
Arc::new(move |line: &str| captured.lock().unwrap().push(line.to_owned()));
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
pump_lines_core(
&b"x\ny\n"[..],
StreamConfig {
encoding: encoding_rs::UTF_8,
handler: Some(handler),
tee: Some(tee_of(VecSink(buf.clone()))),
terminator: LineTerminator::Newline,
},
sink.clone(),
)
.await;
assert_eq!(*seen.lock().unwrap(), vec!["x", "y"], "handler fired");
assert_eq!(
String::from_utf8(buf.lock().unwrap().clone()).unwrap(),
"x\ny\n",
"tee fired"
);
}
mod proptests {
use super::*;
use proptest::prelude::*;
fn arb_line_content() -> impl Strategy<Value = String> {
prop::collection::vec(
any::<char>().prop_filter("no CR/LF/BOM in line content", |c| {
!matches!(*c, '\n' | '\r' | '\u{feff}')
}),
0..12,
)
.prop_map(|chars| chars.into_iter().collect())
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(64))]
#[test]
fn pump_preserves_lines_and_counts_across_arbitrary_chunking(
lines in prop::collection::vec(arb_line_content(), 0..12),
tail in prop::option::of(
arb_line_content().prop_filter("tail must be non-empty", |s| !s.is_empty())
),
chunk_sizes in prop::collection::vec(1usize..=7, 1..20),
) {
let mut text = String::new();
for line in &lines {
text.push_str(line);
text.push('\n');
}
if let Some(t) = &tail {
text.push_str(t);
}
let mut expected = lines.clone();
if let Some(t) = &tail {
expected.push(t.clone());
}
let bytes = text.into_bytes();
let chunks = to_chunks(&bytes, &chunk_sizes);
let reader = ChunkedReader::new(chunks);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.expect("current-thread runtime");
rt.block_on(pump_lines(reader, encoding_rs::UTF_8, None, sink.clone()));
let expected_bytes: usize = expected.iter().map(String::len).sum();
prop_assert_eq!(sink.count(), expected.len(), "no line lost or fabricated");
prop_assert_eq!(sink.seen_bytes(), expected_bytes, "byte counter is exact");
prop_assert_eq!(sink.dropped(), 0, "unbounded policy drops nothing");
prop_assert_eq!(sink.drain(), expected, "every line reassembled correctly");
}
#[test]
fn pump_never_panics_on_arbitrary_bytes_under_any_chunking(
raw in prop::collection::vec(any::<u8>(), 0..512),
chunk_sizes in prop::collection::vec(1usize..=9, 1..20),
encoding_idx in 0usize..4,
) {
const ENCODINGS: [&encoding_rs::Encoding; 4] = [
encoding_rs::UTF_8,
encoding_rs::SHIFT_JIS,
encoding_rs::UTF_16LE,
encoding_rs::WINDOWS_1252,
];
let encoding = ENCODINGS[encoding_idx];
let chunks = to_chunks(&raw, &chunk_sizes);
let reader = ChunkedReader::new(chunks);
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.expect("current-thread runtime");
rt.block_on(pump_lines(reader, encoding, None, sink.clone()));
let lines = sink.drain();
prop_assert!(lines.len() <= sink.count());
}
}
}
}