use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
#[cfg(feature = "async-process")]
use tokio::sync::Notify;
use crate::StreamKind;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutputRecord {
pub sequence: u64,
pub stream: StreamKind,
pub bytes: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CursorRead {
Gap {
from: u64,
to: u64,
},
Record(OutputRecord),
Eof,
}
#[derive(Debug)]
pub struct OutputLog {
capacity_bytes: usize,
retained_bytes: usize,
next_sequence: u64,
records: VecDeque<OutputRecord>,
}
impl OutputLog {
pub fn new(capacity_bytes: usize) -> Self {
Self {
capacity_bytes,
retained_bytes: 0,
next_sequence: 0,
records: VecDeque::new(),
}
}
pub fn append(&mut self, stream: StreamKind, bytes: impl Into<Vec<u8>>) -> u64 {
let sequence = self.next_sequence;
self.next_sequence = self.next_sequence.saturating_add(1);
let bytes = bytes.into();
if bytes.len() > self.capacity_bytes {
self.records.clear();
self.retained_bytes = 0;
return sequence;
}
self.retained_bytes = self.retained_bytes.saturating_add(bytes.len());
self.records.push_back(OutputRecord {
sequence,
stream,
bytes,
});
while self.retained_bytes > self.capacity_bytes {
if let Some(record) = self.records.pop_front() {
self.retained_bytes = self.retained_bytes.saturating_sub(record.bytes.len());
}
}
sequence
}
pub fn first_sequence(&self) -> u64 {
self.records
.front()
.map_or(self.next_sequence, |record| record.sequence)
}
pub fn next_sequence(&self) -> u64 {
self.next_sequence
}
pub fn retained_bytes(&self) -> usize {
self.retained_bytes
}
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
pub fn cursor(&self) -> OutputCursor {
OutputCursor {
next_sequence: self.first_sequence(),
}
}
pub fn cursor_from(&self, sequence: u64) -> OutputCursor {
OutputCursor {
next_sequence: sequence,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OutputCursor {
next_sequence: u64,
}
#[derive(Debug, Clone)]
pub struct SharedOutputLog {
inner: Arc<SharedOutputState>,
}
#[derive(Debug)]
struct SharedOutputState {
log: Mutex<OutputLog>,
#[cfg(feature = "async-process")]
notify: Notify,
closed: AtomicBool,
}
impl SharedOutputLog {
pub fn new(capacity_bytes: usize) -> Self {
Self {
inner: Arc::new(SharedOutputState {
log: Mutex::new(OutputLog::new(capacity_bytes)),
#[cfg(feature = "async-process")]
notify: Notify::new(),
closed: AtomicBool::new(false),
}),
}
}
pub fn append(&self, stream: StreamKind, bytes: impl Into<Vec<u8>>) -> u64 {
let sequence = self
.inner
.log
.lock()
.expect("output log lock is not poisoned")
.append(stream, bytes);
#[cfg(feature = "async-process")]
self.inner.notify.notify_waiters();
sequence
}
pub fn close(&self) {
self.inner.closed.store(true, Ordering::Release);
#[cfg(feature = "async-process")]
self.inner.notify.notify_waiters();
}
pub fn cursor(&self) -> SharedOutputCursor {
let cursor = self
.inner
.log
.lock()
.expect("output log lock is not poisoned")
.cursor();
SharedOutputCursor {
inner: Arc::clone(&self.inner),
cursor,
}
}
pub fn retained_bytes(&self) -> usize {
self.inner
.log
.lock()
.expect("output log lock is not poisoned")
.retained_bytes()
}
}
#[derive(Debug, Clone)]
pub struct SharedOutputCursor {
inner: Arc<SharedOutputState>,
cursor: OutputCursor,
}
impl SharedOutputCursor {
pub fn read_next(&mut self) -> CursorRead {
self.cursor.read_next(
&self
.inner
.log
.lock()
.expect("output log lock is not poisoned"),
)
}
#[cfg(feature = "async-process")]
pub async fn read_next_async(&mut self) -> CursorRead {
self.read_next_async_inner(|| {}).await
}
#[cfg(all(test, feature = "async-process"))]
async fn read_next_async_after_empty<F>(&mut self, after_empty: F) -> CursorRead
where
F: FnMut(),
{
self.read_next_async_inner(after_empty).await
}
#[cfg(feature = "async-process")]
async fn read_next_async_inner<F>(&mut self, mut after_empty: F) -> CursorRead
where
F: FnMut(),
{
loop {
let state = Arc::clone(&self.inner);
let notified = state.notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
match self.read_next() {
CursorRead::Eof if self.inner.closed.load(Ordering::Acquire) => {
return CursorRead::Eof;
}
CursorRead::Eof => {
after_empty();
notified.as_mut().await;
}
result => return result,
}
}
}
pub fn position(&self) -> u64 {
self.cursor.position()
}
pub fn is_closed(&self) -> bool {
self.inner.closed.load(Ordering::Acquire)
}
}
impl OutputCursor {
pub fn position(&self) -> u64 {
self.next_sequence
}
pub fn read_next(&mut self, log: &OutputLog) -> CursorRead {
let first = log.first_sequence();
if self.next_sequence < first {
let from = self.next_sequence;
self.next_sequence = first;
return CursorRead::Gap {
from,
to: first.saturating_sub(1),
};
}
let Some(record) = log
.records
.iter()
.find(|record| record.sequence == self.next_sequence)
else {
return CursorRead::Eof;
};
self.next_sequence = self.next_sequence.saturating_add(1);
CursorRead::Record(record.clone())
}
}
#[cfg(test)]
mod tests {
use super::{CursorRead, OutputLog, SharedOutputLog};
use crate::StreamKind;
#[test]
fn bounded_retention_reports_gaps_to_lagging_cursors() {
let mut log = OutputLog::new(4);
let mut cursor = log.cursor();
assert_eq!(log.append(StreamKind::Stdout, b"aa"), 0);
assert!(
matches!(cursor.read_next(&log), CursorRead::Record(record) if record.sequence == 0)
);
log.append(StreamKind::Stderr, b"bb");
log.append(StreamKind::Stdout, b"cc");
log.append(StreamKind::Stderr, b"dd");
assert_eq!(cursor.read_next(&log), CursorRead::Gap { from: 1, to: 1 });
assert!(
matches!(cursor.read_next(&log), CursorRead::Record(record) if record.sequence == 2)
);
}
#[test]
fn cursors_are_independent_and_oversized_records_are_explicitly_lost() {
let mut log = OutputLog::new(3);
log.append(StreamKind::Stdout, b"1234");
let mut first = log.cursor_from(0);
let mut second = log.cursor_from(0);
log.append(StreamKind::Stderr, b"ok");
assert!(matches!(
first.read_next(&log),
CursorRead::Gap { from: 0, to: 0 }
));
assert!(matches!(
second.read_next(&log),
CursorRead::Gap { from: 0, to: 0 }
));
assert!(
matches!(first.read_next(&log), CursorRead::Record(record) if record.bytes == b"ok")
);
assert!(
matches!(second.read_next(&log), CursorRead::Record(record) if record.bytes == b"ok")
);
assert_eq!(log.retained_bytes(), 2);
}
#[test]
fn shared_log_keeps_cursor_positions_independent() {
let log = SharedOutputLog::new(8);
let mut first = log.cursor();
let mut second = log.cursor();
log.append(StreamKind::Stdout, b"one");
assert!(matches!(first.read_next(), CursorRead::Record(_)));
assert!(matches!(second.read_next(), CursorRead::Record(_)));
assert_eq!(first.position(), second.position());
assert_eq!(log.retained_bytes(), 3);
}
#[cfg(feature = "async-process")]
#[tokio::test]
async fn async_cursor_returns_terminal_eof_after_close() {
let log = SharedOutputLog::new(8);
let mut cursor = log.cursor();
log.close();
assert_eq!(cursor.read_next_async().await, CursorRead::Eof);
assert!(cursor.is_closed());
}
#[cfg(feature = "async-process")]
#[tokio::test]
async fn async_cursor_close_between_empty_check_and_wait_is_not_lost() {
let log = SharedOutputLog::new(8);
let close_log = log.clone();
let mut cursor = log.cursor();
let result = tokio::time::timeout(
std::time::Duration::from_secs(1),
cursor.read_next_async_after_empty(move || close_log.close()),
)
.await
.expect("close wakes a cursor that has no record");
assert_eq!(result, CursorRead::Eof);
assert!(cursor.is_closed());
}
}