use std::path::PathBuf;
use std::thread::JoinHandle;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
#[derive(Debug, Clone)]
pub struct HostOptions {
pub thread_name: &'static str,
pub storage_dir: PathBuf,
pub command_backlog: usize,
}
#[async_trait::async_trait]
pub trait Body: Send + 'static {
type Cmd: Send + 'static;
async fn open(&mut self, ctx: &commonware_runtime::tokio::Context) -> Result<(), String>;
async fn handle(&mut self, ctx: &commonware_runtime::tokio::Context, cmd: Self::Cmd);
async fn on_drain_complete(&mut self, ctx: &commonware_runtime::tokio::Context) {
let _ = ctx;
}
}
#[derive(Debug, thiserror::Error)]
pub enum SpawnError {
#[error("dedicated host thread failed before it could report readiness")]
RuntimeStart,
#[error("host body failed to open: {0}")]
Open(String),
}
#[derive(Debug, thiserror::Error)]
#[error("host is shut down")]
pub struct Closed;
pub async fn call<Cmd, T>(
tx: &mpsc::Sender<Cmd>,
build: impl FnOnce(oneshot::Sender<T>) -> Cmd,
) -> Result<T, Closed> {
let (ack, ack_rx) = oneshot::channel();
tx.send(build(ack)).await.map_err(|_| Closed)?;
ack_rx.await.map_err(|_| Closed)
}
#[derive(Debug)]
pub struct HostHandle<Cmd> {
tx: Option<mpsc::Sender<Cmd>>,
thread: Option<JoinHandle<()>>,
}
impl<Cmd> HostHandle<Cmd> {
#[must_use]
pub const fn sender(&self) -> Option<&mpsc::Sender<Cmd>> {
self.tx.as_ref()
}
}
impl<Cmd> Drop for HostHandle<Cmd> {
fn drop(&mut self) {
drop(self.tx.take());
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
}
}
pub fn spawn_host<B: Body>(
opts: HostOptions,
shutdown: CancellationToken,
body: B,
) -> Result<HostHandle<B::Cmd>, SpawnError> {
let (tx, rx) = mpsc::channel::<B::Cmd>(opts.command_backlog);
spawn_host_with_channel(opts, shutdown, body, tx, rx)
}
pub fn spawn_host_with_channel<B: Body>(
opts: HostOptions,
shutdown: CancellationToken,
body: B,
tx: mpsc::Sender<B::Cmd>,
rx: mpsc::Receiver<B::Cmd>,
) -> Result<HostHandle<B::Cmd>, SpawnError> {
let HostOptions {
thread_name,
storage_dir,
command_backlog: _,
} = opts;
let (ready_tx, ready_rx) = std::sync::mpsc::channel::<Result<(), String>>();
let thread = std::thread::Builder::new()
.name(thread_name.to_owned())
.spawn(move || run_dedicated(storage_dir, rx, &ready_tx, &shutdown, body))
.expect("spawn dedicated host thread");
match ready_rx.recv() {
Ok(Ok(())) => Ok(HostHandle {
tx: Some(tx),
thread: Some(thread),
}),
Ok(Err(reason)) => {
let _ = thread.join();
Err(SpawnError::Open(reason))
}
Err(_) => {
let _ = thread.join();
Err(SpawnError::RuntimeStart)
}
}
}
fn run_dedicated<B: Body>(
storage_dir: PathBuf,
rx: mpsc::Receiver<B::Cmd>,
ready_tx: &std::sync::mpsc::Sender<Result<(), String>>,
shutdown: &CancellationToken,
mut body: B,
) {
use commonware_runtime::Runner as _;
let cfg = commonware_runtime::tokio::Config::default().with_storage_directory(storage_dir);
let runner = commonware_runtime::tokio::Runner::new(cfg);
runner.start(|context| async move {
if let Err(reason) = body.open(&context).await {
tracing::warn!(error = %reason, "dedicated host thread's body failed to open");
let _ = ready_tx.send(Err(reason));
return;
}
let _ = ready_tx.send(Ok(()));
let mut rx = rx;
loop {
let cmd = commonware_macros::select! {
() = shutdown.cancelled() => break,
maybe = rx.recv() => match maybe {
Some(cmd) => cmd,
None => break,
},
};
body.handle(&context, cmd).await;
}
let mut drained = 0usize;
while let Ok(cmd) = rx.try_recv() {
body.handle(&context, cmd).await;
drained += 1;
}
if drained > 0 {
tracing::debug!(drained, "ran queued commands after shutdown");
}
body.on_drain_complete(&context).await;
});
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs, clippy::unwrap_used)]
use super::*;
use tokio_util::sync::CancellationToken;
fn scratch_dir(label: &str) -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!(
"polyc-host-{label}-{}-{}",
std::process::id(),
uuid::Uuid::new_v4()
));
let _ = std::fs::remove_dir_all(&dir);
dir
}
fn opts(label: &str, command_backlog: usize) -> HostOptions {
let thread_name: &'static str =
Box::leak(format!("polyc-host-test-{label}").into_boxed_str());
HostOptions {
thread_name,
storage_dir: scratch_dir(label),
command_backlog,
}
}
struct EagerFailBody;
#[async_trait::async_trait]
impl Body for EagerFailBody {
type Cmd = ();
async fn open(&mut self, _ctx: &commonware_runtime::tokio::Context) -> Result<(), String> {
Err("the store is unopenable".to_owned())
}
async fn handle(&mut self, _ctx: &commonware_runtime::tokio::Context, (): Self::Cmd) {}
}
#[test]
fn an_eager_open_failure_surfaces_through_spawn_host() {
let opts = HostOptions {
thread_name: "polyc-host-test-eager-fail",
storage_dir: scratch_dir("eager-fail"),
command_backlog: 8,
};
let err = spawn_host(opts, CancellationToken::new(), EagerFailBody)
.expect_err("an eager Body::open failure must fail spawn_host");
match err {
SpawnError::Open(reason) => assert_eq!(reason, "the store is unopenable"),
SpawnError::RuntimeStart => panic!("expected Open, got RuntimeStart"),
}
}
enum LazyCmd {
Ping {
id: u32,
ack: tokio::sync::oneshot::Sender<u32>,
},
Fail {
ack: tokio::sync::oneshot::Sender<Result<u32, String>>,
},
}
struct LazyBody;
#[async_trait::async_trait]
impl Body for LazyBody {
type Cmd = LazyCmd;
async fn open(&mut self, _ctx: &commonware_runtime::tokio::Context) -> Result<(), String> {
Ok(())
}
async fn handle(&mut self, _ctx: &commonware_runtime::tokio::Context, cmd: Self::Cmd) {
match cmd {
LazyCmd::Ping { id, ack } => {
let _ = ack.send(id);
}
LazyCmd::Fail { ack } => {
let _ = ack.send(Err("simulated per-command failure".to_owned()));
}
}
}
}
#[tokio::test]
async fn dropping_the_handle_never_deadlocks() {
let handle = spawn_host(opts("drop-order", 8), CancellationToken::new(), LazyBody)
.expect("lazy open always succeeds");
let (ack, ack_rx) = tokio::sync::oneshot::channel();
handle
.sender()
.expect("handle is fresh")
.send(LazyCmd::Ping { id: 42, ack })
.await
.expect("channel is open");
assert_eq!(ack_rx.await.expect("the loop answers"), 42);
drop(handle);
}
#[tokio::test]
async fn a_command_that_fails_inside_handle_does_not_crash_the_loop() {
let handle = spawn_host(
opts("handle-failure", 8),
CancellationToken::new(),
LazyBody,
)
.expect("lazy open always succeeds");
let (fail_ack, fail_ack_rx) = tokio::sync::oneshot::channel();
handle
.sender()
.expect("handle is fresh")
.send(LazyCmd::Fail { ack: fail_ack })
.await
.expect("channel is open");
assert_eq!(
fail_ack_rx.await.expect("the loop still answers"),
Err("simulated per-command failure".to_owned())
);
let (ack, ack_rx) = tokio::sync::oneshot::channel();
handle
.sender()
.expect("handle is fresh")
.send(LazyCmd::Ping { id: 7, ack })
.await
.expect("channel is open");
assert_eq!(ack_rx.await.expect("the loop answers"), 7);
drop(handle);
}
#[tokio::test]
async fn a_command_queued_before_shutdown_cancel_fires_is_still_processed() {
let shutdown = CancellationToken::new();
let handle = spawn_host(opts("post-cancel-drain", 8), shutdown.clone(), LazyBody)
.expect("lazy open always succeeds");
let extra_tx = handle.sender().expect("handle is fresh").clone();
let (ack, ack_rx) = tokio::sync::oneshot::channel();
extra_tx
.send(LazyCmd::Ping { id: 9, ack })
.await
.expect("channel is open");
shutdown.cancel();
assert_eq!(ack_rx.await.expect("a queued command must still run"), 9);
drop(handle);
drop(extra_tx);
}
struct GatedBody {
gate: std::sync::Arc<tokio::sync::Notify>,
order: std::sync::Arc<std::sync::Mutex<Vec<&'static str>>>,
}
enum GatedCmd {
Slow {
ack: tokio::sync::oneshot::Sender<()>,
},
Ping {
id: u32,
ack: tokio::sync::oneshot::Sender<u32>,
},
}
#[async_trait::async_trait]
impl Body for GatedBody {
type Cmd = GatedCmd;
async fn open(&mut self, _ctx: &commonware_runtime::tokio::Context) -> Result<(), String> {
Ok(())
}
async fn handle(&mut self, _ctx: &commonware_runtime::tokio::Context, cmd: Self::Cmd) {
match cmd {
GatedCmd::Slow { ack } => {
self.gate.notified().await;
self.order.lock().unwrap().push("slow");
let _ = ack.send(());
}
GatedCmd::Ping { id, ack } => {
self.order.lock().unwrap().push("ping");
let _ = ack.send(id);
}
}
}
async fn on_drain_complete(&mut self, _ctx: &commonware_runtime::tokio::Context) {
self.order.lock().unwrap().push("drained");
}
}
#[tokio::test]
async fn on_drain_complete_fires_after_the_drain_has_run_every_command() {
let gate = std::sync::Arc::new(tokio::sync::Notify::new());
let order = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let shutdown = CancellationToken::new();
let body = GatedBody {
gate: gate.clone(),
order: order.clone(),
};
let handle = spawn_host(opts("drain-complete", 8), shutdown.clone(), body)
.expect("lazy open always succeeds");
let (slow_ack, slow_ack_rx) = tokio::sync::oneshot::channel();
handle
.sender()
.expect("handle is fresh")
.send(GatedCmd::Slow { ack: slow_ack })
.await
.expect("channel is open");
let (ack_a, ack_a_rx) = tokio::sync::oneshot::channel();
let (ack_b, ack_b_rx) = tokio::sync::oneshot::channel();
handle
.sender()
.expect("handle is fresh")
.send(GatedCmd::Ping { id: 1, ack: ack_a })
.await
.expect("channel is open");
handle
.sender()
.expect("handle is fresh")
.send(GatedCmd::Ping { id: 2, ack: ack_b })
.await
.expect("channel is open");
shutdown.cancel();
gate.notify_one();
slow_ack_rx.await.expect("slow command completes");
ack_a_rx.await.expect("drained");
ack_b_rx.await.expect("drained");
drop(handle);
let seen = order.lock().unwrap().clone();
assert_eq!(
seen,
vec!["slow", "ping", "ping", "drained"],
"on_drain_complete must fire strictly after every drained command"
);
}
#[tokio::test]
async fn call_round_trips_and_reports_closed_once_the_host_is_gone() {
let shutdown = CancellationToken::new();
let handle = spawn_host(opts("call-helper", 8), shutdown.clone(), LazyBody)
.expect("lazy open always succeeds");
let tx = handle.sender().expect("handle is fresh").clone();
let pong = call(&tx, |ack| LazyCmd::Ping { id: 5, ack })
.await
.expect("the channel is open");
assert_eq!(pong, 5);
shutdown.cancel();
drop(handle);
let closed = call(&tx, |ack| LazyCmd::Ping { id: 6, ack }).await;
assert!(matches!(closed, Err(Closed)));
}
#[tokio::test]
async fn a_caller_supplied_channel_works_identically_to_an_internal_one() {
let (tx, rx) = tokio::sync::mpsc::channel::<LazyCmd>(8);
let peer_clone = tx.clone();
let handle = spawn_host_with_channel(
opts("with-channel", 8),
CancellationToken::new(),
LazyBody,
tx,
rx,
)
.expect("lazy open always succeeds");
let (ack, ack_rx) = tokio::sync::oneshot::channel();
peer_clone
.send(LazyCmd::Ping { id: 11, ack })
.await
.expect("channel is open");
assert_eq!(ack_rx.await.expect("the loop answers"), 11);
drop(peer_clone);
drop(handle);
}
}