Skip to main content

commonware_runtime/tokio/
runtime.rs

1#[cfg(feature = "external")]
2use crate::Pacer;
3use crate::{
4    BlobLayout, BlobVersion, Clock, Error, Execution, Handle, METRICS_PREFIX, Name, SinkOf,
5    StreamOf, child_label,
6    network::{
7        metered::Network as MeteredNetwork,
8        tokio::{Config as TokioNetworkConfig, Network as TokioNetwork},
9    },
10    prefixed_name,
11    process::metered::Metrics as MeteredProcess,
12    signal::Signal,
13    storage::{
14        metered::Storage as MeteredStorage,
15        tokio::{Config as TokioStorageConfig, Storage as TokioStorage},
16    },
17    telemetry::metrics::{
18        CounterFamily, GaugeFamily, Metric, Register, Registered, Registry, add_attribute, raw,
19        task::Label, validate_label,
20    },
21    utils::{self, FactoryGuard, Panicker, signal::Stopper, supervision::Tree},
22};
23use commonware_macros::{select, stability};
24#[stability(BETA)]
25use commonware_parallel::{Rayon, rayon::ThreadPoolBuilder};
26use commonware_utils::{sync::Mutex, sys_rng};
27use governor::clock::{Clock as GClock, ReasonablyRealtime};
28use rand_core::{Rng, TryCryptoRng, TryRng};
29use std::{
30    convert::Infallible,
31    env,
32    future::{Future, poll_fn},
33    net::{IpAddr, SocketAddr},
34    num::NonZeroUsize,
35    ops::RangeInclusive,
36    panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
37    path::PathBuf,
38    pin::pin,
39    sync::Arc,
40    task::Poll,
41    time::{Duration, SystemTime},
42};
43use tokio::{
44    runtime::{Builder, Handle as RuntimeHandle},
45    sync::Notify,
46};
47
48#[derive(Debug)]
49struct Metrics {
50    tasks_spawned: CounterFamily<Label>,
51    tasks_running: GaugeFamily<Label>,
52}
53
54impl Metrics {
55    pub fn init(registry: &mut impl Register) -> Self {
56        Self {
57            tasks_spawned: registry.register(
58                "tasks_spawned",
59                "Total number of tasks spawned",
60                raw::Family::default(),
61            ),
62            tasks_running: registry.register(
63                "tasks_running",
64                "Number of tasks currently running",
65                raw::Family::default(),
66            ),
67        }
68    }
69}
70
71#[derive(Clone, Debug)]
72pub struct NetworkConfig {
73    /// If Some, explicitly sets TCP_NODELAY on the socket.
74    /// Otherwise uses system default.
75    ///
76    /// Defaults to `Some(true)`.
77    tcp_nodelay: Option<bool>,
78
79    /// Whether to set `SO_LINGER` to zero on the socket.
80    ///
81    /// When enabled, causes an immediate RST on close, avoiding
82    /// `TIME_WAIT` state. This is useful in adversarial environments to
83    /// reclaim socket resources immediately when closing connections to
84    /// misbehaving peers.
85    ///
86    /// Defaults to `true`.
87    zero_linger: bool,
88
89    /// Timeout for establishing an outbound TCP connection.
90    ///
91    /// Defaults to 10 seconds.
92    connect_timeout: Duration,
93
94    /// Read/write timeout for network operations.
95    ///
96    /// Bounds the full `Sink::send` and `Stream::recv` calls rather than each
97    /// individual socket syscall. Larger batched writes may therefore require a
98    /// larger timeout.
99    ///
100    /// Defaults to 60 seconds.
101    read_write_timeout: Duration,
102}
103
104impl Default for NetworkConfig {
105    fn default() -> Self {
106        Self {
107            tcp_nodelay: Some(true),
108            zero_linger: true,
109            connect_timeout: Duration::from_secs(10),
110            read_write_timeout: Duration::from_secs(60),
111        }
112    }
113}
114
115/// Configuration for the `tokio` runtime.
116#[derive(Clone)]
117pub struct Config {
118    /// Number of threads to use for handling async tasks.
119    ///
120    /// Worker threads are always active (waiting for work).
121    ///
122    /// Tokio sets the default value to the number of logical CPUs.
123    worker_threads: usize,
124
125    /// Number of scheduler ticks between global queue polls.
126    ///
127    /// When unset, Tokio uses its default behavior for the multi-thread
128    /// scheduler. Smaller values reduce the delay before tasks woken from
129    /// outside a worker, such as io_uring completion notifications, are polled
130    /// from the global queue again.
131    global_queue_interval: Option<u32>,
132
133    /// Maximum number of threads to use for blocking tasks.
134    ///
135    /// Unlike worker threads, blocking threads are created as needed and
136    /// exit if left idle for too long.
137    ///
138    /// Tokio sets the default value to 512 to avoid hanging on lower-level
139    /// operations that require blocking (like `fs` and writing to `Stdout`).
140    ///
141    /// Tasks spawned with [`crate::Spawner::shared`] in blocking mode share this pool with the
142    /// blocking work behind storage operations and name resolution. A blocking task that waits
143    /// on them holds a thread they need, so enough such tasks exhaust the pool and deadlock.
144    max_blocking_threads: usize,
145
146    /// Stack size to use for runtime-owned threads.
147    ///
148    /// Defaults to the system stack size when the current platform exposes it,
149    /// and otherwise falls back to Rust's default spawned-thread stack size.
150    ///
151    /// See [utils::thread::system_thread_stack_size].
152    thread_stack_size: usize,
153
154    /// Whether or not to catch panics.
155    catch_panics: bool,
156
157    /// Base directory for all storage operations, created at start if missing
158    /// and held for the run.
159    storage_directory: PathBuf,
160
161    /// Blob layouts accepted by storage.
162    ///
163    /// Defaults to [BlobLayout::ALL] and must be non-empty. New blobs use the latest layout in
164    /// the range. Existing blobs outside it fail to open with [Error::BlobLayoutMismatch], so
165    /// restrict the range to what a rollback target can read before the upgraded binary first
166    /// opens storage.
167    storage_blob_layouts: RangeInclusive<BlobLayout>,
168
169    /// Network configuration.
170    network_cfg: NetworkConfig,
171}
172
173impl Config {
174    /// Returns a new [Config] with default values.
175    pub fn new() -> Self {
176        let rng = sys_rng().next_u64();
177        let storage_directory = env::temp_dir().join(format!("commonware_tokio_runtime_{rng}"));
178        Self {
179            worker_threads: 2,
180            global_queue_interval: None,
181            max_blocking_threads: 512,
182            thread_stack_size: utils::thread::system_thread_stack_size(),
183            catch_panics: false,
184            storage_directory,
185            storage_blob_layouts: BlobLayout::ALL,
186            network_cfg: NetworkConfig::default(),
187        }
188    }
189
190    // Setters
191    /// See [Config]
192    pub const fn with_worker_threads(mut self, n: usize) -> Self {
193        self.worker_threads = n;
194        self
195    }
196    /// See [Config]
197    pub const fn with_global_queue_interval(mut self, n: u32) -> Self {
198        self.global_queue_interval = Some(n);
199        self
200    }
201    /// See [Config]
202    pub const fn with_max_blocking_threads(mut self, n: usize) -> Self {
203        self.max_blocking_threads = n;
204        self
205    }
206    /// See [Config]
207    pub const fn with_thread_stack_size(mut self, n: usize) -> Self {
208        self.thread_stack_size = n;
209        self
210    }
211    /// See [Config]
212    pub const fn with_catch_panics(mut self, b: bool) -> Self {
213        self.catch_panics = b;
214        self
215    }
216    /// See [Config]
217    pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
218        self.network_cfg.connect_timeout = timeout;
219        self
220    }
221    /// See [Config]
222    pub const fn with_read_write_timeout(mut self, d: Duration) -> Self {
223        self.network_cfg.read_write_timeout = d;
224        self
225    }
226    /// See [Config]
227    pub const fn with_tcp_nodelay(mut self, n: Option<bool>) -> Self {
228        self.network_cfg.tcp_nodelay = n;
229        self
230    }
231    /// See [Config]
232    pub const fn with_zero_linger(mut self, l: bool) -> Self {
233        self.network_cfg.zero_linger = l;
234        self
235    }
236    /// See [Config]
237    pub fn with_storage_directory(mut self, p: impl Into<PathBuf>) -> Self {
238        self.storage_directory = p.into();
239        self
240    }
241    /// See [Config]
242    ///
243    /// # Panics
244    ///
245    /// Panics if `layouts` is empty.
246    pub fn with_storage_blob_layouts(mut self, layouts: RangeInclusive<BlobLayout>) -> Self {
247        assert!(
248            !layouts.is_empty(),
249            "storage blob layouts must be non-empty"
250        );
251        self.storage_blob_layouts = layouts;
252        self
253    }
254
255    // Getters
256    /// See [Config]
257    pub const fn worker_threads(&self) -> usize {
258        self.worker_threads
259    }
260    /// See [Config]
261    pub const fn global_queue_interval(&self) -> Option<u32> {
262        self.global_queue_interval
263    }
264    /// See [Config]
265    pub const fn max_blocking_threads(&self) -> usize {
266        self.max_blocking_threads
267    }
268    /// See [Config]
269    pub const fn thread_stack_size(&self) -> usize {
270        self.thread_stack_size
271    }
272    /// See [Config]
273    pub const fn catch_panics(&self) -> bool {
274        self.catch_panics
275    }
276    /// See [Config]
277    pub const fn connect_timeout(&self) -> Duration {
278        self.network_cfg.connect_timeout
279    }
280    /// See [Config]
281    pub const fn read_write_timeout(&self) -> Duration {
282        self.network_cfg.read_write_timeout
283    }
284    /// See [Config]
285    pub const fn tcp_nodelay(&self) -> Option<bool> {
286        self.network_cfg.tcp_nodelay
287    }
288    /// See [Config]
289    pub const fn zero_linger(&self) -> bool {
290        self.network_cfg.zero_linger
291    }
292    /// See [Config]
293    pub const fn storage_directory(&self) -> &PathBuf {
294        &self.storage_directory
295    }
296    /// See [Config]
297    pub const fn storage_blob_layouts(&self) -> &RangeInclusive<BlobLayout> {
298        &self.storage_blob_layouts
299    }
300}
301
302impl Default for Config {
303    fn default() -> Self {
304        Self::new()
305    }
306}
307
308/// Runtime based on [Tokio](https://tokio.rs).
309pub struct Executor {
310    registry: Registry,
311    metrics: Arc<Metrics>,
312    runtime: RuntimeHandle,
313    tasks: Arc<TaskTracker>,
314    shutdown: Mutex<Stopper>,
315    panicker: Panicker,
316    thread_stack_size: usize,
317}
318
319/// Closes task admission and tracks wrappers through user-future drop and descendant cleanup.
320#[derive(Default)]
321struct TaskTracker {
322    state: Mutex<TaskTrackerState>,
323    idle: Notify,
324}
325
326#[derive(Default)]
327struct TaskTrackerState {
328    active: usize,
329    closed: bool,
330}
331
332impl TaskTracker {
333    fn admit(self: &Arc<Self>) -> Option<TaskGuard> {
334        let mut state = self.state.lock();
335        if state.closed {
336            return None;
337        }
338        state.active = state.active.checked_add(1).expect("active task overflow");
339        Some(TaskGuard(Arc::clone(self)))
340    }
341
342    fn close(&self) {
343        self.state.lock().closed = true;
344    }
345
346    async fn wait(&self) {
347        loop {
348            // Subscribe before checking so the final task cannot notify between the check and wait.
349            let idle = self.idle.notified();
350            if self.state.lock().active == 0 {
351                return;
352            }
353            idle.await;
354        }
355    }
356}
357
358struct TaskGuard(Arc<TaskTracker>);
359
360impl Drop for TaskGuard {
361    fn drop(&mut self) {
362        let mut state = self.0.state.lock();
363        state.active = state.active.checked_sub(1).expect("active task underflow");
364        if state.active == 0 {
365            drop(state);
366            self.0.idle.notify_one();
367        }
368    }
369}
370
371/// Implementation of [crate::Runner] for the `tokio` runtime.
372pub struct Runner {
373    cfg: Config,
374}
375
376impl Default for Runner {
377    fn default() -> Self {
378        Self::new(Config::default())
379    }
380}
381
382impl Runner {
383    /// Initialize a new `tokio` runtime with the given number of threads.
384    pub const fn new(cfg: Config) -> Self {
385        Self { cfg }
386    }
387}
388
389impl crate::Runner for Runner {
390    type Context = Context;
391
392    fn start<F, Fut>(self, f: F) -> Fut::Output
393    where
394        F: FnOnce(Self::Context) -> Fut,
395        Fut: Future,
396    {
397        // Create a new registry
398        let mut registry = Registry::new();
399        let mut runtime_registry = registry.sub_registry(METRICS_PREFIX);
400
401        // Initialize runtime
402        let metrics = Arc::new(Metrics::init(&mut runtime_registry));
403        let mut builder = Builder::new_multi_thread();
404        builder
405            .worker_threads(self.cfg.worker_threads)
406            .max_blocking_threads(self.cfg.max_blocking_threads)
407            .thread_stack_size(self.cfg.thread_stack_size)
408            .enable_all();
409        if let Some(global_queue_interval) = self.cfg.global_queue_interval {
410            builder.global_queue_interval(global_queue_interval);
411        }
412        let runtime = builder.build().expect("failed to create Tokio runtime");
413
414        // Initialize panicker
415        let (panicker, panicked) = Panicker::new(self.cfg.catch_panics);
416
417        // Collect process metrics.
418        //
419        // We prefer to collect process metrics outside of `Context` because
420        // we are using `runtime_registry` rather than the one provided by `Context`.
421        let process = MeteredProcess::init(&mut runtime_registry);
422        runtime.spawn(process.collect(tokio::time::sleep));
423
424        // Initialize storage
425        let storage = MeteredStorage::new(
426            TokioStorage::new(TokioStorageConfig::new(
427                self.cfg.storage_directory.clone(),
428                self.cfg.storage_blob_layouts.clone(),
429            )),
430            &mut runtime_registry,
431        );
432
433        // Under the hold, a prior run's operations have finished. Linux flushes the filesystem.
434        // Other platforms sync the storage directory here so inherited partition removals are
435        // durable, then partition directories and existing blob contents on first access.
436        if let Err(e) = crate::storage::sync(&self.cfg.storage_directory) {
437            panic!(
438                "failed to sync storage at startup ({}): {e}",
439                self.cfg.storage_directory.display()
440            );
441        }
442
443        // Initialize network
444        let config = TokioNetworkConfig::default()
445            .with_connect_timeout(self.cfg.network_cfg.connect_timeout)
446            .with_read_timeout(self.cfg.network_cfg.read_write_timeout)
447            .with_write_timeout(self.cfg.network_cfg.read_write_timeout)
448            .with_tcp_nodelay(self.cfg.network_cfg.tcp_nodelay)
449            .with_zero_linger(self.cfg.network_cfg.zero_linger);
450        let network = MeteredNetwork::new(TokioNetwork::new(config), &mut runtime_registry);
451
452        // Initialize executor
453        let executor = Arc::new(Executor {
454            registry,
455            metrics,
456            runtime: runtime.handle().clone(),
457            tasks: Arc::new(TaskTracker::default()),
458            shutdown: Mutex::new(Stopper::default()),
459            panicker,
460            thread_stack_size: self.cfg.thread_stack_size,
461        });
462
463        // Get metrics
464        let label = Label::root();
465        executor.metrics.tasks_spawned.get_or_create(&label).inc();
466        let gauge = executor.metrics.tasks_running.get_or_create(&label).clone();
467
468        // Run the future
469        let tree = Tree::root();
470        let context = Context {
471            storage,
472            name: label.name().to_string(),
473            attributes: Vec::new(),
474            executor: executor.clone(),
475            network,
476            tree: Arc::clone(&tree),
477            execution: Execution::default(),
478        };
479
480        // Keep observing task panics until every admitted task has drained. The boxed root
481        // is destroyed when it completes, while tasks can still be admitted.
482        let output = catch_unwind(AssertUnwindSafe(|| {
483            // The factory runs on this thread outside block_on, so give it the runtime context
484            // that tokio primitives such as timers require.
485            let root = {
486                let _enter = runtime.enter();
487                Box::pin(f(context))
488            };
489            runtime.block_on(panicked.interrupt(async {
490                let output = root.await;
491                executor.tasks.close();
492                tree.abort();
493                executor.tasks.wait().await;
494                output
495            }))
496        }));
497
498        // Finish cleanup if the root or a task panic unwound first.
499        executor.tasks.close();
500        tree.abort();
501        runtime.block_on(executor.tasks.wait());
502        gauge.dec();
503
504        match output {
505            Ok(output) => output,
506            Err(panic) => resume_unwind(panic),
507        }
508    }
509}
510
511type Storage = MeteredStorage<TokioStorage>;
512type Network = MeteredNetwork<TokioNetwork>;
513
514/// Implementation of [crate::Spawner], [crate::Clock],
515/// [crate::Network], and [crate::Storage] for the `tokio`
516/// runtime.
517pub struct Context {
518    name: String,
519    attributes: Vec<(String, String)>,
520    executor: Arc<Executor>,
521    storage: Storage,
522    network: Network,
523    tree: Arc<Tree>,
524    execution: Execution,
525}
526
527impl Context {
528    /// Access the [Metrics] of the runtime.
529    fn metrics(&self) -> &Metrics {
530        &self.executor.metrics
531    }
532}
533
534impl crate::Spawner for Context {
535    fn dedicated(mut self) -> Self {
536        self.execution = Execution::Dedicated;
537        self
538    }
539
540    fn shared(mut self, blocking: bool) -> Self {
541        self.execution = Execution::Shared(blocking);
542        self
543    }
544
545    fn spawn<F, Fut, T>(mut self, f: F) -> Handle<T>
546    where
547        F: FnOnce(Self) -> Fut + Send + 'static,
548        Fut: Future<Output = T> + Send + 'static,
549        T: Send + 'static,
550    {
551        // Get metrics
552        let (_, metric) = spawn_metrics!(self);
553
554        // Track supervision before resetting configuration
555        let parent = Arc::clone(&self.tree);
556        let past = self.execution;
557        self.execution = Execution::default();
558        let (child, aborted) = Tree::child(&parent);
559        if aborted {
560            return Handle::closed(metric);
561        }
562        self.tree = child;
563
564        // Spawn the task
565        let executor = self.executor.clone();
566        let Some(task_guard) = executor.tasks.admit() else {
567            return Handle::closed(metric);
568        };
569
570        // The factory runs user code. If it unwinds, the guard finishes the
571        // running gauge and closes the supervision node the wrapper never received.
572        let guard = FactoryGuard::new(&parent, metric);
573        let future = f(self);
574        let (f, handle) = Handle::init(
575            future,
576            guard.disarm(),
577            executor.panicker.clone(),
578            Arc::clone(&parent),
579        );
580
581        // Destroying a finished or cancelled future, or an undelivered output, can
582        // unwind outside the user-poll boundary. Report it before releasing the tracker.
583        let panicker = executor.panicker.clone();
584        let f = async move {
585            let _task_guard = task_guard;
586            let mut f = pin!(f);
587            poll_fn(|cx| {
588                catch_unwind(AssertUnwindSafe(|| f.as_mut().poll(cx))).unwrap_or_else(|panic| {
589                    panicker.notify(panic);
590                    Poll::Ready(())
591                })
592            })
593            .await;
594        };
595
596        if matches!(past, Execution::Dedicated) {
597            utils::thread::spawn(executor.thread_stack_size, {
598                // Ensure the task can access the tokio runtime
599                let handle = executor.runtime.clone();
600                move || {
601                    handle.block_on(f);
602                }
603            });
604        } else if matches!(past, Execution::Shared(true)) {
605            executor.runtime.spawn_blocking({
606                // Ensure the task can access the tokio runtime
607                let handle = executor.runtime.clone();
608                move || {
609                    handle.block_on(f);
610                }
611            });
612        } else {
613            executor.runtime.spawn(f);
614        }
615
616        handle
617    }
618
619    async fn stop(self, value: i32, timeout: Option<Duration>) -> Result<(), Error> {
620        let stop_resolved = {
621            let mut shutdown = self.executor.shutdown.lock();
622            shutdown.stop(value)
623        };
624
625        // Wait for all tasks to complete or the timeout to fire
626        let timeout_future = timeout.map_or_else(
627            || futures::future::Either::Right(futures::future::pending()),
628            |duration| futures::future::Either::Left(self.sleep(duration)),
629        );
630        select! {
631            result = stop_resolved => {
632                result.map_err(|_| Error::Closed)?;
633                Ok(())
634            },
635            _ = timeout_future => Err(Error::Timeout),
636        }
637    }
638
639    fn stopped(&self) -> Signal {
640        self.executor.shutdown.lock().stopped()
641    }
642}
643
644#[stability(BETA)]
645impl crate::Strategizer for Context {
646    fn strategy(&self, parallelism: NonZeroUsize) -> Rayon {
647        let pool = ThreadPoolBuilder::new()
648            .num_threads(parallelism.get())
649            .stack_size(self.executor.thread_stack_size)
650            .build()
651            .expect("failed to create Tokio Rayon thread pool");
652        Rayon::with_pool(Arc::new(pool))
653    }
654}
655
656impl crate::Supervisor for Context {
657    fn child(&self, label: &'static str) -> Self {
658        let (tree, _) = Tree::child(&self.tree);
659        Self {
660            name: child_label(&self.name, label),
661            attributes: self.attributes.clone(),
662            executor: self.executor.clone(),
663            storage: self.storage.clone(),
664            network: self.network.clone(),
665            tree,
666            execution: Execution::default(),
667        }
668    }
669
670    fn with_attribute(mut self, key: &'static str, value: impl std::fmt::Display) -> Self {
671        // Validate label format (must match [a-zA-Z][a-zA-Z0-9_]*)
672        validate_label(key);
673
674        // Add the attribute to the list of attributes
675        add_attribute(&mut self.attributes, key, value);
676        self
677    }
678
679    fn name(&self) -> Name {
680        Name {
681            label: self.name.clone(),
682            attributes: self.attributes.clone(),
683        }
684    }
685}
686
687impl crate::Metrics for Context {
688    fn register<N: Into<String>, H: Into<String>, M: Metric>(
689        &self,
690        name: N,
691        help: H,
692        metric: M,
693    ) -> Registered<M> {
694        let name = name.into();
695        let help = help.into();
696        let metric = Arc::new(metric);
697        self.executor.registry.register(
698            prefixed_name(&self.name, &name),
699            help,
700            self.attributes.clone(),
701            metric,
702        )
703    }
704
705    fn encode(&self) -> String {
706        self.executor.registry.encode()
707    }
708}
709
710impl Clock for Context {
711    fn current(&self) -> SystemTime {
712        SystemTime::now()
713    }
714
715    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static + use<> {
716        tokio::time::sleep(duration)
717    }
718
719    fn sleep_until(
720        &self,
721        deadline: SystemTime,
722    ) -> impl Future<Output = ()> + Send + 'static + use<> {
723        let duration_until_deadline = deadline.duration_since(self.current()).unwrap_or_default();
724        tokio::time::sleep(duration_until_deadline)
725    }
726}
727
728#[cfg(feature = "external")]
729impl Pacer for Context {
730    fn pace<'a, F, T>(
731        &'a self,
732        _latency: Duration,
733        future: F,
734    ) -> impl Future<Output = T> + Send + 'a
735    where
736        F: Future<Output = T> + Send + 'a,
737        T: Send + 'a,
738    {
739        // Execute the future immediately
740        future
741    }
742}
743
744impl GClock for Context {
745    type Instant = SystemTime;
746
747    fn now(&self) -> Self::Instant {
748        self.current()
749    }
750}
751
752impl ReasonablyRealtime for Context {}
753
754impl crate::Network for Context {
755    type Listener = <Network as crate::Network>::Listener;
756
757    async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
758        self.network.bind(socket).await
759    }
760
761    async fn dial(&self, socket: SocketAddr) -> Result<(SinkOf<Self>, StreamOf<Self>), Error> {
762        self.network.dial(socket).await
763    }
764}
765
766impl crate::Resolver for Context {
767    async fn resolve(&self, host: &str) -> Result<Vec<IpAddr>, Error> {
768        // Uses the host's DNS configuration (e.g. /etc/resolv.conf). This delegates to the
769        // system's libc resolver.
770        //
771        // The `:0` port is required by lookup_host's API but is not used
772        // for DNS resolution.
773        let addrs = tokio::net::lookup_host(format!("{host}:0"))
774            .await
775            .map_err(|e| Error::ResolveFailed(e.to_string()))?;
776        Ok(addrs.map(|addr| addr.ip()).collect())
777    }
778}
779
780impl TryRng for Context {
781    type Error = Infallible;
782
783    fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
784        Ok(sys_rng().next_u32())
785    }
786
787    fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
788        Ok(sys_rng().next_u64())
789    }
790
791    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
792        sys_rng().fill_bytes(dest);
793        Ok(())
794    }
795}
796
797impl TryCryptoRng for Context {}
798
799impl crate::Storage for Context {
800    type Blob = <Storage as crate::Storage>::Blob;
801
802    async fn open_versioned(
803        &self,
804        partition: &str,
805        name: &[u8],
806        versions: std::ops::RangeInclusive<BlobVersion>,
807    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
808        self.storage.open_versioned(partition, name, versions).await
809    }
810
811    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
812        self.storage.remove(partition, name).await
813    }
814
815    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
816        self.storage.scan(partition).await
817    }
818}
819
820#[cfg(test)]
821#[allow(deprecated)]
822mod tests {
823    use super::*;
824    use crate::{
825        AbortOnDrop, Blob as _, Metrics, Network, Resolver, Runner as _, Sink, Spawner as _,
826        Storage as _, Strategizer as _, Stream, Supervisor as _, telemetry::metrics::raw::Counter,
827        tokio::telemetry,
828    };
829    use bytes::Bytes;
830    use commonware_parallel::Strategy as _;
831    use commonware_utils::{NZUsize, channel::oneshot};
832    use rstest::rstest;
833    use std::{
834        self,
835        collections::HashMap,
836        net::{IpAddr, Ipv4Addr, Ipv6Addr},
837        str::FromStr,
838        sync::atomic::{AtomicBool, Ordering},
839    };
840    use tracing::{Level, error};
841
842    struct TaskDropGate {
843        entered: std::sync::mpsc::Sender<()>,
844        release: std::sync::mpsc::Receiver<()>,
845    }
846
847    impl Drop for TaskDropGate {
848        fn drop(&mut self) {
849            let _ = self.entered.send(());
850            let _ = self.release.recv();
851        }
852    }
853
854    /// Signals when dropped.
855    struct NotifyDrop(std::sync::mpsc::Sender<()>);
856
857    impl Drop for NotifyDrop {
858        fn drop(&mut self) {
859            let _ = self.0.send(());
860        }
861    }
862
863    /// Panics when destroyed.
864    struct PanicOnDrop;
865
866    impl Drop for PanicOnDrop {
867        fn drop(&mut self) {
868            panic!("disposal panic");
869        }
870    }
871
872    #[derive(Clone, Copy, Debug)]
873    enum Disposal {
874        /// The task returns a value after its handle was dropped.
875        Output,
876        /// The handle aborts a pending task.
877        Cancel,
878        /// The task's future completes.
879        Ready,
880    }
881
882    fn placed(context: Context, execution: Execution) -> Context {
883        match execution {
884            Execution::Dedicated => context.dedicated(),
885            Execution::Shared(blocking) => context.shared(blocking),
886        }
887    }
888
889    #[derive(Clone, Copy, Debug)]
890    enum RootExit {
891        Return,
892        FuturePanic,
893        ConstructorPanic,
894    }
895
896    #[derive(Clone, Copy, Debug)]
897    enum JoinRelease {
898        Task,
899        Completion,
900        Output,
901    }
902
903    fn spawn_drop_gated_task(
904        context: Context,
905        execution: Execution,
906        drop_gate: TaskDropGate,
907        ready: Option<oneshot::Sender<()>>,
908    ) {
909        placed(context, execution).spawn(move |context| async move {
910            let _context = context;
911            let _drop_gate = drop_gate;
912            if let Some(ready) = ready {
913                ready.send(()).unwrap();
914            }
915            futures::future::pending::<()>().await;
916        });
917    }
918
919    fn assert_runner_drains_spawned_task(execution: Execution, root_exit: RootExit) {
920        let cfg = Config::new();
921        let storage_directory = cfg.storage_directory().clone();
922        let (ready_tx, ready_rx) = oneshot::channel();
923        let (drop_entered_tx, drop_entered_rx) = std::sync::mpsc::channel();
924        let (drop_release_tx, drop_release_rx) = std::sync::mpsc::channel();
925        let (runner_done_tx, runner_done_rx) = std::sync::mpsc::channel();
926        let runner = std::thread::spawn(move || {
927            let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
928                let drop_gate = TaskDropGate {
929                    entered: drop_entered_tx,
930                    release: drop_release_rx,
931                };
932                match root_exit {
933                    RootExit::ConstructorPanic => {
934                        Runner::new(cfg).start(move |context| -> futures::future::Pending<()> {
935                            spawn_drop_gated_task(context, execution, drop_gate, None);
936                            panic!("root constructor panic after spawning child");
937                        })
938                    }
939                    RootExit::Return | RootExit::FuturePanic => {
940                        Runner::new(cfg).start(move |context| async move {
941                            spawn_drop_gated_task(context, execution, drop_gate, Some(ready_tx));
942                            ready_rx.await.unwrap();
943                            assert!(
944                                matches!(root_exit, RootExit::Return),
945                                "root future panic after spawning child"
946                            );
947                        })
948                    }
949                }
950            }));
951            runner_done_tx.send(result.is_err()).unwrap();
952        });
953
954        drop_entered_rx
955            .recv_timeout(Duration::from_secs(5))
956            .expect("spawned task was not canceled after the root returned");
957        let early = runner_done_rx.recv_timeout(Duration::from_millis(250));
958        let returned_before_cleanup = match early {
959            Ok(panicked) => Some(panicked),
960            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => None,
961            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
962                panic!("Runner::start exited without reporting its result")
963            }
964        };
965        drop_release_tx.send(()).unwrap();
966        let panicked = returned_before_cleanup.unwrap_or_else(|| {
967            runner_done_rx
968                .recv_timeout(Duration::from_secs(5))
969                .expect("Runner::start did not return after task cleanup completed")
970        });
971        runner.join().unwrap();
972        let _ = std::fs::remove_dir_all(storage_directory);
973        assert!(
974            returned_before_cleanup.is_none(),
975            "Runner::start returned before {execution:?} task cleanup completed"
976        );
977        assert_eq!(panicked, !matches!(root_exit, RootExit::Return));
978    }
979
980    fn run_with_returned_strategy(retain: bool) -> Option<Rayon> {
981        let cfg = Config::new();
982        let storage_directory = cfg.storage_directory().clone();
983        let (strategy_tx, strategy_rx) = std::sync::mpsc::channel();
984        let runner = std::thread::spawn(move || {
985            let strategy = Runner::new(cfg).start(move |context| async move {
986                let strategy = context.strategy(NZUsize!(2));
987                strategy.spawn(1, |_| ()).await;
988                retain.then_some(strategy)
989            });
990            strategy_tx.send(strategy).unwrap();
991        });
992
993        let strategy = strategy_rx
994            .recv_timeout(Duration::from_secs(5))
995            .expect("Runner::start did not return after the strategy completed work");
996        runner.join().unwrap();
997        let _ = std::fs::remove_dir_all(storage_directory);
998        strategy
999    }
1000
1001    #[test]
1002    fn test_strategy_spawn_keeps_sibling_futures_polling() {
1003        // One runtime worker must be able to poll the branch that releases the offloaded job.
1004        let cfg = Config::new().with_worker_threads(1);
1005        let storage_directory = cfg.storage_directory().clone();
1006        let ready_branch_ran = Runner::new(cfg).start(|context| async move {
1007            let strategy = context.strategy(NZUsize!(2));
1008            let (gate_tx, gate_rx) = std::sync::mpsc::channel::<()>();
1009            let mut job = Box::pin(strategy.spawn(1, move |_| {
1010                gate_rx.recv_timeout(Duration::from_secs(5)).is_ok()
1011            }));
1012            select! {
1013                _ = job.as_mut() => false,
1014                _ = futures::future::ready(()) => {
1015                    gate_tx.send(()).unwrap();
1016                    assert!(job.await, "the gate must open before the job gives up");
1017                    true
1018                },
1019            }
1020        });
1021
1022        // Finish filesystem cleanup before asserting progress.
1023        let _ = std::fs::remove_dir_all(storage_directory);
1024        assert!(
1025            ready_branch_ran,
1026            "the sibling branch must run while the job is in flight"
1027        );
1028    }
1029
1030    #[test]
1031    fn test_storage_blob_layout_restriction() {
1032        // The default policy accepts legacy V0 blobs while selecting V1 for new blobs.
1033        let cfg = Config::new();
1034        let storage_directory = cfg.storage_directory().clone();
1035        assert_eq!(cfg.storage_blob_layouts(), &BlobLayout::ALL);
1036
1037        let partition = "layout_restriction";
1038        let v0_name = b"v0";
1039        let partition_directory = storage_directory.join(partition);
1040        std::fs::create_dir_all(&partition_directory).unwrap();
1041        let v0_path = partition_directory.join(commonware_formatting::hex(v0_name));
1042        let v0_bytes = crate::storage::tests::v0_blob_bytes(0, b"payload");
1043        std::fs::write(&v0_path, &v0_bytes).unwrap();
1044
1045        Runner::new(cfg).start(|context| async move {
1046            let (blob, size) = context.open(partition, v0_name).await.unwrap();
1047            assert_eq!(size, 7);
1048            let payload = blob
1049                .read_at(0, 7, crate::ReadOptions::default())
1050                .await
1051                .unwrap();
1052            assert_eq!(payload.coalesce(), b"payload".as_slice());
1053        });
1054
1055        // A V1-only policy rejects V0 without modifying it and creates new blobs as V1.
1056        let cfg = Config::new()
1057            .with_storage_directory(storage_directory.clone())
1058            .with_storage_blob_layouts(BlobLayout::V1..=BlobLayout::V1);
1059        Runner::new(cfg).start(|context| async move {
1060            let result = context.open(partition, v0_name).await;
1061            assert!(matches!(
1062                result,
1063                Err(Error::BlobLayoutMismatch { expected, found })
1064                    if expected == (BlobLayout::V1..=BlobLayout::V1)
1065                        && found == BlobLayout::V0
1066            ));
1067
1068            context.open(partition, b"v1").await.unwrap();
1069        });
1070
1071        assert_eq!(std::fs::read(&v0_path).unwrap(), v0_bytes);
1072        let v1_path = partition_directory.join(commonware_formatting::hex(b"v1"));
1073        let v1 = std::fs::read(&v1_path).unwrap();
1074        assert_eq!(&v1[..4], &BlobLayout::V1.magic());
1075
1076        // A V0-only policy rejects the intact header-only V1 blob without healing it and
1077        // creates new blobs the rollback target can read and write.
1078        let cfg = Config::new()
1079            .with_storage_directory(storage_directory.clone())
1080            .with_storage_blob_layouts(BlobLayout::V0..=BlobLayout::V0);
1081        Runner::new(cfg).start(|context| async move {
1082            let result = context.open(partition, b"v1").await;
1083            assert!(matches!(
1084                result,
1085                Err(Error::BlobLayoutMismatch { expected, found })
1086                    if expected == (BlobLayout::V0..=BlobLayout::V0)
1087                        && found == BlobLayout::V1
1088            ));
1089
1090            let (blob, size) = context.open(partition, b"rollback").await.unwrap();
1091            assert_eq!(size, 0);
1092            blob.write_at(0, b"rollback!".as_slice(), crate::WriteOptions::SYNC)
1093                .await
1094                .unwrap();
1095            drop(blob);
1096            let (blob, size) = context.open(partition, b"rollback").await.unwrap();
1097            assert_eq!(size, 9);
1098            let payload = blob
1099                .read_at(0, 9, crate::ReadOptions::default())
1100                .await
1101                .unwrap();
1102            assert_eq!(payload.coalesce(), b"rollback!".as_slice());
1103        });
1104
1105        assert_eq!(std::fs::read(&v1_path).unwrap(), v1);
1106        let rollback_path = partition_directory.join(commonware_formatting::hex(b"rollback"));
1107        let rollback = std::fs::read(rollback_path).unwrap();
1108        assert_eq!(&rollback[..4], &BlobLayout::V0.magic());
1109        assert_eq!(&rollback[8..], b"rollback!");
1110        let _ = std::fs::remove_dir_all(storage_directory);
1111    }
1112
1113    #[test]
1114    #[should_panic(expected = "non-empty")]
1115    fn test_storage_blob_layout_restriction_rejects_empty_range() {
1116        let _ = Config::new().with_storage_blob_layouts(BlobLayout::V1..=BlobLayout::V0);
1117    }
1118
1119    #[test]
1120    fn test_default_thread_stack_size_uses_system_default() {
1121        let cfg = Config::new();
1122        assert_eq!(
1123            cfg.thread_stack_size(),
1124            utils::thread::system_thread_stack_size()
1125        );
1126    }
1127
1128    /// The root factory runs on the caller's thread before the root is polled, yet it holds the
1129    /// context, so futures it builds eagerly from runtime primitives must be usable.
1130    #[test]
1131    fn test_root_factory_builds_runtime_futures() {
1132        let cfg = Config::new().with_worker_threads(2);
1133        Runner::new(cfg).start(|context| {
1134            let sleep = context.sleep(Duration::from_millis(1));
1135            let deadline = context.sleep_until(context.current() + Duration::from_millis(1));
1136            async move {
1137                sleep.await;
1138                deadline.await;
1139            }
1140        });
1141    }
1142
1143    #[test]
1144    fn test_runner_waits_for_spawned_task_cancellation() {
1145        for execution in [
1146            Execution::Shared(false),
1147            Execution::Shared(true),
1148            Execution::Dedicated,
1149        ] {
1150            for root_exit in [
1151                RootExit::Return,
1152                RootExit::FuturePanic,
1153                RootExit::ConstructorPanic,
1154            ] {
1155                assert_runner_drains_spawned_task(execution, root_exit);
1156            }
1157        }
1158    }
1159
1160    /// A task panic raised while the runner drains tasks after the root returns
1161    /// fails the runner.
1162    #[test]
1163    fn test_runner_propagates_panic_during_drain() {
1164        let cfg = Config::new();
1165        let storage_directory = cfg.storage_directory().clone();
1166        let result = catch_unwind(AssertUnwindSafe(|| {
1167            Runner::new(cfg).start(|context| async move {
1168                // Only the runner's tree abort after the root returns drops the sentinel.
1169                let (sentinel_started_tx, sentinel_started_rx) = oneshot::channel();
1170                let (sentinel_dropped_tx, sentinel_dropped_rx) = std::sync::mpsc::channel();
1171                context.child("sentinel").spawn(move |_| async move {
1172                    let _notify = NotifyDrop(sentinel_dropped_tx);
1173                    sentinel_started_tx.send(()).unwrap();
1174                    futures::future::pending::<()>().await;
1175                });
1176                sentinel_started_rx.await.unwrap();
1177
1178                // The child blocks its poll until that abort, then panics.
1179                let (started_tx, started_rx) = oneshot::channel();
1180                context
1181                    .child("late")
1182                    .dedicated()
1183                    .spawn(move |_| async move {
1184                        started_tx.send(()).unwrap();
1185                        sentinel_dropped_rx.recv().unwrap();
1186                        panic!("drain panic");
1187                    });
1188                started_rx.await.unwrap();
1189            })
1190        }));
1191        let _ = std::fs::remove_dir_all(storage_directory);
1192        let panic = result.expect_err("Runner::start returned after a task panicked during drain");
1193        assert_eq!(utils::extract_panic_message(&*panic), "drain panic");
1194    }
1195
1196    /// Panics raised while destroying a task's future or undelivered output are
1197    /// task panics. A handle still waiting on the task resolves to closed.
1198    #[rstest]
1199    fn test_disposal_panic_is_task_panic(
1200        #[values(Disposal::Output, Disposal::Cancel, Disposal::Ready)] disposal: Disposal,
1201        #[values(
1202            Execution::Shared(false),
1203            Execution::Shared(true),
1204            Execution::Dedicated
1205        )]
1206        execution: Execution,
1207        #[values(false, true)] catch: bool,
1208    ) {
1209        let cfg = Config::new().with_catch_panics(catch);
1210        let storage_directory = cfg.storage_directory().clone();
1211        let result = catch_unwind(AssertUnwindSafe(|| {
1212            Runner::new(cfg).start(|context| async move {
1213                let child = placed(context.child("disposal"), execution);
1214                match disposal {
1215                    Disposal::Output => {
1216                        // Drop the handle while the task runs so its output
1217                        // cannot be delivered.
1218                        let (entered_tx, entered_rx) = std::sync::mpsc::channel();
1219                        let (release_tx, release_rx) = std::sync::mpsc::channel();
1220                        let handle = child.spawn(move |_| async move {
1221                            entered_tx.send(()).unwrap();
1222                            release_rx.recv().unwrap();
1223                            PanicOnDrop
1224                        });
1225                        entered_rx.recv().unwrap();
1226                        drop(handle);
1227                        release_tx.send(()).unwrap();
1228                    }
1229                    Disposal::Cancel => {
1230                        // Abort the task once it holds a value that panics on drop.
1231                        let (ready_tx, ready_rx) = oneshot::channel();
1232                        let handle = child.spawn(move |_| async move {
1233                            let _guard = PanicOnDrop;
1234                            ready_tx.send(()).unwrap();
1235                            futures::future::pending::<()>().await;
1236                        });
1237                        ready_rx.await.unwrap();
1238                        handle.abort();
1239                        assert!(matches!(handle.await, Err(Error::Closed)));
1240                    }
1241                    Disposal::Ready => {
1242                        // The completed future panics when destroyed, before
1243                        // its result is published.
1244                        let guard = PanicOnDrop;
1245                        let result = child
1246                            .spawn(move |_| {
1247                                poll_fn(move |_| {
1248                                    let _guard = &guard;
1249                                    Poll::Ready(())
1250                                })
1251                            })
1252                            .await;
1253                        assert!(matches!(result, Err(Error::Closed)));
1254                    }
1255                }
1256
1257                // A caught disposal panic leaves the runner usable.
1258                context.child("survivor").spawn(|_| async {}).await.unwrap();
1259            })
1260        }));
1261        let _ = std::fs::remove_dir_all(storage_directory);
1262        if catch {
1263            result.expect("caught disposal panic failed the runner");
1264        } else {
1265            let panic = result.expect_err("disposal panic did not fail the runner");
1266            assert_eq!(utils::extract_panic_message(&*panic), "disposal panic");
1267        }
1268    }
1269
1270    /// Records the tracker's active count for each error logged on the installing thread.
1271    struct ActiveOnError {
1272        /// Tracker of the runner under test.
1273        tasks: Arc<TaskTracker>,
1274        /// Active counts observed when errors were logged.
1275        active: Arc<Mutex<Vec<usize>>>,
1276    }
1277
1278    impl tracing::Subscriber for ActiveOnError {
1279        fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
1280            *metadata.level() == Level::ERROR
1281        }
1282
1283        fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
1284            tracing::span::Id::from_u64(1)
1285        }
1286
1287        fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
1288
1289        fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
1290
1291        fn event(&self, _: &tracing::Event<'_>) {
1292            self.active.lock().push(self.tasks.state.lock().active);
1293        }
1294
1295        fn enter(&self, _: &tracing::span::Id) {}
1296
1297        fn exit(&self, _: &tracing::span::Id) {}
1298    }
1299
1300    /// A disposal panic is reported while its task is still tracked, so the runner cannot
1301    /// finish draining before the panic is published.
1302    #[test]
1303    fn test_disposal_panic_reported_before_tracker_release() {
1304        let cfg = Config::new().with_catch_panics(true);
1305        let storage_directory = cfg.storage_directory().clone();
1306        let active = Arc::new(Mutex::new(Vec::new()));
1307        let observed = active.clone();
1308        Runner::new(cfg).start(|context| async move {
1309            // The only tracked task records the tracker when its panic is logged.
1310            let subscriber = ActiveOnError {
1311                tasks: context.executor.tasks.clone(),
1312                active: observed,
1313            };
1314            let (ready_tx, ready_rx) = oneshot::channel();
1315            let handle = context
1316                .child("disposal")
1317                .dedicated()
1318                .spawn(move |_| async move {
1319                    std::mem::forget(tracing::subscriber::set_default(subscriber));
1320                    let _guard = PanicOnDrop;
1321                    ready_tx.send(()).unwrap();
1322                    futures::future::pending::<()>().await;
1323                });
1324            ready_rx.await.unwrap();
1325
1326            // Aborting destroys the guard, whose panic is logged on the task's thread.
1327            handle.abort();
1328            assert!(matches!(handle.await, Err(Error::Closed)));
1329        });
1330        let _ = std::fs::remove_dir_all(storage_directory);
1331        assert_eq!(*active.lock(), vec![1]);
1332    }
1333
1334    /// A root destroyed on completion can still spawn work, which shutdown drains.
1335    #[test]
1336    fn test_root_destruction_can_spawn_work_before_shutdown() {
1337        /// Ready root that spawns a task while its destructor still owns a context.
1338        struct Root {
1339            context: Option<Context>,
1340            invoked: Arc<AtomicBool>,
1341            dropped: std::sync::mpsc::Sender<()>,
1342        }
1343
1344        impl Future for Root {
1345            type Output = ();
1346
1347            fn poll(self: std::pin::Pin<&mut Self>, _: &mut std::task::Context<'_>) -> Poll<()> {
1348                Poll::Ready(())
1349            }
1350        }
1351
1352        impl Drop for Root {
1353            fn drop(&mut self) {
1354                let invoked = self.invoked.clone();
1355                let payload = NotifyDrop(self.dropped.clone());
1356                self.context.take().unwrap().spawn(move |_| {
1357                    invoked.store(true, Ordering::SeqCst);
1358                    async move {
1359                        let _payload = payload;
1360                        futures::future::pending::<()>().await;
1361                    }
1362                });
1363            }
1364        }
1365
1366        for execution in [
1367            Execution::Shared(false),
1368            Execution::Shared(true),
1369            Execution::Dedicated,
1370        ] {
1371            let cfg = Config::new();
1372            let storage_directory = cfg.storage_directory().clone();
1373            let invoked = Arc::new(AtomicBool::new(false));
1374            let (dropped_tx, dropped_rx) = std::sync::mpsc::channel();
1375            Runner::new(cfg).start(|context| Root {
1376                context: Some(placed(context.child("root"), execution)),
1377                invoked: invoked.clone(),
1378                dropped: dropped_tx,
1379            });
1380            let _ = std::fs::remove_dir_all(storage_directory);
1381
1382            // The root can still spawn during disposal, and shutdown drains its child.
1383            assert!(
1384                invoked.load(Ordering::SeqCst),
1385                "{execution:?} spawn from root destruction was rejected"
1386            );
1387            assert!(
1388                dropped_rx.try_recv().is_ok(),
1389                "{execution:?} child spawned by root destruction outlived the runner"
1390            );
1391        }
1392    }
1393
1394    fn assert_join_all_releases_pending_work(release_owner: JoinRelease) {
1395        let cfg = Config::new();
1396        let storage_directory = cfg.storage_directory().clone();
1397        let (started, started_rx) = oneshot::channel();
1398        let (release, release_rx) = std::sync::mpsc::channel();
1399        let (joining, joining_rx) = std::sync::mpsc::channel();
1400        let (done, done_rx) = std::sync::mpsc::channel();
1401        let (exited, exited_rx) = std::sync::mpsc::channel();
1402        let (drop_release, drop_release_rx) = std::sync::mpsc::channel();
1403        drop(drop_release);
1404        let release_on_drop = TaskDropGate {
1405            entered: release.clone(),
1406            release: drop_release_rx,
1407        };
1408
1409        let runner = std::thread::spawn(move || {
1410            Runner::new(cfg).start(move |context| async move {
1411                let blocked = context
1412                    .child("blocked")
1413                    .dedicated()
1414                    .spawn(move |_| async move {
1415                        started.send(()).unwrap();
1416                        release_rx.recv().unwrap();
1417                        exited.send(()).unwrap();
1418                        Ok::<_, Error>(None::<TaskDropGate>)
1419                    })
1420                    .abort_on_drop();
1421                let failed = Handle::ready(Ok(Err(Error::Closed))).abort_on_drop();
1422                let mut guards = vec![failed, blocked];
1423                if matches!(release_owner, JoinRelease::Output) {
1424                    guards.insert(
1425                        0,
1426                        Handle::ready(Ok(Ok(Some(release_on_drop)))).abort_on_drop(),
1427                    );
1428                } else {
1429                    let pending = async move {
1430                        let _release_on_drop = release_on_drop;
1431                        futures::future::pending::<Result<Option<TaskDropGate>, Error>>().await
1432                    };
1433                    let pending = if matches!(release_owner, JoinRelease::Task) {
1434                        context.child("pending").spawn(move |_| pending)
1435                    } else {
1436                        Handle::from_future(async move { Ok(pending.await) })
1437                    };
1438                    guards.push(pending.abort_on_drop());
1439                }
1440
1441                started_rx.await.unwrap();
1442                joining.send(()).unwrap();
1443                let result = AbortOnDrop::join_all::<Error>(guards).await;
1444                assert!(done.send((result, exited_rx.try_recv().is_ok())).is_ok());
1445            });
1446        });
1447
1448        joining_rx
1449            .recv_timeout(Duration::from_secs(5))
1450            .expect("blocking task did not start");
1451        let early = done_rx.recv_timeout(Duration::from_secs(5));
1452        let completed = early.is_ok();
1453
1454        // Release the blocking poll even when cancellation stalls, so the runner can shut down.
1455        let _ = release.send(());
1456        let (result, exited) = early.unwrap_or_else(|_| {
1457            done_rx
1458                .recv_timeout(Duration::from_secs(5))
1459                .expect("joining did not finish after releasing the blocking task")
1460        });
1461        runner.join().unwrap();
1462        let _ = std::fs::remove_dir_all(storage_directory);
1463        assert!(
1464            completed,
1465            "joining stalled with release owned by {release_owner:?}"
1466        );
1467        assert!(exited, "joining returned before the blocking task exited");
1468        assert!(matches!(result, Err(Error::Closed)));
1469    }
1470
1471    #[test]
1472    fn test_join_all_aborts_all_tasks_before_waiting() {
1473        assert_join_all_releases_pending_work(JoinRelease::Task);
1474    }
1475
1476    #[test]
1477    fn test_join_all_drains_completions_concurrently() {
1478        assert_join_all_releases_pending_work(JoinRelease::Completion);
1479    }
1480
1481    #[test]
1482    fn test_join_all_drops_outputs_before_waiting() {
1483        assert_join_all_releases_pending_work(JoinRelease::Output);
1484    }
1485
1486    #[test]
1487    fn test_runner_start_waits_for_previous_run() {
1488        let cfg = Config::new();
1489        let storage_directory = cfg.storage_directory().clone();
1490
1491        // The first run keeps its context, and with it the storage directory,
1492        // until it is released.
1493        let (started, first_started) = std::sync::mpsc::channel();
1494        let (release, released) = futures::channel::oneshot::channel();
1495        let first_cfg = cfg.clone();
1496        let first = std::thread::spawn(move || {
1497            Runner::new(first_cfg).start(|context| async move {
1498                started.send(()).unwrap();
1499                released.await.unwrap();
1500                drop(context);
1501            });
1502        });
1503        first_started.recv_timeout(Duration::from_secs(10)).unwrap();
1504
1505        // A second run on the same directory cannot start until the first has
1506        // returned.
1507        let (started, second_started) = std::sync::mpsc::channel();
1508        let second = std::thread::spawn(move || {
1509            Runner::new(cfg).start(|_| async move {
1510                started.send(()).unwrap();
1511            });
1512        });
1513        match second_started.recv_timeout(Duration::from_millis(200)) {
1514            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
1515            other => panic!("second run started while the first held the directory: {other:?}"),
1516        }
1517        release.send(()).unwrap();
1518        first.join().unwrap();
1519        second_started
1520            .recv_timeout(Duration::from_secs(10))
1521            .expect("second run did not start after the first returned");
1522        second.join().unwrap();
1523        let _ = std::fs::remove_dir_all(storage_directory);
1524    }
1525
1526    #[test]
1527    fn test_runner_owns_runtime_when_context_escapes() {
1528        let cfg = Config::new();
1529        let storage_directory = cfg.storage_directory().clone();
1530        let (ready_tx, ready_rx) = oneshot::channel();
1531        let (drop_entered_tx, drop_entered_rx) = std::sync::mpsc::channel();
1532        let (drop_release_tx, drop_release_rx) = std::sync::mpsc::channel();
1533        let (runner_returned_tx, runner_returned_rx) = std::sync::mpsc::channel();
1534        let (context_release_tx, context_release_rx) = std::sync::mpsc::channel();
1535        let runner = std::thread::spawn(move || {
1536            let context = Runner::new(cfg).start(move |context| async move {
1537                context.executor.runtime.spawn(async move {
1538                    let _drop_gate = TaskDropGate {
1539                        entered: drop_entered_tx,
1540                        release: drop_release_rx,
1541                    };
1542                    ready_tx.send(()).unwrap();
1543                    futures::future::pending::<()>().await;
1544                });
1545                ready_rx.await.unwrap();
1546                context
1547            });
1548            runner_returned_tx.send(()).unwrap();
1549            context_release_rx.recv().unwrap();
1550            drop(context);
1551        });
1552
1553        let returned_early = runner_returned_rx
1554            .recv_timeout(Duration::from_millis(500))
1555            .is_ok();
1556        if returned_early {
1557            drop_release_tx.send(()).unwrap();
1558            context_release_tx.send(()).unwrap();
1559            drop_entered_rx
1560                .recv_timeout(Duration::from_secs(5))
1561                .expect("escaped Context did not retain the raw runtime task");
1562        } else {
1563            drop_entered_rx
1564                .recv_timeout(Duration::from_secs(5))
1565                .expect("Runner did not cancel its raw runtime task");
1566            drop_release_tx.send(()).unwrap();
1567            runner_returned_rx
1568                .recv_timeout(Duration::from_secs(5))
1569                .expect("Runner did not return after raw task cleanup");
1570            context_release_tx.send(()).unwrap();
1571        }
1572        runner.join().unwrap();
1573        let _ = std::fs::remove_dir_all(storage_directory);
1574        assert!(
1575            !returned_early,
1576            "a returned Context kept the Tokio runtime alive after Runner::start"
1577        );
1578    }
1579
1580    #[test]
1581    fn test_runner_returns_strategy_after_pool_work() {
1582        assert!(run_with_returned_strategy(false).is_none());
1583
1584        let strategy = run_with_returned_strategy(true).unwrap();
1585        assert_eq!(futures::executor::block_on(strategy.spawn(1, |_| 42)), 42);
1586    }
1587
1588    #[test]
1589    fn test_runner_resumes_strategy_panic_payload_after_pool_work() {
1590        let cfg = Config::new();
1591        let storage_directory = cfg.storage_directory().clone();
1592        let (strategy_tx, strategy_rx) = std::sync::mpsc::channel();
1593        let runner = std::thread::spawn(move || {
1594            let result: std::thread::Result<()> =
1595                std::panic::catch_unwind(AssertUnwindSafe(|| {
1596                    Runner::new(cfg).start(move |context| async move {
1597                        let strategy = context.strategy(NZUsize!(2));
1598                        strategy.spawn(1, |_| ()).await;
1599                        std::panic::panic_any(strategy);
1600                    });
1601                }));
1602            let strategy = result
1603                .expect_err("Runner::start did not resume the root panic")
1604                .downcast::<Rayon>()
1605                .expect("Runner::start changed the root panic payload");
1606            strategy_tx.send(*strategy).unwrap();
1607        });
1608
1609        let strategy = strategy_rx
1610            .recv_timeout(Duration::from_secs(5))
1611            .expect("Runner::start did not resume the strategy panic payload");
1612        runner.join().unwrap();
1613        let _ = std::fs::remove_dir_all(storage_directory);
1614        assert_eq!(futures::executor::block_on(strategy.spawn(1, |_| 42)), 42);
1615    }
1616
1617    #[test]
1618    fn test_thread_stack_size_override() {
1619        let cfg = Config::new().with_thread_stack_size(4 * 1024 * 1024);
1620        assert_eq!(cfg.thread_stack_size(), 4 * 1024 * 1024);
1621    }
1622
1623    #[test]
1624    fn test_process_rss_metric() {
1625        let executor = Runner::default();
1626        executor.start(|context| async move {
1627            loop {
1628                // Wait for RSS metric to be available
1629                let metrics = context.encode();
1630                if !metrics.contains("runtime_process_rss") {
1631                    context.sleep(Duration::from_millis(100)).await;
1632                    continue;
1633                }
1634
1635                // Verify the RSS value is eventually populated (greater than 0)
1636                for line in metrics.lines() {
1637                    if line.starts_with("runtime_process_rss")
1638                        && !line.starts_with("runtime_process_rss{")
1639                    {
1640                        let parts: Vec<&str> = line.split_whitespace().collect();
1641                        if parts.len() >= 2 {
1642                            let rss_value: i64 =
1643                                parts[1].parse().expect("Failed to parse RSS value");
1644                            if rss_value > 0 {
1645                                return;
1646                            }
1647                        }
1648                    }
1649                }
1650            }
1651        });
1652    }
1653
1654    #[test]
1655    fn test_telemetry() {
1656        let executor = Runner::default();
1657        executor.start(|context| async move {
1658            // Define the server address
1659            let address = SocketAddr::from_str("127.0.0.1:8000").unwrap();
1660
1661            // Configure telemetry
1662            telemetry::init(
1663                context.child("metrics"),
1664                telemetry::Logs {
1665                    level: Level::INFO,
1666                    json: false,
1667                },
1668                Some(address),
1669                None,
1670            );
1671
1672            // Register a test metric
1673            let counter: Counter<u64> = Counter::default();
1674            let _registered = context.register("test_counter", "Test counter", counter.clone());
1675            counter.inc();
1676
1677            // Helper functions to parse HTTP response
1678            async fn read_line<St: Stream>(stream: &mut St) -> Result<String, Error> {
1679                let mut line = Vec::new();
1680                loop {
1681                    let received = stream.recv(1).await?;
1682                    let byte = received.coalesce().as_ref()[0];
1683                    if byte == b'\n' {
1684                        if line.last() == Some(&b'\r') {
1685                            line.pop(); // Remove trailing \r
1686                        }
1687                        break;
1688                    }
1689                    line.push(byte);
1690                }
1691                String::from_utf8(line).map_err(|_| Error::ReadFailed)
1692            }
1693
1694            async fn read_headers<St: Stream>(
1695                stream: &mut St,
1696            ) -> Result<HashMap<String, String>, Error> {
1697                let mut headers = HashMap::new();
1698                loop {
1699                    let line = read_line(stream).await?;
1700                    if line.is_empty() {
1701                        break;
1702                    }
1703                    let parts: Vec<&str> = line.splitn(2, ": ").collect();
1704                    if parts.len() == 2 {
1705                        headers.insert(parts[0].to_string(), parts[1].to_string());
1706                    }
1707                }
1708                Ok(headers)
1709            }
1710
1711            async fn read_body<St: Stream>(
1712                stream: &mut St,
1713                content_length: usize,
1714            ) -> Result<String, Error> {
1715                let received = stream.recv(content_length).await?;
1716                String::from_utf8(received.coalesce().into()).map_err(|_| Error::ReadFailed)
1717            }
1718
1719            // Simulate a client connecting to the server
1720            let client_handle = context.child("client").spawn(move |context| async move {
1721                let (mut sink, mut stream) = loop {
1722                    match context.dial(address).await {
1723                        Ok((sink, stream)) => break (sink, stream),
1724                        Err(e) => {
1725                            // The client may be polled before the server is ready, that's alright!
1726                            error!(err =?e, "failed to connect");
1727                            context.sleep(Duration::from_millis(10)).await;
1728                        }
1729                    }
1730                };
1731
1732                // Send a GET request to the server
1733                let request = format!(
1734                    "GET /metrics HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n"
1735                );
1736                sink.send(Bytes::from(request)).await.unwrap();
1737
1738                // Read and verify the HTTP status line
1739                let status_line = read_line(&mut stream).await.unwrap();
1740                assert_eq!(status_line, "HTTP/1.1 200 OK");
1741
1742                // Read and parse headers
1743                let headers = read_headers(&mut stream).await.unwrap();
1744                println!("Headers: {headers:?}");
1745                let content_length = headers
1746                    .get("content-length")
1747                    .unwrap()
1748                    .parse::<usize>()
1749                    .unwrap();
1750
1751                // Read and verify the body
1752                let body = read_body(&mut stream, content_length).await.unwrap();
1753                assert!(body.contains("test_counter_total 1"));
1754            });
1755
1756            // Wait for the client task to complete
1757            client_handle.await.unwrap();
1758        });
1759    }
1760
1761    #[test]
1762    fn test_resolver() {
1763        let executor = Runner::default();
1764        executor.start(|context| async move {
1765            let addrs = context.resolve("localhost").await.unwrap();
1766            assert!(!addrs.is_empty());
1767            for addr in addrs {
1768                assert!(
1769                    addr == IpAddr::V4(Ipv4Addr::LOCALHOST)
1770                        || addr == IpAddr::V6(Ipv6Addr::LOCALHOST)
1771                );
1772            }
1773        });
1774    }
1775}