use std::{
env, io,
io::{BufRead, BufReader, Write},
process::{Child, Command, Stdio},
sync::{
Arc, Condvar, Mutex,
atomic::{AtomicBool, Ordering},
},
thread,
time::Duration,
};
use anyhow::{Context as _, Result, bail};
use tracing::{debug, error, info, trace};
use crate::{Config, EvalError, Event, WORKER_ENV};
const CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(200);
struct SharedState {
todo: Vec<Vec<String>>,
active: usize,
error: Option<String>,
}
pub fn run<F>(config: &Config, cancel: &Arc<AtomicBool>, sink: Arc<Mutex<F>>) -> Result<()>
where
F: FnMut(&Event) -> Result<()> + Send + 'static,
{
let shared = Arc::new((
Mutex::new(SharedState {
todo: vec![vec![]],
active: 0,
error: None,
}),
Condvar::new(),
));
let n = config.workers.max(1);
let mut handles = Vec::with_capacity(n);
for _ in 0..n {
let shared = Arc::clone(&shared);
let sink = Arc::clone(&sink);
let cancel = Arc::clone(cancel);
let config = config.clone();
handles.push(thread::spawn(move || {
collector(&config, &cancel, shared, sink)
}));
}
for h in handles {
h.join()
.map_err(|_| anyhow::anyhow!("collector thread panicked"))??;
}
let (lock, _) = &*shared;
if let Some(e) = &lock.lock().unwrap().error {
bail!("{e}");
}
Ok(())
}
fn collector<F>(
config: &Config,
cancel: &Arc<AtomicBool>,
shared: Arc<(Mutex<SharedState>, Condvar)>,
sink: Arc<Mutex<F>>,
) -> Result<()>
where
F: FnMut(&Event) -> Result<()>,
{
let (lock, cvar) = &*shared;
let (mut proc, mut child_stdin, mut reader) = spawn_worker_pair(config)?;
loop {
let path = {
let mut s = lock.lock().unwrap();
loop {
if cancel.load(Ordering::Relaxed) {
writeln!(child_stdin, "exit")?;
child_stdin.flush()?;
let _ = proc.wait();
info!("cancellation requested, collector exiting");
return Ok(());
}
if let Some(ref e) = s.error {
let msg = e.clone();
writeln!(child_stdin, "exit")?;
child_stdin.flush()?;
error!(error = %msg, "stopping collector due to fatal error");
bail!("{msg}");
}
if !s.todo.is_empty() {
let p = s.todo.remove(0);
s.active += 1;
debug!(attr = %p.join("."), active = s.active, pending = s.todo.len(), "dispatched attribute");
break Some(p);
}
if s.active == 0 {
writeln!(child_stdin, "exit")?;
child_stdin.flush()?;
let _ = proc.wait();
info!("evaluation queue empty, exiting worker");
return Ok(());
}
let (guard, _timeout) = cvar.wait_timeout(s, CANCEL_POLL_INTERVAL).unwrap();
s = guard;
}
};
let Some(path) = path else {
continue;
};
let attr = path.join(".");
trace!(attr = %attr, "sending work to worker");
writeln!(child_stdin, "do {}", serde_json::to_string(&path)?)?;
child_stdin.flush()?;
let event = read_event(&mut reader, &path).with_context(|| {
format!(
"reading event for {attr}; worker status: {}",
worker_status(&mut proc)
)
})?;
{
let mut s = lock.lock().unwrap();
s.active -= 1;
if let Event::AttrSet { attrs, .. } = &event {
debug!(attr = %attr, new_attrs = attrs.len(), "expanded attrset");
for name in attrs {
let mut child = path.clone();
child.push(name.clone());
s.todo.push(child);
}
} else {
if let Event::Error(EvalError {
fatal: true, error, ..
}) = &event
{
error!(attr = %attr, error = %error, "fatal evaluation error");
s.error = Some(error.clone());
}
sink.lock().unwrap()(&event).context("event sink returned an error")?;
}
cvar.notify_all();
}
let status = read_line(&mut reader).with_context(|| {
format!(
"reading worker status for {attr}; worker status: {}",
worker_status(&mut proc)
)
})?;
trace!(attr = %attr, status = %status, "received worker status");
match status.as_str() {
"ready" => {}
"restart" => {
info!("restarting worker due to memory limit");
if let Ok(status) = proc.wait() {
debug!(?status, "previous worker exited");
}
(proc, child_stdin, reader) = spawn_worker_pair(config)?;
}
other => {
if let Ok(Event::Error(EvalError { error, .. })) =
serde_json::from_str::<Event>(other)
{
error!(error = %error, "worker reported error");
let mut s = lock.lock().unwrap();
s.error = Some(error.clone());
cvar.notify_all();
bail!("{error}");
}
bail!("unexpected worker message: {other}");
}
}
}
}
fn read_event(reader: &mut BufReader<impl io::Read>, path: &[String]) -> Result<Event> {
let resp = read_line(reader)?;
serde_json::from_str(&resp)
.with_context(|| format!("parsing worker response for {path:?}: {resp}"))
}
fn read_line(reader: &mut BufReader<impl io::Read>) -> Result<String> {
let mut line = String::new();
if reader.read_line(&mut line)? == 0 {
bail!("worker closed stdout unexpectedly");
}
Ok(line.trim_matches(['\n', '\r', ' ']).to_string())
}
fn spawn_worker_pair(config: &Config) -> Result<(Child, impl Write, BufReader<impl io::Read>)> {
let exe = env::current_exe().context("resolving current exe")?;
debug!("spawning worker process");
let mut child = Command::new(&exe)
.env(WORKER_ENV, "1")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.with_context(|| format!("spawning worker process from {exe:?}"))?;
let mut stdin = child.stdin.take().context("worker stdin")?;
writeln!(stdin, "{}", serde_json::to_string(config)?)?;
stdin.flush()?;
let mut reader = BufReader::new(child.stdout.take().context("worker stdout")?);
read_ready(&mut reader)?;
info!("worker ready");
Ok((child, stdin, reader))
}
fn read_ready(reader: &mut BufReader<impl io::Read>) -> Result<()> {
let line = read_line(reader)?;
if line != "ready" {
bail!("unexpected worker handshake: {line:?}");
}
Ok(())
}
fn worker_status(proc: &mut Child) -> String {
proc.try_wait()
.ok()
.flatten()
.map_or_else(|| "unknown".to_string(), |s| s.to_string())
}