use std::collections::HashMap;
use std::io::{Read, Write};
use std::process::{Child, Output};
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use crate::selector::Situation;
use crate::types::{Decision, EvaluatorSource};
#[derive(Debug, thiserror::Error)]
pub enum ExternalError {
#[error("evaluator kind '{0}' not registered")]
NotRegistered(String),
#[error("source resolution failed: {0}")]
SourceResolution(String),
#[error("evaluator failed: {0}")]
Execution(String),
#[error(transparent)]
Other(#[from] Box<dyn std::error::Error + Send + Sync>),
}
impl ExternalError {
pub fn code(&self) -> &'static str {
match self {
ExternalError::NotRegistered(_) => "evaluator_not_registered",
ExternalError::SourceResolution(_) => "evaluator_source_unresolved",
ExternalError::Execution(_) => "evaluator_failed",
ExternalError::Other(e) if e.is::<TimedOut>() => "evaluator_timeout",
ExternalError::Other(_) => "evaluator_error",
}
}
}
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(Debug, thiserror::Error)]
#[error("evaluator timed out after {0:?}")]
pub struct TimedOut(pub Duration);
pub fn wait_with_timeout(
mut child: Child,
input: Option<Vec<u8>>,
timeout: Duration,
) -> Result<Output, ExternalError> {
let started = Instant::now();
let stdin = child.stdin.take();
let writer = input.and_then(|input| {
let mut stdin = stdin?;
Some(thread::spawn(move || stdin.write_all(&input)))
});
let stdout = child.stdout.take().map(read_on_thread);
let stderr = child.stderr.take().map(read_on_thread);
let mut pause = Duration::from_millis(1);
let status = loop {
let exited = child
.try_wait()
.map_err(|e| ExternalError::Execution(format!("wait: {e}")))?;
if let Some(status) = exited {
break status;
}
let elapsed = started.elapsed();
if elapsed >= timeout {
let _ = child.kill();
let _ = child.wait();
return Err(ExternalError::Other(Box::new(TimedOut(timeout))));
}
thread::sleep(pause.min(timeout - elapsed));
pause = (pause * 2).min(Duration::from_millis(50));
};
if let Some(Ok(Err(e))) = writer.map(JoinHandle::join) {
if e.kind() != std::io::ErrorKind::BrokenPipe {
return Err(ExternalError::Execution(format!("write stdin: {e}")));
}
}
let collect = |pipe: Option<JoinHandle<Vec<u8>>>| {
pipe.and_then(|reader| reader.join().ok())
.unwrap_or_default()
};
Ok(Output {
status,
stdout: collect(stdout),
stderr: collect(stderr),
})
}
fn read_on_thread(mut pipe: impl Read + Send + 'static) -> JoinHandle<Vec<u8>> {
thread::spawn(move || {
let mut out = Vec::new();
let _ = pipe.read_to_end(&mut out);
out
})
}
pub trait ExternalEvaluator: Send + Sync {
fn kind(&self) -> &'static str;
fn evaluate(
&self,
source: &EvaluatorSource,
situation: &Situation,
action: &str,
agent_id: &str,
) -> Result<Decision, ExternalError>;
}
#[derive(Default)]
pub struct ExternalEvaluatorRegistry {
impls: HashMap<&'static str, Arc<dyn ExternalEvaluator>>,
}
impl ExternalEvaluatorRegistry {
pub fn new() -> Self {
Self {
impls: HashMap::new(),
}
}
pub fn register(&mut self, eval: Arc<dyn ExternalEvaluator>) {
self.impls.insert(eval.kind(), eval);
}
pub fn get(&self, kind: &str) -> Option<&Arc<dyn ExternalEvaluator>> {
self.impls.get(kind)
}
pub fn kinds(&self) -> Vec<&'static str> {
self.impls.keys().copied().collect()
}
pub fn len(&self) -> usize {
self.impls.len()
}
pub fn is_empty(&self) -> bool {
self.impls.is_empty()
}
}
impl std::fmt::Debug for ExternalEvaluatorRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExternalEvaluatorRegistry")
.field("kinds", &self.kinds())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockRunner;
impl ExternalEvaluator for MockRunner {
fn kind(&self) -> &'static str {
"mock"
}
fn evaluate(
&self,
_source: &EvaluatorSource,
_situation: &Situation,
_action: &str,
_agent_id: &str,
) -> Result<Decision, ExternalError> {
Ok(Decision::NoPolicyMatch)
}
}
#[cfg(unix)]
#[test]
fn a_process_past_its_timeout_is_killed() {
use std::process::{Command, Stdio};
let child = Command::new("sleep")
.arg("5")
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap();
let started = Instant::now();
let error = wait_with_timeout(child, None, Duration::from_millis(100)).unwrap_err();
assert!(started.elapsed() < Duration::from_secs(4));
assert_eq!(error.code(), "evaluator_timeout");
}
#[cfg(unix)]
#[test]
fn a_process_gets_its_input_and_returns_its_output() {
use std::process::{Command, Stdio};
let child = Command::new("cat")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap();
let input = vec![b'x'; 1 << 20];
let output = wait_with_timeout(child, Some(input.clone()), DEFAULT_TIMEOUT).unwrap();
assert!(output.status.success());
assert_eq!(output.stdout, input);
}
#[test]
fn registry_register_and_get() {
let mut reg = ExternalEvaluatorRegistry::new();
assert!(reg.is_empty());
reg.register(Arc::new(MockRunner));
assert_eq!(reg.len(), 1);
assert!(reg.get("mock").is_some());
assert!(reg.get("missing").is_none());
assert_eq!(reg.kinds(), vec!["mock"]);
}
}