hopr_transport_session/
utils.rs

1use std::time::Duration;
2
3use futures::{FutureExt, SinkExt, StreamExt, TryStreamExt};
4use hopr_api::types::internal::routing::DestinationRouting;
5use hopr_protocol_app::prelude::{ApplicationData, ApplicationDataOut};
6use hopr_protocol_start::{KeepAliveFlag, KeepAliveMessage};
7/// Bidirectionally copies data between a [`HoprSession`](crate::HoprSession) and an arbitrary
8/// async IO stream.
9///
10/// Re-exported from `hopr-utils`; the published generic version accepts any two
11/// `AsyncRead + AsyncWrite` streams. Only available with Tokio.
12#[cfg(feature = "runtime-tokio")]
13pub use hopr_utils::network_types::utils::transfer_session;
14use hopr_utils::runtime::AbortHandle;
15use tracing::{Instrument, debug, error, instrument};
16
17use crate::{
18    AtomicSurbFlowEstimator, SessionId,
19    balancer::{BalancerStateValues, RateController, RateLimitStreamExt, SurbFlowEstimator},
20    errors::TransportSessionError,
21    types::HoprStartProtocol,
22};
23
24/// This function will use the given generator to generate an initial seeding key.
25/// It will check whether the given cache already contains a value for that key, and if not,
26/// calls the generator (with the previous value) to generate a new seeding key and retry.
27/// The function either finds a suitable free slot, inserting value generated by `value_fn` and returns the found key,
28/// or terminates with `None` when `gen` returns the initial seed again.
29pub(crate) fn insert_into_next_slot<F, K, U, V>(
30    cache: &moka::sync::Cache<K, V>,
31    mut generator: F,
32    value_fn: U,
33    max_capacity: Option<u64>,
34) -> Option<(K, V)>
35where
36    F: FnMut(Option<K>) -> K,
37    K: Copy + std::hash::Hash + Eq + Send + Sync + 'static,
38    U: FnOnce(K) -> V,
39    V: Clone + Send + Sync + 'static,
40{
41    cache.run_pending_tasks();
42
43    // Reject when the cache is already at capacity to avoid Moka evicting an
44    // existing entry before we can insert the new one.
45    if let Some(max) = max_capacity
46        && cache.entry_count() >= max
47    {
48        return None;
49    }
50
51    // Wrap the FnOnce so we can "consume" it exactly once,
52    // but only when we actually insert into a free slot.
53    let value_fn = std::sync::Arc::new(parking_lot::Mutex::new(Some(value_fn)));
54
55    let initial = generator(None);
56    let mut next = initial;
57    loop {
58        let value_fn = value_fn.clone();
59        let insertion_result = cache.entry(next).and_compute_with(move |e| {
60            if e.is_none() {
61                let f = value_fn
62                    .lock()
63                    .take()
64                    .expect("impossible: value_fn was already consumed");
65
66                moka::ops::compute::Op::Put(f(next))
67            } else {
68                moka::ops::compute::Op::Nop
69            }
70        });
71
72        // If we inserted successfully, break the loop and return the insertion key
73        if let moka::ops::compute::CompResult::Inserted(val) = insertion_result {
74            return Some((next, val.into_value()));
75        }
76
77        // Otherwise, generate the next key
78        next = generator(Some(next));
79
80        // If generated keys made it to full loop, return failure
81        if next == initial {
82            return None;
83        }
84    }
85}
86
87/// Indicates whether the [keep-alive stream](spawn_keep_alive_stream) should notify the Session counterparty
88/// about the SURB target (Entry) or SURB level (Exit).
89#[derive(Debug, Clone)]
90pub(crate) enum SurbNotificationMode {
91    /// No keep-alive messages are sent to the Session counterparty.
92    DoNotNotify,
93    /// Session initiator notifies the Session recipient about the desired SURB target level.
94    Target,
95    /// Session recipient notifies the Session initiator about the current SURB level.
96    Level(AtomicSurbFlowEstimator),
97}
98
99/// Spawns a task for a rate-limited stream of Keep-Alive messages to the Session counterparty.
100#[instrument(level = "debug", skip(sender, routing, notification_mode, cfg))]
101pub(crate) fn spawn_keep_alive_stream<S>(
102    session_id: SessionId,
103    sender: S,
104    routing: DestinationRouting,
105    notification_mode: SurbNotificationMode,
106    cfg: std::sync::Arc<BalancerStateValues>,
107) -> (RateController, AbortHandle)
108where
109    S: futures::Sink<(DestinationRouting, ApplicationDataOut)> + Clone + Send + Sync + Unpin + 'static,
110    S::Error: std::error::Error + Send + Sync + 'static,
111{
112    // The stream is suspended until the caller sets a rate via the Controller
113    let controller = RateController::new(0, Duration::from_secs(1));
114
115    // DropAbortable not needed because the stream only generates items when polled
116    let (ka_stream, abort_handle) = futures::stream::abortable(
117        futures::stream::repeat_with(move || match &notification_mode {
118            SurbNotificationMode::Target => HoprStartProtocol::KeepAlive(KeepAliveMessage {
119                session_id,
120                flags: KeepAliveFlag::BalancerTarget.into(),
121                additional_data: cfg.target_surb_buffer_size.load(std::sync::atomic::Ordering::Relaxed),
122            }),
123            SurbNotificationMode::Level(estimator) => HoprStartProtocol::KeepAlive(KeepAliveMessage {
124                session_id,
125                flags: KeepAliveFlag::BalancerState.into(),
126                additional_data: estimator.saturating_diff(),
127            }),
128            SurbNotificationMode::DoNotNotify => HoprStartProtocol::KeepAlive(KeepAliveMessage {
129                session_id,
130                flags: None.into(),
131                additional_data: 0,
132            }),
133        })
134        .rate_limit_with_controller(&controller),
135    );
136
137    let sender_clone = sender.clone();
138    let fwd_routing_clone = routing.clone();
139
140    // This task will automatically terminate once the returned abort handle is used.
141    debug!(%session_id, "spawning keep-alive stream");
142    let keep_alive_diag = hopr_utils::runtime::diagnostics::ConcurrentDiagnostics::new(
143        "session_keep_alive_try_for_each_concurrent",
144        module_path!(),
145        file!(),
146        line!(),
147    );
148    hopr_utils::runtime::prelude::spawn(hopr_utils::runtime::diagnostics::instrument(
149        ka_stream
150            .map(move |msg| {
151                ApplicationData::try_from(msg)
152                    .map(|data| (fwd_routing_clone.clone(), ApplicationDataOut::with_no_packet_info(data)))
153            })
154            .map_err(TransportSessionError::from)
155            .try_for_each_concurrent(None, move |msg| {
156                let mut sender_clone = sender_clone.clone();
157                let keep_alive_diag = keep_alive_diag.clone();
158                keep_alive_diag.wrap(|| async move {
159                    sender_clone
160                        .send(msg)
161                        .await
162                        .map_err(TransportSessionError::packet_sending)
163                })
164            })
165            .then(move |res| {
166                match res {
167                    Ok(_) => tracing::debug!(
168                        component = "session",
169                        %session_id,
170                        task = "session keepalive",
171                        "background task finished"
172                    ),
173                    Err(error) => error!(%session_id, %error, "keep-alive stream failed"),
174                }
175                futures::future::ready(())
176            })
177            .in_current_span(),
178        "session_keep_alive",
179        module_path!(),
180        file!(),
181        line!(),
182    ));
183
184    (controller, abort_handle)
185}
186
187#[cfg(test)]
188mod tests {
189    use anyhow::anyhow;
190
191    use super::*;
192
193    /// Generator that cycles through 0..4, wrapping at 5 back to 0.
194    fn cycling_generator(prev: Option<u8>) -> u8 {
195        prev.map(|v| (v + 1) % 5).unwrap_or(0)
196    }
197
198    /// Tests sequential insertion into an empty cache: each call fills the next slot.
199    #[tokio::test]
200    async fn test_insert_into_next_slot_sequential() -> anyhow::Result<()> {
201        let cache = moka::sync::Cache::new(10);
202
203        for i in 0..5 {
204            let (k, v) = insert_into_next_slot(&cache, cycling_generator, |k| format!("foo_{k}"), Some(10u64))
205                .ok_or(anyhow!("should insert into slot {i}"))?;
206            assert_eq!(k, i);
207            assert_eq!(format!("foo_{i}"), v);
208            assert_eq!(Some(v), cache.get(&i));
209        }
210
211        Ok(())
212    }
213
214    /// Tests that insertion returns `None` when all slots are occupied and the generator cycles back.
215    #[tokio::test]
216    async fn test_insert_into_next_slot_returns_none_when_full() -> anyhow::Result<()> {
217        let cache = moka::sync::Cache::new(10);
218
219        for _ in 0..5 {
220            insert_into_next_slot(&cache, cycling_generator, |k| format!("foo_{k}"), Some(10u64))
221                .ok_or(anyhow!("precondition: should insert"))?;
222        }
223
224        assert!(
225            insert_into_next_slot(&cache, cycling_generator, |_| "foo".to_string(), Some(10u64)).is_none(),
226            "must not find slot when full"
227        );
228
229        Ok(())
230    }
231
232    /// Tests that a cache with max capacity of 1 rejects a second distinct key.
233    #[tokio::test]
234    async fn test_insert_into_next_slot_capacity_one_rejects_second_key() -> anyhow::Result<()> {
235        let unit_cache = moka::sync::Cache::new(1);
236
237        let (k0, _v0) = insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64))
238            .ok_or(anyhow!("first insertion must succeed"))?;
239        assert_eq!(k0, 0);
240
241        assert!(
242            insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64)).is_none(),
243            "second distinct key must be rejected when cache capacity is 1"
244        );
245
246        Ok(())
247    }
248
249    /// Tests that a rejected insertion does not evict the existing entry.
250    #[tokio::test]
251    async fn test_insert_into_next_slot_rejected_insertion_does_not_evict() -> anyhow::Result<()> {
252        let unit_cache = moka::sync::Cache::new(1);
253
254        let (k0, v0) = insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64))
255            .ok_or(anyhow!("first insertion must succeed"))?;
256
257        insert_into_next_slot(&unit_cache, |prev| prev.map(|v| v + 1).unwrap_or(0), |k| k, Some(1u64));
258
259        assert_eq!(
260            Some(v0),
261            unit_cache.get(&k0),
262            "first entry must still be present after rejection"
263        );
264
265        Ok(())
266    }
267}
hopr_transport_session/utils.rs

hopr_transport_session/
utils.rs
