hopr_utils/
runtime.rs

1//! Executor API for HOPR which exposes the necessary async functions depending on the enabled
2//! runtime.
3
4use std::{
5    future::Future,
6    hash::Hash,
7    pin::Pin,
8    sync::{
9        Arc, OnceLock,
10        atomic::{AtomicBool, AtomicU64, Ordering},
11    },
12    task::{Context, Poll},
13    time::{Duration, Instant},
14};
15
16pub use futures::future::AbortHandle;
17
18// Both features could be enabled during testing; therefore, we only use tokio when it's
19// exclusively enabled.
20pub mod prelude {
21    #[cfg(feature = "async-lock")]
22    pub use async_lock::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
23    pub use futures::future::{AbortHandle, abortable};
24    #[cfg(all(feature = "runtime-tokio", not(feature = "async-lock")))]
25    pub use tokio::sync::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
26    #[cfg(feature = "runtime-tokio")]
27    pub use tokio::{
28        task::{JoinError, JoinHandle, spawn, spawn_blocking, spawn_local, yield_now},
29        time::{sleep, timeout as timeout_fut},
30    };
31}
32
33/// Passive diagnostics for attributing high-frequency async polling.
34///
35/// Diagnostics are disabled unless `HOPR_RUNTIME_POLL_DIAGNOSTICS=1` is set.
36/// When enabled, warnings are rate-limited by `HOPR_RUNTIME_POLL_WARN_EVERY`
37/// and start after `HOPR_RUNTIME_POLL_WARN_AFTER` polls or child futures.
38pub mod diagnostics {
39    use super::*;
40
41    const ENV_ENABLED: &str = "HOPR_RUNTIME_POLL_DIAGNOSTICS";
42    const ENV_WARN_AFTER: &str = "HOPR_RUNTIME_POLL_WARN_AFTER";
43    const ENV_WARN_EVERY: &str = "HOPR_RUNTIME_POLL_WARN_EVERY";
44    const DEFAULT_WARN_AFTER: u64 = 100_000;
45    const DEFAULT_WARN_EVERY: Duration = Duration::from_secs(5);
46
47    #[derive(Debug)]
48    struct DiagnosticsConfig {
49        enabled: bool,
50        warn_after: u64,
51        warn_every: Duration,
52    }
53
54    static CONFIG: OnceLock<DiagnosticsConfig> = OnceLock::new();
55
56    fn parse_bool(value: &str) -> bool {
57        matches!(value.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on")
58    }
59
60    fn parse_duration(value: &str) -> Option<Duration> {
61        let trimmed = value.trim();
62        if let Some(stripped) = trimmed.strip_suffix("ms") {
63            stripped.trim().parse::<u64>().ok().map(Duration::from_millis)
64        } else if let Some(stripped) = trimmed.strip_suffix('s') {
65            stripped.trim().parse::<u64>().ok().map(Duration::from_secs)
66        } else {
67            trimmed.parse::<u64>().ok().map(Duration::from_secs)
68        }
69    }
70
71    fn config() -> &'static DiagnosticsConfig {
72        CONFIG.get_or_init(|| DiagnosticsConfig {
73            enabled: std::env::var(ENV_ENABLED).is_ok_and(|v| parse_bool(&v)),
74            warn_after: std::env::var(ENV_WARN_AFTER)
75                .ok()
76                .and_then(|v| v.trim().parse::<u64>().ok())
77                .filter(|&v| v > 0)
78                .unwrap_or(DEFAULT_WARN_AFTER),
79            warn_every: std::env::var(ENV_WARN_EVERY)
80                .ok()
81                .and_then(|v| parse_duration(&v))
82                .filter(|&v| v > Duration::ZERO)
83                .unwrap_or(DEFAULT_WARN_EVERY),
84        })
85    }
86
87    /// Returns whether runtime poll diagnostics are enabled for this process.
88    pub fn enabled() -> bool {
89        config().enabled
90    }
91
92    fn location(module_path: &'static str, file: &'static str, line: u32) -> String {
93        format!("{module_path} at {file}:{line}")
94    }
95
96    fn should_warn(last_warn_nanos: &AtomicU64, elapsed: Duration, warn_every: Duration) -> bool {
97        let elapsed_nanos = elapsed.as_nanos().min(u64::MAX as u128) as u64;
98        let warn_every_nanos = warn_every.as_nanos().min(u64::MAX as u128) as u64;
99
100        loop {
101            let last = last_warn_nanos.load(Ordering::Relaxed);
102            if elapsed_nanos.saturating_sub(last) < warn_every_nanos {
103                return false;
104            }
105
106            if last_warn_nanos
107                .compare_exchange(last, elapsed_nanos, Ordering::Relaxed, Ordering::Relaxed)
108                .is_ok()
109            {
110                return true;
111            }
112        }
113    }
114
115    fn warn_high_poll_rate(name: &str, location: &str, total_polls: u64, elapsed: Duration, polls_per_sec: f64) {
116        #[cfg(feature = "runtime-tokio")]
117        tracing::warn!(
118            task = name,
119            location,
120            total_polls,
121            elapsed_ms = elapsed.as_millis() as u64,
122            polls_per_sec,
123            "high-frequency polling detected"
124        );
125
126        #[cfg(not(feature = "runtime-tokio"))]
127        eprintln!(
128            "high-frequency polling detected: task={name} location={location} total_polls={total_polls} elapsed_ms={} \
129             polls_per_sec={polls_per_sec}",
130            elapsed.as_millis()
131        );
132    }
133
134    #[allow(clippy::too_many_arguments)]
135    fn warn_high_child_rate(
136        name: &str,
137        location: &str,
138        started: u64,
139        completed: u64,
140        dropped: u64,
141        in_flight: u64,
142        elapsed: Duration,
143        completed_per_sec: f64,
144    ) {
145        #[cfg(feature = "runtime-tokio")]
146        tracing::warn!(
147            task = name,
148            location,
149            started,
150            completed,
151            dropped,
152            in_flight,
153            elapsed_ms = elapsed.as_millis() as u64,
154            completed_per_sec,
155            "high-frequency concurrent child churn detected"
156        );
157
158        #[cfg(not(feature = "runtime-tokio"))]
159        eprintln!(
160            "high-frequency concurrent child churn detected: task={name} location={location} started={started} \
161             completed={completed} dropped={dropped} in_flight={in_flight} elapsed_ms={} \
162             completed_per_sec={completed_per_sec}",
163            elapsed.as_millis()
164        );
165    }
166
167    struct PollDiagnosticState {
168        name: String,
169        location: String,
170        started_at: Instant,
171        total_polls: AtomicU64,
172        last_warn_nanos: AtomicU64,
173    }
174
175    /// Future wrapper that counts polls and emits rate-limited warnings.
176    #[must_use = "futures do nothing unless polled or awaited"]
177    pub struct PollDiagnosticFuture<F> {
178        inner: F,
179        state: Option<Arc<PollDiagnosticState>>,
180    }
181
182    impl<F> PollDiagnosticFuture<F> {
183        fn new(inner: F, name: String, module_path: &'static str, file: &'static str, line: u32) -> Self {
184            let state = enabled().then(|| {
185                Arc::new(PollDiagnosticState {
186                    name,
187                    location: location(module_path, file, line),
188                    started_at: Instant::now(),
189                    total_polls: AtomicU64::new(0),
190                    last_warn_nanos: AtomicU64::new(0),
191                })
192            });
193
194            Self { inner, state }
195        }
196    }
197
198    impl<F: Future> Future for PollDiagnosticFuture<F> {
199        type Output = F::Output;
200
201        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
202            // SAFETY: `inner` is pinned in place as part of `self`; we never move it.
203            let this = unsafe { self.get_unchecked_mut() };
204
205            if let Some(state) = &this.state {
206                let total_polls = state.total_polls.fetch_add(1, Ordering::Relaxed) + 1;
207                let cfg = config();
208                if total_polls >= cfg.warn_after {
209                    let elapsed = state.started_at.elapsed();
210                    if should_warn(&state.last_warn_nanos, elapsed, cfg.warn_every) {
211                        let polls_per_sec = total_polls as f64 / elapsed.as_secs_f64().max(0.001);
212                        warn_high_poll_rate(&state.name, &state.location, total_polls, elapsed, polls_per_sec);
213                    }
214                }
215            }
216
217            // SAFETY: see above; projecting only the pinned `inner` field.
218            unsafe { Pin::new_unchecked(&mut this.inner) }.poll(cx)
219        }
220    }
221
222    /// Instruments a future using callsite metadata.
223    pub fn instrument<F>(
224        inner: F,
225        name: &'static str,
226        module_path: &'static str,
227        file: &'static str,
228        line: u32,
229    ) -> PollDiagnosticFuture<F> {
230        PollDiagnosticFuture::new(inner, name.to_owned(), module_path, file, line)
231    }
232
233    struct ConcurrentDiagnosticState {
234        name: String,
235        location: String,
236        started_at: Instant,
237        started: AtomicU64,
238        completed: AtomicU64,
239        dropped: AtomicU64,
240        in_flight: AtomicU64,
241        last_warn_nanos: AtomicU64,
242    }
243
244    /// Shared diagnostics for `for_each_concurrent`/`try_for_each_concurrent`
245    /// child futures.
246    #[derive(Clone)]
247    pub struct ConcurrentDiagnostics {
248        state: Option<Arc<ConcurrentDiagnosticState>>,
249    }
250
251    impl ConcurrentDiagnostics {
252        /// Creates a diagnostics handle for one concurrent stream operator.
253        pub fn new(name: &'static str, module_path: &'static str, file: &'static str, line: u32) -> Self {
254            let state = enabled().then(|| {
255                Arc::new(ConcurrentDiagnosticState {
256                    name: name.to_owned(),
257                    location: location(module_path, file, line),
258                    started_at: Instant::now(),
259                    started: AtomicU64::new(0),
260                    completed: AtomicU64::new(0),
261                    dropped: AtomicU64::new(0),
262                    in_flight: AtomicU64::new(0),
263                    last_warn_nanos: AtomicU64::new(0),
264                })
265            });
266
267            Self { state }
268        }
269
270        /// Wraps one child future produced by a concurrent stream operator.
271        pub fn wrap<F: FnOnce() -> T, T>(&self, inner: F) -> ConcurrentDiagnosticFuture<T> {
272            if let Some(state) = &self.state {
273                state.started.fetch_add(1, Ordering::Relaxed);
274                state.in_flight.fetch_add(1, Ordering::Relaxed);
275            }
276
277            ConcurrentDiagnosticFuture {
278                inner: inner(),
279                state: self.state.clone(),
280                finished: AtomicBool::new(false),
281            }
282        }
283    }
284
285    /// Child-future wrapper for concurrent stream diagnostics.
286    #[must_use = "futures do nothing unless polled or awaited"]
287    pub struct ConcurrentDiagnosticFuture<F> {
288        inner: F,
289        state: Option<Arc<ConcurrentDiagnosticState>>,
290        finished: AtomicBool,
291    }
292
293    impl<F> ConcurrentDiagnosticFuture<F> {
294        fn mark_finished(&self, completed: bool) {
295            if self.finished.swap(true, Ordering::Relaxed) {
296                return;
297            }
298
299            let Some(state) = &self.state else {
300                return;
301            };
302
303            state.in_flight.fetch_sub(1, Ordering::Relaxed);
304            if completed {
305                let completed = state.completed.fetch_add(1, Ordering::Relaxed) + 1;
306                let cfg = config();
307                if completed >= cfg.warn_after {
308                    let elapsed = state.started_at.elapsed();
309                    if should_warn(&state.last_warn_nanos, elapsed, cfg.warn_every) {
310                        let started = state.started.load(Ordering::Relaxed);
311                        let dropped = state.dropped.load(Ordering::Relaxed);
312                        let in_flight = state.in_flight.load(Ordering::Relaxed);
313                        let completed_per_sec = completed as f64 / elapsed.as_secs_f64().max(0.001);
314                        warn_high_child_rate(
315                            &state.name,
316                            &state.location,
317                            started,
318                            completed,
319                            dropped,
320                            in_flight,
321                            elapsed,
322                            completed_per_sec,
323                        );
324                    }
325                }
326            } else {
327                state.dropped.fetch_add(1, Ordering::Relaxed);
328            }
329        }
330    }
331
332    impl<F: Future> Future for ConcurrentDiagnosticFuture<F> {
333        type Output = F::Output;
334
335        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
336            // SAFETY: `inner` is pinned in place as part of `self`; we never move it.
337            let this = unsafe { self.get_unchecked_mut() };
338
339            // SAFETY: see above; projecting only the pinned `inner` field.
340            match unsafe { Pin::new_unchecked(&mut this.inner) }.poll(cx) {
341                Poll::Ready(output) => {
342                    this.mark_finished(true);
343                    Poll::Ready(output)
344                }
345                Poll::Pending => Poll::Pending,
346            }
347        }
348    }
349
350    impl<F> Drop for ConcurrentDiagnosticFuture<F> {
351        fn drop(&mut self) {
352            self.mark_finished(false);
353        }
354    }
355}
356
357#[macro_export]
358macro_rules! spawn_as_abortable {
359    ($expr:expr $(,)?) => {{
360        let (proc, abort_handle) = $crate::runtime::prelude::abortable($expr);
361        let _jh = $crate::runtime::prelude::spawn(proc);
362        abort_handle
363    }};
364}
365
366#[macro_export]
367macro_rules! spawn_as_abortable_named {
368    ($name:expr, $expr:expr $(,)?) => {{
369        let proc = $crate::runtime::diagnostics::instrument($expr, $name, module_path!(), file!(), line!());
370        let (proc, abort_handle) = $crate::runtime::prelude::abortable(proc);
371        let _jh = $crate::runtime::prelude::spawn(proc);
372        abort_handle
373    }};
374}
375
376/// Wrapper around [`futures::stream::Abortable`] that also automatically drops the inner stream when the abort handle
377/// is fired.
378///
379/// This is mostly useful to make MPSC channels close the channel when the AbortHandle is fired.
380/// If plain `futures::stream::Abortable` is used, the inner stream will not be dropped when the `AbortHandle` is fired,
381/// and therefore the senders can still send items, but the aborted receivers cannot process them.
382///
383/// For unbounded channels, this may cause the channels to keep memory growing indefinitely.
384/// For bounded channels, the senders will continue sending data until the channel is full and the senders will then
385/// become stuck indefinitely.
386///
387/// For streams that do not have active senders and only produce items when polled, this wrapper is not needed and
388/// the standard ` futures::stream::Abortable ` will work just fine.
389pub struct DropAbortable<St> {
390    inner: Option<futures::stream::Abortable<St>>,
391}
392
393impl<St> DropAbortable<St> {
394    pub fn new(stream: St) -> (Self, AbortHandle) {
395        let (abort_handle, reg) = futures::stream::AbortHandle::new_pair();
396        (Self::new_with_registration(stream, reg), abort_handle)
397    }
398
399    pub fn new_with_registration(stream: St, reg: futures::stream::AbortRegistration) -> Self {
400        Self {
401            inner: Some(futures::stream::Abortable::new(stream, reg)),
402        }
403    }
404}
405
406impl<St: futures::Stream + Unpin> futures::Stream for DropAbortable<St> {
407    type Item = St::Item;
408
409    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> {
410        let this = self.get_mut();
411        let Some(inner) = this.inner.as_mut() else {
412            return Poll::Ready(None);
413        };
414        match Pin::new(inner).poll_next(cx) {
415            Poll::Ready(None) => {
416                this.inner = None; // drops Abortable -> drops the receiver NOW
417                Poll::Ready(None)
418            }
419            other => other,
420        }
421    }
422}
423
424/// Abstraction over tasks that can be aborted (such as join or abort handles).
425#[auto_impl::auto_impl(&, Box, Arc)]
426pub trait Abortable {
427    /// Notifies the task that it should abort.
428    ///
429    /// Must be idempotent and not panic if it was already called before, due to implementation-specific
430    /// semantics of [`Abortable::was_aborted`].
431    fn abort_task(&self);
432
433    /// Returns `true` if [`abort_task`](Abortable::abort_task) was already called or the task has finished.
434    ///
435    /// It is implementation-specific whether `true` actually means that the task has been finished.
436    /// The [`Abortable::abort_task`] therefore can be also called if `true` is returned without a consequence.
437    fn was_aborted(&self) -> bool;
438}
439
440impl Abortable for AbortHandle {
441    fn abort_task(&self) {
442        self.abort();
443    }
444
445    fn was_aborted(&self) -> bool {
446        self.is_aborted()
447    }
448}
449
450#[cfg(feature = "runtime-tokio")]
451impl Abortable for tokio::task::JoinHandle<()> {
452    fn abort_task(&self) {
453        self.abort();
454    }
455
456    fn was_aborted(&self) -> bool {
457        self.is_finished()
458    }
459}
460
461/// List of [`Abortable`] tasks with each task identified by a unique key of type `T`.
462///
463/// Abortable objects, such as join or abort handles, do not by design abort when dropped.
464/// Sometimes this behavior is not desirable, and spawned run-away tasks may still continue to live
465/// e.g.: after an error is raised.
466///
467/// This object allows safely managing abortable tasks and will terminate all the tasks in reverse insertion order once
468/// dropped.
469///
470/// Additionally, this object also implements [`Abortable`] allowing it to be arbitrarily nested.
471pub struct AbortableList<T>(indexmap::IndexMap<T, Box<dyn Abortable + Send + Sync>>);
472
473impl<T> Default for AbortableList<T> {
474    fn default() -> Self {
475        Self(indexmap::IndexMap::new())
476    }
477}
478
479impl<T: std::fmt::Debug> std::fmt::Debug for AbortableList<T> {
480    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
481        f.debug_list().entries(self.0.keys()).finish()
482    }
483}
484
485impl<T: Hash + Eq> AbortableList<T> {
486    /// Appends a new [`abortable task`](Abortable) to the end of this list.
487    pub fn insert<A: Abortable + Send + Sync + 'static>(&mut self, process: T, task: A) {
488        self.0.insert(process, Box::new(task));
489    }
490
491    /// Checks if the list contains a task with the given key.
492    pub fn contains(&self, process: &T) -> bool {
493        self.0.contains_key(process)
494    }
495
496    /// Looks up a task by its key, removes it and aborts it.
497    ///
498    /// Returns `true` if the task was aborted and removed.
499    /// Otherwise, returns `false` (including a situation when the task was present but already aborted).
500    pub fn abort_one(&mut self, process: &T) -> bool {
501        if let Some(item) = self.0.shift_remove(process).filter(|t| !t.was_aborted()) {
502            item.abort_task();
503            true
504        } else {
505            false
506        }
507    }
508
509    /// Extends this list by appending `other`.
510    ///
511    /// The tasks from `other` are moved to this list without aborting them.
512    /// Afterward, `other` will be empty.
513    pub fn extend_from(&mut self, mut other: AbortableList<T>) {
514        self.0.extend(other.0.drain(..));
515    }
516
517    /// Extends this list by appending `other` while mapping its keys to the ones in this list.
518    ///
519    /// The tasks from `other` are moved to this list without aborting them.
520    /// Afterward, `other` will be empty.
521    pub fn flat_map_extend_from<U>(&mut self, mut other: AbortableList<U>, key_map: impl Fn(U) -> T) {
522        self.0.extend(other.0.drain(..).map(|(k, v)| (key_map(k), v)));
523    }
524}
525impl<T> AbortableList<T> {
526    /// Checks if the list is empty.
527    pub fn is_empty(&self) -> bool {
528        self.0.is_empty()
529    }
530
531    /// Returns the number of abortable tasks in the list.
532    pub fn size(&self) -> usize {
533        self.0.len()
534    }
535
536    /// Returns an iterator over the task names in the insertion order.
537    pub fn iter_names(&self) -> impl Iterator<Item = &T> {
538        self.0.keys()
539    }
540
541    /// Aborts all tasks in this list in the reverse insertion order.
542    ///
543    /// Skips tasks which were [already aborted](Abortable::was_aborted).
544    pub fn abort_all(&self) {
545        for (_, task) in self.0.iter().rev().filter(|(_, task)| !task.was_aborted()) {
546            task.abort_task();
547        }
548    }
549}
550
551impl<T> Abortable for AbortableList<T> {
552    fn abort_task(&self) {
553        self.abort_all();
554    }
555
556    fn was_aborted(&self) -> bool {
557        self.0.iter().all(|(_, task)| task.was_aborted())
558    }
559}
560
561impl<T> Drop for AbortableList<T> {
562    fn drop(&mut self) {
563        self.abort_all();
564        self.0.clear();
565    }
566}
567
568#[cfg(test)]
569mod tests {
570    use std::sync::{
571        Arc,
572        atomic::{AtomicBool, Ordering},
573    };
574
575    use super::*;
576
577    #[derive(Default)]
578    struct MockTask {
579        aborted: AtomicBool,
580    }
581
582    impl Abortable for MockTask {
583        fn abort_task(&self) {
584            self.aborted.store(true, Ordering::SeqCst);
585        }
586
587        fn was_aborted(&self) -> bool {
588            self.aborted.load(Ordering::SeqCst)
589        }
590    }
591
592    #[test]
593    fn test_insert_and_contains() {
594        let mut list = AbortableList::default();
595        let task1 = Arc::new(MockTask::default());
596        let task2 = Arc::new(MockTask::default());
597
598        list.insert("task1", task1.clone());
599        list.insert("task2", task2.clone());
600
601        assert!(list.contains(&"task1"));
602        assert!(list.contains(&"task2"));
603        assert!(!list.contains(&"task3"));
604        assert_eq!(list.size(), 2);
605        assert!(!list.is_empty());
606    }
607
608    #[test]
609    fn test_abort_one() {
610        let mut list = AbortableList::default();
611        let task1 = Arc::new(MockTask::default());
612
613        list.insert("task1", task1.clone());
614        assert!(list.abort_one(&"task1"));
615        assert!(task1.was_aborted());
616        assert!(!list.contains(&"task1"));
617        assert_eq!(list.size(), 0);
618
619        // Aborting already removed task
620        assert!(!list.abort_one(&"task1"));
621    }
622
623    #[test]
624    fn test_abort_one_already_aborted() {
625        let mut list = AbortableList::default();
626        let task1 = Arc::new(MockTask::default());
627        task1.abort_task();
628
629        list.insert("task1", task1.clone());
630        // abort_one returns false if already aborted
631        assert!(!list.abort_one(&"task1"));
632        // Check that it was still removed from the list even if already aborted
633        assert!(!list.contains(&"task1"));
634    }
635
636    #[test]
637    fn test_debug_impl() {
638        let mut list = AbortableList::default();
639        list.insert("task1", MockTask::default());
640        list.insert("task2", MockTask::default());
641        let debug_str = format!("{:?}", list);
642        assert!(debug_str.contains("task1"));
643        assert!(debug_str.contains("task2"));
644    }
645
646    #[test]
647    fn test_abort_all() {
648        let mut list = AbortableList::default();
649        let task1 = Arc::new(MockTask::default());
650        let task2 = Arc::new(MockTask::default());
651
652        list.insert(1, task1.clone());
653        list.insert(2, task2.clone());
654
655        list.abort_all();
656
657        assert!(task1.was_aborted());
658        assert!(task2.was_aborted());
659        // abort_all doesn't remove from list
660        assert_eq!(list.size(), 2);
661    }
662
663    #[test]
664    fn test_drop_aborts_all() {
665        let task1 = Arc::new(MockTask::default());
666        let task2 = Arc::new(MockTask::default());
667
668        {
669            let mut list = AbortableList::default();
670            list.insert(1, task1.clone());
671            list.insert(2, task2.clone());
672        }
673
674        assert!(task1.was_aborted());
675        assert!(task2.was_aborted());
676    }
677
678    #[test]
679    fn test_extend_from() {
680        let mut list1 = AbortableList::default();
681        let mut list2 = AbortableList::default();
682
683        let task1 = Arc::new(MockTask::default());
684        let task2 = Arc::new(MockTask::default());
685
686        list1.insert(1, task1.clone());
687        list2.insert(2, task2.clone());
688
689        list1.extend_from(list2);
690
691        assert_eq!(list1.size(), 2);
692        assert!(list1.contains(&1));
693        assert!(list1.contains(&2));
694
695        // Ensure task2 was not aborted during extend
696        assert!(!task2.was_aborted());
697    }
698
699    #[test]
700    fn test_flat_map_extend_from() {
701        let mut list1 = AbortableList::default();
702        let mut list2 = AbortableList::default();
703
704        let task1 = Arc::new(MockTask::default());
705        let task2 = Arc::new(MockTask::default());
706
707        list1.insert("a", task1.clone());
708        list2.insert(1, task2.clone());
709
710        list1.flat_map_extend_from(list2, |k| if k == 1 { "b" } else { "c" });
711
712        assert_eq!(list1.size(), 2);
713        assert!(list1.contains(&"a"));
714        assert!(list1.contains(&"b"));
715    }
716
717    #[test]
718    fn test_nested_abortable_list() {
719        let mut outer = AbortableList::default();
720        let mut inner = AbortableList::default();
721
722        let task1 = Arc::new(MockTask::default());
723        inner.insert(1, task1.clone());
724
725        outer.insert("inner", inner);
726
727        outer.abort_all();
728        assert!(task1.was_aborted());
729    }
730
731    #[test]
732    fn test_was_aborted_all() {
733        let mut list = AbortableList::default();
734        let task1 = Arc::new(MockTask::default());
735        let task2 = Arc::new(MockTask::default());
736
737        list.insert(1, task1.clone());
738        list.insert(2, task2.clone());
739
740        assert!(!list.was_aborted());
741
742        task1.abort_task();
743        assert!(!list.was_aborted());
744
745        task2.abort_task();
746        assert!(list.was_aborted());
747    }
748
749    #[test]
750    fn test_iter_names() {
751        let mut list = AbortableList::default();
752        list.insert("a", MockTask::default());
753        list.insert("b", MockTask::default());
754        list.insert("c", MockTask::default());
755
756        let names: Vec<&&str> = list.iter_names().collect();
757        assert_eq!(names, vec![&"a", &"b", &"c"]);
758    }
759
760    #[test]
761    fn test_reverse_insertion_order_on_abort() {
762        use std::sync::Mutex;
763        let abort_order = Arc::new(Mutex::new(Vec::new()));
764
765        struct OrderedMockTask {
766            id: i32,
767            order: Arc<Mutex<Vec<i32>>>,
768        }
769
770        impl Abortable for OrderedMockTask {
771            fn abort_task(&self) {
772                self.order.lock().unwrap().push(self.id);
773            }
774
775            fn was_aborted(&self) -> bool {
776                self.order.lock().unwrap().contains(&self.id)
777            }
778        }
779
780        let mut list = AbortableList::default();
781        list.insert(
782            1,
783            OrderedMockTask {
784                id: 1,
785                order: abort_order.clone(),
786            },
787        );
788        list.insert(
789            2,
790            OrderedMockTask {
791                id: 2,
792                order: abort_order.clone(),
793            },
794        );
795        list.insert(
796            3,
797            OrderedMockTask {
798                id: 3,
799                order: abort_order.clone(),
800            },
801        );
802
803        list.abort_all();
804
805        let order = abort_order.lock().unwrap();
806        assert_eq!(*order, vec![3, 2, 1]);
807    }
808
809    #[tokio::test]
810    async fn test_drop_abortable_drops_inner_stream_on_abort() {
811        use std::{
812            sync::atomic::{AtomicBool, Ordering},
813            time::Duration,
814        };
815
816        use futures::{channel::mpsc, stream::StreamExt};
817
818        // Track if the receiver was dropped
819        let receiver_dropped = Arc::new(AtomicBool::new(false));
820
821        let (tx, rx) = mpsc::channel::<i32>(16);
822
823        // Clone for the spawn
824        let receiver_dropped_clone = receiver_dropped.clone();
825
826        // Create DropAbortable wrapping a channel stream
827        let (drop_abortable, abort_handle) = DropAbortable::new(rx);
828
829        // Spawn task that processes items from the stream
830        let handle = tokio::spawn(async move {
831            let mut rx = drop_abortable;
832            while let Some(_item) = rx.next().await {
833                // Process items
834            }
835            // Stream ended - this is only reached when stream is fully consumed or dropped
836            receiver_dropped_clone.store(true, Ordering::SeqCst);
837        });
838
839        // Send some items using the original sender (not clone)
840        let mut tx = tx;
841        tx.try_send(1).unwrap();
842        tx.try_send(2).unwrap();
843        tx.try_send(3).unwrap();
844
845        // Give some time for processing
846        tokio::time::sleep(Duration::from_millis(50)).await;
847
848        // Now abort - this should drop the inner stream immediately
849        abort_handle.abort();
850
851        // Drop tx after abort to ensure we're testing abort-triggered behavior
852        drop(tx);
853
854        // Wait for the task to finish
855        let _ = handle.await;
856
857        // Verify the receiver was dropped when abort was triggered
858        assert!(
859            receiver_dropped.load(Ordering::SeqCst),
860            "Receiver should be dropped when abort handle is fired"
861        );
862    }
863
864    #[tokio::test]
865    async fn test_drop_abortable_closes_channel_on_abort() {
866        use std::time::Duration;
867
868        use futures::{channel::mpsc, stream::StreamExt};
869
870        // Small buffer
871        let (tx, rx) = mpsc::channel::<i32>(2);
872
873        let (drop_abortable, abort_handle) = DropAbortable::new(rx);
874
875        // Spawn task that processes items very slowly (never actually)
876        let handle = tokio::spawn(async move {
877            let mut rx = drop_abortable;
878            // Don't process any items - just wait to be aborted
879            tokio::time::sleep(Duration::from_secs(10)).await;
880            while let Some(_item) = rx.next().await {
881                tokio::time::sleep(Duration::from_secs(10)).await;
882            }
883        });
884
885        // Send some items
886        let mut tx = tx;
887        tx.try_send(1).unwrap();
888        tx.try_send(2).unwrap();
889
890        // Abort - this should drop the receiver, closing the channel
891        abort_handle.abort();
892
893        // Wait for the task to finish
894        let _ = handle.await;
895
896        // Now try_send should fail because receiver is gone (channel closed)
897        // Must fail specifically due to disconnection, not fullness
898        let err = tx.try_send(4).expect_err("expected channel closure after abort");
899        assert!(err.is_disconnected(), "expected disconnected sender after abort");
900    }
901}
hopr_utils/runtime.rs

hopr_utils/
runtime.rs
