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 tcp_nodelay: Option<bool>,
78
79 zero_linger: bool,
88
89 connect_timeout: Duration,
93
94 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#[derive(Clone)]
117pub struct Config {
118 worker_threads: usize,
124
125 global_queue_interval: Option<u32>,
132
133 max_blocking_threads: usize,
145
146 thread_stack_size: usize,
153
154 catch_panics: bool,
156
157 storage_directory: PathBuf,
160
161 storage_blob_layouts: RangeInclusive<BlobLayout>,
168
169 network_cfg: NetworkConfig,
171}
172
173impl Config {
174 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 pub const fn with_worker_threads(mut self, n: usize) -> Self {
193 self.worker_threads = n;
194 self
195 }
196 pub const fn with_global_queue_interval(mut self, n: u32) -> Self {
198 self.global_queue_interval = Some(n);
199 self
200 }
201 pub const fn with_max_blocking_threads(mut self, n: usize) -> Self {
203 self.max_blocking_threads = n;
204 self
205 }
206 pub const fn with_thread_stack_size(mut self, n: usize) -> Self {
208 self.thread_stack_size = n;
209 self
210 }
211 pub const fn with_catch_panics(mut self, b: bool) -> Self {
213 self.catch_panics = b;
214 self
215 }
216 pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
218 self.network_cfg.connect_timeout = timeout;
219 self
220 }
221 pub const fn with_read_write_timeout(mut self, d: Duration) -> Self {
223 self.network_cfg.read_write_timeout = d;
224 self
225 }
226 pub const fn with_tcp_nodelay(mut self, n: Option<bool>) -> Self {
228 self.network_cfg.tcp_nodelay = n;
229 self
230 }
231 pub const fn with_zero_linger(mut self, l: bool) -> Self {
233 self.network_cfg.zero_linger = l;
234 self
235 }
236 pub fn with_storage_directory(mut self, p: impl Into<PathBuf>) -> Self {
238 self.storage_directory = p.into();
239 self
240 }
241 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 pub const fn worker_threads(&self) -> usize {
258 self.worker_threads
259 }
260 pub const fn global_queue_interval(&self) -> Option<u32> {
262 self.global_queue_interval
263 }
264 pub const fn max_blocking_threads(&self) -> usize {
266 self.max_blocking_threads
267 }
268 pub const fn thread_stack_size(&self) -> usize {
270 self.thread_stack_size
271 }
272 pub const fn catch_panics(&self) -> bool {
274 self.catch_panics
275 }
276 pub const fn connect_timeout(&self) -> Duration {
278 self.network_cfg.connect_timeout
279 }
280 pub const fn read_write_timeout(&self) -> Duration {
282 self.network_cfg.read_write_timeout
283 }
284 pub const fn tcp_nodelay(&self) -> Option<bool> {
286 self.network_cfg.tcp_nodelay
287 }
288 pub const fn zero_linger(&self) -> bool {
290 self.network_cfg.zero_linger
291 }
292 pub const fn storage_directory(&self) -> &PathBuf {
294 &self.storage_directory
295 }
296 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
308pub 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#[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 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
371pub 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 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 let mut registry = Registry::new();
399 let mut runtime_registry = registry.sub_registry(METRICS_PREFIX);
400
401 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 let (panicker, panicked) = Panicker::new(self.cfg.catch_panics);
416
417 let process = MeteredProcess::init(&mut runtime_registry);
422 runtime.spawn(process.collect(tokio::time::sleep));
423
424 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 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 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 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 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 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 let output = catch_unwind(AssertUnwindSafe(|| {
483 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 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
514pub 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 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 let (_, metric) = spawn_metrics!(self);
553
554 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 let executor = self.executor.clone();
566 let Some(task_guard) = executor.tasks.admit() else {
567 return Handle::closed(metric);
568 };
569
570 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 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 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 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 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(key);
673
674 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 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 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 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 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 Output,
876 Cancel,
878 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 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 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 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 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 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 #[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 #[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 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 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 #[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 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 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 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 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 struct ActiveOnError {
1272 tasks: Arc<TaskTracker>,
1274 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 #[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 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 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 #[test]
1336 fn test_root_destruction_can_spawn_work_before_shutdown() {
1337 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 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 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 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 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 let metrics = context.encode();
1630 if !metrics.contains("runtime_process_rss") {
1631 context.sleep(Duration::from_millis(100)).await;
1632 continue;
1633 }
1634
1635 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 let address = SocketAddr::from_str("127.0.0.1:8000").unwrap();
1660
1661 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 let counter: Counter<u64> = Counter::default();
1674 let _registered = context.register("test_counter", "Test counter", counter.clone());
1675 counter.inc();
1676
1677 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(); }
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 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 error!(err =?e, "failed to connect");
1727 context.sleep(Duration::from_millis(10)).await;
1728 }
1729 }
1730 };
1731
1732 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 let status_line = read_line(&mut stream).await.unwrap();
1740 assert_eq!(status_line, "HTTP/1.1 200 OK");
1741
1742 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 let body = read_body(&mut stream, content_length).await.unwrap();
1753 assert!(body.contains("test_counter_total 1"));
1754 });
1755
1756 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}