#[cfg(feature = "external")]
use crate::Pacer;
use crate::{
BlobLayout, BlobVersion, Clock, Error, Execution, Handle, METRICS_PREFIX, Name, SinkOf,
StreamOf, child_label,
network::{
metered::Network as MeteredNetwork,
tokio::{Config as TokioNetworkConfig, Network as TokioNetwork},
},
prefixed_name,
process::metered::Metrics as MeteredProcess,
signal::Signal,
storage::{
metered::Storage as MeteredStorage,
tokio::{Config as TokioStorageConfig, Storage as TokioStorage},
},
telemetry::metrics::{
CounterFamily, GaugeFamily, Metric, Register, Registered, Registry, add_attribute, raw,
task::Label, validate_label,
},
utils::{self, FactoryGuard, Panicker, signal::Stopper, supervision::Tree},
};
use commonware_macros::{select, stability};
#[stability(BETA)]
use commonware_parallel::{Rayon, rayon::ThreadPoolBuilder};
use commonware_utils::{sync::Mutex, sys_rng};
use governor::clock::{Clock as GClock, ReasonablyRealtime};
use rand_core::{Rng, TryCryptoRng, TryRng};
use std::{
convert::Infallible,
env,
future::{Future, poll_fn},
net::{IpAddr, SocketAddr},
num::NonZeroUsize,
ops::RangeInclusive,
panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
path::PathBuf,
pin::pin,
sync::Arc,
task::Poll,
time::{Duration, SystemTime},
};
use tokio::{
runtime::{Builder, Handle as RuntimeHandle},
sync::Notify,
};
#[derive(Debug)]
struct Metrics {
tasks_spawned: CounterFamily<Label>,
tasks_running: GaugeFamily<Label>,
}
impl Metrics {
pub fn init(registry: &mut impl Register) -> Self {
Self {
tasks_spawned: registry.register(
"tasks_spawned",
"Total number of tasks spawned",
raw::Family::default(),
),
tasks_running: registry.register(
"tasks_running",
"Number of tasks currently running",
raw::Family::default(),
),
}
}
}
#[derive(Clone, Debug)]
pub struct NetworkConfig {
tcp_nodelay: Option<bool>,
zero_linger: bool,
connect_timeout: Duration,
read_write_timeout: Duration,
}
impl Default for NetworkConfig {
fn default() -> Self {
Self {
tcp_nodelay: Some(true),
zero_linger: true,
connect_timeout: Duration::from_secs(10),
read_write_timeout: Duration::from_secs(60),
}
}
}
#[derive(Clone)]
pub struct Config {
worker_threads: usize,
global_queue_interval: Option<u32>,
max_blocking_threads: usize,
thread_stack_size: usize,
catch_panics: bool,
storage_directory: PathBuf,
storage_blob_layouts: RangeInclusive<BlobLayout>,
network_cfg: NetworkConfig,
}
impl Config {
pub fn new() -> Self {
let rng = sys_rng().next_u64();
let storage_directory = env::temp_dir().join(format!("commonware_tokio_runtime_{rng}"));
Self {
worker_threads: 2,
global_queue_interval: None,
max_blocking_threads: 512,
thread_stack_size: utils::thread::system_thread_stack_size(),
catch_panics: false,
storage_directory,
storage_blob_layouts: BlobLayout::ALL,
network_cfg: NetworkConfig::default(),
}
}
pub const fn with_worker_threads(mut self, n: usize) -> Self {
self.worker_threads = n;
self
}
pub const fn with_global_queue_interval(mut self, n: u32) -> Self {
self.global_queue_interval = Some(n);
self
}
pub const fn with_max_blocking_threads(mut self, n: usize) -> Self {
self.max_blocking_threads = n;
self
}
pub const fn with_thread_stack_size(mut self, n: usize) -> Self {
self.thread_stack_size = n;
self
}
pub const fn with_catch_panics(mut self, b: bool) -> Self {
self.catch_panics = b;
self
}
pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
self.network_cfg.connect_timeout = timeout;
self
}
pub const fn with_read_write_timeout(mut self, d: Duration) -> Self {
self.network_cfg.read_write_timeout = d;
self
}
pub const fn with_tcp_nodelay(mut self, n: Option<bool>) -> Self {
self.network_cfg.tcp_nodelay = n;
self
}
pub const fn with_zero_linger(mut self, l: bool) -> Self {
self.network_cfg.zero_linger = l;
self
}
pub fn with_storage_directory(mut self, p: impl Into<PathBuf>) -> Self {
self.storage_directory = p.into();
self
}
pub fn with_storage_blob_layouts(mut self, layouts: RangeInclusive<BlobLayout>) -> Self {
assert!(
!layouts.is_empty(),
"storage blob layouts must be non-empty"
);
self.storage_blob_layouts = layouts;
self
}
pub const fn worker_threads(&self) -> usize {
self.worker_threads
}
pub const fn global_queue_interval(&self) -> Option<u32> {
self.global_queue_interval
}
pub const fn max_blocking_threads(&self) -> usize {
self.max_blocking_threads
}
pub const fn thread_stack_size(&self) -> usize {
self.thread_stack_size
}
pub const fn catch_panics(&self) -> bool {
self.catch_panics
}
pub const fn connect_timeout(&self) -> Duration {
self.network_cfg.connect_timeout
}
pub const fn read_write_timeout(&self) -> Duration {
self.network_cfg.read_write_timeout
}
pub const fn tcp_nodelay(&self) -> Option<bool> {
self.network_cfg.tcp_nodelay
}
pub const fn zero_linger(&self) -> bool {
self.network_cfg.zero_linger
}
pub const fn storage_directory(&self) -> &PathBuf {
&self.storage_directory
}
pub const fn storage_blob_layouts(&self) -> &RangeInclusive<BlobLayout> {
&self.storage_blob_layouts
}
}
impl Default for Config {
fn default() -> Self {
Self::new()
}
}
pub struct Executor {
registry: Registry,
metrics: Arc<Metrics>,
runtime: RuntimeHandle,
tasks: Arc<TaskTracker>,
shutdown: Mutex<Stopper>,
panicker: Panicker,
thread_stack_size: usize,
}
#[derive(Default)]
struct TaskTracker {
state: Mutex<TaskTrackerState>,
idle: Notify,
}
#[derive(Default)]
struct TaskTrackerState {
active: usize,
closed: bool,
}
impl TaskTracker {
fn admit(self: &Arc<Self>) -> Option<TaskGuard> {
let mut state = self.state.lock();
if state.closed {
return None;
}
state.active = state.active.checked_add(1).expect("active task overflow");
Some(TaskGuard(Arc::clone(self)))
}
fn close(&self) {
self.state.lock().closed = true;
}
async fn wait(&self) {
loop {
let idle = self.idle.notified();
if self.state.lock().active == 0 {
return;
}
idle.await;
}
}
}
struct TaskGuard(Arc<TaskTracker>);
impl Drop for TaskGuard {
fn drop(&mut self) {
let mut state = self.0.state.lock();
state.active = state.active.checked_sub(1).expect("active task underflow");
if state.active == 0 {
drop(state);
self.0.idle.notify_one();
}
}
}
pub struct Runner {
cfg: Config,
}
impl Default for Runner {
fn default() -> Self {
Self::new(Config::default())
}
}
impl Runner {
pub const fn new(cfg: Config) -> Self {
Self { cfg }
}
}
impl crate::Runner for Runner {
type Context = Context;
fn start<F, Fut>(self, f: F) -> Fut::Output
where
F: FnOnce(Self::Context) -> Fut,
Fut: Future,
{
let mut registry = Registry::new();
let mut runtime_registry = registry.sub_registry(METRICS_PREFIX);
let metrics = Arc::new(Metrics::init(&mut runtime_registry));
let mut builder = Builder::new_multi_thread();
builder
.worker_threads(self.cfg.worker_threads)
.max_blocking_threads(self.cfg.max_blocking_threads)
.thread_stack_size(self.cfg.thread_stack_size)
.enable_all();
if let Some(global_queue_interval) = self.cfg.global_queue_interval {
builder.global_queue_interval(global_queue_interval);
}
let runtime = builder.build().expect("failed to create Tokio runtime");
let (panicker, panicked) = Panicker::new(self.cfg.catch_panics);
let process = MeteredProcess::init(&mut runtime_registry);
runtime.spawn(process.collect(tokio::time::sleep));
let storage = MeteredStorage::new(
TokioStorage::new(TokioStorageConfig::new(
self.cfg.storage_directory.clone(),
self.cfg.storage_blob_layouts.clone(),
)),
&mut runtime_registry,
);
if let Err(e) = crate::storage::sync(&self.cfg.storage_directory) {
panic!(
"failed to sync storage at startup ({}): {e}",
self.cfg.storage_directory.display()
);
}
let config = TokioNetworkConfig::default()
.with_connect_timeout(self.cfg.network_cfg.connect_timeout)
.with_read_timeout(self.cfg.network_cfg.read_write_timeout)
.with_write_timeout(self.cfg.network_cfg.read_write_timeout)
.with_tcp_nodelay(self.cfg.network_cfg.tcp_nodelay)
.with_zero_linger(self.cfg.network_cfg.zero_linger);
let network = MeteredNetwork::new(TokioNetwork::new(config), &mut runtime_registry);
let executor = Arc::new(Executor {
registry,
metrics,
runtime: runtime.handle().clone(),
tasks: Arc::new(TaskTracker::default()),
shutdown: Mutex::new(Stopper::default()),
panicker,
thread_stack_size: self.cfg.thread_stack_size,
});
let label = Label::root();
executor.metrics.tasks_spawned.get_or_create(&label).inc();
let gauge = executor.metrics.tasks_running.get_or_create(&label).clone();
let tree = Tree::root();
let context = Context {
storage,
name: label.name().to_string(),
attributes: Vec::new(),
executor: executor.clone(),
network,
tree: Arc::clone(&tree),
execution: Execution::default(),
};
let output = catch_unwind(AssertUnwindSafe(|| {
let root = {
let _enter = runtime.enter();
Box::pin(f(context))
};
runtime.block_on(panicked.interrupt(async {
let output = root.await;
executor.tasks.close();
tree.abort();
executor.tasks.wait().await;
output
}))
}));
executor.tasks.close();
tree.abort();
runtime.block_on(executor.tasks.wait());
gauge.dec();
match output {
Ok(output) => output,
Err(panic) => resume_unwind(panic),
}
}
}
type Storage = MeteredStorage<TokioStorage>;
type Network = MeteredNetwork<TokioNetwork>;
pub struct Context {
name: String,
attributes: Vec<(String, String)>,
executor: Arc<Executor>,
storage: Storage,
network: Network,
tree: Arc<Tree>,
execution: Execution,
}
impl Context {
fn metrics(&self) -> &Metrics {
&self.executor.metrics
}
}
impl crate::Spawner for Context {
fn dedicated(mut self) -> Self {
self.execution = Execution::Dedicated;
self
}
fn shared(mut self, blocking: bool) -> Self {
self.execution = Execution::Shared(blocking);
self
}
fn spawn<F, Fut, T>(mut self, f: F) -> Handle<T>
where
F: FnOnce(Self) -> Fut + Send + 'static,
Fut: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let (_, metric) = spawn_metrics!(self);
let parent = Arc::clone(&self.tree);
let past = self.execution;
self.execution = Execution::default();
let (child, aborted) = Tree::child(&parent);
if aborted {
return Handle::closed(metric);
}
self.tree = child;
let executor = self.executor.clone();
let Some(task_guard) = executor.tasks.admit() else {
return Handle::closed(metric);
};
let guard = FactoryGuard::new(&parent, metric);
let future = f(self);
let (f, handle) = Handle::init(
future,
guard.disarm(),
executor.panicker.clone(),
Arc::clone(&parent),
);
let panicker = executor.panicker.clone();
let f = async move {
let _task_guard = task_guard;
let mut f = pin!(f);
poll_fn(|cx| {
catch_unwind(AssertUnwindSafe(|| f.as_mut().poll(cx))).unwrap_or_else(|panic| {
panicker.notify(panic);
Poll::Ready(())
})
})
.await;
};
if matches!(past, Execution::Dedicated) {
utils::thread::spawn(executor.thread_stack_size, {
let handle = executor.runtime.clone();
move || {
handle.block_on(f);
}
});
} else if matches!(past, Execution::Shared(true)) {
executor.runtime.spawn_blocking({
let handle = executor.runtime.clone();
move || {
handle.block_on(f);
}
});
} else {
executor.runtime.spawn(f);
}
handle
}
async fn stop(self, value: i32, timeout: Option<Duration>) -> Result<(), Error> {
let stop_resolved = {
let mut shutdown = self.executor.shutdown.lock();
shutdown.stop(value)
};
let timeout_future = timeout.map_or_else(
|| futures::future::Either::Right(futures::future::pending()),
|duration| futures::future::Either::Left(self.sleep(duration)),
);
select! {
result = stop_resolved => {
result.map_err(|_| Error::Closed)?;
Ok(())
},
_ = timeout_future => Err(Error::Timeout),
}
}
fn stopped(&self) -> Signal {
self.executor.shutdown.lock().stopped()
}
}
#[stability(BETA)]
impl crate::Strategizer for Context {
fn strategy(&self, parallelism: NonZeroUsize) -> Rayon {
let pool = ThreadPoolBuilder::new()
.num_threads(parallelism.get())
.stack_size(self.executor.thread_stack_size)
.build()
.expect("failed to create Tokio Rayon thread pool");
Rayon::with_pool(Arc::new(pool))
}
}
impl crate::Supervisor for Context {
fn child(&self, label: &'static str) -> Self {
let (tree, _) = Tree::child(&self.tree);
Self {
name: child_label(&self.name, label),
attributes: self.attributes.clone(),
executor: self.executor.clone(),
storage: self.storage.clone(),
network: self.network.clone(),
tree,
execution: Execution::default(),
}
}
fn with_attribute(mut self, key: &'static str, value: impl std::fmt::Display) -> Self {
validate_label(key);
add_attribute(&mut self.attributes, key, value);
self
}
fn name(&self) -> Name {
Name {
label: self.name.clone(),
attributes: self.attributes.clone(),
}
}
}
impl crate::Metrics for Context {
fn register<N: Into<String>, H: Into<String>, M: Metric>(
&self,
name: N,
help: H,
metric: M,
) -> Registered<M> {
let name = name.into();
let help = help.into();
let metric = Arc::new(metric);
self.executor.registry.register(
prefixed_name(&self.name, &name),
help,
self.attributes.clone(),
metric,
)
}
fn encode(&self) -> String {
self.executor.registry.encode()
}
}
impl Clock for Context {
fn current(&self) -> SystemTime {
SystemTime::now()
}
fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static + use<> {
tokio::time::sleep(duration)
}
fn sleep_until(
&self,
deadline: SystemTime,
) -> impl Future<Output = ()> + Send + 'static + use<> {
let duration_until_deadline = deadline.duration_since(self.current()).unwrap_or_default();
tokio::time::sleep(duration_until_deadline)
}
}
#[cfg(feature = "external")]
impl Pacer for Context {
fn pace<'a, F, T>(
&'a self,
_latency: Duration,
future: F,
) -> impl Future<Output = T> + Send + 'a
where
F: Future<Output = T> + Send + 'a,
T: Send + 'a,
{
future
}
}
impl GClock for Context {
type Instant = SystemTime;
fn now(&self) -> Self::Instant {
self.current()
}
}
impl ReasonablyRealtime for Context {}
impl crate::Network for Context {
type Listener = <Network as crate::Network>::Listener;
async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
self.network.bind(socket).await
}
async fn dial(&self, socket: SocketAddr) -> Result<(SinkOf<Self>, StreamOf<Self>), Error> {
self.network.dial(socket).await
}
}
impl crate::Resolver for Context {
async fn resolve(&self, host: &str) -> Result<Vec<IpAddr>, Error> {
let addrs = tokio::net::lookup_host(format!("{host}:0"))
.await
.map_err(|e| Error::ResolveFailed(e.to_string()))?;
Ok(addrs.map(|addr| addr.ip()).collect())
}
}
impl TryRng for Context {
type Error = Infallible;
fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
Ok(sys_rng().next_u32())
}
fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
Ok(sys_rng().next_u64())
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
sys_rng().fill_bytes(dest);
Ok(())
}
}
impl TryCryptoRng for Context {}
impl crate::Storage for Context {
type Blob = <Storage as crate::Storage>::Blob;
async fn open_versioned(
&self,
partition: &str,
name: &[u8],
versions: std::ops::RangeInclusive<BlobVersion>,
) -> Result<(Self::Blob, u64, BlobVersion), Error> {
self.storage.open_versioned(partition, name, versions).await
}
async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
self.storage.remove(partition, name).await
}
async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
self.storage.scan(partition).await
}
}
#[cfg(test)]
#[allow(deprecated)]
mod tests {
use super::*;
use crate::{
AbortOnDrop, Blob as _, Metrics, Network, Resolver, Runner as _, Sink, Spawner as _,
Storage as _, Strategizer as _, Stream, Supervisor as _, telemetry::metrics::raw::Counter,
tokio::telemetry,
};
use bytes::Bytes;
use commonware_parallel::Strategy as _;
use commonware_utils::{NZUsize, channel::oneshot};
use rstest::rstest;
use std::{
self,
collections::HashMap,
net::{IpAddr, Ipv4Addr, Ipv6Addr},
str::FromStr,
sync::atomic::{AtomicBool, Ordering},
};
use tracing::{Level, error};
struct TaskDropGate {
entered: std::sync::mpsc::Sender<()>,
release: std::sync::mpsc::Receiver<()>,
}
impl Drop for TaskDropGate {
fn drop(&mut self) {
let _ = self.entered.send(());
let _ = self.release.recv();
}
}
struct NotifyDrop(std::sync::mpsc::Sender<()>);
impl Drop for NotifyDrop {
fn drop(&mut self) {
let _ = self.0.send(());
}
}
struct PanicOnDrop;
impl Drop for PanicOnDrop {
fn drop(&mut self) {
panic!("disposal panic");
}
}
#[derive(Clone, Copy, Debug)]
enum Disposal {
Output,
Cancel,
Ready,
}
fn placed(context: Context, execution: Execution) -> Context {
match execution {
Execution::Dedicated => context.dedicated(),
Execution::Shared(blocking) => context.shared(blocking),
}
}
#[derive(Clone, Copy, Debug)]
enum RootExit {
Return,
FuturePanic,
ConstructorPanic,
}
#[derive(Clone, Copy, Debug)]
enum JoinRelease {
Task,
Completion,
Output,
}
fn spawn_drop_gated_task(
context: Context,
execution: Execution,
drop_gate: TaskDropGate,
ready: Option<oneshot::Sender<()>>,
) {
placed(context, execution).spawn(move |context| async move {
let _context = context;
let _drop_gate = drop_gate;
if let Some(ready) = ready {
ready.send(()).unwrap();
}
futures::future::pending::<()>().await;
});
}
fn assert_runner_drains_spawned_task(execution: Execution, root_exit: RootExit) {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (ready_tx, ready_rx) = oneshot::channel();
let (drop_entered_tx, drop_entered_rx) = std::sync::mpsc::channel();
let (drop_release_tx, drop_release_rx) = std::sync::mpsc::channel();
let (runner_done_tx, runner_done_rx) = std::sync::mpsc::channel();
let runner = std::thread::spawn(move || {
let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
let drop_gate = TaskDropGate {
entered: drop_entered_tx,
release: drop_release_rx,
};
match root_exit {
RootExit::ConstructorPanic => {
Runner::new(cfg).start(move |context| -> futures::future::Pending<()> {
spawn_drop_gated_task(context, execution, drop_gate, None);
panic!("root constructor panic after spawning child");
})
}
RootExit::Return | RootExit::FuturePanic => {
Runner::new(cfg).start(move |context| async move {
spawn_drop_gated_task(context, execution, drop_gate, Some(ready_tx));
ready_rx.await.unwrap();
assert!(
matches!(root_exit, RootExit::Return),
"root future panic after spawning child"
);
})
}
}
}));
runner_done_tx.send(result.is_err()).unwrap();
});
drop_entered_rx
.recv_timeout(Duration::from_secs(5))
.expect("spawned task was not canceled after the root returned");
let early = runner_done_rx.recv_timeout(Duration::from_millis(250));
let returned_before_cleanup = match early {
Ok(panicked) => Some(panicked),
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => None,
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
panic!("Runner::start exited without reporting its result")
}
};
drop_release_tx.send(()).unwrap();
let panicked = returned_before_cleanup.unwrap_or_else(|| {
runner_done_rx
.recv_timeout(Duration::from_secs(5))
.expect("Runner::start did not return after task cleanup completed")
});
runner.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
assert!(
returned_before_cleanup.is_none(),
"Runner::start returned before {execution:?} task cleanup completed"
);
assert_eq!(panicked, !matches!(root_exit, RootExit::Return));
}
fn run_with_returned_strategy(retain: bool) -> Option<Rayon> {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (strategy_tx, strategy_rx) = std::sync::mpsc::channel();
let runner = std::thread::spawn(move || {
let strategy = Runner::new(cfg).start(move |context| async move {
let strategy = context.strategy(NZUsize!(2));
strategy.spawn(1, |_| ()).await;
retain.then_some(strategy)
});
strategy_tx.send(strategy).unwrap();
});
let strategy = strategy_rx
.recv_timeout(Duration::from_secs(5))
.expect("Runner::start did not return after the strategy completed work");
runner.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
strategy
}
#[test]
fn test_strategy_spawn_keeps_sibling_futures_polling() {
let cfg = Config::new().with_worker_threads(1);
let storage_directory = cfg.storage_directory().clone();
let ready_branch_ran = Runner::new(cfg).start(|context| async move {
let strategy = context.strategy(NZUsize!(2));
let (gate_tx, gate_rx) = std::sync::mpsc::channel::<()>();
let mut job = Box::pin(strategy.spawn(1, move |_| {
gate_rx.recv_timeout(Duration::from_secs(5)).is_ok()
}));
select! {
_ = job.as_mut() => false,
_ = futures::future::ready(()) => {
gate_tx.send(()).unwrap();
assert!(job.await, "the gate must open before the job gives up");
true
},
}
});
let _ = std::fs::remove_dir_all(storage_directory);
assert!(
ready_branch_ran,
"the sibling branch must run while the job is in flight"
);
}
#[test]
fn test_storage_blob_layout_restriction() {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
assert_eq!(cfg.storage_blob_layouts(), &BlobLayout::ALL);
let partition = "layout_restriction";
let v0_name = b"v0";
let partition_directory = storage_directory.join(partition);
std::fs::create_dir_all(&partition_directory).unwrap();
let v0_path = partition_directory.join(commonware_formatting::hex(v0_name));
let v0_bytes = crate::storage::tests::v0_blob_bytes(0, b"payload");
std::fs::write(&v0_path, &v0_bytes).unwrap();
Runner::new(cfg).start(|context| async move {
let (blob, size) = context.open(partition, v0_name).await.unwrap();
assert_eq!(size, 7);
let payload = blob
.read_at(0, 7, crate::ReadOptions::default())
.await
.unwrap();
assert_eq!(payload.coalesce(), b"payload".as_slice());
});
let cfg = Config::new()
.with_storage_directory(storage_directory.clone())
.with_storage_blob_layouts(BlobLayout::V1..=BlobLayout::V1);
Runner::new(cfg).start(|context| async move {
let result = context.open(partition, v0_name).await;
assert!(matches!(
result,
Err(Error::BlobLayoutMismatch { expected, found })
if expected == (BlobLayout::V1..=BlobLayout::V1)
&& found == BlobLayout::V0
));
context.open(partition, b"v1").await.unwrap();
});
assert_eq!(std::fs::read(&v0_path).unwrap(), v0_bytes);
let v1_path = partition_directory.join(commonware_formatting::hex(b"v1"));
let v1 = std::fs::read(&v1_path).unwrap();
assert_eq!(&v1[..4], &BlobLayout::V1.magic());
let cfg = Config::new()
.with_storage_directory(storage_directory.clone())
.with_storage_blob_layouts(BlobLayout::V0..=BlobLayout::V0);
Runner::new(cfg).start(|context| async move {
let result = context.open(partition, b"v1").await;
assert!(matches!(
result,
Err(Error::BlobLayoutMismatch { expected, found })
if expected == (BlobLayout::V0..=BlobLayout::V0)
&& found == BlobLayout::V1
));
let (blob, size) = context.open(partition, b"rollback").await.unwrap();
assert_eq!(size, 0);
blob.write_at(0, b"rollback!".as_slice(), crate::WriteOptions::SYNC)
.await
.unwrap();
drop(blob);
let (blob, size) = context.open(partition, b"rollback").await.unwrap();
assert_eq!(size, 9);
let payload = blob
.read_at(0, 9, crate::ReadOptions::default())
.await
.unwrap();
assert_eq!(payload.coalesce(), b"rollback!".as_slice());
});
assert_eq!(std::fs::read(&v1_path).unwrap(), v1);
let rollback_path = partition_directory.join(commonware_formatting::hex(b"rollback"));
let rollback = std::fs::read(rollback_path).unwrap();
assert_eq!(&rollback[..4], &BlobLayout::V0.magic());
assert_eq!(&rollback[8..], b"rollback!");
let _ = std::fs::remove_dir_all(storage_directory);
}
#[test]
#[should_panic(expected = "non-empty")]
fn test_storage_blob_layout_restriction_rejects_empty_range() {
let _ = Config::new().with_storage_blob_layouts(BlobLayout::V1..=BlobLayout::V0);
}
#[test]
fn test_default_thread_stack_size_uses_system_default() {
let cfg = Config::new();
assert_eq!(
cfg.thread_stack_size(),
utils::thread::system_thread_stack_size()
);
}
#[test]
fn test_root_factory_builds_runtime_futures() {
let cfg = Config::new().with_worker_threads(2);
Runner::new(cfg).start(|context| {
let sleep = context.sleep(Duration::from_millis(1));
let deadline = context.sleep_until(context.current() + Duration::from_millis(1));
async move {
sleep.await;
deadline.await;
}
});
}
#[test]
fn test_runner_waits_for_spawned_task_cancellation() {
for execution in [
Execution::Shared(false),
Execution::Shared(true),
Execution::Dedicated,
] {
for root_exit in [
RootExit::Return,
RootExit::FuturePanic,
RootExit::ConstructorPanic,
] {
assert_runner_drains_spawned_task(execution, root_exit);
}
}
}
#[test]
fn test_runner_propagates_panic_during_drain() {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let result = catch_unwind(AssertUnwindSafe(|| {
Runner::new(cfg).start(|context| async move {
let (sentinel_started_tx, sentinel_started_rx) = oneshot::channel();
let (sentinel_dropped_tx, sentinel_dropped_rx) = std::sync::mpsc::channel();
context.child("sentinel").spawn(move |_| async move {
let _notify = NotifyDrop(sentinel_dropped_tx);
sentinel_started_tx.send(()).unwrap();
futures::future::pending::<()>().await;
});
sentinel_started_rx.await.unwrap();
let (started_tx, started_rx) = oneshot::channel();
context
.child("late")
.dedicated()
.spawn(move |_| async move {
started_tx.send(()).unwrap();
sentinel_dropped_rx.recv().unwrap();
panic!("drain panic");
});
started_rx.await.unwrap();
})
}));
let _ = std::fs::remove_dir_all(storage_directory);
let panic = result.expect_err("Runner::start returned after a task panicked during drain");
assert_eq!(utils::extract_panic_message(&*panic), "drain panic");
}
#[rstest]
fn test_disposal_panic_is_task_panic(
#[values(Disposal::Output, Disposal::Cancel, Disposal::Ready)] disposal: Disposal,
#[values(
Execution::Shared(false),
Execution::Shared(true),
Execution::Dedicated
)]
execution: Execution,
#[values(false, true)] catch: bool,
) {
let cfg = Config::new().with_catch_panics(catch);
let storage_directory = cfg.storage_directory().clone();
let result = catch_unwind(AssertUnwindSafe(|| {
Runner::new(cfg).start(|context| async move {
let child = placed(context.child("disposal"), execution);
match disposal {
Disposal::Output => {
let (entered_tx, entered_rx) = std::sync::mpsc::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let handle = child.spawn(move |_| async move {
entered_tx.send(()).unwrap();
release_rx.recv().unwrap();
PanicOnDrop
});
entered_rx.recv().unwrap();
drop(handle);
release_tx.send(()).unwrap();
}
Disposal::Cancel => {
let (ready_tx, ready_rx) = oneshot::channel();
let handle = child.spawn(move |_| async move {
let _guard = PanicOnDrop;
ready_tx.send(()).unwrap();
futures::future::pending::<()>().await;
});
ready_rx.await.unwrap();
handle.abort();
assert!(matches!(handle.await, Err(Error::Closed)));
}
Disposal::Ready => {
let guard = PanicOnDrop;
let result = child
.spawn(move |_| {
poll_fn(move |_| {
let _guard = &guard;
Poll::Ready(())
})
})
.await;
assert!(matches!(result, Err(Error::Closed)));
}
}
context.child("survivor").spawn(|_| async {}).await.unwrap();
})
}));
let _ = std::fs::remove_dir_all(storage_directory);
if catch {
result.expect("caught disposal panic failed the runner");
} else {
let panic = result.expect_err("disposal panic did not fail the runner");
assert_eq!(utils::extract_panic_message(&*panic), "disposal panic");
}
}
struct ActiveOnError {
tasks: Arc<TaskTracker>,
active: Arc<Mutex<Vec<usize>>>,
}
impl tracing::Subscriber for ActiveOnError {
fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
*metadata.level() == Level::ERROR
}
fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
tracing::span::Id::from_u64(1)
}
fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
fn event(&self, _: &tracing::Event<'_>) {
self.active.lock().push(self.tasks.state.lock().active);
}
fn enter(&self, _: &tracing::span::Id) {}
fn exit(&self, _: &tracing::span::Id) {}
}
#[test]
fn test_disposal_panic_reported_before_tracker_release() {
let cfg = Config::new().with_catch_panics(true);
let storage_directory = cfg.storage_directory().clone();
let active = Arc::new(Mutex::new(Vec::new()));
let observed = active.clone();
Runner::new(cfg).start(|context| async move {
let subscriber = ActiveOnError {
tasks: context.executor.tasks.clone(),
active: observed,
};
let (ready_tx, ready_rx) = oneshot::channel();
let handle = context
.child("disposal")
.dedicated()
.spawn(move |_| async move {
std::mem::forget(tracing::subscriber::set_default(subscriber));
let _guard = PanicOnDrop;
ready_tx.send(()).unwrap();
futures::future::pending::<()>().await;
});
ready_rx.await.unwrap();
handle.abort();
assert!(matches!(handle.await, Err(Error::Closed)));
});
let _ = std::fs::remove_dir_all(storage_directory);
assert_eq!(*active.lock(), vec![1]);
}
#[test]
fn test_root_destruction_can_spawn_work_before_shutdown() {
struct Root {
context: Option<Context>,
invoked: Arc<AtomicBool>,
dropped: std::sync::mpsc::Sender<()>,
}
impl Future for Root {
type Output = ();
fn poll(self: std::pin::Pin<&mut Self>, _: &mut std::task::Context<'_>) -> Poll<()> {
Poll::Ready(())
}
}
impl Drop for Root {
fn drop(&mut self) {
let invoked = self.invoked.clone();
let payload = NotifyDrop(self.dropped.clone());
self.context.take().unwrap().spawn(move |_| {
invoked.store(true, Ordering::SeqCst);
async move {
let _payload = payload;
futures::future::pending::<()>().await;
}
});
}
}
for execution in [
Execution::Shared(false),
Execution::Shared(true),
Execution::Dedicated,
] {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let invoked = Arc::new(AtomicBool::new(false));
let (dropped_tx, dropped_rx) = std::sync::mpsc::channel();
Runner::new(cfg).start(|context| Root {
context: Some(placed(context.child("root"), execution)),
invoked: invoked.clone(),
dropped: dropped_tx,
});
let _ = std::fs::remove_dir_all(storage_directory);
assert!(
invoked.load(Ordering::SeqCst),
"{execution:?} spawn from root destruction was rejected"
);
assert!(
dropped_rx.try_recv().is_ok(),
"{execution:?} child spawned by root destruction outlived the runner"
);
}
}
fn assert_join_all_releases_pending_work(release_owner: JoinRelease) {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (started, started_rx) = oneshot::channel();
let (release, release_rx) = std::sync::mpsc::channel();
let (joining, joining_rx) = std::sync::mpsc::channel();
let (done, done_rx) = std::sync::mpsc::channel();
let (exited, exited_rx) = std::sync::mpsc::channel();
let (drop_release, drop_release_rx) = std::sync::mpsc::channel();
drop(drop_release);
let release_on_drop = TaskDropGate {
entered: release.clone(),
release: drop_release_rx,
};
let runner = std::thread::spawn(move || {
Runner::new(cfg).start(move |context| async move {
let blocked = context
.child("blocked")
.dedicated()
.spawn(move |_| async move {
started.send(()).unwrap();
release_rx.recv().unwrap();
exited.send(()).unwrap();
Ok::<_, Error>(None::<TaskDropGate>)
})
.abort_on_drop();
let failed = Handle::ready(Ok(Err(Error::Closed))).abort_on_drop();
let mut guards = vec![failed, blocked];
if matches!(release_owner, JoinRelease::Output) {
guards.insert(
0,
Handle::ready(Ok(Ok(Some(release_on_drop)))).abort_on_drop(),
);
} else {
let pending = async move {
let _release_on_drop = release_on_drop;
futures::future::pending::<Result<Option<TaskDropGate>, Error>>().await
};
let pending = if matches!(release_owner, JoinRelease::Task) {
context.child("pending").spawn(move |_| pending)
} else {
Handle::from_future(async move { Ok(pending.await) })
};
guards.push(pending.abort_on_drop());
}
started_rx.await.unwrap();
joining.send(()).unwrap();
let result = AbortOnDrop::join_all::<Error>(guards).await;
assert!(done.send((result, exited_rx.try_recv().is_ok())).is_ok());
});
});
joining_rx
.recv_timeout(Duration::from_secs(5))
.expect("blocking task did not start");
let early = done_rx.recv_timeout(Duration::from_secs(5));
let completed = early.is_ok();
let _ = release.send(());
let (result, exited) = early.unwrap_or_else(|_| {
done_rx
.recv_timeout(Duration::from_secs(5))
.expect("joining did not finish after releasing the blocking task")
});
runner.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
assert!(
completed,
"joining stalled with release owned by {release_owner:?}"
);
assert!(exited, "joining returned before the blocking task exited");
assert!(matches!(result, Err(Error::Closed)));
}
#[test]
fn test_join_all_aborts_all_tasks_before_waiting() {
assert_join_all_releases_pending_work(JoinRelease::Task);
}
#[test]
fn test_join_all_drains_completions_concurrently() {
assert_join_all_releases_pending_work(JoinRelease::Completion);
}
#[test]
fn test_join_all_drops_outputs_before_waiting() {
assert_join_all_releases_pending_work(JoinRelease::Output);
}
#[test]
fn test_runner_start_waits_for_previous_run() {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (started, first_started) = std::sync::mpsc::channel();
let (release, released) = futures::channel::oneshot::channel();
let first_cfg = cfg.clone();
let first = std::thread::spawn(move || {
Runner::new(first_cfg).start(|context| async move {
started.send(()).unwrap();
released.await.unwrap();
drop(context);
});
});
first_started.recv_timeout(Duration::from_secs(10)).unwrap();
let (started, second_started) = std::sync::mpsc::channel();
let second = std::thread::spawn(move || {
Runner::new(cfg).start(|_| async move {
started.send(()).unwrap();
});
});
match second_started.recv_timeout(Duration::from_millis(200)) {
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
other => panic!("second run started while the first held the directory: {other:?}"),
}
release.send(()).unwrap();
first.join().unwrap();
second_started
.recv_timeout(Duration::from_secs(10))
.expect("second run did not start after the first returned");
second.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
}
#[test]
fn test_runner_owns_runtime_when_context_escapes() {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (ready_tx, ready_rx) = oneshot::channel();
let (drop_entered_tx, drop_entered_rx) = std::sync::mpsc::channel();
let (drop_release_tx, drop_release_rx) = std::sync::mpsc::channel();
let (runner_returned_tx, runner_returned_rx) = std::sync::mpsc::channel();
let (context_release_tx, context_release_rx) = std::sync::mpsc::channel();
let runner = std::thread::spawn(move || {
let context = Runner::new(cfg).start(move |context| async move {
context.executor.runtime.spawn(async move {
let _drop_gate = TaskDropGate {
entered: drop_entered_tx,
release: drop_release_rx,
};
ready_tx.send(()).unwrap();
futures::future::pending::<()>().await;
});
ready_rx.await.unwrap();
context
});
runner_returned_tx.send(()).unwrap();
context_release_rx.recv().unwrap();
drop(context);
});
let returned_early = runner_returned_rx
.recv_timeout(Duration::from_millis(500))
.is_ok();
if returned_early {
drop_release_tx.send(()).unwrap();
context_release_tx.send(()).unwrap();
drop_entered_rx
.recv_timeout(Duration::from_secs(5))
.expect("escaped Context did not retain the raw runtime task");
} else {
drop_entered_rx
.recv_timeout(Duration::from_secs(5))
.expect("Runner did not cancel its raw runtime task");
drop_release_tx.send(()).unwrap();
runner_returned_rx
.recv_timeout(Duration::from_secs(5))
.expect("Runner did not return after raw task cleanup");
context_release_tx.send(()).unwrap();
}
runner.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
assert!(
!returned_early,
"a returned Context kept the Tokio runtime alive after Runner::start"
);
}
#[test]
fn test_runner_returns_strategy_after_pool_work() {
assert!(run_with_returned_strategy(false).is_none());
let strategy = run_with_returned_strategy(true).unwrap();
assert_eq!(futures::executor::block_on(strategy.spawn(1, |_| 42)), 42);
}
#[test]
fn test_runner_resumes_strategy_panic_payload_after_pool_work() {
let cfg = Config::new();
let storage_directory = cfg.storage_directory().clone();
let (strategy_tx, strategy_rx) = std::sync::mpsc::channel();
let runner = std::thread::spawn(move || {
let result: std::thread::Result<()> =
std::panic::catch_unwind(AssertUnwindSafe(|| {
Runner::new(cfg).start(move |context| async move {
let strategy = context.strategy(NZUsize!(2));
strategy.spawn(1, |_| ()).await;
std::panic::panic_any(strategy);
});
}));
let strategy = result
.expect_err("Runner::start did not resume the root panic")
.downcast::<Rayon>()
.expect("Runner::start changed the root panic payload");
strategy_tx.send(*strategy).unwrap();
});
let strategy = strategy_rx
.recv_timeout(Duration::from_secs(5))
.expect("Runner::start did not resume the strategy panic payload");
runner.join().unwrap();
let _ = std::fs::remove_dir_all(storage_directory);
assert_eq!(futures::executor::block_on(strategy.spawn(1, |_| 42)), 42);
}
#[test]
fn test_thread_stack_size_override() {
let cfg = Config::new().with_thread_stack_size(4 * 1024 * 1024);
assert_eq!(cfg.thread_stack_size(), 4 * 1024 * 1024);
}
#[test]
fn test_process_rss_metric() {
let executor = Runner::default();
executor.start(|context| async move {
loop {
let metrics = context.encode();
if !metrics.contains("runtime_process_rss") {
context.sleep(Duration::from_millis(100)).await;
continue;
}
for line in metrics.lines() {
if line.starts_with("runtime_process_rss")
&& !line.starts_with("runtime_process_rss{")
{
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 2 {
let rss_value: i64 =
parts[1].parse().expect("Failed to parse RSS value");
if rss_value > 0 {
return;
}
}
}
}
}
});
}
#[test]
fn test_telemetry() {
let executor = Runner::default();
executor.start(|context| async move {
let address = SocketAddr::from_str("127.0.0.1:8000").unwrap();
telemetry::init(
context.child("metrics"),
telemetry::Logs {
level: Level::INFO,
json: false,
},
Some(address),
None,
);
let counter: Counter<u64> = Counter::default();
let _registered = context.register("test_counter", "Test counter", counter.clone());
counter.inc();
async fn read_line<St: Stream>(stream: &mut St) -> Result<String, Error> {
let mut line = Vec::new();
loop {
let received = stream.recv(1).await?;
let byte = received.coalesce().as_ref()[0];
if byte == b'\n' {
if line.last() == Some(&b'\r') {
line.pop(); }
break;
}
line.push(byte);
}
String::from_utf8(line).map_err(|_| Error::ReadFailed)
}
async fn read_headers<St: Stream>(
stream: &mut St,
) -> Result<HashMap<String, String>, Error> {
let mut headers = HashMap::new();
loop {
let line = read_line(stream).await?;
if line.is_empty() {
break;
}
let parts: Vec<&str> = line.splitn(2, ": ").collect();
if parts.len() == 2 {
headers.insert(parts[0].to_string(), parts[1].to_string());
}
}
Ok(headers)
}
async fn read_body<St: Stream>(
stream: &mut St,
content_length: usize,
) -> Result<String, Error> {
let received = stream.recv(content_length).await?;
String::from_utf8(received.coalesce().into()).map_err(|_| Error::ReadFailed)
}
let client_handle = context.child("client").spawn(move |context| async move {
let (mut sink, mut stream) = loop {
match context.dial(address).await {
Ok((sink, stream)) => break (sink, stream),
Err(e) => {
error!(err =?e, "failed to connect");
context.sleep(Duration::from_millis(10)).await;
}
}
};
let request = format!(
"GET /metrics HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n"
);
sink.send(Bytes::from(request)).await.unwrap();
let status_line = read_line(&mut stream).await.unwrap();
assert_eq!(status_line, "HTTP/1.1 200 OK");
let headers = read_headers(&mut stream).await.unwrap();
println!("Headers: {headers:?}");
let content_length = headers
.get("content-length")
.unwrap()
.parse::<usize>()
.unwrap();
let body = read_body(&mut stream, content_length).await.unwrap();
assert!(body.contains("test_counter_total 1"));
});
client_handle.await.unwrap();
});
}
#[test]
fn test_resolver() {
let executor = Runner::default();
executor.start(|context| async move {
let addrs = context.resolve("localhost").await.unwrap();
assert!(!addrs.is_empty());
for addr in addrs {
assert!(
addr == IpAddr::V4(Ipv4Addr::LOCALHOST)
|| addr == IpAddr::V6(Ipv6Addr::LOCALHOST)
);
}
});
}
}