1#![doc(
20 html_logo_url = "https://commonware.xyz/imgs/rustdoc_logo.svg",
21 html_favicon_url = "https://commonware.xyz/favicon.ico"
22)]
23
24use commonware_macros::stability_scope;
25
26#[macro_use]
27mod macros;
28
29mod network;
30mod process;
31mod storage;
32
33stability_scope!(ALPHA {
34 #[cfg(feature = "arbitrary")]
35 pub mod conformance;
36 pub mod deterministic;
37 pub mod mocks;
38});
39stability_scope!(ALPHA, cfg(not(target_arch = "wasm32")) {
40 pub mod benchmarks;
41});
42stability_scope!(ALPHA, cfg(any(feature = "iouring-storage", feature = "iouring-network")) {
43 mod iouring;
44});
45stability_scope!(BETA, cfg(not(target_arch = "wasm32")) {
46 pub mod tokio;
47});
48stability_scope!(BETA {
49 pub use bytes::{Buf, BufMut};
51 use commonware_macros::select;
52 use commonware_parallel::Rayon;
53 pub use governor::Quota;
55 use iobuf::PoolError;
56 use std::{
57 future::Future,
58 io::Error as IoError,
59 net::SocketAddr,
60 num::NonZeroUsize,
61 sync::Arc,
62 time::{Duration, SystemTime},
63 };
64 pub(crate) use telemetry::metrics::{METRICS_PREFIX, child_label, prefixed_name};
65 use thiserror::Error;
66
67 pub mod iobuf;
68 pub use iobuf::{
69 BufferPool, BufferPoolClassConfig, BufferPoolConfig, BufferPoolThreadCache,
70 Builder as IoBufsBuilder, IoBuf, IoBufMut, IoBufs, IoBufsMut, cache_line_size, page_size,
71 };
72
73 pub mod utils;
74 pub use utils::*;
75
76 pub mod telemetry;
77
78 #[repr(u16)]
86 #[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
87 pub enum BlobLayout {
88 #[deprecated(note = "unaligned pages can degrade performance")]
95 V0 = 0,
96 V1 = 1,
98 }
99
100 pub const DEFAULT_BLOB_LAYOUT: BlobLayout = BlobLayout::V1;
103
104 impl BlobLayout {
105 #[allow(deprecated)]
107 pub const ALL: std::ops::RangeInclusive<Self> = Self::V0..=DEFAULT_BLOB_LAYOUT;
108 }
109
110 #[derive(Clone, Copy, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
114 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
115 pub struct BlobVersion(u16);
116
117 impl BlobVersion {
118 pub const fn new(version: u16) -> Self {
120 Self(version)
121 }
122
123 pub const fn get(self) -> u16 {
125 self.0
126 }
127 }
128
129 impl std::fmt::Display for BlobVersion {
130 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131 write!(f, "{}", self.0)
132 }
133 }
134
135 impl std::fmt::Debug for BlobVersion {
136 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
137 write!(f, "{}", self.0)
138 }
139 }
140
141 pub const DEFAULT_BLOB_VERSION: BlobVersion = BlobVersion::new(0);
143
144 #[derive(Error, Debug, Clone)]
146 pub enum Error {
147 #[error("exited")]
148 Exited,
149 #[error("closed")]
150 Closed,
151 #[error("aborted")]
152 Aborted,
153 #[error("timeout")]
154 Timeout,
155 #[error("bind failed")]
156 BindFailed,
157 #[error("connection failed")]
158 ConnectionFailed,
159 #[error("write failed")]
160 WriteFailed,
161 #[error("read failed")]
162 ReadFailed,
163 #[error("send failed")]
164 SendFailed,
165 #[error("recv failed")]
166 RecvFailed,
167 #[error("dns resolution failed: {0}")]
168 ResolveFailed(String),
169 #[error(
170 "partition name invalid, must only contain alphanumeric, dash ('-'), or underscore ('_') characters: {0}"
171 )]
172 PartitionNameInvalid(String),
173 #[error("partition creation failed: {0}")]
174 PartitionCreationFailed(String),
175 #[error("partition missing: {0}")]
176 PartitionMissing(String),
177 #[error("partition corrupt: {0}")]
178 PartitionCorrupt(String),
179 #[error("blob open failed: {0}/{1} error: {2}")]
180 BlobOpenFailed(String, String, Arc<IoError>),
181 #[error("blob missing: {0}/{1}")]
182 BlobMissing(String, String),
183 #[error("blob resize failed: {0}/{1} error: {2}")]
184 BlobResizeFailed(String, String, Arc<IoError>),
185 #[error("blob sync failed: {0}/{1} error: {2}")]
186 BlobSyncFailed(String, String, Arc<IoError>),
187 #[error("blob insufficient length")]
188 BlobInsufficientLength,
189 #[error("blob corrupt: {0}/{1} reason: {2}")]
190 BlobCorrupt(String, String, String),
191 #[error("blob layout mismatch: expected one of {expected:?}, found {found:?}")]
192 BlobLayoutMismatch {
193 expected: std::ops::RangeInclusive<BlobLayout>,
194 found: BlobLayout,
195 },
196 #[error("blob version mismatch: expected one of {expected:?}, found {found}")]
197 BlobVersionMismatch {
198 expected: std::ops::RangeInclusive<BlobVersion>,
199 found: BlobVersion,
200 },
201 #[error("invalid or missing checksum")]
202 InvalidChecksum,
203 #[error("offset overflow")]
204 OffsetOverflow,
205 #[error("io error: {0}")]
206 Io(Arc<IoError>),
207 #[error("buffer pool: {0}")]
208 Pool(#[from] PoolError),
209 }
210
211 impl From<IoError> for Error {
212 fn from(err: IoError) -> Self {
213 Self::Io(Arc::new(err))
214 }
215 }
216
217 pub trait Runner {
220 type Context;
222
223 fn start<F, Fut>(self, f: F) -> Fut::Output
229 where
230 F: FnOnce(Self::Context) -> Fut,
231 Fut: Future;
232 }
233
234 #[derive(Clone, Debug, Default)]
236 pub struct Name {
237 pub label: String,
239 pub attributes: Vec<(String, String)>,
241 }
242
243 pub trait Supervisor: Send + Sync + 'static {
245 fn name(&self) -> Name;
247
248 #[must_use]
258 fn child(&self, label: &'static str) -> Self;
259
260 #[must_use]
325 fn with_attribute(self, key: &'static str, value: impl std::fmt::Display) -> Self;
326 }
327
328 pub trait Spawner: Supervisor {
330 #[must_use]
339 fn shared(self, blocking: bool) -> Self;
340
341 #[must_use]
348 fn dedicated(self) -> Self;
349
350 fn spawn<F, Fut, T>(self, f: F) -> Handle<T>
385 where
386 Self: Sized,
387 F: FnOnce(Self) -> Fut + Send + 'static,
388 Fut: Future<Output = T> + Send + 'static,
389 T: Send + 'static;
390
391 fn stop(
411 self,
412 value: i32,
413 timeout: Option<Duration>,
414 ) -> impl Future<Output = Result<(), Error>> + Send;
415
416 fn stopped(&self) -> signal::Signal;
423 }
424
425 pub trait Strategizer: Spawner {
427 fn strategy(&self, parallelism: NonZeroUsize) -> Rayon;
436 }
437
438 pub trait Metrics: Supervisor {
440 fn register<N: Into<String>, H: Into<String>, M: telemetry::metrics::Metric>(
457 &self,
458 name: N,
459 help: H,
460 metric: M,
461 ) -> telemetry::metrics::Registered<M>;
462
463 fn encode(&self) -> String;
465 }
466
467 pub type RateLimiter<C> = governor::RateLimiter<
472 governor::state::NotKeyed,
473 governor::state::InMemoryState,
474 C,
475 governor::middleware::NoOpMiddleware<<C as governor::clock::Clock>::Instant>,
476 >;
477
478 pub type KeyedRateLimiter<K, C> = governor::RateLimiter<
485 K,
486 governor::state::keyed::HashMapStateStore<K>,
487 C,
488 governor::middleware::NoOpMiddleware<<C as governor::clock::Clock>::Instant>,
489 >;
490
491 pub trait Clock:
497 governor::clock::Clock<Instant = SystemTime>
498 + governor::clock::ReasonablyRealtime
499 + Send
500 + Sync
501 + 'static
502 {
503 fn current(&self) -> SystemTime;
505
506 fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static;
508
509 fn sleep_until(&self, deadline: SystemTime) -> impl Future<Output = ()> + Send + 'static;
511
512 fn timeout<F, T>(
533 &self,
534 duration: Duration,
535 future: F,
536 ) -> impl Future<Output = Result<T, Error>> + Send + '_
537 where
538 F: Future<Output = T> + Send + 'static,
539 T: Send + 'static,
540 {
541 async move {
542 select! {
543 result = future => Ok(result),
544 _ = self.sleep(duration) => Err(Error::Timeout),
545 }
546 }
547 }
548 }
549
550 pub type SinkOf<N> = <<N as Network>::Listener as Listener>::Sink;
552
553 pub type StreamOf<N> = <<N as Network>::Listener as Listener>::Stream;
555
556 pub type ListenerOf<N> = <N as crate::Network>::Listener;
558
559 pub trait Network: Send + Sync + 'static {
562 type Listener: Listener;
566
567 fn bind(
569 &self,
570 socket: SocketAddr,
571 ) -> impl Future<Output = Result<Self::Listener, Error>> + Send;
572
573 fn dial(
575 &self,
576 socket: SocketAddr,
577 ) -> impl Future<Output = Result<(SinkOf<Self>, StreamOf<Self>), Error>> + Send;
578 }
579
580 pub trait Resolver: Send + Sync + 'static {
582 fn resolve(
586 &self,
587 host: &str,
588 ) -> impl Future<Output = Result<Vec<std::net::IpAddr>, Error>> + Send;
589 }
590
591 pub trait Listener: Sync + Send + 'static {
594 type Sink: Sink;
597 type Stream: Stream;
600
601 fn accept(
603 &mut self,
604 ) -> impl Future<Output = Result<(SocketAddr, Self::Sink, Self::Stream), Error>> + Send;
605
606 fn local_addr(&self) -> Result<SocketAddr, std::io::Error>;
608 }
609
610 pub trait Sink: Sync + Send + 'static {
613 fn send(
624 &mut self,
625 bufs: impl Into<IoBufs> + Send,
626 ) -> impl Future<Output = Result<(), Error>> + Send;
627 }
628
629 pub trait Stream: Sync + Send + 'static {
632 fn recv(&mut self, len: usize) -> impl Future<Output = Result<IoBufs, Error>> + Send;
645
646 fn peek(&self, max_len: usize) -> &[u8];
654 }
655
656 pub trait Storage: Send + Sync + 'static {
685 type Blob: Blob;
687
688 fn open(
691 &self,
692 partition: &str,
693 name: &[u8],
694 ) -> impl Future<Output = Result<(Self::Blob, u64), Error>> + Send {
695 async move {
696 let (blob, size, _) = self
697 .open_versioned(partition, name, DEFAULT_BLOB_VERSION..=DEFAULT_BLOB_VERSION)
698 .await?;
699 Ok((blob, size))
700 }
701 }
702
703 fn open_versioned(
727 &self,
728 partition: &str,
729 name: &[u8],
730 versions: std::ops::RangeInclusive<BlobVersion>,
731 ) -> impl Future<Output = Result<(Self::Blob, u64, BlobVersion), Error>> + Send;
732
733 fn remove(
752 &self,
753 partition: &str,
754 name: Option<&[u8]>,
755 ) -> impl Future<Output = Result<(), Error>> + Send;
756
757 fn scan(&self, partition: &str)
759 -> impl Future<Output = Result<Vec<Vec<u8>>, Error>> + Send;
760 }
761
762 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
766 pub struct ReadOptions(u8);
767
768 impl ReadOptions {
769 pub const DONT_CACHE: Self = Self(1 << 0);
775
776 #[must_use]
778 pub const fn contains(self, options: Self) -> bool {
779 self.0 & options.0 == options.0
780 }
781
782 #[must_use]
784 pub const fn without(self, options: Self) -> Self {
785 Self(self.0 & !options.0)
786 }
787 }
788
789 impl std::ops::BitOr for ReadOptions {
790 type Output = Self;
791
792 fn bitor(self, rhs: Self) -> Self::Output {
793 Self(self.0 | rhs.0)
794 }
795 }
796
797 impl std::ops::BitOrAssign for ReadOptions {
798 fn bitor_assign(&mut self, rhs: Self) {
799 self.0 |= rhs.0;
800 }
801 }
802
803 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
808 pub struct WriteOptions(u8);
809
810 impl WriteOptions {
811 pub const SYNC: Self = Self(1 << 0);
816
817 pub const DONT_CACHE: Self = Self(1 << 1);
823
824 #[must_use]
826 pub const fn contains(self, options: Self) -> bool {
827 self.0 & options.0 == options.0
828 }
829
830 #[must_use]
832 pub const fn without(self, options: Self) -> Self {
833 Self(self.0 & !options.0)
834 }
835 }
836
837 impl std::ops::BitOr for WriteOptions {
838 type Output = Self;
839
840 fn bitor(self, rhs: Self) -> Self::Output {
841 Self(self.0 | rhs.0)
842 }
843 }
844
845 impl std::ops::BitOrAssign for WriteOptions {
846 fn bitor_assign(&mut self, rhs: Self) {
847 self.0 |= rhs.0;
848 }
849 }
850
851 #[allow(clippy::len_without_is_empty)]
872 pub trait Blob: Clone + Send + Sync + 'static {
873 fn read_at_buf(
881 &self,
882 offset: u64,
883 len: usize,
884 bufs: impl Into<IoBufsMut> + Send,
885 options: ReadOptions,
886 ) -> impl Future<Output = Result<IoBufsMut, Error>> + Send;
887
888 fn read_at(
892 &self,
893 offset: u64,
894 len: usize,
895 options: ReadOptions,
896 ) -> impl Future<Output = Result<IoBufsMut, Error>> + Send;
897
898 fn write_at(
902 &self,
903 offset: u64,
904 bufs: impl Into<IoBufs> + Send,
905 options: WriteOptions,
906 ) -> impl Future<Output = Result<(), Error>> + Send;
907
908 fn resize(&self, len: u64) -> impl Future<Output = Result<(), Error>> + Send;
913
914 fn sync(&self) -> impl Future<Output = Result<(), Error>> + Send;
916
917 fn start_sync(&self) -> impl Future<Output = Handle<()>> + Send;
922 }
923
924 pub trait BufferPooler: Send + Sync + 'static {
926 fn network_buffer_pool(&self) -> &BufferPool;
928
929 fn storage_buffer_pool(&self) -> &BufferPool;
931 }
932});
933stability_scope!(BETA, cfg(feature = "external") {
934 pub trait Pacer: Clock + Send + Sync + 'static {
936 fn pace<'a, F, T>(
956 &'a self,
957 latency: Duration,
958 future: F,
959 ) -> impl Future<Output = T> + Send + 'a
960 where
961 F: Future<Output = T> + Send + 'a,
962 T: Send + 'a;
963 }
964
965 pub trait FutureExt: Future + Send + Sized {
970 fn pace<'a, E>(
972 self,
973 pacer: &'a E,
974 latency: Duration,
975 ) -> impl Future<Output = Self::Output> + Send + 'a
976 where
977 E: Pacer + 'a,
978 Self: Send + 'a,
979 Self::Output: Send + 'a,
980 {
981 pacer.pace(latency, self)
982 }
983 }
984
985 impl<F> FutureExt for F where F: Future + Send {}
986});
987
988#[cfg(test)]
989mod tests {
990 use super::*;
991 use crate::telemetry::metrics::{
992 EncodeLabelKey, EncodeLabelSetTrait as EncodeLabelSet,
993 EncodeLabelValueTrait as EncodeLabelValue, LabelSetEncoder, count_running_tasks,
994 raw::{Counter, Family},
995 };
996 use commonware_macros::select;
997 use commonware_parallel::Strategy as _;
998 use commonware_utils::{
999 NZU32, NZUsize, SystemTimeExt,
1000 channel::{mpsc, oneshot},
1001 futures::Pool as FuturesPool,
1002 sync::Mutex,
1003 };
1004 use futures::{
1005 FutureExt,
1006 future::{pending, ready},
1007 join, pin_mut,
1008 };
1009 use rstest::rstest;
1010 use std::{
1011 pin::Pin,
1012 sync::{
1013 Arc,
1014 atomic::{AtomicU32, Ordering},
1015 },
1016 task::{Context as TContext, Poll, Waker},
1017 };
1018 use utils::reschedule;
1019
1020 #[test]
1021 fn test_blob_version() {
1022 let version = BlobVersion::new(7);
1023 assert_eq!(version.get(), 7);
1024 assert_eq!(version.to_string(), "7");
1025 assert_eq!(format!("{version:?}"), "7");
1026 assert_eq!(BlobVersion::default(), DEFAULT_BLOB_VERSION);
1027 assert!((BlobVersion::new(3)..=BlobVersion::new(7)).contains(&version));
1028 }
1029
1030 #[test]
1031 fn test_read_options_compose() {
1032 let options = ReadOptions::default() | ReadOptions::DONT_CACHE;
1034 let mut assigned = ReadOptions::default();
1035 assigned |= ReadOptions::DONT_CACHE;
1036 assert!(options.contains(ReadOptions::DONT_CACHE));
1037 assert_eq!(assigned, options);
1038 assert_eq!(
1039 options.without(ReadOptions::DONT_CACHE),
1040 ReadOptions::default()
1041 );
1042 assert!(!ReadOptions::default().contains(ReadOptions::DONT_CACHE));
1043 }
1044
1045 #[test]
1046 fn test_write_options_compose() {
1047 let options = WriteOptions::SYNC | WriteOptions::DONT_CACHE;
1048 let mut assigned = WriteOptions::SYNC;
1049 assigned |= WriteOptions::DONT_CACHE;
1050 assert!(options.contains(WriteOptions::SYNC));
1051 assert!(options.contains(WriteOptions::DONT_CACHE));
1052 assert_eq!(assigned, options);
1053 assert_eq!(
1054 options.without(WriteOptions::SYNC),
1055 WriteOptions::DONT_CACHE
1056 );
1057 let default = WriteOptions::default();
1058 assert!(!default.contains(WriteOptions::SYNC));
1059 assert!(!default.contains(WriteOptions::DONT_CACHE));
1060 }
1061
1062 #[rstest]
1063 #[case::deterministic(deterministic::Runner::default())]
1064 #[case::tokio(tokio::Runner::default())]
1065 fn test_error_future<R: Runner>(#[case] runner: R) {
1066 #[allow(clippy::unused_async)]
1067 async fn error_future() -> Result<&'static str, &'static str> {
1068 Err("An error occurred")
1069 }
1070 let result = runner.start(|_| error_future());
1071 assert_eq!(result, Err("An error occurred"));
1072 }
1073
1074 #[rstest]
1075 #[case::deterministic(deterministic::Runner::default())]
1076 #[case::tokio(tokio::Runner::default())]
1077 fn test_handle_can_use_futures_pool<R: Runner>(#[case] runner: R) {
1078 runner.start(|_| async move {
1079 let mut pool = FuturesPool::<Result<(), Error>>::default();
1080 pool.push(Handle::ready(Ok(())));
1081 assert!(pool.next_completed().await.is_ok());
1082 });
1083 }
1084
1085 #[rstest]
1086 #[case::deterministic(deterministic::Runner::default())]
1087 #[case::tokio(tokio::Runner::default())]
1088 fn test_clock_sleep<R: Runner>(#[case] runner: R)
1089 where
1090 R::Context: Spawner + Clock,
1091 {
1092 runner.start(|context| async move {
1093 let start = context.current();
1095 let sleep_duration = Duration::from_millis(10);
1096 context.sleep(sleep_duration).await;
1097
1098 let end = context.current();
1100 assert!(end.duration_since(start).unwrap() >= sleep_duration);
1101 });
1102 }
1103
1104 #[rstest]
1105 #[case::deterministic(deterministic::Runner::default())]
1106 #[case::tokio(tokio::Runner::default())]
1107 fn test_clock_sleep_until<R: Runner>(#[case] runner: R)
1108 where
1109 R::Context: Spawner + Clock + Metrics,
1110 {
1111 runner.start(|context| async move {
1112 let now = context.current();
1114 context.sleep_until(now + Duration::from_millis(100)).await;
1115
1116 let elapsed = now.elapsed().unwrap();
1118 assert!(elapsed >= Duration::from_millis(100));
1119 });
1120 }
1121
1122 #[rstest]
1123 #[case::deterministic(deterministic::Runner::default())]
1124 #[case::tokio(tokio::Runner::default())]
1125 fn test_clock_sleep_until_far_future<R: Runner>(#[case] runner: R)
1126 where
1127 R::Context: Spawner + Clock,
1128 {
1129 runner.start(|context| async move {
1130 let sleep = context.sleep_until(SystemTime::limit());
1131 let result = context.timeout(Duration::from_millis(1), sleep).await;
1132 assert!(matches!(result, Err(Error::Timeout)));
1133 });
1134 }
1135
1136 #[rstest]
1137 #[case::deterministic(deterministic::Runner::default())]
1138 #[case::tokio(tokio::Runner::default())]
1139 fn test_clock_timeout<R: Runner>(#[case] runner: R)
1140 where
1141 R::Context: Spawner + Clock,
1142 {
1143 runner.start(|context| async move {
1144 let result = context
1146 .timeout(Duration::from_millis(100), async { "success" })
1147 .await;
1148 assert_eq!(result.unwrap(), "success");
1149
1150 let result = context
1152 .timeout(Duration::from_millis(50), pending::<()>())
1153 .await;
1154 assert!(matches!(result, Err(Error::Timeout)));
1155
1156 let result = context
1158 .timeout(
1159 Duration::from_millis(100),
1160 context.sleep(Duration::from_millis(50)),
1161 )
1162 .await;
1163 assert!(result.is_ok());
1164 });
1165 }
1166
1167 #[rstest]
1168 #[case::deterministic(deterministic::Runner::default())]
1169 #[case::tokio(tokio::Runner::default())]
1170 fn test_root_finishes<R: Runner>(#[case] runner: R)
1171 where
1172 R::Context: Spawner,
1173 {
1174 runner.start(|context| async move {
1175 context.spawn(|_| async move {
1176 loop {
1177 reschedule().await;
1178 }
1179 });
1180 });
1181 }
1182
1183 #[rstest]
1184 #[case::deterministic(deterministic::Runner::default())]
1185 #[case::tokio(tokio::Runner::default())]
1186 fn test_spawn_after_abort<R>(#[case] runner: R)
1187 where
1188 R: Runner,
1189 R::Context: Spawner,
1190 {
1191 runner.start(|context| async move {
1192 let child = context.child("child");
1194
1195 let parent_handle = context.spawn(move |_| async move {
1197 pending::<()>().await;
1198 });
1199 parent_handle.abort();
1200
1201 let child_handle = child.spawn(move |_| async move {
1203 pending::<()>().await;
1204 });
1205 assert!(matches!(child_handle.await, Err(Error::Closed)));
1206 });
1207 }
1208
1209 #[rstest]
1210 #[case::deterministic(deterministic::Runner::default())]
1211 #[case::tokio(tokio::Runner::default())]
1212 fn test_spawn_abort<R: Runner>(
1213 #[case] runner: R,
1214 #[values(
1215 Execution::Shared(false),
1216 Execution::Shared(true),
1217 Execution::Dedicated
1218 )]
1219 execution: Execution,
1220 ) where
1221 R::Context: Spawner,
1222 {
1223 runner.start(|context| async move {
1224 let context = match execution {
1225 Execution::Dedicated => context.dedicated(),
1226 Execution::Shared(blocking) => context.shared(blocking),
1227 };
1228
1229 let handle = context.spawn(|_| async move {
1230 loop {
1231 reschedule().await;
1232 }
1233 });
1234 handle.abort();
1235 assert!(matches!(handle.await, Err(Error::Closed)));
1236 });
1237 }
1238
1239 #[rstest]
1240 #[case::deterministic(deterministic::Runner::default())]
1241 #[case::deterministic_caught(deterministic::Runner::new(
1242 deterministic::Config::default().with_catch_panics(true)
1243 ))]
1244 #[case::tokio(tokio::Runner::default())]
1245 #[case::tokio_caught(tokio::Runner::new(tokio::Config::default().with_catch_panics(true)))]
1246 #[should_panic(expected = "blah")]
1247 fn test_panic_aborts_root<R: Runner>(#[case] runner: R) {
1248 let result: Result<(), Error> = runner.start(|_| async move {
1249 panic!("blah");
1250 });
1251 result.unwrap_err();
1252 }
1253
1254 #[rstest]
1255 #[case::deterministic(deterministic::Runner::default())]
1256 #[case::tokio(tokio::Runner::default())]
1257 #[should_panic(expected = "blah")]
1258 fn test_panic_aborts_spawn<R: Runner>(#[case] runner: R)
1259 where
1260 R::Context: Spawner + Clock,
1261 {
1262 runner.start(|context| async move {
1263 context.child("panic").spawn(|_| async move {
1264 panic!("blah");
1265 });
1266
1267 loop {
1269 context.sleep(Duration::from_millis(100)).await;
1270 }
1271 });
1272 }
1273
1274 #[rstest]
1275 #[case::deterministic(deterministic::Runner::new(
1276 deterministic::Config::default().with_catch_panics(true)
1277 ))]
1278 #[case::tokio(tokio::Runner::new(tokio::Config::default().with_catch_panics(true)))]
1279 fn test_panic_aborts_spawn_caught<R: Runner>(#[case] runner: R)
1280 where
1281 R::Context: Spawner + Clock,
1282 {
1283 let result: Result<(), Error> = runner.start(|context| async move {
1284 let result = context.child("panic").spawn(|_| async move {
1285 panic!("blah");
1286 });
1287 result.await
1288 });
1289 assert!(matches!(result, Err(Error::Exited)));
1290 }
1291
1292 #[rstest]
1293 #[case::deterministic(deterministic::Runner::default())]
1294 #[case::tokio(tokio::Runner::default())]
1295 #[should_panic(expected = "boom")]
1296 fn test_multiple_panics<R: Runner>(#[case] runner: R)
1297 where
1298 R::Context: Spawner + Clock,
1299 {
1300 runner.start(|context| async move {
1301 context.child("panic").spawn(|_| async move {
1302 panic!("boom 1");
1303 });
1304 context.child("panic").spawn(|_| async move {
1305 panic!("boom 2");
1306 });
1307 context.child("panic").spawn(|_| async move {
1308 panic!("boom 3");
1309 });
1310
1311 loop {
1313 context.sleep(Duration::from_millis(100)).await;
1314 }
1315 });
1316 }
1317
1318 #[rstest]
1319 #[case::deterministic(deterministic::Runner::new(
1320 deterministic::Config::default().with_catch_panics(true)
1321 ))]
1322 #[case::tokio(tokio::Runner::new(tokio::Config::default().with_catch_panics(true)))]
1323 fn test_multiple_panics_caught<R: Runner>(#[case] runner: R)
1324 where
1325 R::Context: Spawner + Clock,
1326 {
1327 let (res1, res2, res3) = runner.start(|context| async move {
1328 let handle1 = context.child("panic").spawn(|_| async move {
1329 panic!("boom 1");
1330 });
1331 let handle2 = context.child("panic").spawn(|_| async move {
1332 panic!("boom 2");
1333 });
1334 let handle3 = context.child("panic").spawn(|_| async move {
1335 panic!("boom 3");
1336 });
1337
1338 join!(handle1, handle2, handle3)
1339 });
1340 assert!(matches!(res1, Err(Error::Exited)));
1341 assert!(matches!(res2, Err(Error::Exited)));
1342 assert!(matches!(res3, Err(Error::Exited)));
1343 }
1344
1345 #[rstest]
1346 #[case::deterministic(deterministic::Runner::default())]
1347 #[case::tokio(tokio::Runner::default())]
1348 fn test_select<R: Runner>(#[case] runner: R) {
1349 runner.start(|_| async move {
1350 let output = Mutex::new(0);
1352 select! {
1353 v1 = ready(1) => {
1354 *output.lock() = v1;
1355 },
1356 v2 = ready(2) => {
1357 *output.lock() = v2;
1358 },
1359 };
1360 assert_eq!(*output.lock(), 1);
1361
1362 select! {
1364 v1 = std::future::pending::<i32>() => {
1365 *output.lock() = v1;
1366 },
1367 v2 = ready(2) => {
1368 *output.lock() = v2;
1369 },
1370 };
1371 assert_eq!(*output.lock(), 2);
1372 });
1373 }
1374
1375 #[rstest]
1377 #[case::deterministic(deterministic::Runner::default())]
1378 #[case::tokio(tokio::Runner::default())]
1379 fn test_select_loop<R: Runner>(#[case] runner: R)
1380 where
1381 R::Context: Clock,
1382 {
1383 runner.start(|context| async move {
1384 let (sender, mut receiver) = mpsc::unbounded_channel();
1386 for _ in 0..2 {
1387 select! {
1388 v = receiver.recv() => {
1389 panic!("unexpected value: {v:?}");
1390 },
1391 _ = context.sleep(Duration::from_millis(100)) => {
1392 continue;
1393 },
1394 };
1395 }
1396
1397 sender.send(0).unwrap();
1399 sender.send(1).unwrap();
1400
1401 select! {
1403 _ = async {} => {
1404 },
1406 v = receiver.recv() => {
1407 panic!("unexpected value: {v:?}");
1408 },
1409 };
1410
1411 for i in 0..2 {
1413 select! {
1414 _ = context.sleep(Duration::from_millis(100)) => {
1415 panic!("timeout");
1416 },
1417 v = receiver.recv() => {
1418 assert_eq!(v.unwrap(), i);
1419 },
1420 };
1421 }
1422 });
1423 }
1424
1425 #[rstest]
1426 #[case::deterministic(deterministic::Runner::default())]
1427 #[case::tokio(tokio::Runner::default())]
1428 fn test_storage_operations<R: Runner>(#[case] runner: R)
1429 where
1430 R::Context: Storage,
1431 {
1432 runner.start(|context| async move {
1433 let partition = "test_partition";
1434 let name = b"test_blob";
1435
1436 let (blob, size) = context
1438 .open(partition, name)
1439 .await
1440 .expect("Failed to open blob");
1441 assert_eq!(size, 0, "new blob should have size 0");
1442
1443 let data = b"Hello, Storage!";
1445 blob.write_at(0, data, WriteOptions::default())
1446 .await
1447 .expect("Failed to write to blob");
1448
1449 blob.sync().await.expect("Failed to sync blob");
1451
1452 let read = blob
1454 .read_at(0, data.len(), ReadOptions::default())
1455 .await
1456 .expect("Failed to read from blob");
1457 assert_eq!(read.coalesce(), data);
1458
1459 blob.sync().await.expect("Failed to sync blob");
1461
1462 let blobs = context
1464 .scan(partition)
1465 .await
1466 .expect("Failed to scan partition");
1467 assert!(blobs.contains(&name.to_vec()));
1468
1469 let (blob, len) = context
1471 .open(partition, name)
1472 .await
1473 .expect("Failed to reopen blob");
1474 assert_eq!(len, data.len() as u64);
1475
1476 let read = blob
1478 .read_at(7, 7, ReadOptions::default())
1479 .await
1480 .expect("Failed to read data");
1481 assert_eq!(read.coalesce(), b"Storage");
1482
1483 blob.sync().await.expect("Failed to sync blob");
1485
1486 context
1488 .remove(partition, Some(name))
1489 .await
1490 .expect("Failed to remove blob");
1491
1492 let blobs = context
1494 .scan(partition)
1495 .await
1496 .expect("Failed to scan partition");
1497 assert!(!blobs.contains(&name.to_vec()));
1498
1499 context
1501 .remove(partition, None)
1502 .await
1503 .expect("Failed to remove partition");
1504
1505 let result = context.scan(partition).await;
1507 assert!(matches!(result, Err(Error::PartitionMissing(_))));
1508 });
1509 }
1510
1511 #[rstest]
1512 #[case::deterministic(deterministic::Runner::default())]
1513 #[case::tokio(tokio::Runner::default())]
1514 fn test_blob_read_write<R: Runner>(#[case] runner: R)
1515 where
1516 R::Context: Storage,
1517 {
1518 runner.start(|context| async move {
1519 let partition = "test_partition";
1520 let name = b"test_blob_rw";
1521
1522 let (blob, _) = context
1524 .open(partition, name)
1525 .await
1526 .expect("Failed to open blob");
1527
1528 let data1 = b"Hello";
1530 let data2 = b"World";
1531 blob.write_at(0, data1, WriteOptions::default())
1532 .await
1533 .expect("Failed to write data1");
1534 blob.write_at(5, data2, WriteOptions::default())
1535 .await
1536 .expect("Failed to write data2");
1537
1538 let read = blob
1540 .read_at(0, 10, ReadOptions::default())
1541 .await
1542 .expect("Failed to read data");
1543 let read = read.coalesce();
1544 assert_eq!(&read.as_ref()[..5], data1);
1545 assert_eq!(&read.as_ref()[5..], data2);
1546
1547 let result = blob.read_at(10, 10, ReadOptions::default()).await;
1549 assert!(result.is_err());
1550
1551 let data3 = b"Store";
1553 blob.write_at(5, data3, WriteOptions::default())
1554 .await
1555 .expect("Failed to write data3");
1556
1557 let read = blob
1559 .read_at(0, 10, ReadOptions::default())
1560 .await
1561 .expect("Failed to read data");
1562 let read = read.coalesce();
1563 assert_eq!(&read.as_ref()[..5], data1);
1564 assert_eq!(&read.as_ref()[5..], data3);
1565
1566 let result = blob.read_at(10, 10, ReadOptions::default()).await;
1568 assert!(result.is_err());
1569 });
1570 }
1571
1572 #[rstest]
1573 #[case::deterministic(deterministic::Runner::default())]
1574 #[case::tokio(tokio::Runner::default())]
1575 fn test_blob_resize<R: Runner>(#[case] runner: R)
1576 where
1577 R::Context: Storage,
1578 {
1579 runner.start(|context| async move {
1580 let partition = "test_partition_resize";
1581 let name = b"test_blob_resize";
1582
1583 let (blob, _) = context
1585 .open(partition, name)
1586 .await
1587 .expect("Failed to open blob");
1588
1589 let data = b"some data";
1590 blob.write_at(0, data.to_vec(), WriteOptions::default())
1591 .await
1592 .expect("Failed to write");
1593 blob.sync().await.expect("Failed to sync after write");
1594
1595 let (blob, len) = context.open(partition, name).await.unwrap();
1597 assert_eq!(len, data.len() as u64);
1598
1599 let new_len = (data.len() as u64) * 2;
1601 blob.resize(new_len)
1602 .await
1603 .expect("Failed to resize to extend");
1604 blob.sync().await.expect("Failed to sync after resize");
1605
1606 let (blob, len) = context.open(partition, name).await.unwrap();
1608 assert_eq!(len, new_len);
1609
1610 let read_buf = blob
1612 .read_at(0, data.len(), ReadOptions::default())
1613 .await
1614 .unwrap();
1615 assert_eq!(read_buf.coalesce(), data);
1616
1617 let extended_part = blob
1619 .read_at(data.len() as u64, data.len(), ReadOptions::default())
1620 .await
1621 .unwrap();
1622 assert_eq!(extended_part.coalesce(), vec![0; data.len()].as_slice());
1623
1624 blob.resize(data.len() as u64).await.unwrap();
1626 blob.sync().await.unwrap();
1627
1628 let (blob, size) = context.open(partition, name).await.unwrap();
1630 assert_eq!(size, data.len() as u64);
1631
1632 let read_buf = blob
1634 .read_at(0, data.len(), ReadOptions::default())
1635 .await
1636 .unwrap();
1637 assert_eq!(read_buf.coalesce(), data);
1638 blob.sync().await.unwrap();
1639 });
1640 }
1641
1642 #[rstest]
1643 #[case::deterministic(deterministic::Runner::default())]
1644 #[case::tokio(tokio::Runner::default())]
1645 fn test_many_partition_read_write<R: Runner>(#[case] runner: R)
1646 where
1647 R::Context: Storage,
1648 {
1649 runner.start(|context| async move {
1650 let partitions = ["partition1", "partition2", "partition3"];
1651 let name = b"test_blob_rw";
1652 let data1 = b"Hello";
1653 let data2 = b"World";
1654
1655 for (additional, partition) in partitions.iter().enumerate() {
1656 let (blob, _) = context
1658 .open(partition, name)
1659 .await
1660 .expect("Failed to open blob");
1661
1662 blob.write_at(0, data1, WriteOptions::default())
1664 .await
1665 .expect("Failed to write data1");
1666 blob.write_at(5 + additional as u64, data2, WriteOptions::default())
1667 .await
1668 .expect("Failed to write data2");
1669
1670 blob.sync().await.expect("Failed to sync blob");
1672 }
1673
1674 for (additional, partition) in partitions.iter().enumerate() {
1675 let (blob, len) = context
1677 .open(partition, name)
1678 .await
1679 .expect("Failed to open blob");
1680 assert_eq!(len, (data1.len() + data2.len() + additional) as u64);
1681
1682 let read = blob
1684 .read_at(0, 10 + additional, ReadOptions::default())
1685 .await
1686 .expect("Failed to read data");
1687 let read = read.coalesce();
1688 assert_eq!(&read.as_ref()[..5], b"Hello");
1689 assert_eq!(&read.as_ref()[5 + additional..], b"World");
1690 }
1691 });
1692 }
1693
1694 #[rstest]
1695 #[case::deterministic(deterministic::Runner::default())]
1696 #[case::tokio(tokio::Runner::default())]
1697 fn test_blob_read_past_length<R: Runner>(#[case] runner: R)
1698 where
1699 R::Context: Storage,
1700 {
1701 runner.start(|context| async move {
1702 let partition = "test_partition";
1703 let name = b"test_blob_rw";
1704
1705 let (blob, _) = context
1707 .open(partition, name)
1708 .await
1709 .expect("Failed to open blob");
1710
1711 let result = blob.read_at(0, 10, ReadOptions::default()).await;
1713 assert!(result.is_err());
1714
1715 let data = b"Hello, Storage!".to_vec();
1717 blob.write_at(0, data, WriteOptions::default())
1718 .await
1719 .expect("Failed to write to blob");
1720
1721 let result = blob.read_at(0, 20, ReadOptions::default()).await;
1723 assert!(result.is_err());
1724 })
1725 }
1726
1727 #[rstest]
1728 #[case::deterministic(deterministic::Runner::default())]
1729 #[case::tokio(tokio::Runner::default())]
1730 fn test_blob_clone_and_concurrent_read<R: Runner>(#[case] runner: R)
1731 where
1732 R::Context: Spawner + Storage + Metrics,
1733 {
1734 runner.start(|context| async move {
1735 let partition = "test_partition";
1736 let name = b"test_blob_rw";
1737
1738 let (blob, _) = context
1740 .open(partition, name)
1741 .await
1742 .expect("Failed to open blob");
1743
1744 let data = b"Hello, Storage!";
1746 blob.write_at(0, data, WriteOptions::default())
1747 .await
1748 .expect("Failed to write to blob");
1749
1750 blob.sync().await.expect("Failed to sync blob");
1752
1753 let check1 = context.child("check1").spawn({
1755 let blob = blob.clone();
1756 let data_len = data.len();
1757 move |_| async move {
1758 let read = blob
1759 .read_at(0, data_len, ReadOptions::default())
1760 .await
1761 .expect("Failed to read from blob");
1762 assert_eq!(read.coalesce(), data);
1763 }
1764 });
1765 let check2 = context.child("check2").spawn({
1766 let blob = blob.clone();
1767 let data_len = data.len();
1768 move |_| async move {
1769 let read = blob
1770 .read_at(0, data_len, ReadOptions::default())
1771 .await
1772 .expect("Failed to read from blob");
1773 assert_eq!(read.coalesce(), data);
1774 }
1775 });
1776
1777 let result = join!(check1, check2);
1779 assert!(result.0.is_ok());
1780 assert!(result.1.is_ok());
1781
1782 let read = blob
1784 .read_at(0, data.len(), ReadOptions::default())
1785 .await
1786 .expect("Failed to read from blob");
1787 assert_eq!(read.coalesce(), data);
1788
1789 drop(blob);
1791
1792 let buffer = context.encode();
1794 assert!(buffer.contains("open_blobs 0"));
1795 });
1796 }
1797
1798 #[rstest]
1799 #[case::deterministic(deterministic::Runner::default())]
1800 #[case::tokio(tokio::Runner::default())]
1801 fn test_shutdown<R: Runner>(#[case] runner: R)
1802 where
1803 R::Context: Spawner + Metrics + Clock,
1804 {
1805 let kill = 9;
1806 runner.start(|context| async move {
1807 let before = context.child("before").spawn(move |context| async move {
1809 let mut signal = context.stopped();
1810 let value = (&mut signal).await.unwrap();
1811 assert_eq!(value, kill);
1812 drop(signal);
1813 });
1814
1815 let result = context.child("stop").stop(kill, None).await;
1817 assert!(result.is_ok());
1818
1819 let after = context.child("after").spawn(move |context| async move {
1821 let value = context.stopped().await.unwrap();
1823 assert_eq!(value, kill);
1824 });
1825
1826 let result = join!(before, after);
1828 assert!(result.0.is_ok());
1829 assert!(result.1.is_ok());
1830 });
1831 }
1832
1833 #[rstest]
1834 #[case::deterministic(deterministic::Runner::default())]
1835 #[case::tokio(tokio::Runner::default())]
1836 fn test_shutdown_multiple_signals<R: Runner>(#[case] runner: R)
1837 where
1838 R::Context: Spawner + Metrics + Clock,
1839 {
1840 let kill = 42;
1841 runner.start(|context| async move {
1842 let (started_tx, mut started_rx) = mpsc::channel(3);
1843 let counter = Arc::new(AtomicU32::new(0));
1844
1845 let task = |context: R::Context, cleanup_duration: Duration| {
1848 let counter = counter.clone();
1849 let started_tx = started_tx.clone();
1850 context.spawn(move |context| async move {
1851 let mut signal = context.stopped();
1853 started_tx.send(()).await.unwrap();
1854
1855 let value = (&mut signal).await.unwrap();
1857 assert_eq!(value, kill);
1858 context.sleep(cleanup_duration).await;
1859 counter.fetch_add(1, Ordering::SeqCst);
1860
1861 drop(signal);
1863 })
1864 };
1865
1866 let task1 = task(context.child("cleanup"), Duration::from_millis(10));
1867 let task2 = task(context.child("cleanup"), Duration::from_millis(20));
1868 let task3 = task(context.child("cleanup"), Duration::from_millis(30));
1869
1870 for _ in 0..3 {
1872 started_rx.recv().await.unwrap();
1873 }
1874
1875 context.stop(kill, None).await.unwrap();
1877 assert_eq!(counter.load(Ordering::SeqCst), 3);
1878
1879 let result = join!(task1, task2, task3);
1881 assert!(result.0.is_ok());
1882 assert!(result.1.is_ok());
1883 assert!(result.2.is_ok());
1884 });
1885 }
1886
1887 #[rstest]
1888 #[case::deterministic(deterministic::Runner::default())]
1889 #[case::tokio(tokio::Runner::default())]
1890 fn test_shutdown_timeout<R: Runner>(#[case] runner: R)
1891 where
1892 R::Context: Spawner + Metrics + Clock,
1893 {
1894 let kill = 42;
1895 runner.start(|context| async move {
1896 let (started_tx, started_rx) = oneshot::channel();
1898
1899 context.child("signal").spawn(move |context| async move {
1901 let signal = context.stopped();
1902 started_tx.send(()).unwrap();
1903 pending::<()>().await;
1904 signal.await.unwrap();
1905 });
1906
1907 started_rx.await.unwrap();
1909 let result = context.stop(kill, Some(Duration::from_millis(100))).await;
1910
1911 assert!(matches!(result, Err(Error::Timeout)));
1913 });
1914 }
1915
1916 #[rstest]
1917 #[case::deterministic(deterministic::Runner::default())]
1918 #[case::tokio(tokio::Runner::default())]
1919 fn test_shutdown_multiple_stop_calls<R: Runner>(#[case] runner: R)
1920 where
1921 R::Context: Spawner + Metrics + Clock,
1922 {
1923 let kill1 = 42;
1924 let kill2 = 43;
1925
1926 runner.start(|context| async move {
1927 let (started_tx, started_rx) = oneshot::channel();
1928 let counter = Arc::new(AtomicU32::new(0));
1929
1930 let task = context.child("blocking_task").spawn({
1932 let counter = counter.clone();
1933 move |context| async move {
1934 let mut signal = context.stopped();
1936 started_tx.send(()).unwrap();
1937
1938 let value = (&mut signal).await.unwrap();
1940 assert_eq!(value, kill1);
1941 context.sleep(Duration::from_millis(50)).await;
1942
1943 counter.fetch_add(1, Ordering::SeqCst);
1945 drop(signal);
1946 }
1947 });
1948
1949 started_rx.await.unwrap();
1951
1952 let stop_task1 = context.child("stop").stop(kill1, None);
1955 pin_mut!(stop_task1);
1956 let stop_task2 = context.child("stop").stop(kill2, None);
1957 pin_mut!(stop_task2);
1958
1959 assert!(stop_task1.as_mut().now_or_never().is_none());
1961 assert!(stop_task2.as_mut().now_or_never().is_none());
1962
1963 assert!(stop_task1.await.is_ok());
1965 assert!(stop_task2.await.is_ok());
1966
1967 let sig = context.stopped().await;
1969 assert_eq!(sig.unwrap(), kill1);
1970
1971 let result = task.await;
1973 assert!(result.is_ok());
1974 assert_eq!(counter.load(Ordering::SeqCst), 1);
1975
1976 assert!(context.stop(kill2, None).now_or_never().unwrap().is_ok());
1978 });
1979 }
1980
1981 #[rstest]
1982 #[case::deterministic(deterministic::Runner::default())]
1983 #[case::tokio(tokio::Runner::default())]
1984 fn test_unfulfilled_shutdown<R: Runner>(#[case] runner: R)
1985 where
1986 R::Context: Spawner + Metrics,
1987 {
1988 runner.start(|context| async move {
1989 context.child("before").spawn(move |context| async move {
1991 let mut signal = context.stopped();
1992 let value = (&mut signal).await.unwrap();
1993
1994 assert_eq!(value, 42);
1996 drop(signal);
1997 });
1998
1999 reschedule().await;
2001 });
2002 }
2003
2004 #[rstest]
2005 #[case::deterministic(deterministic::Runner::default())]
2006 #[case::tokio(tokio::Runner::default())]
2007 fn test_spawn_dedicated<R: Runner>(#[case] runner: R)
2008 where
2009 R::Context: Spawner,
2010 {
2011 runner.start(|context| async move {
2012 let handle = context.dedicated().spawn(|_| async move { 42 });
2013 assert!(matches!(handle.await, Ok(42)));
2014 });
2015 }
2016
2017 #[rstest]
2018 #[case::deterministic(deterministic::Runner::default())]
2019 #[case::tokio(tokio::Runner::default())]
2020 fn test_spawn<R: Runner>(#[case] runner: R)
2021 where
2022 R::Context: Spawner + Clock,
2023 {
2024 runner.start(|context| async move {
2025 let child_handle = Arc::new(Mutex::new(None));
2026 let child_handle2 = child_handle.clone();
2027
2028 let (parent_initialized_tx, parent_initialized_rx) = oneshot::channel();
2029 let (parent_complete_tx, parent_complete_rx) = oneshot::channel();
2030 let parent_handle = context.spawn(move |context| async move {
2031 let handle = context.spawn(|_| async {});
2033
2034 *child_handle2.lock() = Some(handle);
2036
2037 parent_initialized_tx.send(()).unwrap();
2038
2039 parent_complete_rx.await.unwrap();
2041 });
2042
2043 parent_initialized_rx.await.unwrap();
2045
2046 let child_handle = child_handle.lock().take().unwrap();
2048 assert!(child_handle.await.is_ok());
2049
2050 parent_complete_tx.send(()).unwrap();
2052
2053 assert!(parent_handle.await.is_ok());
2055 });
2056 }
2057
2058 #[rstest]
2059 #[case::deterministic(deterministic::Runner::default())]
2060 #[case::tokio(tokio::Runner::default())]
2061 fn test_spawn_abort_on_parent_abort<R: Runner>(#[case] runner: R)
2062 where
2063 R::Context: Spawner + Clock,
2064 {
2065 runner.start(|context| async move {
2066 let child_handle = Arc::new(Mutex::new(None));
2067 let child_handle2 = child_handle.clone();
2068
2069 let (parent_initialized_tx, parent_initialized_rx) = oneshot::channel();
2070 let parent_handle = context.spawn(move |context| async move {
2071 let handle = context.spawn(|_| pending::<()>());
2073
2074 *child_handle2.lock() = Some(handle);
2076
2077 parent_initialized_tx.send(()).unwrap();
2078
2079 pending::<()>().await
2081 });
2082
2083 parent_initialized_rx.await.unwrap();
2085
2086 parent_handle.abort();
2088 assert!(matches!(parent_handle.await, Err(Error::Closed)));
2089
2090 let child_handle = child_handle.lock().take().unwrap();
2092 assert!(matches!(child_handle.await, Err(Error::Closed)));
2093 });
2094 }
2095
2096 #[rstest]
2097 #[case::deterministic(deterministic::Runner::default())]
2098 #[case::tokio(tokio::Runner::default())]
2099 fn test_spawn_abort_on_parent_completion<R: Runner>(#[case] runner: R)
2100 where
2101 R::Context: Spawner + Clock,
2102 {
2103 runner.start(|context| async move {
2104 let child_handle = Arc::new(Mutex::new(None));
2105 let child_handle2 = child_handle.clone();
2106
2107 let (parent_complete_tx, parent_complete_rx) = oneshot::channel();
2108 let parent_handle = context.spawn(move |context| async move {
2109 let handle = context.spawn(|_| pending::<()>());
2111
2112 *child_handle2.lock() = Some(handle);
2114
2115 parent_complete_rx.await.unwrap();
2117 });
2118
2119 parent_complete_tx.send(()).unwrap();
2121
2122 assert!(parent_handle.await.is_ok());
2124
2125 let child_handle = child_handle.lock().take().unwrap();
2127 assert!(matches!(child_handle.await, Err(Error::Closed)));
2128 });
2129 }
2130
2131 #[rstest]
2132 #[case::deterministic(deterministic::Runner::default())]
2133 #[case::tokio(tokio::Runner::default())]
2134 fn test_spawn_cascading_abort<R: Runner>(#[case] runner: R)
2135 where
2136 R::Context: Spawner + Clock,
2137 {
2138 runner.start(|context| async move {
2139 let c0 = context.child("c0");
2149 let g0 = c0.child("g0");
2150 let g1 = c0.child("g1");
2151 let c1 = context.child("c1");
2152 let g2 = c1.child("g2");
2153 let g3 = c1.child("g3");
2154 let c2 = context.child("c2");
2155 let g4 = c2.child("g4");
2156 let g5 = c2.child("g5");
2157
2158 let handles = Arc::new(Mutex::new(Vec::new()));
2160 let (initialized_tx, mut initialized_rx) = mpsc::channel(9);
2161 let root_task = context.spawn({
2162 let handles = handles.clone();
2163 move |_| async move {
2164 for (context, grandchildren) in [(c0, [g0, g1]), (c1, [g2, g3]), (c2, [g4, g5])]
2165 {
2166 let handle = context.spawn({
2167 let handles = handles.clone();
2168 let initialized_tx = initialized_tx.clone();
2169 move |_| async move {
2170 for grandchild in grandchildren {
2171 let handle = grandchild.spawn(|_| async {
2172 pending::<()>().await;
2173 });
2174 handles.lock().push(handle);
2175 initialized_tx.send(()).await.unwrap();
2176 }
2177
2178 pending::<()>().await;
2179 }
2180 });
2181 handles.lock().push(handle);
2182 initialized_tx.send(()).await.unwrap();
2183 }
2184
2185 pending::<()>().await;
2186 }
2187 });
2188
2189 for _ in 0..9 {
2191 initialized_rx.recv().await.unwrap();
2192 }
2193
2194 assert_eq!(handles.lock().len(), 9);
2196
2197 root_task.abort();
2199 assert!(matches!(root_task.await, Err(Error::Closed)));
2200
2201 let handles = handles.lock().drain(..).collect::<Vec<_>>();
2203 for handle in handles {
2204 assert!(matches!(handle.await, Err(Error::Closed)));
2205 }
2206 });
2207 }
2208
2209 #[rstest]
2210 #[case::deterministic(deterministic::Runner::default())]
2211 #[case::tokio(tokio::Runner::default())]
2212 fn test_child_survives_sibling_completion<R: Runner>(#[case] runner: R)
2213 where
2214 R::Context: Spawner + Clock,
2215 {
2216 runner.start(|context| async move {
2217 let (child_started_tx, child_started_rx) = oneshot::channel();
2218 let (child_complete_tx, child_complete_rx) = oneshot::channel();
2219 let (child_handle_tx, child_handle_rx) = oneshot::channel();
2220 let (sibling_started_tx, sibling_started_rx) = oneshot::channel();
2221 let (sibling_complete_tx, sibling_complete_rx) = oneshot::channel();
2222 let (sibling_handle_tx, sibling_handle_rx) = oneshot::channel();
2223 let (parent_complete_tx, parent_complete_rx) = oneshot::channel();
2224
2225 let parent = context.spawn(move |context| async move {
2226 let child_handle = context.child("child").spawn(|_| async move {
2228 child_started_tx.send(()).unwrap();
2229 child_complete_rx.await.unwrap();
2231 });
2232 assert!(
2233 child_handle_tx.send(child_handle).is_ok(),
2234 "child handle receiver dropped"
2235 );
2236
2237 let sibling_handle = context.child("sibling").spawn(move |_| async move {
2239 sibling_started_tx.send(()).unwrap();
2240 sibling_complete_rx.await.unwrap();
2242 });
2243 assert!(
2244 sibling_handle_tx.send(sibling_handle).is_ok(),
2245 "sibling handle receiver dropped"
2246 );
2247
2248 parent_complete_rx.await.unwrap();
2250 });
2251
2252 child_started_rx.await.unwrap();
2254 sibling_started_rx.await.unwrap();
2255
2256 sibling_complete_tx.send(()).unwrap();
2258 assert!(sibling_handle_rx.await.is_ok());
2259
2260 child_complete_tx.send(()).unwrap();
2262 assert!(child_handle_rx.await.is_ok());
2263
2264 parent_complete_tx.send(()).unwrap();
2266 assert!(parent.await.is_ok());
2267 });
2268 }
2269
2270 #[rstest]
2271 #[case::deterministic(deterministic::Runner::default())]
2272 #[case::tokio(tokio::Runner::default())]
2273 fn test_spawn_clone_chain<R: Runner>(#[case] runner: R)
2274 where
2275 R::Context: Spawner + Clock,
2276 {
2277 runner.start(|context| async move {
2278 let (parent_started_tx, parent_started_rx) = oneshot::channel();
2279 let (child_started_tx, child_started_rx) = oneshot::channel();
2280 let (grandchild_started_tx, grandchild_started_rx) = oneshot::channel();
2281 let (child_handle_tx, child_handle_rx) = oneshot::channel();
2282 let (grandchild_handle_tx, grandchild_handle_rx) = oneshot::channel();
2283
2284 let parent = context.child("parent").spawn({
2285 move |context| async move {
2286 let child = context.child("child").spawn({
2287 move |context| async move {
2288 let grandchild = context.child("grandchild").spawn({
2289 move |_| async move {
2290 grandchild_started_tx.send(()).unwrap();
2291 pending::<()>().await;
2292 }
2293 });
2294 assert!(
2295 grandchild_handle_tx.send(grandchild).is_ok(),
2296 "grandchild handle receiver dropped"
2297 );
2298 child_started_tx.send(()).unwrap();
2299 pending::<()>().await;
2300 }
2301 });
2302 assert!(
2303 child_handle_tx.send(child).is_ok(),
2304 "child handle receiver dropped"
2305 );
2306 parent_started_tx.send(()).unwrap();
2307 pending::<()>().await;
2308 }
2309 });
2310
2311 parent_started_rx.await.unwrap();
2312 child_started_rx.await.unwrap();
2313 grandchild_started_rx.await.unwrap();
2314
2315 let child_handle = child_handle_rx.await.unwrap();
2316 let grandchild_handle = grandchild_handle_rx.await.unwrap();
2317
2318 parent.abort();
2319 assert!(parent.await.is_err());
2320
2321 assert!(child_handle.await.is_err());
2322 assert!(grandchild_handle.await.is_err());
2323 });
2324 }
2325
2326 #[rstest]
2327 #[case::deterministic(deterministic::Runner::default())]
2328 #[case::tokio(tokio::Runner::default())]
2329 fn test_spawn_sparse_clone_chain<R: Runner>(#[case] runner: R)
2330 where
2331 R::Context: Spawner + Clock,
2332 {
2333 runner.start(|context| async move {
2334 let (leaf_started_tx, leaf_started_rx) = oneshot::channel();
2335 let (leaf_handle_tx, leaf_handle_rx) = oneshot::channel();
2336
2337 let parent = context.child("parent").spawn({
2338 move |context| async move {
2339 let clone1 = context.child("clone1");
2340 let clone2 = clone1.child("clone2");
2341 let clone3 = clone2.child("clone3");
2342
2343 let leaf = clone3.spawn({
2344 move |_| async move {
2345 leaf_started_tx.send(()).unwrap();
2346 pending::<()>().await;
2347 }
2348 });
2349
2350 leaf_handle_tx
2351 .send(leaf)
2352 .unwrap_or_else(|_| panic!("leaf handle receiver dropped"));
2353 pending::<()>().await;
2354 }
2355 });
2356
2357 leaf_started_rx.await.unwrap();
2358 let leaf_handle = leaf_handle_rx.await.unwrap();
2359
2360 parent.abort();
2361 assert!(parent.await.is_err());
2362 assert!(leaf_handle.await.is_err());
2363 });
2364 }
2365
2366 #[rstest]
2367 #[case::deterministic(deterministic::Runner::default())]
2368 #[case::tokio(tokio::Runner::default())]
2369 fn test_spawn_blocking<R: Runner>(
2370 #[case] runner: R,
2371 #[values(Execution::Shared(true), Execution::Dedicated)] execution: Execution,
2372 ) where
2373 R::Context: Spawner,
2374 {
2375 runner.start(|context| async move {
2376 let context = match execution {
2377 Execution::Dedicated => context.dedicated(),
2378 Execution::Shared(blocking) => context.shared(blocking),
2379 };
2380
2381 let handle = context.spawn(|_| async move { 42 });
2382 let result = handle.await;
2383 assert!(matches!(result, Ok(42)));
2384 });
2385 }
2386
2387 #[rstest]
2388 #[case::deterministic(deterministic::Runner::default())]
2389 #[case::tokio(tokio::Runner::default())]
2390 #[should_panic(expected = "blocking task panicked")]
2391 fn test_spawn_blocking_panic<R: Runner>(
2392 #[case] runner: R,
2393 #[values(Execution::Shared(true), Execution::Dedicated)] execution: Execution,
2394 ) where
2395 R::Context: Spawner + Clock,
2396 {
2397 runner.start(|context| async move {
2398 let spawner = match execution {
2399 Execution::Dedicated => context.child("blocking").dedicated(),
2400 Execution::Shared(blocking) => context.child("blocking").shared(blocking),
2401 };
2402 spawner.spawn(|_| async move {
2403 panic!("blocking task panicked");
2404 });
2405
2406 loop {
2408 context.sleep(Duration::from_millis(100)).await;
2409 }
2410 });
2411 }
2412
2413 #[rstest]
2414 #[case::deterministic(deterministic::Runner::new(
2415 deterministic::Config::default().with_catch_panics(true)
2416 ))]
2417 #[case::tokio(tokio::Runner::new(tokio::Config::default().with_catch_panics(true)))]
2418 fn test_spawn_blocking_panic_caught<R: Runner>(
2419 #[case] runner: R,
2420 #[values(Execution::Shared(true), Execution::Dedicated)] execution: Execution,
2421 ) where
2422 R::Context: Spawner + Clock,
2423 {
2424 let result: Result<(), Error> = runner.start(|context| async move {
2425 let spawner = match execution {
2426 Execution::Dedicated => context.child("blocking").dedicated(),
2427 Execution::Shared(blocking) => context.child("blocking").shared(blocking),
2428 };
2429 let handle = spawner.spawn(|_| async move {
2430 panic!("blocking task panicked");
2431 });
2432 handle.await
2433 });
2434 assert!(matches!(result, Err(Error::Exited)));
2435 }
2436
2437 #[rstest]
2438 #[case::deterministic(deterministic::Runner::default())]
2439 #[case::tokio(tokio::Runner::default())]
2440 fn test_circular_reference_prevents_cleanup<R: Runner>(#[case] runner: R) {
2441 runner.start(|_| async move {
2442 let dropper = Arc::new(());
2444 let executor = deterministic::Runner::default();
2445 executor.start({
2446 let dropper = dropper.clone();
2447 move |context| async move {
2448 let (setup_tx, mut setup_rx) = mpsc::unbounded_channel::<()>();
2450 let (tx1, mut rx1) = mpsc::unbounded_channel::<()>();
2451 let (tx2, mut rx2) = mpsc::unbounded_channel::<()>();
2452
2453 context.child("task1").spawn({
2455 let setup_tx = setup_tx.clone();
2456 let dropper = dropper.clone();
2457 move |_| async move {
2458 tx2.send(()).unwrap();
2460 rx1.recv().await.unwrap();
2461 setup_tx.send(()).unwrap();
2462
2463 while rx1.recv().await.is_some() {}
2465 drop(tx2);
2466 drop(dropper);
2467 }
2468 });
2469
2470 context.child("task2").spawn(move |_| async move {
2472 tx1.send(()).unwrap();
2474 rx2.recv().await.unwrap();
2475 setup_tx.send(()).unwrap();
2476
2477 while rx2.recv().await.is_some() {}
2479 drop(tx1);
2480 drop(dropper);
2481 });
2482
2483 setup_rx.recv().await.unwrap();
2485 setup_rx.recv().await.unwrap();
2486 }
2487 });
2488
2489 Arc::try_unwrap(dropper).expect("references remaining");
2491 });
2492 }
2493
2494 #[rstest]
2495 #[case::deterministic(deterministic::Runner::default())]
2496 #[case::tokio(tokio::Runner::default())]
2497 fn test_late_waker<R: Runner>(#[case] runner: R)
2498 where
2499 R::Context: Metrics + Spawner,
2500 {
2501 struct CaptureWaker {
2504 tx: Option<oneshot::Sender<Waker>>,
2505 sent: bool,
2506 }
2507 impl Future for CaptureWaker {
2508 type Output = ();
2509 fn poll(mut self: Pin<&mut Self>, cx: &mut TContext<'_>) -> Poll<Self::Output> {
2510 if !self.sent {
2511 if let Some(tx) = self.tx.take() {
2512 let _ = tx.send(cx.waker().clone());
2514 }
2515 self.sent = true;
2516 }
2517 Poll::Pending
2518 }
2519 }
2520
2521 struct WakeOnDrop(Option<Waker>);
2523 impl Drop for WakeOnDrop {
2524 fn drop(&mut self) {
2525 if let Some(w) = self.0.take() {
2526 w.wake_by_ref();
2527 }
2528 }
2529 }
2530
2531 let holder = runner.start(|context| async move {
2533 let (tx, rx) = oneshot::channel::<Waker>();
2535
2536 context.child("capture_waker").spawn(move |_| async move {
2538 CaptureWaker {
2539 tx: Some(tx),
2540 sent: false,
2541 }
2542 .await;
2543 });
2544
2545 utils::reschedule().await;
2547
2548 let waker = rx.await.expect("waker not received");
2550
2551 WakeOnDrop(Some(waker))
2553 });
2554
2555 drop(holder);
2558 }
2559
2560 #[rstest]
2561 #[case::deterministic(deterministic::Runner::default())]
2562 #[case::tokio(tokio::Runner::default())]
2563 fn test_metrics<R: Runner>(#[case] runner: R)
2564 where
2565 R::Context: Metrics,
2566 {
2567 runner.start(|context| async move {
2568 assert_eq!(context.name().label, "");
2570
2571 let counter = Counter::<u64>::default();
2573 let _registered = context.register("test", "test", counter.clone());
2574
2575 counter.inc();
2577
2578 let buffer = context.encode();
2580 assert!(buffer.contains("test_total 1"));
2581
2582 let context = context.child("nested");
2584 let nested_counter = Counter::<u64>::default();
2585 let _nested_registered = context.register("test", "test", nested_counter.clone());
2586
2587 nested_counter.inc();
2589
2590 let buffer = context.encode();
2592 assert!(buffer.contains("nested_test_total 1"));
2593 assert!(buffer.contains("test_total 1"));
2594 });
2595 }
2596
2597 #[rstest]
2598 #[case::deterministic(deterministic::Runner::default())]
2599 #[case::tokio(tokio::Runner::default())]
2600 fn test_metrics_with_attribute<R: Runner>(#[case] runner: R)
2601 where
2602 R::Context: Metrics,
2603 {
2604 runner.start(|context| async move {
2605 let ctx_epoch5 = context.child("consensus").with_attribute("epoch", "e5");
2607
2608 let counter = Counter::<u64>::default();
2610 let _epoch5 = ctx_epoch5.register("votes", "vote count", counter.clone());
2611 counter.inc();
2612
2613 let buffer = context.encode();
2615 assert!(
2616 buffer.contains("consensus_votes_total{epoch=\"e5\"} 1"),
2617 "Expected metric with epoch attribute, got: {}",
2618 buffer
2619 );
2620
2621 let ctx_epoch6 = context.child("consensus").with_attribute("epoch", "e6");
2623 let counter2 = Counter::<u64>::default();
2624 let _epoch6 = ctx_epoch6.register("votes", "vote count", counter2.clone());
2625 counter2.inc();
2626 counter2.inc();
2627
2628 let buffer = context.encode();
2630 assert!(
2631 buffer.contains("consensus_votes_total{epoch=\"e5\"} 1"),
2632 "Expected metric with epoch=e5, got: {}",
2633 buffer
2634 );
2635 assert!(
2636 buffer.contains("consensus_votes_total{epoch=\"e6\"} 2"),
2637 "Expected metric with epoch=e6, got: {}",
2638 buffer
2639 );
2640
2641 assert_eq!(
2643 buffer.matches("# HELP consensus_votes").count(),
2644 1,
2645 "HELP should appear exactly once, got: {}",
2646 buffer
2647 );
2648 assert_eq!(
2649 buffer.matches("# TYPE consensus_votes").count(),
2650 1,
2651 "TYPE should appear exactly once, got: {}",
2652 buffer
2653 );
2654
2655 let ctx_multi = context
2657 .child("engine")
2658 .with_attribute("region", "us")
2659 .with_attribute("instance", "i1");
2660 let counter3 = Counter::<u64>::default();
2661 let _multi = ctx_multi.register("requests", "request count", counter3.clone());
2662 counter3.inc();
2663
2664 let buffer = context.encode();
2665 assert!(
2666 buffer.contains("engine_requests_total{instance=\"i1\",region=\"us\"} 1"),
2667 "Expected metric with sorted attributes, got: {}",
2668 buffer
2669 );
2670 });
2671 }
2672
2673 #[rstest]
2674 #[case::deterministic(deterministic::Runner::default())]
2675 #[case::tokio(tokio::Runner::default())]
2676 fn test_metrics_attribute_with_nested_label<R: Runner>(#[case] runner: R)
2677 where
2678 R::Context: Metrics,
2679 {
2680 runner.start(|context| async move {
2681 let ctx = context
2683 .child("orchestrator")
2684 .with_attribute("epoch", "e5")
2685 .child("engine");
2686
2687 let counter = Counter::<u64>::default();
2689 let _registered = ctx.register("votes", "vote count", counter.clone());
2690 counter.inc();
2691
2692 let buffer = context.encode();
2694 assert!(
2695 buffer.contains("orchestrator_engine_votes_total{epoch=\"e5\"} 1"),
2696 "Expected metric with preserved epoch attribute, got: {}",
2697 buffer
2698 );
2699
2700 let ctx2 = context
2702 .child("outer")
2703 .with_attribute("region", "us")
2704 .child("middle")
2705 .with_attribute("az", "east")
2706 .child("inner");
2707
2708 let counter2 = Counter::<u64>::default();
2709 let _registered2 = ctx2.register("requests", "request count", counter2.clone());
2710 counter2.inc();
2711 counter2.inc();
2712
2713 let buffer = context.encode();
2714 assert!(
2715 buffer.contains("outer_middle_inner_requests_total{az=\"east\",region=\"us\"} 2"),
2716 "Expected metric with all attributes preserved and sorted, got: {}",
2717 buffer
2718 );
2719 });
2720 }
2721
2722 #[rstest]
2723 #[case::deterministic(deterministic::Runner::default())]
2724 #[case::tokio(tokio::Runner::default())]
2725 fn test_metrics_attributes_isolated_between_contexts<R: Runner>(#[case] runner: R)
2726 where
2727 R::Context: Metrics,
2728 {
2729 runner.start(|context| async move {
2730 let ctx_a = context.child("component_a").with_attribute("epoch", 1);
2732 let ctx_b = context.child("component_b").with_attribute("epoch", 2);
2733
2734 let c1 = Counter::<u64>::default();
2736 let _ctx_a_requests = ctx_a.register("requests", "help", c1);
2737
2738 let c2 = Counter::<u64>::default();
2740 let _ctx_b_requests = ctx_b.register("requests", "help", c2);
2741
2742 let c3 = Counter::<u64>::default();
2744 let _ctx_a_errors = ctx_a.register("errors", "help", c3);
2745
2746 let output = context.encode();
2747
2748 assert!(
2750 output.contains("component_a_requests_total{epoch=\"1\"} 0"),
2751 "ctx_a requests should have epoch=1: {output}"
2752 );
2753 assert!(
2754 output.contains("component_a_errors_total{epoch=\"1\"} 0"),
2755 "ctx_a errors should have epoch=1: {output}"
2756 );
2757 assert!(
2758 !output.contains("component_a_requests_total{epoch=\"2\"}"),
2759 "ctx_a requests should not have epoch=2: {output}"
2760 );
2761
2762 assert!(
2764 output.contains("component_b_requests_total{epoch=\"2\"} 0"),
2765 "ctx_b should have epoch=2: {output}"
2766 );
2767 assert!(
2768 !output.contains("component_b_requests_total{epoch=\"1\"}"),
2769 "ctx_b should not have epoch=1: {output}"
2770 );
2771 });
2772 }
2773
2774 #[rstest]
2781 #[case::deterministic(deterministic::Runner::default())]
2782 #[case::tokio(tokio::Runner::default())]
2783 fn test_metrics_spawn_attribute_cardinality<R: Runner>(#[case] runner: R)
2784 where
2785 R::Context: Spawner + Metrics + Clock,
2786 {
2787 runner.start(|context| async move {
2788 const ROUNDS: u64 = 128;
2789
2790 let mut handles = Vec::with_capacity(ROUNDS as usize);
2791 for round in 0..ROUNDS {
2792 let handle = context
2793 .child("deferred_verify")
2794 .with_attribute("round", round)
2795 .spawn(move |_| async move { round });
2796 handles.push(handle);
2797 }
2798 for (expected, handle) in handles.into_iter().enumerate() {
2799 assert_eq!(handle.await.expect("task failed"), expected as u64);
2800 }
2801
2802 while count_running_tasks(&context, "deferred_verify") > 0 {
2806 context.sleep(Duration::from_millis(10)).await;
2807 }
2808 let buffer = context.encode();
2809
2810 let spawned_lines = buffer
2814 .lines()
2815 .filter(|line| {
2816 line.starts_with("runtime_tasks_spawned_total{")
2817 && line.contains("name=\"deferred_verify\"")
2818 })
2819 .count();
2820 let running_lines = buffer
2821 .lines()
2822 .filter(|line| {
2823 line.starts_with("runtime_tasks_running{")
2824 && line.contains("name=\"deferred_verify\"")
2825 })
2826 .count();
2827 assert_eq!(
2828 spawned_lines, 1,
2829 "expected exactly 1 runtime_tasks_spawned entry for deferred_verify, got {spawned_lines}: {buffer}",
2830 );
2831 assert_eq!(
2832 running_lines, 1,
2833 "expected exactly 1 runtime_tasks_running entry for deferred_verify, got {running_lines}: {buffer}",
2834 );
2835
2836 let spawned_value = format!(
2838 "runtime_tasks_spawned_total{{name=\"deferred_verify\",kind=\"Task\",execution=\"Shared\"}} {ROUNDS}"
2839 );
2840 assert!(
2841 buffer.contains(&spawned_value),
2842 "expected accumulated spawned counter `{spawned_value}`, got: {buffer}",
2843 );
2844 let running_value = "runtime_tasks_running{name=\"deferred_verify\",kind=\"Task\",execution=\"Shared\"} 0";
2845 assert!(
2846 buffer.contains(running_value),
2847 "expected running gauge to return to 0, got: {buffer}",
2848 );
2849
2850 assert!(
2853 !buffer
2854 .lines()
2855 .any(|line| line.starts_with("runtime_tasks_")
2856 && line.contains("round=")),
2857 "task metrics must not carry `round` attribute: {buffer}",
2858 );
2859 });
2860 }
2861
2862 #[rstest]
2863 #[case::deterministic(deterministic::Runner::default())]
2864 #[case::tokio(tokio::Runner::default())]
2865 fn test_metrics_attributes_sorted_deterministically<R: Runner>(#[case] runner: R)
2866 where
2867 R::Context: Metrics,
2868 {
2869 runner.start(|context| async move {
2870 let ctx_ab = context
2872 .child("service")
2873 .with_attribute("region", "us")
2874 .with_attribute("env", "prod");
2875
2876 let ctx_ba = context
2877 .child("service")
2878 .with_attribute("env", "prod")
2879 .with_attribute("region", "us");
2880
2881 let c1 = Counter::<u64>::default();
2883 let _requests = ctx_ab.register("requests", "help", c1.clone());
2884 c1.inc();
2885
2886 let c2 = Counter::<u64>::default();
2888 let _errors = ctx_ba.register("errors", "help", c2.clone());
2889 c2.inc();
2890 c2.inc();
2891
2892 let output = context.encode();
2893
2894 assert!(
2896 output.contains("service_requests_total{env=\"prod\",region=\"us\"} 1"),
2897 "requests should have sorted labels: {output}"
2898 );
2899 assert!(
2900 output.contains("service_errors_total{env=\"prod\",region=\"us\"} 2"),
2901 "errors should have sorted labels: {output}"
2902 );
2903
2904 assert!(
2906 !output.contains("region=\"us\",env=\"prod\""),
2907 "should not have unsorted label order: {output}"
2908 );
2909 });
2910 }
2911
2912 #[rstest]
2913 #[case::deterministic(deterministic::Runner::default())]
2914 #[case::tokio(tokio::Runner::default())]
2915 fn test_metrics_nested_labels_with_attributes<R: Runner>(#[case] runner: R)
2916 where
2917 R::Context: Metrics,
2918 {
2919 runner.start(|context| async move {
2920 let svc_a = context.child("service_a");
2922
2923 let svc_a_v2 = context.child("service_a").with_attribute("version", 2);
2925
2926 let svc_b_worker = context.child("service_b").child("worker");
2928
2929 let svc_b_worker_shard = context
2931 .child("service_b")
2932 .child("worker")
2933 .with_attribute("shard", 99);
2934
2935 let svc_b_manager = context.child("service_b").child("manager");
2937
2938 let svc_c = context.child("service_c");
2940
2941 let c1 = Counter::<u64>::default();
2943 let _svc_a = svc_a.register("requests", "help", c1);
2944
2945 let c2 = Counter::<u64>::default();
2946 let _svc_a_v2 = svc_a_v2.register("requests", "help", c2);
2947
2948 let c3 = Counter::<u64>::default();
2949 let _svc_b_worker = svc_b_worker.register("tasks", "help", c3);
2950
2951 let c4 = Counter::<u64>::default();
2952 let _svc_b_worker_shard = svc_b_worker_shard.register("tasks", "help", c4);
2953
2954 let c5 = Counter::<u64>::default();
2955 let _svc_b_manager = svc_b_manager.register("decisions", "help", c5);
2956
2957 let c6 = Counter::<u64>::default();
2958 let _svc_c = svc_c.register("requests", "help", c6);
2959
2960 let output = context.encode();
2961
2962 assert!(
2964 output.contains("service_a_requests_total 0"),
2965 "svc_a plain should exist: {output}"
2966 );
2967 assert!(
2968 output.contains("service_a_requests_total{version=\"2\"} 0"),
2969 "svc_a_v2 should have version=2: {output}"
2970 );
2971
2972 assert!(
2974 output.contains("service_b_worker_tasks_total 0"),
2975 "svc_b_worker plain should exist: {output}"
2976 );
2977 assert!(
2978 output.contains("service_b_worker_tasks_total{shard=\"99\"} 0"),
2979 "svc_b_worker_shard should have shard=99: {output}"
2980 );
2981
2982 assert!(
2984 output.contains("service_b_manager_decisions_total 0"),
2985 "svc_b_manager should have no attributes: {output}"
2986 );
2987 assert!(
2988 !output.contains("service_b_manager_decisions_total{"),
2989 "svc_b_manager should have no attributes at all: {output}"
2990 );
2991
2992 assert!(
2994 output.contains("service_c_requests_total 0"),
2995 "svc_c should have no attributes: {output}"
2996 );
2997 assert!(
2998 !output.contains("service_c_requests_total{"),
2999 "svc_c should have no attributes at all: {output}"
3000 );
3001
3002 assert!(
3004 !output.contains("service_b_manager_decisions_total{shard="),
3005 "svc_b_manager should not have shard: {output}"
3006 );
3007 assert!(
3008 !output.contains("service_a_requests_total{shard="),
3009 "svc_a should not have shard: {output}"
3010 );
3011 assert!(
3012 !output.contains("service_c_requests_total{version="),
3013 "svc_c should not have version: {output}"
3014 );
3015 });
3016 }
3017
3018 #[rstest]
3019 #[case::deterministic(deterministic::Runner::default())]
3020 #[case::tokio(tokio::Runner::default())]
3021 fn test_metrics_family_with_attributes<R: Runner>(#[case] runner: R)
3022 where
3023 R::Context: Metrics,
3024 {
3025 runner.start(|context| async move {
3026 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
3027 struct RequestLabels {
3028 method: String,
3029 status: u16,
3030 }
3031
3032 let ctx = context
3034 .child("api")
3035 .with_attribute("region", "us_east")
3036 .with_attribute("env", "prod");
3037
3038 let requests: Family<RequestLabels, Counter<u64>> = Family::default();
3040 let _requests = ctx.register("requests", "HTTP requests", requests.clone());
3041
3042 requests
3044 .get_or_create(&RequestLabels {
3045 method: "GET".to_string(),
3046 status: 200,
3047 })
3048 .inc();
3049 requests
3050 .get_or_create(&RequestLabels {
3051 method: "POST".to_string(),
3052 status: 201,
3053 })
3054 .inc();
3055 requests
3056 .get_or_create(&RequestLabels {
3057 method: "GET".to_string(),
3058 status: 404,
3059 })
3060 .inc();
3061
3062 let output = context.encode();
3063
3064 assert!(
3068 output.contains(
3069 "api_requests_total{env=\"prod\",region=\"us_east\",method=\"GET\",status=\"200\"} 1"
3070 ),
3071 "GET 200 should have merged labels: {output}"
3072 );
3073 assert!(
3074 output.contains(
3075 "api_requests_total{env=\"prod\",region=\"us_east\",method=\"POST\",status=\"201\"} 1"
3076 ),
3077 "POST 201 should have merged labels: {output}"
3078 );
3079 assert!(
3080 output.contains(
3081 "api_requests_total{env=\"prod\",region=\"us_east\",method=\"GET\",status=\"404\"} 1"
3082 ),
3083 "GET 404 should have merged labels: {output}"
3084 );
3085
3086 let ctx_plain = context.child("api_plain");
3088 let plain_requests: Family<RequestLabels, Counter<u64>> = Family::default();
3089 let _plain_requests =
3090 ctx_plain.register("requests", "HTTP requests", plain_requests.clone());
3091
3092 plain_requests
3093 .get_or_create(&RequestLabels {
3094 method: "DELETE".to_string(),
3095 status: 204,
3096 })
3097 .inc();
3098
3099 let output = context.encode();
3100
3101 assert!(
3103 output.contains("api_plain_requests_total{method=\"DELETE\",status=\"204\"} 1"),
3104 "plain DELETE should have only family labels: {output}"
3105 );
3106 assert!(
3107 !output.contains("api_plain_requests_total{env="),
3108 "plain should not have env attribute: {output}"
3109 );
3110 assert!(
3111 !output.contains("api_plain_requests_total{region="),
3112 "plain should not have region attribute: {output}"
3113 );
3114 });
3115 }
3116
3117 #[rstest]
3118 #[case::deterministic(deterministic::Runner::default())]
3119 #[case::tokio(tokio::Runner::default())]
3120 fn test_register_and_encode<R: Runner>(#[case] runner: R)
3121 where
3122 R::Context: Metrics,
3123 {
3124 runner.start(|context| async move {
3125 let counter =
3126 context
3127 .child("engine")
3128 .register("votes", "vote count", Counter::<u64>::default());
3129 counter.inc();
3130
3131 let buffer = context.encode();
3132 assert!(
3133 buffer.contains("engine_votes_total 1"),
3134 "registered metric should appear in encode: {buffer}"
3135 );
3136 });
3137 }
3138
3139 #[rstest]
3140 #[case::deterministic(deterministic::Runner::default())]
3141 #[case::tokio(tokio::Runner::default())]
3142 fn test_register_drop_removes_metrics<R: Runner>(#[case] runner: R)
3143 where
3144 R::Context: Metrics,
3145 {
3146 runner.start(|context| async move {
3147 let permanent = context.child("permanent").register(
3148 "counter",
3149 "permanent counter",
3150 Counter::<u64>::default(),
3151 );
3152 permanent.inc();
3153
3154 let counter =
3155 context
3156 .child("engine")
3157 .register("votes", "vote count", Counter::<u64>::default());
3158 counter.inc();
3159
3160 let buffer = context.encode();
3161 assert!(buffer.contains("permanent_counter_total 1"));
3162 assert!(buffer.contains("engine_votes_total 1"));
3163
3164 drop(counter);
3165
3166 let buffer = context.encode();
3167 assert!(
3168 buffer.contains("permanent_counter_total 1"),
3169 "other registered metrics should survive handle drop: {buffer}"
3170 );
3171 assert!(
3172 !buffer.contains("engine_votes"),
3173 "metric should be removed after handle drop: {buffer}"
3174 );
3175 });
3176 }
3177
3178 #[rstest]
3179 #[case::deterministic(deterministic::Runner::default())]
3180 #[case::tokio(tokio::Runner::default())]
3181 fn test_register_with_attributes<R: Runner>(#[case] runner: R)
3182 where
3183 R::Context: Metrics,
3184 {
3185 runner.start(|context| async move {
3186 let epoch1 = context.child("engine").with_attribute("epoch", 1).register(
3187 "votes",
3188 "vote count",
3189 Counter::<u64>::default(),
3190 );
3191 epoch1.inc();
3192
3193 let epoch2 = context.child("engine").with_attribute("epoch", 2).register(
3194 "votes",
3195 "vote count",
3196 Counter::<u64>::default(),
3197 );
3198 epoch2.inc();
3199 epoch2.inc();
3200
3201 let buffer = context.encode();
3202 assert!(buffer.contains("engine_votes_total{epoch=\"1\"} 1"));
3203 assert!(buffer.contains("engine_votes_total{epoch=\"2\"} 2"));
3204
3205 assert_eq!(
3206 buffer.matches("# HELP engine_votes").count(),
3207 1,
3208 "HELP should appear once: {buffer}"
3209 );
3210 assert_eq!(
3211 buffer.matches("# TYPE engine_votes").count(),
3212 1,
3213 "TYPE should appear once: {buffer}"
3214 );
3215
3216 drop(epoch1);
3217 let buffer = context.encode();
3218 assert!(
3219 !buffer.contains("epoch=\"1\""),
3220 "epoch 1 should be gone: {buffer}"
3221 );
3222 assert!(buffer.contains("engine_votes_total{epoch=\"2\"} 2"));
3223
3224 drop(epoch2);
3225 let buffer = context.encode();
3226 assert!(
3227 !buffer.contains("engine_votes"),
3228 "all epoch metrics should be gone: {buffer}"
3229 );
3230 });
3231 }
3232
3233 #[rstest]
3234 #[case::deterministic(deterministic::Runner::default())]
3235 #[case::tokio(tokio::Runner::default())]
3236 fn test_reregister_after_drop<R: Runner>(#[case] runner: R)
3237 where
3238 R::Context: Metrics,
3239 {
3240 runner.start(|context| async move {
3241 let votes = context.child("engine").with_attribute("epoch", 1).register(
3242 "votes",
3243 "vote count",
3244 Counter::<u64>::default(),
3245 );
3246 drop(votes);
3247
3248 let replacement = context.child("engine").with_attribute("epoch", 1).register(
3249 "votes",
3250 "vote count",
3251 Counter::<u64>::default(),
3252 );
3253 drop(replacement);
3254 });
3255 }
3256
3257 #[rstest]
3258 #[case::deterministic(deterministic::Runner::default())]
3259 #[case::tokio(tokio::Runner::default())]
3260 fn test_register_clone_keeps_metric_alive<R: Runner>(#[case] runner: R)
3261 where
3262 R::Context: Metrics,
3263 {
3264 runner.start(|context| async move {
3265 let registered =
3266 context
3267 .child("engine")
3268 .register("votes", "vote count", Counter::<u64>::default());
3269 registered.inc();
3270 let clone = registered.clone();
3271
3272 let buffer = context.encode();
3273 assert!(
3274 buffer.contains("engine_votes_total 1"),
3275 "metric should remain registered while any handle exists: {buffer}"
3276 );
3277
3278 drop(registered);
3279 let buffer = context.encode();
3280 assert!(
3281 buffer.contains("engine_votes_total 1"),
3282 "metric should survive while clone is retained: {buffer}"
3283 );
3284
3285 drop(clone);
3286 let buffer = context.encode();
3287 assert!(
3288 !buffer.contains("engine_votes"),
3289 "metric should be removed when all handle clones are dropped: {buffer}"
3290 );
3291 });
3292 }
3293
3294 #[rstest]
3295 #[case::deterministic(deterministic::Runner::default())]
3296 #[case::tokio(tokio::Runner::default())]
3297 fn test_encode_single_eof<R: Runner>(#[case] runner: R)
3298 where
3299 R::Context: Metrics,
3300 {
3301 runner.start(|context| async move {
3302 let root_counter = context.register("root", "root metric", Counter::<u64>::default());
3303 root_counter.inc();
3304
3305 let child =
3306 context
3307 .child("engine")
3308 .register("ops", "child metric", Counter::<u64>::default());
3309 child.inc();
3310
3311 let buffer = context.encode();
3312 assert!(
3313 buffer.contains("root_total 1"),
3314 "root metric missing: {buffer}"
3315 );
3316 assert!(
3317 buffer.contains("engine_ops_total 1"),
3318 "child metric missing: {buffer}"
3319 );
3320 assert_eq!(
3321 buffer.matches("# EOF").count(),
3322 1,
3323 "expected exactly one EOF marker: {buffer}"
3324 );
3325 assert!(
3326 buffer.ends_with("# EOF\n"),
3327 "EOF must be the last line: {buffer}"
3328 );
3329 });
3330 }
3331
3332 #[rstest]
3333 #[case::deterministic(deterministic::Runner::default())]
3334 #[case::tokio(tokio::Runner::default())]
3335 fn test_family_with_attributes<R: Runner>(#[case] runner: R)
3336 where
3337 R::Context: Metrics,
3338 {
3339 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
3340 struct Peer {
3341 name: String,
3342 }
3343 impl EncodeLabelSet for Peer {
3344 fn encode(&self, encoder: &mut LabelSetEncoder<'_>) -> Result<(), std::fmt::Error> {
3345 let mut label = encoder.encode_label();
3346 let mut key = label.encode_label_key()?;
3347 EncodeLabelKey::encode(&"peer", &mut key)?;
3348 let mut value = key.encode_label_value()?;
3349 EncodeLabelValue::encode(&self.name.as_str(), &mut value)?;
3350 value.finish()
3351 }
3352 }
3353
3354 runner.start(|context| async move {
3355 let family = context
3356 .child("batcher")
3357 .with_attribute("epoch", 1)
3358 .register(
3359 "votes",
3360 "votes per peer",
3361 Family::<Peer, Counter>::default(),
3362 );
3363 family
3364 .get_or_create(&Peer {
3365 name: "alice".into(),
3366 })
3367 .inc();
3368 family.get_or_create(&Peer { name: "bob".into() }).inc();
3369
3370 let buffer = context.encode();
3371 assert!(
3372 buffer.contains("batcher_votes_total{epoch=\"1\",peer=\"alice\"} 1"),
3373 "family with attributes should combine labels: {buffer}"
3374 );
3375 assert!(
3376 buffer.contains("batcher_votes_total{epoch=\"1\",peer=\"bob\"} 1"),
3377 "family with attributes should combine labels: {buffer}"
3378 );
3379
3380 drop(family);
3381 let buffer = context.encode();
3382 assert!(
3383 !buffer.contains("batcher_votes"),
3384 "family metrics should be removed: {buffer}"
3385 );
3386 });
3387 }
3388
3389 #[rstest]
3390 #[case::deterministic(deterministic::Runner::default())]
3391 #[case::tokio(tokio::Runner::default())]
3392 fn test_strategy<R: Runner>(#[case] runner: R)
3393 where
3394 R::Context: Strategizer + Metrics,
3395 {
3396 runner.start(|context| async move {
3397 let strategy = context.child("pool").strategy(NZUsize!(4));
3399 assert_eq!(strategy.manual().parallelism(), 4);
3400
3401 let sum = strategy.fold(0..10000, || 0i32, |acc, n| acc + n, |a, b| a + b);
3403 assert_eq!(sum, 10000 * 9999 / 2);
3404 });
3405 }
3406
3407 #[rstest]
3408 #[case::deterministic(deterministic::Runner::default())]
3409 #[case::tokio(tokio::Runner::default())]
3410 fn test_nested_strategy_runs_inline<R: Runner>(#[case] runner: R)
3411 where
3412 R::Context: Strategizer + Metrics,
3413 {
3414 runner.start(|context| async move {
3415 let strategy = context.child("pool").strategy(NZUsize!(1)).manual();
3416
3417 let output = strategy
3418 .spawn(2, |strategy| strategy.map_collect_vec(0..2, |i| i + 1))
3419 .await;
3420
3421 assert_eq!(output, vec![1, 2]);
3422 });
3423 }
3424
3425 #[rstest]
3426 #[case::deterministic(deterministic::Runner::default(), 4096, 64)]
3427 #[case::deterministic_custom(
3428 deterministic::Runner::new(
3429 deterministic::Config::default()
3430 .with_network_buffer_pool_config(
3431 BufferPoolConfig::for_network().with_max_per_class(NZU32!(64)),
3432 )
3433 .with_storage_buffer_pool_config(
3434 BufferPoolConfig::for_storage().with_max_per_class(NZU32!(8)),
3435 ),
3436 ),
3437 64,
3438 8
3439 )]
3440 #[case::tokio(tokio::Runner::default(), 4096, 64)]
3441 #[case::tokio_custom(
3442 tokio::Runner::new(
3443 tokio::Config::default()
3444 .with_network_buffer_pool_config(
3445 BufferPoolConfig::for_network().with_max_per_class(NZU32!(64)),
3446 )
3447 .with_storage_buffer_pool_config(
3448 BufferPoolConfig::for_storage().with_max_per_class(NZU32!(8)),
3449 ),
3450 ),
3451 64,
3452 8
3453 )]
3454 fn test_buffer_pooler<R: Runner>(
3455 #[case] runner: R,
3456 #[case] expected_network_max_per_class: u32,
3457 #[case] expected_storage_max_per_class: u32,
3458 ) where
3459 R::Context: BufferPooler,
3460 {
3461 runner.start(|context| async move {
3462 let net_buf = context.network_buffer_pool().try_alloc(1024).unwrap();
3464 assert!(net_buf.capacity() >= 1024);
3465
3466 let storage_buf = context.storage_buffer_pool().try_alloc(1024).unwrap();
3468 assert!(storage_buf.capacity() >= 4096);
3469
3470 assert!(
3472 context
3473 .network_buffer_pool()
3474 .config()
3475 .size_classes()
3476 .all(|class| class.max_buffers.get() == expected_network_max_per_class)
3477 );
3478 assert!(
3479 context
3480 .storage_buffer_pool()
3481 .config()
3482 .size_classes()
3483 .all(|class| class.max_buffers.get() == expected_storage_max_per_class)
3484 );
3485 });
3486 }
3487}