#![allow(
clippy::missing_docs_in_private_items,
clippy::pattern_type_mismatch,
clippy::multiple_crate_versions
)]
#![cfg_attr(doctest, doc = include_str!("../README.md"))]
use std::any;
use std::env;
use std::fmt;
use std::io::{self, Read, Write};
use std::num::NonZeroUsize;
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};
use postcard::accumulator::{CobsAccumulator, FeedResult};
const WORKER_ENV_PREFIX: &str = "__TARNISH_WORKER_";
#[doc(hidden)]
#[must_use]
pub fn is_worker_process() -> bool {
env::vars().any(|(key, _)| key.starts_with(WORKER_ENV_PREFIX))
}
const GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
fn worker_env_name<T: 'static>() -> String {
let type_name = any::type_name::<T>()
.replace("::", "_")
.replace(['<', '>'], "_");
format!("{WORKER_ENV_PREFIX}{type_name}__")
}
#[derive(Debug, serde::Serialize, serde::Deserialize)]
enum Message {
Request(Vec<u8>),
Response(Vec<u8>),
Error(String),
Shutdown,
}
impl Message {
fn encode(&self) -> io::Result<Vec<u8>> {
postcard::to_allocvec_cobs(self).map_err(|e| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("Message encoding failed: {e}"),
)
})
}
}
#[derive(Debug)]
pub enum ProcessError {
SpawnError(io::Error),
ExecutablePathError(io::Error),
CommunicationError(io::Error),
ProcessTerminated,
ProcessPanicked(String),
TaskError(String),
ProtocolError(String),
}
impl fmt::Display for ProcessError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::SpawnError(e) => write!(f, "Failed to spawn process: {e}"),
Self::ExecutablePathError(e) => {
write!(f, "Failed to get executable path: {e}")
}
Self::CommunicationError(e) => write!(f, "Communication error: {e}"),
Self::ProcessTerminated => write!(f, "Process terminated unexpectedly"),
Self::ProcessPanicked(msg) => write!(f, "Process panicked: {msg}"),
Self::TaskError(msg) => write!(f, "Task error: {msg}"),
Self::ProtocolError(msg) => write!(f, "Protocol error: {msg}"),
}
}
}
impl std::error::Error for ProcessError {}
impl From<io::Error> for ProcessError {
fn from(err: io::Error) -> Self {
Self::CommunicationError(err)
}
}
pub type Result<T> = std::result::Result<T, ProcessError>;
pub trait MessageEncode {
fn encode(&self) -> std::result::Result<Vec<u8>, String>;
}
pub trait MessageDecode: Sized {
fn decode(bytes: &[u8]) -> std::result::Result<Self, String>;
}
impl<T: serde::Serialize> MessageEncode for T {
fn encode(&self) -> std::result::Result<Vec<u8>, String> {
postcard::to_allocvec_cobs(self).map_err(|e| format!("COBS encoding error: {e}"))
}
}
impl<T: for<'de> serde::Deserialize<'de>> MessageDecode for T {
fn decode(bytes: &[u8]) -> std::result::Result<Self, String> {
let mut buf = bytes.to_vec();
postcard::from_bytes_cobs(&mut buf).map_err(|e| format!("COBS decoding error: {e}"))
}
}
pub trait Task: Default + 'static {
type Input: MessageEncode + MessageDecode;
type Output: MessageEncode + MessageDecode;
type Error: fmt::Display;
fn run(&mut self, input: Self::Input) -> std::result::Result<Self::Output, Self::Error>;
}
pub struct Process<T: Task> {
child: Child,
stdin: std::process::ChildStdin,
stdout: std::process::ChildStdout,
_phantom: std::marker::PhantomData<T>,
}
impl<T: Task> Process<T> {
pub fn spawn() -> Result<Self> {
Self::spawn_internal()
}
fn spawn_internal() -> Result<Self> {
let exe_path = env::current_exe().map_err(ProcessError::ExecutablePathError)?;
let env_name = worker_env_name::<T>();
let mut child = Command::new(exe_path)
.env(&env_name, "1")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.map_err(ProcessError::SpawnError)?;
let stdin = child.stdin.take().ok_or_else(|| {
ProcessError::SpawnError(io::Error::new(
io::ErrorKind::BrokenPipe,
"Failed to capture child stdin",
))
})?;
let stdout = child.stdout.take().ok_or_else(|| {
ProcessError::SpawnError(io::Error::new(
io::ErrorKind::BrokenPipe,
"Failed to capture child stdout",
))
})?;
Ok(Self {
child,
stdin,
stdout,
_phantom: std::marker::PhantomData,
})
}
#[allow(clippy::needless_pass_by_value)] pub fn call(&mut self, input: T::Input) -> Result<T::Output> {
let encoded_input = input
.encode()
.map_err(|e| ProcessError::ProtocolError(format!("Failed to encode input: {e}")))?;
if let Err(e) = self.send_message(&Message::Request(encoded_input)) {
self.restart()?;
return Err(e);
}
match self.receive_message() {
Ok(Message::Response(encoded_output)) => {
T::Output::decode(&encoded_output).map_err(|e| {
ProcessError::ProtocolError(format!("Failed to decode output: {e}"))
})
}
Ok(Message::Error(err)) => Err(ProcessError::TaskError(err)),
Ok(msg) => {
self.restart()?;
Err(ProcessError::ProtocolError(format!(
"Unexpected message: {msg:?}"
)))
}
Err(e) => {
self.restart()?;
Err(e)
}
}
}
fn send_message(&mut self, msg: &Message) -> Result<()> {
let bytes = msg.encode()?;
self.stdin.write_all(&bytes)?;
Ok(())
}
fn receive_message(&mut self) -> Result<Message> {
let mut raw_buf = [0_u8; 256];
let mut cobs_buf: CobsAccumulator<1024> = CobsAccumulator::new();
loop {
let bytes_read = self.stdout.read(&mut raw_buf)?;
if bytes_read == 0 {
return Err(ProcessError::ProcessTerminated);
}
let mut window = raw_buf.get(..bytes_read).ok_or_else(|| {
ProcessError::ProtocolError(format!(
"Read returned invalid byte count: {bytes_read} > {}",
raw_buf.len()
))
})?;
while !window.is_empty() {
window = match cobs_buf.feed::<Message>(window) {
FeedResult::Consumed => break,
FeedResult::OverFull(remaining) => remaining,
FeedResult::DeserError(_remaining) => {
return Err(ProcessError::ProtocolError(
"COBS deserialization error".to_owned(),
));
}
FeedResult::Success { data, .. } => {
return Ok(data);
}
};
}
}
}
fn restart(&mut self) -> Result<()> {
drop(self.child.kill());
drop(self.child.wait());
let new_handle = Self::spawn_internal()?;
*self = new_handle;
Ok(())
}
pub fn is_running(&mut self) -> Result<bool> {
match self.child.try_wait() {
Ok(Some(_)) => Ok(false),
Ok(None) => Ok(true),
Err(e) => Err(ProcessError::CommunicationError(e)),
}
}
}
impl<T: Task> Drop for Process<T> {
fn drop(&mut self) {
if self.send_message(&Message::Shutdown).is_ok() {
let start = Instant::now();
while start.elapsed() < GRACEFUL_SHUTDOWN_TIMEOUT {
if let Ok(Some(_)) = self.child.try_wait() {
return; }
std::thread::sleep(Duration::from_millis(100));
}
}
drop(self.child.kill());
drop(self.child.wait());
}
}
pub struct ProcessPool<T: Task> {
workers: Vec<Process<T>>,
next_worker: std::sync::atomic::AtomicUsize,
}
impl<T: Task> ProcessPool<T> {
pub fn new(size: NonZeroUsize) -> Result<Self> {
let workers = (0..size.get())
.map(|_| Process::<T>::spawn())
.collect::<Result<Vec<_>>>()?;
Ok(Self {
workers,
next_worker: std::sync::atomic::AtomicUsize::new(0),
})
}
#[allow(clippy::arithmetic_side_effects, clippy::indexing_slicing)]
pub fn call(&mut self, input: T::Input) -> Result<T::Output> {
let idx = self
.next_worker
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
% self.workers.len();
self.workers
.get_mut(idx)
.ok_or_else(|| ProcessError::ProtocolError(format!("Invalid worker index: {idx}")))?
.call(input)
}
#[must_use]
pub const fn size(&self) -> usize {
self.workers.len()
}
}
pub fn main<T: Task>(parent_main: fn()) {
let env_name = worker_env_name::<T>();
if env::var(&env_name).is_ok() {
let exit_code = run_worker_loop::<T>();
std::process::exit(exit_code);
}
parent_main();
}
#[must_use]
pub fn worker_main<T: Task>() -> Option<i32> {
let env_name = worker_env_name::<T>();
if env::var(&env_name).is_err() {
return None; }
let exit_code = run_worker_loop::<T>();
Some(exit_code)
}
#[allow(clippy::print_stderr)] fn run_worker_loop<T: Task>() -> i32 {
let mut stdin = io::stdin();
let mut stdout = io::stdout();
let mut worker = T::default();
let mut raw_buf = [0_u8; 256];
let mut cobs_buf: CobsAccumulator<1024> = CobsAccumulator::new();
loop {
let bytes_read = match stdin.read(&mut raw_buf) {
Ok(n) => n,
Err(e) => {
eprintln!("[CHILD] Failed to read from parent: {e}");
return 1;
}
};
if bytes_read == 0 {
return 0;
}
let Some(mut window) = raw_buf.get(..bytes_read) else {
eprintln!(
"[CHILD] Read returned invalid byte count: {bytes_read} > {}",
raw_buf.len()
);
return 1;
};
while !window.is_empty() {
let message = match cobs_buf.feed::<Message>(window) {
FeedResult::Consumed => break,
FeedResult::OverFull(remaining) => {
window = remaining;
continue;
}
FeedResult::DeserError(_remaining) => {
eprintln!("[CHILD] COBS deserialization error");
return 1;
}
FeedResult::Success { data, remaining } => {
window = remaining;
data
}
};
match message {
Message::Shutdown => {
return 0;
}
Message::Request(encoded_input) => {
let input = match T::Input::decode(&encoded_input) {
Ok(inp) => inp,
Err(e) => {
eprintln!("[WORKER] Failed to decode input: {e}");
let err_msg = Message::Error(format!("Decode error: {e}"));
if send_message(&mut stdout, &err_msg).is_err() {
return 1;
}
continue;
}
};
let response = match worker.run(input) {
Ok(output) => match output.encode() {
Ok(bytes) => Message::Response(bytes),
Err(e) => Message::Error(format!("Encoding error: {e}")),
},
Err(err) => Message::Error(err.to_string()),
};
if send_message(&mut stdout, &response).is_err() {
return 1;
}
}
Message::Response(_) | Message::Error(_) => {
eprintln!("[CHILD] Unexpected message type");
return 1;
}
}
}
}
}
fn send_message(stdout: &mut io::Stdout, msg: &Message) -> io::Result<()> {
let bytes = msg.encode()?;
stdout.write_all(&bytes)?;
stdout.flush()?;
Ok(())
}
#[macro_export]
macro_rules! task {
($name:ident: || $body:block) => {{
paste::paste! {
#[derive(Default)]
struct [<__TarnishTask $name:camel>];
impl $crate::Task for [<__TarnishTask $name:camel>] {
type Input = ();
type Output = ();
type Error = $crate::ProcessError;
fn run(&mut self, _input: ()) -> ::std::result::Result<Self::Output, Self::Error> {
(|| $body)()
}
}
if let ::std::option::Option::Some(exit_code) = $crate::worker_main::<[<__TarnishTask $name:camel>]>() {
::std::process::exit(exit_code);
}
if !$crate::is_worker_process() {
(|| -> $crate::Result<()> {
let mut process = $crate::Process::<[<__TarnishTask $name:camel>]>::spawn()?;
process.call(())
})()
} else {
::std::result::Result::Err($crate::ProcessError::ProtocolError(
"Worker process did not find its task".to_owned()
))
}
}
}};
($name:ident: || -> Result<$ok:ty, $err:ty> $body:block) => {{
paste::paste! {
#[derive(Default)]
struct [<__TarnishTask $name:camel>];
impl $crate::Task for [<__TarnishTask $name:camel>] {
type Input = ();
type Output = $ok;
type Error = $err;
fn run(&mut self, _input: ()) -> ::std::result::Result<Self::Output, Self::Error> {
(|| $body)()
}
}
if let ::std::option::Option::Some(exit_code) = $crate::worker_main::<[<__TarnishTask $name:camel>]>() {
::std::process::exit(exit_code);
}
if !$crate::is_worker_process() {
(|| -> $crate::Result<$ok> {
let mut process = $crate::Process::<[<__TarnishTask $name:camel>]>::spawn()?;
process.call(())
})()
} else {
::std::result::Result::Err($crate::ProcessError::ProtocolError(
"Worker process did not find its task".to_owned()
))
}
}
}};
}