use std::{
thread::{self, JoinHandle},
time::Duration,
};
use thiserror::Error;
use tokio::runtime::{Builder, Runtime};
use super::{
ManagedClient, ManagedClientConfig, ManagedClientHandle, ManagedClientTask, ManagedConfigError,
ManagedEventSubscription, ManagedOutcome, ManagedSubscribeError, ManagedTargetConfig,
};
const BLOCKING_RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(100);
pub(super) struct WorkerRuntime {
runtime: Option<Runtime>,
}
impl WorkerRuntime {
pub(super) fn new(runtime: Runtime) -> Self {
Self {
runtime: Some(runtime),
}
}
fn block_on<F: std::future::Future>(&self, task: F) -> F::Output {
self.runtime
.as_ref()
.expect("worker runtime is available until shutdown")
.block_on(task)
}
pub(super) fn shutdown(mut self) {
self.runtime
.take()
.expect("worker runtime is available until shutdown")
.shutdown_timeout(BLOCKING_RUNTIME_SHUTDOWN_TIMEOUT);
}
}
impl Drop for WorkerRuntime {
fn drop(&mut self) {
if let Some(runtime) = self.runtime.take() {
runtime.shutdown_background();
}
}
}
#[derive(Debug, Error)]
pub enum BlockingManagedStartError {
#[error(transparent)]
Config(#[from] ManagedConfigError),
#[error(transparent)]
Subscribe(#[from] ManagedSubscribeError),
#[error("failed to create blocking managed Tokio runtime: {source}")]
Runtime {
#[source]
source: std::io::Error,
},
#[error("failed to spawn blocking managed worker thread: {source}")]
ThreadSpawn {
#[source]
source: std::io::Error,
},
}
#[derive(Debug, Error, Eq, PartialEq)]
pub enum BlockingManagedJoinError {
#[error("blocking managed worker thread panicked")]
WorkerPanicked,
}
#[must_use = "dropping the blocking managed owner requests stop and joins its worker"]
pub struct BlockingManagedClient {
handle: ManagedClientHandle,
worker: Option<JoinHandle<ManagedOutcome>>,
}
impl BlockingManagedClient {
pub fn start(
config: ManagedClientConfig,
targets: Vec<ManagedTargetConfig>,
) -> Result<Self, BlockingManagedStartError> {
let (task, handle) = ManagedClient::task(config, targets)?;
Self::start_task(task, handle)
}
pub fn start_with_subscription(
config: ManagedClientConfig,
targets: Vec<ManagedTargetConfig>,
) -> Result<(Self, ManagedEventSubscription), BlockingManagedStartError> {
let (task, handle) = ManagedClient::task(config, targets)?;
let subscription = handle.subscribe()?;
let owner = Self::start_task(task, handle)?;
Ok((owner, subscription))
}
fn start_task(
task: ManagedClientTask,
handle: ManagedClientHandle,
) -> Result<Self, BlockingManagedStartError> {
let runtime = Builder::new_current_thread()
.enable_io()
.enable_time()
.build()
.map_err(|source| BlockingManagedStartError::Runtime { source })?;
let runtime = WorkerRuntime::new(runtime);
let worker = thread::Builder::new()
.name("irtt-managed".to_owned())
.spawn(move || {
let outcome = runtime.block_on(task);
runtime.shutdown();
outcome
})
.map_err(|source| BlockingManagedStartError::ThreadSpawn { source })?;
Ok(Self {
handle,
worker: Some(worker),
})
}
pub fn handle(&self) -> ManagedClientHandle {
self.handle.clone()
}
pub fn join(mut self) -> Result<ManagedOutcome, BlockingManagedJoinError> {
self.join_worker()
}
fn join_worker(&mut self) -> Result<ManagedOutcome, BlockingManagedJoinError> {
self.worker
.take()
.expect("blocking managed worker is joined at most once")
.join()
.map_err(|_| BlockingManagedJoinError::WorkerPanicked)
}
}
impl Drop for BlockingManagedClient {
fn drop(&mut self) {
if self.worker.is_some() {
let _receipt = self.handle.stop();
let _ = self.join_worker();
}
}
}