hopr_transport_session/balancer/
controller.rs

1use std::{
2    sync::{
3        Arc,
4        atomic::{AtomicU8, AtomicU64},
5    },
6    time::Duration,
7};
8
9use futures::{StreamExt, pin_mut};
10use hopr_utils::runtime::AbortHandle;
11use tracing::{Instrument, instrument};
12
13use super::{
14    BalancerControllerBounds, MIN_BALANCER_SAMPLING_INTERVAL, SimpleSurbFlowEstimator, SurbBalancerController,
15    SurbFlowController, SurbFlowEstimator,
16};
17use crate::SessionId;
18
19#[cfg(all(feature = "telemetry", not(test)))]
20lazy_static::lazy_static! {
21    static ref METRIC_TARGET_ERROR_ESTIMATE: hopr_api::types::telemetry::MultiGauge =
22        hopr_api::types::telemetry::MultiGauge::new(
23            "hopr_surb_balancer_target_error_estimate",
24            "Target error estimation by the SURB balancer",
25            &["session_id"]
26    ).unwrap();
27    static ref METRIC_CONTROL_OUTPUT: hopr_api::types::telemetry::MultiGauge =
28        hopr_api::types::telemetry::MultiGauge::new(
29            "hopr_surb_balancer_control_output",
30            "Control output of the SURB balancer",
31            &["session_id"]
32    ).unwrap();
33    static ref METRIC_CURRENT_BUFFER: hopr_api::types::telemetry::MultiGauge =
34        hopr_api::types::telemetry::MultiGauge::new(
35            "hopr_surb_balancer_current_buffer_estimate",
36            "Estimated number of SURBs in the buffer",
37            &["session_id"]
38    ).unwrap();
39    static ref METRIC_CURRENT_TARGET: hopr_api::types::telemetry::MultiGauge =
40        hopr_api::types::telemetry::MultiGauge::new(
41            "hopr_surb_balancer_current_buffer_target",
42            "Current target (setpoint) number of SURBs in the buffer",
43            &["session_id"]
44    ).unwrap();
45    static ref METRIC_SURB_RATE: hopr_api::types::telemetry::MultiGauge =
46        hopr_api::types::telemetry::MultiGauge::new(
47            "hopr_surb_balancer_surbs_rate",
48            "Estimation of SURB rate per second (positive is buffer surplus, negative is buffer loss)",
49            &["session_id"]
50    ).unwrap();
51}
52
53/// Configuration for the `SurbBalancer`.
54#[derive(Clone, Copy, Debug, PartialEq, smart_default::SmartDefault)]
55pub struct SurbBalancerConfig {
56    /// The desired number of SURBs to be always kept as a buffer locally or at the Session counterparty.
57    ///
58    /// The `SurbBalancer` will try to maintain approximately this number of SURBs
59    /// locally or remotely (at the counterparty) at all times.
60    ///
61    /// The local buffer is maintained by regulating (`SurbFlowController`) the egress from the Session.
62    /// The remote buffer (at session counterparty) is maintained by regulating the flow of non-organic SURBs via
63    /// keep-alive messages.
64    ///
65    /// It does not make sense to set this value higher than the [`max_surb_buffer_size`](crate::SessionManagerConfig)
66    /// configuration at the counterparty.
67    ///
68    /// Default is 7000 SURBs.
69    #[default(7_000)]
70    pub target_surb_buffer_size: u64,
71    /// Maximum outflow of SURBs.
72    ///
73    /// - In the context of the local SURB buffer (Entry), this is the maximum egress Session traffic (= SURB
74    ///   consumption).
75    /// - In the context of the remote SURB buffer (Exit), this is the maximum egress of keep-alive messages to the
76    ///   counterparty (= artificial SURB production).
77    ///
78    /// The default is 5000 (which is 2500 packets/second currently)
79    #[default(5_000)]
80    pub max_surbs_per_sec: u64,
81
82    /// Sets what percentage of the target buffer size should be discarded at each window.
83    ///
84    /// The `SurbBalancer` will discard the given percentage of `target_surb_buffer_size` at each
85    /// window with the given `Duration`.
86    ///
87    /// The default is `(60, 0.05)` (5% of the target buffer size is discarded every 60 seconds).
88    #[default(_code = "Some((Duration::from_secs(60), 0.05))")]
89    pub surb_decay: Option<(Duration, f64)>,
90}
91
92impl SurbBalancerConfig {
93    /// Convenience function to convert the [`SurbBalancerConfig`] into `BalancerControllerBounds`.
94    #[inline]
95    pub fn as_controller_bounds(&self) -> BalancerControllerBounds {
96        BalancerControllerBounds::new(self.target_surb_buffer_size, self.max_surbs_per_sec)
97    }
98}
99
100/// Runtime state of the `SurbBalancer`.
101#[derive(Debug, Default)]
102pub struct BalancerStateValues {
103    pub target_surb_buffer_size: AtomicU64,
104    pub max_surbs_per_sec: AtomicU64,
105    pub decay_duration_msec: AtomicU64,
106    pub decay_volume_pct: AtomicU8,
107    pub buffer_level: AtomicU64,
108}
109
110impl BalancerStateValues {
111    /// Constructor from a [`SurbBalancerConfig`].
112    pub fn new(cfg: SurbBalancerConfig) -> Self {
113        let state = Self::default();
114        state.update(&cfg);
115        state
116    }
117
118    /// Performs update of the [`BalancerStateValues`] from the [`SurbBalancerConfig`] and
119    /// enables it.
120    pub fn update(&self, cfg: &SurbBalancerConfig) {
121        self.target_surb_buffer_size
122            .store(cfg.target_surb_buffer_size, std::sync::atomic::Ordering::Relaxed);
123        self.max_surbs_per_sec
124            .store(cfg.max_surbs_per_sec, std::sync::atomic::Ordering::Relaxed);
125        self.decay_duration_msec.store(
126            cfg.surb_decay
127                .map(|(d, _)| d.as_millis().min(u64::MAX as u128) as u64)
128                .unwrap_or_default(),
129            std::sync::atomic::Ordering::Relaxed,
130        );
131        self.decay_volume_pct.store(
132            cfg.surb_decay
133                .map(|(_, p)| (p.clamp(0.0, 1.0) * 100.0).round() as u8)
134                .unwrap_or_default(),
135            std::sync::atomic::Ordering::Relaxed,
136        );
137    }
138
139    /// Extracts the [`SurbBalancerConfig`] from the [`BalancerStateValues`].
140    pub fn as_config(&self) -> SurbBalancerConfig {
141        SurbBalancerConfig {
142            target_surb_buffer_size: self.target_surb_buffer_size.load(std::sync::atomic::Ordering::Relaxed),
143            max_surbs_per_sec: self.max_surbs_per_sec.load(std::sync::atomic::Ordering::Relaxed),
144            surb_decay: self.surb_decay(),
145        }
146    }
147
148    /// Checks if SURB balancing is disabled (no target buffer size set).
149    pub fn is_disabled(&self) -> bool {
150        self.target_surb_buffer_size.load(std::sync::atomic::Ordering::Relaxed) == 0
151    }
152
153    /// Extracts the SURB decay configuration from the [`BalancerStateValues`].
154    pub fn surb_decay(&self) -> Option<(Duration, f64)> {
155        Some((
156            self.decay_duration_msec.load(std::sync::atomic::Ordering::Relaxed),
157            self.decay_volume_pct.load(std::sync::atomic::Ordering::Relaxed),
158        ))
159        .filter(|&(d, p)| d > 0 && p > 0)
160        .map(|(d, p)| (Duration::from_millis(d), p as f64 / 100.0))
161    }
162
163    /// Gets the current estimated SURB buffer level.
164    #[inline]
165    pub fn buffer_level(&self) -> u64 {
166        self.buffer_level.load(std::sync::atomic::Ordering::Relaxed)
167    }
168
169    /// Returns the current `BalancerControllerBounds` from the [`BalancerStateValues`].
170    #[inline]
171    pub fn controller_bounds(&self) -> BalancerControllerBounds {
172        BalancerControllerBounds::new(
173            self.target_surb_buffer_size.load(std::sync::atomic::Ordering::Relaxed),
174            self.max_surbs_per_sec.load(std::sync::atomic::Ordering::Relaxed),
175        )
176    }
177}
178
179impl From<SurbBalancerConfig> for BalancerStateValues {
180    fn from(cfg: SurbBalancerConfig) -> Self {
181        Self::new(cfg)
182    }
183}
184
185/// Runs a continuous process that attempts to [evaluate](SurbFlowEstimator) and
186/// [regulate](SurbFlowController) the flow of SURBs to the Session counterparty,
187/// to keep the number of SURBs locally or at the counterparty at a certain level.
188///
189/// Internally, the Balancer uses an implementation of [`SurbBalancerController`] to
190/// control the rate of SURBs consumed or sent to the counterparty
191/// each time the [`update`](SurbBalancer::update) method is called:
192///
193/// 1. The size of the SURB buffer at locally or at the counterparty is estimated using [`SurbFlowEstimator`].
194/// 2. Error against a set-point given in [`SurbBalancerConfig`] is evaluated in the `SurbBalancerController`.
195/// 3. The `SurbBalancerController` applies a new SURB flow rate value using the [`SurbFlowController`].
196///
197/// In the local context, the `SurbFlowController` might simply regulate the egress traffic from the
198/// Session, slowing it down to avoid fast SURB drainage.
199///
200/// In the remote context, the `SurbFlowController` might regulate the flow of non-organic SURBs via
201/// Start protocol's `KeepAlive` messages to deliver additional
202/// SURBs to the counterparty.
203pub struct SurbBalancer<C, E, F> {
204    session_id: SessionId,
205    controller: C,
206    surb_estimator: E,
207    flow_control: F,
208    state: Arc<BalancerStateValues>,
209    last_estimator_state: SimpleSurbFlowEstimator,
210    last_update: std::time::Instant,
211    last_decay: std::time::Instant,
212    was_below_target: bool,
213}
214
215impl<C, E, F> SurbBalancer<C, E, F>
216where
217    C: SurbBalancerController + Send + Sync + 'static,
218    E: SurbFlowEstimator + Send + Sync + 'static,
219    F: SurbFlowController + Send + Sync + 'static,
220{
221    pub fn new(
222        session_id: SessionId,
223        mut controller: C,
224        surb_estimator: E,
225        flow_control: F,
226        state: Arc<BalancerStateValues>,
227    ) -> Self {
228        #[cfg(all(feature = "telemetry", not(test)))]
229        {
230            let sid: &str = session_id.as_ref();
231            METRIC_TARGET_ERROR_ESTIMATE.set(&[sid], 0.0);
232            METRIC_CONTROL_OUTPUT.set(&[sid], 0.0);
233        }
234
235        controller.set_target_and_limit(state.controller_bounds());
236
237        Self {
238            surb_estimator,
239            flow_control,
240            controller,
241            session_id,
242            state,
243            last_estimator_state: Default::default(),
244            last_update: std::time::Instant::now(),
245            last_decay: std::time::Instant::now(),
246            was_below_target: true,
247        }
248    }
249
250    /// Computes the next control update and adjusts the [`SurbFlowController`] rate accordingly.
251    #[tracing::instrument(level = "trace", skip_all)]
252    fn update(&mut self) -> u64 {
253        let dt = self.last_update.elapsed();
254
255        // Load the updated current buffer level
256        let mut current = self.state.buffer_level.load(std::sync::atomic::Ordering::Acquire);
257
258        if dt < Duration::from_millis(10) {
259            tracing::debug!("time elapsed since last update is too short, skipping update");
260            return current;
261        }
262
263        self.last_update = std::time::Instant::now();
264
265        // Take a snapshot of the active SURB estimator and calculate the balance change
266        let snapshot = SimpleSurbFlowEstimator::from(&self.surb_estimator);
267        let Some(target_buffer_change) = snapshot.estimated_surb_buffer_change(&self.last_estimator_state) else {
268            tracing::error!("non-monotonic change in SURB estimators");
269            return current;
270        };
271
272        self.last_estimator_state = snapshot;
273        current = current.saturating_add_signed(target_buffer_change);
274
275        // If SURB decaying is enabled, check if the decay window has elapsed
276        // and calculate the number of SURBs that will be discarded
277        if let Some(num_decayed_surbs) = self
278            .state
279            .surb_decay()
280            .filter(|(decay_window, _)| &self.last_decay.elapsed() >= decay_window)
281            .map(|(_, decay_coeff)| (self.controller.bounds().target() as f64 * decay_coeff).round() as u64)
282        {
283            current = current.saturating_sub(num_decayed_surbs);
284            self.last_decay = std::time::Instant::now();
285            tracing::trace!(num_decayed_surbs, "SURBs were discarded due to automatic decay");
286        }
287
288        self.state
289            .buffer_level
290            .store(current, std::sync::atomic::Ordering::Release);
291
292        // Error from the desired target SURB buffer size at counterparty
293        let error = current as i64 - self.controller.bounds().target() as i64;
294
295        if self.was_below_target && error >= 0 {
296            tracing::trace!(current, "reached target SURB buffer size");
297            self.was_below_target = false;
298        } else if !self.was_below_target && error < 0 {
299            tracing::trace!(current, "SURB buffer size is below target");
300            self.was_below_target = true;
301        }
302
303        tracing::trace!(
304            ?dt,
305            delta = target_buffer_change,
306            rate = target_buffer_change as f64 / dt.as_secs_f64(),
307            current,
308            error,
309            "estimated SURB buffer change"
310        );
311
312        let output = self.controller.next_control_output(current);
313        tracing::trace!(output, "next balancer control output for session");
314
315        self.flow_control.adjust_surb_flow(output as usize);
316
317        #[cfg(all(feature = "telemetry", not(test)))]
318        {
319            let sid: &str = self.session_id.as_ref();
320            METRIC_CURRENT_BUFFER.set(&[sid], current as f64);
321            METRIC_CURRENT_TARGET.set(&[sid], self.controller.bounds().target() as f64);
322            METRIC_TARGET_ERROR_ESTIMATE.set(&[sid], error as f64);
323            METRIC_CONTROL_OUTPUT.set(&[sid], output as f64);
324            METRIC_SURB_RATE.set(&[sid], target_buffer_change as f64 / dt.as_secs_f64());
325        }
326
327        current
328    }
329
330    /// Spawns a new task that performs updates of the given [`SurbBalancer`] at the given `sampling_interval`.
331    ///
332    /// If `cfg_feedback` is given, [`SurbBalancerConfig`] can be queried for updates and also updated
333    /// if the underlying [`SurbBalancerController`] also does target updates.
334    ///
335    /// Returns a stream of current estimated buffer levels, and also an `AbortHandle`
336    /// to terminate the loop. If `abort_reg` was given, the returned `AbortHandle` corresponds
337    /// to it.
338    #[instrument(level = "debug", skip(self), fields(session_id = %self.session_id))]
339    pub fn start_control_loop(
340        mut self,
341        sampling_interval: Duration,
342    ) -> (impl futures::Stream<Item = u64>, AbortHandle) {
343        let (abort_handle, abort_reg) = AbortHandle::new_pair();
344
345        // Start an interval stream at which the balancer will sample and perform updates
346
347        // DropAbortable not needed because the stream only generates items when polled
348        let sampling_stream = futures::stream::Abortable::new(
349            futures_time::stream::interval(sampling_interval.max(MIN_BALANCER_SAMPLING_INTERVAL).into()),
350            abort_reg,
351        );
352
353        let balancer_level_capacity = std::env::var("HOPR_INTERNAL_SESSION_BALANCER_LEVEL_CAPACITY")
354            .ok()
355            .and_then(|s| s.trim().parse::<usize>().ok())
356            .filter(|&c| c > 0)
357            .unwrap_or(32_768);
358
359        tracing::debug!(
360            capacity = balancer_level_capacity,
361            "Creating session balancer level channel"
362        );
363        let (mut level_tx, level_rx) = futures::channel::mpsc::channel(balancer_level_capacity);
364        hopr_utils::runtime::prelude::spawn(
365            async move {
366                pin_mut!(sampling_stream);
367                while sampling_stream.next().await.is_some() {
368                    // Check if the balancer controller needs to be reconfigured
369                    let current_bounds = self.state.controller_bounds();
370                    if current_bounds != self.controller.bounds() {
371                        self.controller.set_target_and_limit(current_bounds);
372                        tracing::debug!(new_cfg = ?self.state.as_config(), "surb balancer has been reconfigured");
373                    }
374
375                    // Perform controller update (this internally samples the SurbFlowEstimator)
376                    // and send an update about the current level to the outgoing stream.
377                    // If the other party has closed the stream, we don't care about the update.
378                    let level = self.update();
379                    if !level_tx.is_closed()
380                        && let Err(error) = level_tx.try_send(level)
381                    {
382                        tracing::error!(%error, "cannot send balancer level update");
383                    }
384                }
385
386                tracing::debug!("balancer done");
387            }
388            .in_current_span(),
389        );
390
391        (level_rx, abort_handle)
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use std::sync::{Arc, atomic::AtomicU64};
398
399    use hopr_api::types::{crypto_random::Randomizable, internal::prelude::HoprPseudonym};
400
401    use super::*;
402    use crate::balancer::{AtomicSurbFlowEstimator, MockSurbFlowController, pid::PidBalancerController};
403
404    #[test]
405    fn surb_balancer_config_should_be_convertible_to_atomics() {
406        let cfg = SurbBalancerConfig::default();
407        let state_data = BalancerStateValues::new(cfg);
408        assert_eq!(cfg, state_data.as_config());
409    }
410
411    #[test]
412    fn surb_balancer_config_default_snapshot() {
413        let cfg = SurbBalancerConfig::default();
414        insta::assert_debug_snapshot!(cfg);
415    }
416
417    #[test]
418    fn surb_balancer_config_as_controller_bounds() {
419        let cfg = SurbBalancerConfig {
420            target_surb_buffer_size: 1000,
421            max_surbs_per_sec: 500,
422            surb_decay: None,
423        };
424        let bounds = cfg.as_controller_bounds();
425        assert_eq!(bounds.target(), 1000);
426        assert_eq!(bounds.output_limit(), 500);
427    }
428
429    #[test]
430    fn balancer_state_values_disabled_when_target_is_zero() {
431        let cfg = SurbBalancerConfig {
432            target_surb_buffer_size: 0,
433            max_surbs_per_sec: 0,
434            surb_decay: None,
435        };
436        let state = BalancerStateValues::new(cfg);
437        assert!(state.is_disabled());
438    }
439
440    #[test]
441    fn balancer_state_values_enabled_when_target_is_nonzero() {
442        let state = BalancerStateValues::new(SurbBalancerConfig::default());
443        assert!(!state.is_disabled());
444    }
445
446    #[test]
447    fn balancer_state_values_update_propagates_all_fields() {
448        let state = BalancerStateValues::default();
449        let cfg = SurbBalancerConfig {
450            target_surb_buffer_size: 3000,
451            max_surbs_per_sec: 1500,
452            surb_decay: Some((Duration::from_secs(30), 0.10)),
453        };
454        state.update(&cfg);
455        assert_eq!(state.as_config(), cfg);
456        assert_eq!(state.controller_bounds(), cfg.as_controller_bounds());
457    }
458
459    #[test]
460    fn balancer_state_values_surb_decay_none_maps_to_none() {
461        let cfg = SurbBalancerConfig {
462            target_surb_buffer_size: 1000,
463            max_surbs_per_sec: 500,
464            surb_decay: None,
465        };
466        let state = BalancerStateValues::new(cfg);
467        assert!(state.surb_decay().is_none());
468    }
469
470    #[test]
471    fn balancer_state_values_buffer_level_default_is_zero() {
472        let state = BalancerStateValues::default();
473        assert_eq!(state.buffer_level(), 0);
474    }
475
476    #[test]
477    fn balancer_state_values_buffer_level_can_be_updated() {
478        let state = BalancerStateValues::default();
479        state.buffer_level.store(42, std::sync::atomic::Ordering::Relaxed);
480        assert_eq!(state.buffer_level(), 42);
481    }
482
483    #[test]
484    fn balancer_state_values_from_config() {
485        let cfg = SurbBalancerConfig {
486            target_surb_buffer_size: 5000,
487            max_surbs_per_sec: 2500,
488            surb_decay: Some((Duration::from_secs(60), 0.05)),
489        };
490        let state: BalancerStateValues = cfg.into();
491        assert_eq!(state.as_config(), cfg);
492    }
493
494    #[test]
495    fn balancer_state_values_decay_zero_duration_should_map_to_none() {
496        let cfg = SurbBalancerConfig {
497            surb_decay: Some((Duration::ZERO, 0.10)),
498            ..Default::default()
499        };
500        let state = BalancerStateValues::new(cfg);
501        assert!(
502            state.surb_decay().is_none(),
503            "zero duration decay should be filtered out"
504        );
505    }
506
507    #[test]
508    fn balancer_state_values_decay_zero_percent_should_map_to_none() {
509        let cfg = SurbBalancerConfig {
510            surb_decay: Some((Duration::from_secs(60), 0.0)),
511            ..Default::default()
512        };
513        let state = BalancerStateValues::new(cfg);
514        assert!(
515            state.surb_decay().is_none(),
516            "zero percent decay should be filtered out"
517        );
518    }
519
520    #[test]
521    fn balancer_state_values_decay_should_clamp_above_one() {
522        let cfg = SurbBalancerConfig {
523            surb_decay: Some((Duration::from_secs(1), 1.5)), // > 1.0 should be clamped
524            ..Default::default()
525        };
526        let state = BalancerStateValues::new(cfg);
527        let (_, pct) = state.surb_decay().expect("decay should be present");
528        assert!((pct - 1.0).abs() < f64::EPSILON, "percentage should be clamped to 1.0");
529    }
530
531    #[test_log::test]
532    fn surb_balancer_should_start_increase_level_when_below_target() {
533        let production_rate = Arc::new(AtomicU64::new(0));
534        let consumption_rate = 100;
535        let steps = 3;
536        let step_duration = std::time::Duration::from_millis(1000);
537
538        let mut controller = MockSurbFlowController::new();
539        let production_rate_clone = production_rate.clone();
540        controller
541            .expect_adjust_surb_flow()
542            .times(steps)
543            .with(mockall::predicate::ge(100))
544            .returning(move |r| {
545                production_rate_clone.store(r as u64, std::sync::atomic::Ordering::Relaxed);
546            });
547
548        let surb_estimator = AtomicSurbFlowEstimator::default();
549        let mut balancer = SurbBalancer::new(
550            HoprPseudonym::random(),
551            PidBalancerController::default(),
552            surb_estimator.clone(),
553            controller,
554            Arc::new(
555                SurbBalancerConfig {
556                    target_surb_buffer_size: 5_000,
557                    max_surbs_per_sec: 2500,
558                    surb_decay: None,
559                }
560                .into(),
561            ),
562        );
563
564        let mut last_update = 0;
565        for i in 0..steps {
566            std::thread::sleep(step_duration);
567            surb_estimator.produced.fetch_add(
568                production_rate.load(std::sync::atomic::Ordering::Relaxed) * step_duration.as_secs(),
569                std::sync::atomic::Ordering::Relaxed,
570            );
571            surb_estimator.consumed.fetch_add(
572                consumption_rate * step_duration.as_secs(),
573                std::sync::atomic::Ordering::Relaxed,
574            );
575
576            let next_update = balancer.update();
577            assert!(
578                i == 0 || next_update > last_update,
579                "{next_update} should be greater than {last_update}"
580            );
581            last_update = next_update;
582        }
583    }
584
585    #[test_log::test]
586    fn surb_balancer_should_start_decrease_level_when_above_target() {
587        let production_rate = Arc::new(AtomicU64::new(11_000));
588        let consumption_rate = 100;
589        let steps = 3;
590        let step_duration = std::time::Duration::from_millis(1000);
591
592        let mut controller = MockSurbFlowController::new();
593        let production_rate_clone = production_rate.clone();
594        controller
595            .expect_adjust_surb_flow()
596            .times(steps)
597            .with(mockall::predicate::ge(0))
598            .returning(move |r| {
599                production_rate_clone.store(r as u64, std::sync::atomic::Ordering::Relaxed);
600            });
601
602        let surb_estimator = AtomicSurbFlowEstimator::default();
603        let mut balancer = SurbBalancer::new(
604            HoprPseudonym::random(),
605            PidBalancerController::default(),
606            surb_estimator.clone(),
607            controller,
608            Arc::new(
609                SurbBalancerConfig {
610                    surb_decay: None,
611                    ..Default::default()
612                }
613                .into(),
614            ),
615        );
616
617        let mut last_update = 0;
618        for i in 0..steps {
619            std::thread::sleep(step_duration);
620            surb_estimator.produced.fetch_add(
621                production_rate.load(std::sync::atomic::Ordering::Relaxed) * step_duration.as_secs(),
622                std::sync::atomic::Ordering::Relaxed,
623            );
624            surb_estimator.consumed.fetch_add(
625                consumption_rate * step_duration.as_secs(),
626                std::sync::atomic::Ordering::Relaxed,
627            );
628
629            let next_update = balancer.update();
630            assert!(
631                i == 0 || next_update < last_update,
632                "{next_update} should be greater than {last_update}"
633            );
634            last_update = next_update;
635        }
636    }
637
638    #[test_log::test(tokio::test)]
639    async fn surb_balancer_should_start_decrease_level_when_above_target_and_decay_enabled() {
640        const NUM_STEPS: usize = 5;
641        let session_id = HoprPseudonym::random();
642        let cfg = SurbBalancerConfig {
643            target_surb_buffer_size: 5_000,
644            max_surbs_per_sec: 2500,
645            surb_decay: Some((Duration::from_millis(200), 0.05)),
646        };
647
648        let mut mock_flow_ctl = MockSurbFlowController::new();
649        mock_flow_ctl
650            .expect_adjust_surb_flow()
651            .times(NUM_STEPS)
652            .returning(|_| ());
653
654        let balancer = SurbBalancer::new(
655            session_id,
656            PidBalancerController::default(),
657            SimpleSurbFlowEstimator::default(),
658            mock_flow_ctl,
659            Arc::new(cfg.into()),
660        );
661
662        balancer
663            .state
664            .buffer_level
665            .store(5000, std::sync::atomic::Ordering::Relaxed);
666
667        let (stream, handle) = balancer.start_control_loop(Duration::from_millis(100));
668        let levels = stream.take(NUM_STEPS).collect::<Vec<_>>().await;
669        handle.abort();
670
671        assert_eq!(levels.len(), NUM_STEPS);
672        assert!(
673            levels.windows(2).all(|w| w[1] <= w[0]),
674            "buffer levels should be monotonic non-increasing: {levels:?}"
675        );
676        assert!(
677            levels.last().is_some_and(|last| *last < 5_000),
678            "expected at least one decay step: {levels:?}"
679        );
680    }
681}
hopr_transport_session/balancer/controller.rs

hopr_transport_session/balancer/
controller.rs
