use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::FailureSchedule;
pub struct ChaosReader<R: Read> {
inner: R,
schedule: FailureSchedule,
attempt: AtomicUsize,
}
impl<R: Read> ChaosReader<R> {
pub fn new(inner: R, schedule: FailureSchedule) -> Self {
Self {
inner,
schedule,
attempt: AtomicUsize::new(0),
}
}
pub fn attempt_count(&self) -> usize {
self.attempt.load(Ordering::Relaxed)
}
pub fn into_inner(self) -> R {
self.inner
}
}
impl<R: Read> Read for ChaosReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let n = self.attempt.fetch_add(1, Ordering::Relaxed) + 1;
if let Err(f) = self.schedule.maybe_fail(n) {
return Err(f.into());
}
self.inner.read(buf)
}
}
pub struct ChaosWriter<W: Write> {
inner: W,
schedule: FailureSchedule,
attempt: AtomicUsize,
}
impl<W: Write> ChaosWriter<W> {
pub fn new(inner: W, schedule: FailureSchedule) -> Self {
Self {
inner,
schedule,
attempt: AtomicUsize::new(0),
}
}
pub fn attempt_count(&self) -> usize {
self.attempt.load(Ordering::Relaxed)
}
pub fn into_inner(self) -> W {
self.inner
}
}
impl<W: Write> Write for ChaosWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let n = self.attempt.fetch_add(1, Ordering::Relaxed) + 1;
if let Err(f) = self.schedule.maybe_fail(n) {
if matches!(f.mode, crate::FailureMode::PartialWrite) && !buf.is_empty() {
let _ = self.inner.write(&buf[..1])?;
}
return Err(f.into());
}
self.inner.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
pub type ChaosFile = ChaosWriter<std::fs::File>;
impl ChaosFile {
pub fn create(
path: impl AsRef<std::path::Path>,
schedule: FailureSchedule,
) -> io::Result<Self> {
let f = std::fs::File::create(path)?;
Ok(Self::new(f, schedule))
}
pub fn append(
path: impl AsRef<std::path::Path>,
schedule: FailureSchedule,
) -> io::Result<Self> {
let f = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)?;
Ok(Self::new(f, schedule))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::FailureMode;
#[test]
fn reader_passes_through_when_schedule_does_not_fire() {
let data: &[u8] = b"hello";
let schedule = FailureSchedule::on_attempts(&[10], FailureMode::IoError);
let mut r = ChaosReader::new(data, schedule);
let mut buf = [0u8; 5];
let n = r.read(&mut buf).unwrap();
assert_eq!(n, 5);
assert_eq!(&buf, b"hello");
}
#[test]
fn reader_fails_when_schedule_fires() {
let data: &[u8] = b"hello";
let schedule = FailureSchedule::on_attempts(&[1], FailureMode::Timeout);
let mut r = ChaosReader::new(data, schedule);
let mut buf = [0u8; 5];
let err = r.read(&mut buf).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::TimedOut);
}
#[test]
fn reader_attempt_count_increments() {
let data: &[u8] = b"abc";
let schedule = FailureSchedule::on_attempts(&[], FailureMode::IoError);
let mut r = ChaosReader::new(data, schedule);
let mut buf = [0u8; 1];
for _ in 0..3 {
let _ = r.read(&mut buf);
}
assert_eq!(r.attempt_count(), 3);
}
#[test]
fn writer_passes_through_bytes() {
let sink: Vec<u8> = Vec::new();
let schedule = FailureSchedule::on_attempts(&[], FailureMode::IoError);
let mut w = ChaosWriter::new(sink, schedule);
w.write_all(b"hello").unwrap();
let sink = w.into_inner();
assert_eq!(sink, b"hello");
}
#[test]
fn writer_fails_on_scheduled_attempt() {
let sink: Vec<u8> = Vec::new();
let schedule = FailureSchedule::on_attempts(&[2], FailureMode::ConnectionReset);
let mut w = ChaosWriter::new(sink, schedule);
w.write_all(b"a").unwrap();
let err = w.write_all(b"b").unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::ConnectionReset);
let sink = w.into_inner();
assert_eq!(sink, b"a");
}
#[test]
fn writer_partial_write_emits_one_byte_then_error() {
let sink: Vec<u8> = Vec::new();
let schedule = FailureSchedule::on_attempts(&[1], FailureMode::PartialWrite);
let mut w = ChaosWriter::new(sink, schedule);
let err = w.write(b"abcd").unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::WriteZero);
let sink = w.into_inner();
assert_eq!(sink, b"a");
}
#[test]
fn chaos_file_writes_and_truncates_on_partial() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("log.txt");
let schedule = FailureSchedule::on_attempts(&[2], FailureMode::PartialWrite);
let mut f = ChaosFile::create(&path, schedule).unwrap();
f.write_all(b"first").unwrap();
let _ = f.write(b"second"); drop(f);
let bytes = std::fs::read(&path).unwrap();
assert_eq!(bytes, b"firsts");
}
#[test]
fn into_inner_returns_underlying() {
let data: &[u8] = b"x";
let schedule = FailureSchedule::on_attempts(&[], FailureMode::IoError);
let r = ChaosReader::new(data, schedule);
let inner = r.into_inner();
assert_eq!(inner, b"x");
}
}