hopr_transport_session/
utils.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::Arc,
5    task::{Context, Poll, Waker},
6};
7
8use futures::{FutureExt, SinkExt, StreamExt, TryStreamExt};
9use hopr_api::types::internal::routing::DestinationRouting;
10use hopr_protocol_app::prelude::{ApplicationData, ApplicationDataOut};
11use hopr_protocol_start::{KeepAliveFlag, KeepAliveMessage};
12/// Bidirectionally copies data between a [`HoprSession`](crate::HoprSession) and an arbitrary
13/// async IO stream.
14///
15/// Re-exported from `hopr-utils`; the published generic version accepts any two
16/// `AsyncRead + AsyncWrite` streams. Only available with Tokio.
17#[cfg(feature = "runtime-tokio")]
18pub use hopr_utils::network_types::utils::{transfer_session, transfer_session_datagram};
19use hopr_utils::runtime::AbortHandle;
20use tracing::{Instrument, debug, error, instrument};
21
22use crate::{
23    AtomicSurbFlowEstimator, SessionId,
24    balancer::{BalancerStateValues, RateController, RateLimitStreamExt},
25    errors::TransportSessionError,
26    types::HoprStartProtocol,
27};
28
29/// Runtime-agnostic multi-waker notification primitive.
30///
31/// Uses a generation counter to detect notification events: [`notify_waiters`](SlotNotify::notify_waiters)
32/// bumps the generation, and [`notified`](SlotNotify::notified) futures compare the generation at
33/// creation time against the current one on each `poll()`. This prevents two
34/// race conditions a simple waker-vector approach cannot handle:
35///
36/// 1. **Latent wake.** A [`notified`](SlotNotify::notified) future registers its waker on first `poll()`. If
37///    [`notify_waiters`](SlotNotify::notify_waiters) fires between the creation of the future and that first `poll`,
38///    the waker-vector is empty and the notification is lost. With a generation counter, `gen_at_creation` already
39///    captures the pre-notification value, so the first `poll()` sees the advanced generation and returns `Ready`.
40///
41/// 2. **Spurious `Ready`.** A second `poll()` of an already-registered future can return `Ready` unconditionally if the
42///    notification check is only done on registration. Generation re-check on every `poll()` prevents this.
43///
44/// Uses unique IDs per waiter so that dropping a future cleanly
45/// removes its waker from the vector — no waker leak on cancellation.
46/// No tokio dependency — works with any async runtime.
47///
48/// Clone is cheap (clones the inner `Arc`).
49#[derive(Clone)]
50pub struct SlotNotify {
51    inner: Arc<parking_lot::Mutex<SlotNotifyInner>>,
52}
53
54struct SlotNotifyInner {
55    wakers: Vec<(u64, Waker)>,
56    next_id: u64,
57    generation: u64,
58}
59
60impl SlotNotify {
61    pub fn new() -> Self {
62        Self {
63            inner: Arc::new(parking_lot::Mutex::new(SlotNotifyInner {
64                wakers: Vec::new(),
65                next_id: 0,
66                generation: 0,
67            })),
68        }
69    }
70
71    /// Wake all parked waiters.
72    ///
73    /// Bumps the generation counter **before** draining wakers so that
74    /// futures created between this call and their first `poll()` see the
75    /// advanced generation and return `Ready` immediately.
76    ///
77    /// The guard is released before any waker runs. `Waker::wake` executes arbitrary executor code,
78    /// and an executor that polls the woken task inline re-enters [`SlotNotifyFuture::poll`] — or its
79    /// `Drop` — both of which lock this same non-reentrant mutex. Waking under the guard is therefore
80    /// a self-deadlock against exactly the runtimes this type advertises support for; tokio only
81    /// enqueues, so the shipped path never hit it. Draining into a local buffer costs one allocation
82    /// per notification that actually has waiters to wake, and nothing at all when it does not.
83    pub fn notify_waiters(&self) {
84        let woken = {
85            let mut inner = self.inner.lock();
86            inner.generation += 1;
87            // `drain` rather than `mem::take`: an empty drain allocates nothing and the inner
88            // `Vec` keeps its capacity, which is the common case on a gate that rarely parks.
89            inner.wakers.drain(..).collect::<Vec<_>>()
90        };
91        for (_, waker) in woken {
92            waker.wake();
93        }
94    }
95
96    /// Return a future that completes the next time `notify_waiters` is called.
97    ///
98    /// The future captures the current generation at creation time. A
99    /// concurrent [`notify_waiters`](SlotNotify::notify_waiters) call that fires before the first
100    /// `poll()` will have bumped the generation, so the future still
101    /// completes — no notification is lost.
102    pub fn notified(&self) -> SlotNotifyFuture {
103        let generation = self.inner.lock().generation;
104        SlotNotifyFuture {
105            inner: self.inner.clone(),
106            waker_id: 0,
107            registered: false,
108            gen_at_creation: generation,
109        }
110    }
111}
112
113/// Future returned by [`SlotNotify::notified`].
114///
115/// On cancellation (drop without completion), the registered waker is
116/// automatically removed from [`SlotNotify`] so stale entries are never
117/// left behind.
118pub struct SlotNotifyFuture {
119    inner: Arc<parking_lot::Mutex<SlotNotifyInner>>,
120    waker_id: u64,
121    registered: bool,
122    gen_at_creation: u64,
123}
124
125impl Drop for SlotNotifyFuture {
126    fn drop(&mut self) {
127        if self.registered {
128            self.inner.lock().wakers.retain(|(id, _)| *id != self.waker_id);
129        }
130    }
131}
132
133impl Future for SlotNotifyFuture {
134    type Output = ();
135
136    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
137        let this = self.get_mut();
138        let mut inner = this.inner.lock();
139
140        // If the generation advanced past our creation snapshot, a
141        // notification happened — we are done.
142        if inner.generation != this.gen_at_creation {
143            return Poll::Ready(());
144        }
145
146        if this.registered {
147            // Already registered but generation hasn't advanced — this is a
148            // spurious wake (or a replaced waker). Update the stored waker
149            // and stay Pending.
150            if let Some((_, w)) = inner.wakers.iter_mut().find(|(id, _)| *id == this.waker_id) {
151                // `will_wake` first: a re-poll of the same task is the common case here — the gate
152                // parks callers that are then polled spuriously — and the clone costs a refcount
153                // bump, or an allocation for a waker that is not `Arc`-backed. Replacing only when
154                // the target actually changed keeps the correctness property that matters, which is
155                // that the *latest* waker is the one stored: `will_wake` returning true means
156                // waking the stored one wakes the same task.
157                if !w.will_wake(cx.waker()) {
158                    *w = cx.waker().clone();
159                }
160            }
161            return Poll::Pending;
162        }
163
164        // First poll: register.
165        this.waker_id = inner.next_id;
166        inner.next_id += 1;
167        inner.wakers.push((this.waker_id, cx.waker().clone()));
168        this.registered = true;
169        Poll::Pending
170    }
171}
172
173/// This function will use the given generator to generate an initial seeding key.
174/// It will check whether the given cache already contains a value for that key, and if not,
175/// calls the generator (with the previous value) to generate a new seeding key and retry.
176/// The function either finds a suitable free slot, inserting value generated by `value_fn` and returns the found key,
177/// or terminates with `None` when `gen` returns the initial seed again.
178pub(crate) fn insert_into_next_slot<F, K, U, V>(
179    cache: &moka::sync::Cache<K, V>,
180    mut generator: F,
181    value_fn: U,
182    max_capacity: Option<u64>,
183) -> Option<(K, V)>
184where
185    F: FnMut(Option<K>) -> K,
186    K: Copy + std::hash::Hash + Eq + Send + Sync + 'static,
187    U: FnOnce(K) -> V,
188    V: Clone + Send + Sync + 'static,
189{
190    cache.run_pending_tasks();
191
192    // Reject when the cache is already at capacity to avoid Moka evicting an
193    // existing entry before we can insert the new one.
194    if let Some(max) = max_capacity
195        && cache.entry_count() >= max
196    {
197        return None;
198    }
199
200    // Wrap the FnOnce so we can "consume" it exactly once,
201    // but only when we actually insert into a free slot.
202    let value_fn = std::sync::Arc::new(parking_lot::Mutex::new(Some(value_fn)));
203
204    let initial = generator(None);
205    let mut next = initial;
206    loop {
207        let value_fn = value_fn.clone();
208        let insertion_result = cache.entry(next).and_compute_with(move |e| {
209            if e.is_none() {
210                let f = value_fn
211                    .lock()
212                    .take()
213                    .expect("impossible: value_fn was already consumed");
214
215                moka::ops::compute::Op::Put(f(next))
216            } else {
217                moka::ops::compute::Op::Nop
218            }
219        });
220
221        // If we inserted successfully, break the loop and return the insertion key
222        if let moka::ops::compute::CompResult::Inserted(val) = insertion_result {
223            return Some((next, val.into_value()));
224        }
225
226        // Otherwise, generate the next key
227        next = generator(Some(next));
228
229        // If generated keys made it to full loop, return failure
230        if next == initial {
231            return None;
232        }
233    }
234}
235
236/// Indicates whether the [keep-alive stream](spawn_keep_alive_stream) should notify the Session counterparty
237/// about the SURB target (Entry) or SURB level (Exit).
238#[derive(Debug, Clone)]
239pub(crate) enum SurbNotificationMode {
240    /// No keep-alive messages are sent to the Session counterparty.
241    DoNotNotify,
242    /// Session initiator notifies the Session recipient about the desired SURB target level.
243    Target,
244    /// Session recipient notifies the Session initiator about the current SURB level.
245    Level(AtomicSurbFlowEstimator),
246}
247
248/// Spawns a task for a rate-limited stream of Keep-Alive messages to the Session counterparty.
249///
250/// `controller` is supplied rather than created here because the Exit shares this one stream between
251/// two independent producers — its SURB-level notification and the PIX filler — and both the rate and
252/// the admission decision below have to be steered from the same handle. The stream is suspended
253/// until the controller is given a non-zero rate.
254///
255/// `admit` is consulted once per released packet and drops it when it answers `false`. It is a
256/// predicate rather than a sink wrapper because the sink here must stay `Clone` — `try_for_each_concurrent`
257/// clones it per message — and the natural wrapper for a sink that may emit nothing,
258/// `SinkExt::with_flat_map`, is not. Filtering after the rate limiter also puts the decision where
259/// the Exit wants it: the rate says how often a packet *may* go, and this says whether this
260/// particular one should.
261#[instrument(level = "debug", skip(sender, routing, notification_mode, cfg, controller, admit))]
262pub(crate) fn spawn_keep_alive_stream<S>(
263    session_id: SessionId,
264    sender: S,
265    routing: DestinationRouting,
266    notification_mode: SurbNotificationMode,
267    cfg: std::sync::Arc<BalancerStateValues>,
268    controller: &RateController,
269    admit: impl Fn() -> bool + Send + 'static,
270) -> AbortHandle
271where
272    S: futures::Sink<(DestinationRouting, ApplicationDataOut)> + Clone + Send + Sync + Unpin + 'static,
273    S::Error: std::error::Error + Send + Sync + 'static,
274{
275    // DropAbortable not needed because the stream only generates items when polled
276    let (ka_stream, abort_handle) = futures::stream::abortable(
277        futures::stream::repeat_with(move || match &notification_mode {
278            SurbNotificationMode::Target => HoprStartProtocol::KeepAlive(KeepAliveMessage {
279                session_id,
280                flags: KeepAliveFlag::BalancerTarget.into(),
281                additional_data: cfg.target_surb_buffer_size.load(std::sync::atomic::Ordering::Relaxed),
282            }),
283            // Report the level this side acts on: the balancer's accumulator, fed by
284            // `produced - consumed - evicted` and clamped to the store capacity on every update. The
285            // raw estimator diff now tracks true occupancy too, but the accumulator additionally
286            // collapses to 0 while the return path is degraded, which is exactly the moment the Entry
287            // should mint at maximum rather than read the counterparty as full. The estimator stays in
288            // the enum for the callers.
289            SurbNotificationMode::Level(_estimator) => HoprStartProtocol::KeepAlive(KeepAliveMessage {
290                session_id,
291                flags: KeepAliveFlag::BalancerState.into(),
292                additional_data: cfg.buffer_level(),
293            }),
294            SurbNotificationMode::DoNotNotify => HoprStartProtocol::KeepAlive(KeepAliveMessage {
295                session_id,
296                flags: None.into(),
297                additional_data: 0,
298            }),
299        })
300        .rate_limit_with_controller(controller)
301        .filter(move |_| futures::future::ready(admit())),
302    );
303
304    let sender_clone = sender.clone();
305    let fwd_routing_clone = routing.clone();
306
307    // This task will automatically terminate once the returned abort handle is used.
308    debug!(%session_id, "spawning keep-alive stream");
309    let keep_alive_diag = hopr_utils::runtime::diagnostics::ConcurrentDiagnostics::new(
310        "session_keep_alive_try_for_each_concurrent",
311        module_path!(),
312        file!(),
313        line!(),
314    );
315    hopr_utils::runtime::prelude::spawn(hopr_utils::runtime::diagnostics::instrument(
316        ka_stream
317            .map(move |msg| {
318                ApplicationData::try_from(msg)
319                    .map(|data| (fwd_routing_clone.clone(), ApplicationDataOut::with_no_packet_info(data)))
320            })
321            .map_err(TransportSessionError::from)
322            .try_for_each_concurrent(None, move |msg| {
323                let mut sender_clone = sender_clone.clone();
324                let keep_alive_diag = keep_alive_diag.clone();
325                keep_alive_diag.wrap(|| async move {
326                    sender_clone
327                        .send(msg)
328                        .await
329                        .map_err(TransportSessionError::packet_sending)
330                })
331            })
332            .then(move |res| {
333                match res {
334                    Ok(_) => tracing::debug!(
335                        component = "session",
336                        %session_id,
337                        task = "session keepalive",
338                        "background task finished"
339                    ),
340                    Err(error) => error!(%session_id, %error, "keep-alive stream failed"),
341                }
342                futures::future::ready(())
343            })
344            .in_current_span(),
345        "session_keep_alive",
346        module_path!(),
347        file!(),
348        line!(),
349    ));
350
351    abort_handle
352}
353
354#[cfg(test)]
355mod tests {
356    use std::time::Duration;
357
358    use anyhow::anyhow;
359    use futures::FutureExt;
360
361    use super::*;
362
363    /// Generator that cycles through 0..4, wrapping at 5 back to 0.
364    fn cycling_generator(prev: Option<u8>) -> u8 {
365        prev.map(|v| (v + 1) % 5).unwrap_or(0)
366    }
367
368    /// Tests sequential insertion into an empty cache: each call fills the next slot.
369    #[tokio::test]
370    async fn test_insert_into_next_slot_sequential() -> anyhow::Result<()> {
371        let cache = moka::sync::Cache::new(10);
372
373        for i in 0..5 {
374            let (k, v) = insert_into_next_slot(&cache, cycling_generator, |k| format!("foo_{k}"), Some(10u64))
375                .ok_or(anyhow!("should insert into slot {i}"))?;
376            assert_eq!(k, i);
377            assert_eq!(format!("foo_{i}"), v);
378            assert_eq!(Some(v), cache.get(&i));
379        }
380
381        Ok(())
382    }
383
384    /// Tests that insertion returns `None` when all slots are occupied and the generator cycles back.
385    #[tokio::test]
386    async fn test_insert_into_next_slot_returns_none_when_full() -> anyhow::Result<()> {
387        let cache = moka::sync::Cache::new(10);
388
389        for _ in 0..5 {
390            insert_into_next_slot(&cache, cycling_generator, |k| format!("foo_{k}"), Some(10u64))
391                .ok_or(anyhow!("precondition: should insert"))?;
392        }
393
394        assert!(
395            insert_into_next_slot(&cache, cycling_generator, |_| "foo".to_string(), Some(10u64)).is_none(),
396            "must not find slot when full"
397        );
398
399        Ok(())
400    }
401
402    /// Tests that a cache with max capacity of 1 rejects a second distinct key.
403    #[tokio::test]
404    async fn test_insert_into_next_slot_capacity_one_rejects_second_key() -> anyhow::Result<()> {
405        let unit_cache = moka::sync::Cache::new(1);
406
407        let (k0, _v0) = insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64))
408            .ok_or(anyhow!("first insertion must succeed"))?;
409        assert_eq!(k0, 0);
410
411        assert!(
412            insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64)).is_none(),
413            "second distinct key must be rejected when cache capacity is 1"
414        );
415
416        Ok(())
417    }
418
419    /// Tests that a rejected insertion does not evict the existing entry.
420    #[tokio::test]
421    async fn test_insert_into_next_slot_rejected_insertion_does_not_evict() -> anyhow::Result<()> {
422        let unit_cache = moka::sync::Cache::new(1);
423
424        let (k0, v0) = insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64))
425            .ok_or(anyhow!("first insertion must succeed"))?;
426
427        insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64));
428
429        assert_eq!(
430            Some(v0),
431            unit_cache.get(&k0),
432            "first entry must still be present after rejection"
433        );
434
435        Ok(())
436    }
437
438    fn noop_waker() -> std::task::Waker {
439        futures::task::noop_waker()
440    }
441
442    // -------------------------------------------------------------------
443    // Generation-counter wake detection
444    // -------------------------------------------------------------------
445
446    #[test]
447    fn notification_after_first_poll_completes_future() {
448        let n = Arc::new(SlotNotify::new());
449        let mut fut = n.notified();
450
451        let waker = noop_waker();
452        let mut cx = std::task::Context::from_waker(&waker);
453
454        // First poll registers and returns Pending.
455        assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending);
456
457        n.notify_waiters();
458
459        // Second poll sees the advanced generation -> Ready.
460        assert_eq!(fut.poll_unpin(&mut cx), Poll::Ready(()));
461    }
462
463    #[test]
464    fn notify_after_creation_before_first_poll_completes() {
465        let n = Arc::new(SlotNotify::new());
466        let mut fut = n.notified();
467
468        // Notify AFTER creation but BEFORE first poll.
469        n.notify_waiters();
470
471        let waker = noop_waker();
472        let mut cx = std::task::Context::from_waker(&waker);
473
474        // First poll should see the advanced generation -> Ready.
475        assert_eq!(fut.poll_unpin(&mut cx), Poll::Ready(()));
476    }
477
478    #[test]
479    fn spurious_repoll_stays_pending_without_notification() {
480        let n = Arc::new(SlotNotify::new());
481        let mut fut = n.notified();
482
483        let waker = noop_waker();
484        let mut cx = std::task::Context::from_waker(&waker);
485
486        // First poll returns Pending.
487        assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending);
488
489        // Second poll without any notify -> spurious wake, stays Pending.
490        assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending);
491    }
492
493    #[test]
494    fn multiple_notifications_bump_generation() {
495        let n = Arc::new(SlotNotify::new());
496        let mut fut = n.notified();
497
498        let waker = noop_waker();
499        let mut cx = std::task::Context::from_waker(&waker);
500
501        assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending);
502
503        n.notify_waiters();
504        assert_eq!(fut.poll_unpin(&mut cx), Poll::Ready(()));
505
506        // A new future should also work for the next notification cycle.
507        let mut fut2 = n.notified();
508        assert_eq!(fut2.poll_unpin(&mut cx), Poll::Pending);
509
510        n.notify_waiters();
511        assert_eq!(fut2.poll_unpin(&mut cx), Poll::Ready(()));
512    }
513
514    #[test]
515    fn drop_uncompleted_future_removes_waker() {
516        let n = Arc::new(SlotNotify::new());
517        {
518            let mut fut = n.notified();
519
520            let waker = noop_waker();
521            let mut cx = std::task::Context::from_waker(&waker);
522
523            assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending);
524            // fut drops here -> waker should be removed.
525        }
526        // After the future is dropped, the inner waker list should be empty.
527        assert!(n.inner.lock().wakers.is_empty());
528    }
529
530    /// `notify_waiters` must have released the lock by the time any waker runs.
531    ///
532    /// `Waker::wake` runs arbitrary executor code, and an executor that polls the woken task inline
533    /// re-enters this type: both [`SlotNotifyFuture::poll`] and its `Drop` take the same
534    /// `parking_lot::Mutex`, which is not reentrant. Waking under the guard self-deadlocks against
535    /// precisely the runtime-agnostic contract this type advertises. Tokio only enqueues, so no
536    /// shipped caller ever exercised it and no test could have caught it.
537    ///
538    /// The waker below stands in for such an executor. Note the failure mode: on a regression this
539    /// test does not fail, it *hangs*, and surfaces as a timeout. That is the only observable a
540    /// deadlock has.
541    #[test]
542    fn notify_waiters_releases_the_lock_before_waking() {
543        struct ReentrantWaker(SlotNotify);
544
545        impl std::task::Wake for ReentrantWaker {
546            fn wake(self: Arc<Self>) {
547                // What an inline-polling executor does: re-enter the notifier from inside `wake`.
548                // `notified` takes the very lock `notify_waiters` is in the middle of.
549                drop(self.0.notified());
550            }
551        }
552
553        let n = SlotNotify::new();
554        let mut fut = n.notified();
555        let waker = std::task::Waker::from(Arc::new(ReentrantWaker(n.clone())));
556        let mut cx = std::task::Context::from_waker(&waker);
557
558        assert_eq!(
559            fut.poll_unpin(&mut cx),
560            Poll::Pending,
561            "the first poll registers the waker"
562        );
563
564        n.notify_waiters();
565
566        assert_eq!(
567            fut.poll_unpin(&mut cx),
568            Poll::Ready(()),
569            "the generation advanced, so the waiter is done"
570        );
571    }
572
573    /// One `notify_waiters` call must wake *every* parked waiter, not just one.
574    ///
575    /// `wakers` is a `Vec` and the type calls itself a multi-waker primitive, but every other test
576    /// here drives a single future, so the vector never holds more than one entry and a `pop()` in
577    /// place of `drain(..)` would pass all of them.
578    ///
579    /// Asserted on the wakes rather than on the futures completing, which is the part that is easy
580    /// to get wrong: the generation counter makes `poll` return `Ready` after any notification
581    /// whether or not *that* waiter was woken, so polling every future to completion would pass with
582    /// a single-waker regression. What a real executor needs is the wake itself — without it the
583    /// future is never polled again and the generation is never read.
584    #[test]
585    fn one_notification_wakes_every_parked_waiter() {
586        struct CountingWaker(std::sync::atomic::AtomicUsize);
587
588        impl std::task::Wake for CountingWaker {
589            fn wake(self: Arc<Self>) {
590                self.wake_by_ref();
591            }
592
593            fn wake_by_ref(self: &Arc<Self>) {
594                self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
595            }
596        }
597
598        const WAITERS: usize = 4;
599
600        let n = SlotNotify::new();
601        let counters: Vec<Arc<CountingWaker>> = (0..WAITERS)
602            .map(|_| Arc::new(CountingWaker(std::sync::atomic::AtomicUsize::new(0))))
603            .collect();
604        let mut futures: Vec<_> = (0..WAITERS).map(|_| n.notified()).collect();
605
606        for (i, (fut, counter)) in futures.iter_mut().zip(&counters).enumerate() {
607            let waker = std::task::Waker::from(counter.clone());
608            let mut cx = std::task::Context::from_waker(&waker);
609            assert_eq!(fut.poll_unpin(&mut cx), Poll::Pending, "waiter {i} must park");
610        }
611        assert_eq!(
612            WAITERS,
613            n.inner.lock().wakers.len(),
614            "every waiter must have registered"
615        );
616
617        n.notify_waiters();
618
619        for (i, counter) in counters.iter().enumerate() {
620            assert_eq!(
621                1,
622                counter.0.load(std::sync::atomic::Ordering::Relaxed),
623                "waiter {i} must be woken exactly once by the single notification"
624            );
625        }
626    }
627
628    /// Re-polling with a *different* waker must replace the stored one.
629    ///
630    /// The `will_wake` guard in `poll` skips the clone when the stored waker would wake the same
631    /// task, which is the common case and the point of the guard. Getting its polarity backwards
632    /// would keep the first waker forever: the future would then be parked against a task that is no
633    /// longer polling it, and nothing would wake the one that is — a hang, not a slowdown. This is
634    /// the half of that guard worth pinning; the skipped clone is unobservable by design.
635    #[test]
636    fn repolling_with_a_new_waker_replaces_the_stored_one() {
637        struct CountingWaker(std::sync::atomic::AtomicUsize);
638
639        impl std::task::Wake for CountingWaker {
640            fn wake(self: Arc<Self>) {
641                self.wake_by_ref();
642            }
643
644            fn wake_by_ref(self: &Arc<Self>) {
645                self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
646            }
647        }
648
649        let n = SlotNotify::new();
650        let mut fut = n.notified();
651
652        let first = Arc::new(CountingWaker(std::sync::atomic::AtomicUsize::new(0)));
653        let second = Arc::new(CountingWaker(std::sync::atomic::AtomicUsize::new(0)));
654
655        let w1 = std::task::Waker::from(first.clone());
656        assert_eq!(
657            fut.poll_unpin(&mut std::task::Context::from_waker(&w1)),
658            Poll::Pending,
659            "the first poll registers"
660        );
661
662        // A distinct task, so `will_wake` is false and the replacement must happen.
663        let w2 = std::task::Waker::from(second.clone());
664        assert_eq!(
665            fut.poll_unpin(&mut std::task::Context::from_waker(&w2)),
666            Poll::Pending,
667            "no notification has happened, so the future stays parked"
668        );
669
670        n.notify_waiters();
671
672        assert_eq!(
673            0,
674            first.0.load(std::sync::atomic::Ordering::Relaxed),
675            "the superseded waker must not be the one woken"
676        );
677        assert_eq!(
678            1,
679            second.0.load(std::sync::atomic::Ordering::Relaxed),
680            "the waker from the most recent poll must be woken exactly once"
681        );
682    }
683
684    // -------------------------------------------------------------------
685    // Async tests (tokio runtime)
686    // -------------------------------------------------------------------
687
688    #[tokio::test]
689    async fn notified_awaited_completes_on_notify() {
690        let n = Arc::new(SlotNotify::new());
691        let n2 = n.clone();
692
693        let handle = tokio::spawn(async move {
694            n2.notified().await;
695            42u32
696        });
697
698        tokio::time::sleep(Duration::from_millis(10)).await;
699        n.notify_waiters();
700
701        let result = tokio::time::timeout(Duration::from_secs(1), handle)
702            .await
703            .expect("timeout")
704            .expect("task panicked");
705        assert_eq!(result, 42);
706    }
707}
hopr_transport_session/utils.rs

hopr_transport_session/
utils.rs
