use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
use std::thread::JoinHandle;
use crate::error::{PathErrorExt, SsgError};
const QUEUE_CAP_PER_WORKER: usize = 32;
const MAX_WRITERS: usize = 4;
#[derive(Debug)]
struct WriteJob {
path: PathBuf,
bytes: Vec<u8>,
}
#[derive(Debug, Default)]
struct StateInner {
pending: usize,
errors: Vec<(PathBuf, io::Error)>,
completed: usize,
}
#[derive(Debug)]
struct PoolState {
inner: Mutex<StateInner>,
all_done: Condvar,
}
impl PoolState {
fn lock(&self) -> MutexGuard<'_, StateInner> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
}
#[derive(Debug)]
pub struct IoPool {
tx: Option<SyncSender<WriteJob>>,
workers: Vec<JoinHandle<()>>,
state: Arc<PoolState>,
}
impl Default for IoPool {
fn default() -> Self {
Self::new()
}
}
impl IoPool {
#[must_use]
pub fn new() -> Self {
Self::with_threads(default_writer_threads())
}
#[must_use]
pub fn with_threads(threads: usize) -> Self {
let threads = threads.clamp(1, MAX_WRITERS);
let (tx, rx) = sync_channel::<WriteJob>(threads * QUEUE_CAP_PER_WORKER);
let rx = Arc::new(Mutex::new(rx));
let state = Arc::new(PoolState {
inner: Mutex::new(StateInner::default()),
all_done: Condvar::new(),
});
let workers = (0..threads)
.map(|i| spawn_writer(i, Arc::clone(&rx), Arc::clone(&state)))
.filter_map(|handle| match handle {
Ok(h) => Some(h),
Err(e) => {
log::error!("io_pool: failed to spawn writer thread: {e}");
None
}
})
.collect::<Vec<_>>();
Self {
tx: if workers.is_empty() { None } else { Some(tx) },
workers,
state,
}
}
pub fn write(
&self,
path: impl Into<PathBuf>,
bytes: Vec<u8>,
) -> Result<(), SsgError> {
let path = path.into();
let Some(tx) = self.tx.as_ref() else {
perform_write(&path, &bytes).with_path(&path)?;
self.state.lock().completed += 1;
return Ok(());
};
self.state.lock().pending += 1;
if let Err(e) = tx.send(WriteJob { path, bytes }) {
let mut inner = self.state.lock();
inner.pending -= 1;
if inner.pending == 0 {
self.state.all_done.notify_all();
}
drop(inner);
let job = e.0;
return Err(SsgError::Io {
path: job.path,
source: io::Error::other(
"io_pool: writer threads terminated; job not enqueued",
),
});
}
Ok(())
}
pub fn flush(&self) -> Result<(), SsgError> {
let mut inner = self.state.lock();
while inner.pending > 0 {
inner = self
.state
.all_done
.wait(inner)
.unwrap_or_else(PoisonError::into_inner);
}
let errors = std::mem::take(&mut inner.errors);
drop(inner);
let mut iter = errors.into_iter();
let Some((first_path, first_err)) = iter.next() else {
return Ok(());
};
for (path, err) in iter {
log::error!(
"io_pool: additional write failure at '{}': {err}",
path.display()
);
}
Err(SsgError::Io {
path: first_path,
source: first_err,
})
}
#[must_use]
pub fn completed_writes(&self) -> usize {
self.state.lock().completed
}
#[must_use]
pub const fn threads(&self) -> usize {
self.workers.len()
}
}
impl Drop for IoPool {
fn drop(&mut self) {
drop(self.tx.take());
for handle in self.workers.drain(..) {
if handle.join().is_err() {
log::error!("io_pool: writer thread panicked");
}
}
let inner = self.state.lock();
for (path, err) in &inner.errors {
log::error!(
"io_pool: write failure at '{}' dropped without flush(): {err}",
path.display()
);
}
}
}
fn spawn_writer(
index: usize,
rx: Arc<Mutex<Receiver<WriteJob>>>,
state: Arc<PoolState>,
) -> io::Result<JoinHandle<()>> {
fail_point!("io_pool::spawn", |_| {
Err(io::Error::other("injected: io_pool::spawn"))
});
std::thread::Builder::new()
.name(format!("ssg-io-writer-{index}"))
.spawn(move || worker_loop(&rx, &state))
}
fn default_writer_threads() -> usize {
let cores = std::thread::available_parallelism()
.map_or(2, std::num::NonZeroUsize::get);
(cores / 2).clamp(1, MAX_WRITERS)
}
fn worker_loop(rx: &Mutex<Receiver<WriteJob>>, state: &PoolState) {
loop {
let job = {
let guard = rx.lock().unwrap_or_else(PoisonError::into_inner);
guard.recv()
};
let Ok(job) = job else {
return; };
let result = perform_write(&job.path, &job.bytes);
let mut inner = state.lock();
match result {
Ok(()) => inner.completed += 1,
Err(e) => inner.errors.push((job.path, e)),
}
inner.pending -= 1;
if inner.pending == 0 {
state.all_done.notify_all();
}
}
}
fn perform_write(path: &Path, bytes: &[u8]) -> io::Result<()> {
fail_point!("io_pool::write", |_| {
Err(io::Error::other("injected: io_pool::write"))
});
fs::write(path, bytes)
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use rayon::prelude::*;
use tempfile::tempdir;
fn io_error_path(err: &SsgError) -> Option<PathBuf> {
match err {
SsgError::Io { path, .. } => Some(path.clone()),
_ => None,
}
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn io_error_path_returns_none_for_non_io() {
let err = SsgError::Validation {
field: "f".to_string(),
message: "m".to_string(),
};
assert!(io_error_path(&err).is_none());
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn default_pool_behaves_like_new() {
let pool = IoPool::default();
assert!(pool.threads() >= 1);
assert!(pool.flush().is_ok());
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn flush_logs_additional_errors_and_returns_first() {
crate::test_support::init_logger();
let dir = tempdir().unwrap();
let bad_a = dir.path().join("no-dir-a").join("a.html");
let bad_b = dir.path().join("no-dir-b").join("b.html");
let pool = IoPool::with_threads(1);
pool.write(&bad_a, b"a".to_vec()).unwrap();
pool.write(&bad_b, b"b".to_vec()).unwrap();
let err = pool.flush().expect_err("flush must fail");
let path = io_error_path(&err).expect("must be an Io error");
assert!(path == bad_a || path == bad_b);
assert!(pool.flush().is_ok());
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn drop_without_flush_logs_unobserved_errors() {
crate::test_support::init_logger();
let dir = tempdir().unwrap();
let bad = dir.path().join("missing-dir").join("x.html");
{
let pool = IoPool::with_threads(1);
pool.write(&bad, b"x".to_vec()).unwrap();
while pool.state.lock().pending > 0 {
std::thread::yield_now();
}
}
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn empty_pool_flush_is_ok() {
let pool = IoPool::new();
assert!(pool.flush().is_ok());
assert_eq!(pool.completed_writes(), 0);
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn thread_count_is_clamped() {
assert_eq!(IoPool::with_threads(0).threads(), 1);
assert_eq!(IoPool::with_threads(100).threads(), MAX_WRITERS);
let d = default_writer_threads();
assert!((1..=MAX_WRITERS).contains(&d));
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn concurrent_rayon_producers_all_bytes_correct() {
let dir = tempdir().unwrap();
let pool = IoPool::with_threads(3);
let n = 200usize;
(0..n)
.into_par_iter()
.try_for_each(|i| {
pool.write(
dir.path().join(format!("f{i}.html")),
format!("<p>page {i}</p>").into_bytes(),
)
})
.unwrap();
pool.flush().unwrap();
assert_eq!(pool.completed_writes(), n);
for i in 0..n {
let got = fs::read_to_string(dir.path().join(format!("f{i}.html")))
.unwrap();
assert_eq!(got, format!("<p>page {i}</p>"));
}
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn write_error_surfaces_at_flush_and_pool_survives() {
let dir = tempdir().unwrap();
let missing = dir.path().join("no-such-subdir").join("x.html");
let pool = IoPool::with_threads(2);
pool.write(&missing, b"x".to_vec()).unwrap(); let err = pool.flush().expect_err("flush must surface the failure");
let path = io_error_path(&err)
.expect("flush error must be SsgError::Io with a path");
assert_eq!(path, missing);
assert!(pool.flush().is_ok());
pool.write(dir.path().join("ok.html"), b"ok".to_vec())
.unwrap();
pool.flush().unwrap();
assert_eq!(
fs::read_to_string(dir.path().join("ok.html")).unwrap(),
"ok"
);
}
#[cfg(unix)]
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn unwritable_dir_error_surfaces_at_flush() {
use std::os::unix::fs::PermissionsExt;
let dir = tempdir().unwrap();
let ro = dir.path().join("readonly");
fs::create_dir(&ro).unwrap();
fs::set_permissions(&ro, fs::Permissions::from_mode(0o555)).unwrap();
let pool = IoPool::with_threads(2);
pool.write(ro.join("blocked.html"), b"x".to_vec()).unwrap();
assert!(pool.flush().is_err());
fs::set_permissions(&ro, fs::Permissions::from_mode(0o755)).unwrap();
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn drop_without_flush_completes_queued_writes() {
let dir = tempdir().unwrap();
{
let pool = IoPool::with_threads(2);
for i in 0..50 {
pool.write(
dir.path().join(format!("d{i}.txt")),
format!("v{i}").into_bytes(),
)
.unwrap();
}
}
for i in 0..50 {
assert_eq!(
fs::read_to_string(dir.path().join(format!("d{i}.txt")))
.unwrap(),
format!("v{i}")
);
}
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn flush_is_reusable_across_batches() {
let dir = tempdir().unwrap();
let pool = IoPool::with_threads(2);
pool.write(dir.path().join("a.txt"), b"a".to_vec()).unwrap();
pool.flush().unwrap();
assert_eq!(pool.completed_writes(), 1);
pool.write(dir.path().join("b.txt"), b"b".to_vec()).unwrap();
pool.flush().unwrap();
assert_eq!(pool.completed_writes(), 2);
assert_eq!(fs::read_to_string(dir.path().join("a.txt")).unwrap(), "a");
assert_eq!(fs::read_to_string(dir.path().join("b.txt")).unwrap(), "b");
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn poisoned_state_lock_recovers_via_into_inner() {
let state = Arc::new(PoolState {
inner: Mutex::new(StateInner::default()),
all_done: Condvar::new(),
});
let poison_state = Arc::clone(&state);
let handle = std::thread::spawn(move || {
let _guard = poison_state.inner.lock().unwrap();
panic!("deliberately poison the mutex for test coverage");
});
assert!(handle.join().is_err());
let guard = state.lock();
assert_eq!(guard.pending, 0);
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn flush_condvar_wait_recovers_from_poison() {
let state = Arc::new(PoolState {
inner: Mutex::new(StateInner {
pending: 1,
..StateInner::default()
}),
all_done: Condvar::new(),
});
let (ready_tx, ready_rx) = std::sync::mpsc::channel::<()>();
let waiter_state = Arc::clone(&state);
let waiter = std::thread::spawn(move || {
let mut inner = waiter_state.lock();
let _ = ready_tx.send(());
while inner.pending > 0 {
inner = waiter_state
.all_done
.wait(inner)
.unwrap_or_else(PoisonError::into_inner);
}
inner.pending
});
ready_rx.recv().unwrap();
while state.inner.try_lock().is_err() {
std::thread::yield_now();
}
let poison_state = Arc::clone(&state);
let poisoner = std::thread::spawn(move || {
let mut inner = poison_state.inner.lock().unwrap();
inner.pending = 0;
poison_state.all_done.notify_all();
panic!("poison the mutex while notifying the waiter");
});
assert!(poisoner.join().is_err());
let pending_after =
waiter.join().expect("waiter must recover, not panic");
assert_eq!(pending_after, 0);
}
#[test]
#[serial_test::parallel(io_pool_failpoints)]
fn worker_loop_receiver_lock_recovers_from_poison() {
let (tx, rx) = sync_channel::<WriteJob>(1);
let rx = Arc::new(Mutex::new(rx));
let poison_rx = Arc::clone(&rx);
let handle = std::thread::spawn(move || {
let _guard = poison_rx.lock().unwrap();
panic!("deliberately poison the receiver mutex");
});
assert!(handle.join().is_err());
drop(tx);
let guard = rx.lock().unwrap_or_else(PoisonError::into_inner);
assert!(guard.recv().is_err());
}
#[cfg(feature = "test-fault-injection")]
mod fault_injection {
use super::*;
use serial_test::serial;
struct FailGuard<'a>(&'a str);
impl Drop for FailGuard<'_> {
fn drop(&mut self) {
let _ = fail::cfg(self.0, "off");
}
}
#[test]
#[serial(io_pool_failpoints)]
fn injected_write_failure_surfaces_at_flush() {
let _guard = FailGuard("io_pool::write");
fail::cfg("io_pool::write", "return").expect("activate failpoint");
let dir = tempdir().unwrap();
let pool = IoPool::with_threads(1);
pool.write(dir.path().join("x.html"), b"x".to_vec())
.unwrap();
let err = pool.flush().expect_err("injected write must fail");
assert!(
format!("{err}").contains("injected: io_pool::write"),
"got: {err}"
);
}
#[test]
#[serial(io_pool_failpoints)]
fn spawn_failure_falls_back_to_inline_writes() {
crate::test_support::init_logger();
let dir = tempdir().unwrap();
let pool = {
let _guard = FailGuard("io_pool::spawn");
fail::cfg("io_pool::spawn", "return")
.expect("activate failpoint");
IoPool::with_threads(2)
};
assert_eq!(pool.threads(), 0);
pool.write(dir.path().join("inline.html"), b"i".to_vec())
.unwrap();
assert_eq!(pool.completed_writes(), 1);
assert_eq!(
fs::read_to_string(dir.path().join("inline.html")).unwrap(),
"i"
);
let err = pool
.write(dir.path().join("no-dir").join("y.html"), b"y".to_vec())
.expect_err("inline write into missing dir must fail");
assert!(io_error_path(&err).is_some());
assert!(pool.flush().is_ok());
}
#[test]
#[serial(io_pool_failpoints)]
fn write_after_workers_died_returns_typed_error() {
crate::test_support::init_logger();
let dir = tempdir().unwrap();
let pool = IoPool::with_threads(1);
{
let _guard = FailGuard("io_pool::write");
fail::cfg("io_pool::write", "panic")
.expect("activate failpoint");
pool.write(dir.path().join("doomed.html"), b"d".to_vec())
.unwrap();
let err = loop {
match pool
.write(dir.path().join("after.html"), b"a".to_vec())
{
Err(e) => break e,
Ok(()) => std::thread::sleep(
std::time::Duration::from_millis(1),
),
}
};
assert!(
format!("{err}").contains("writer threads terminated"),
"got: {err}"
);
}
drop(pool);
}
}
}