hopr_transport_session/balancer/mod.rs
1mod controller;
2/// Contains implementation of the [`SurbBalancerController`] trait using a Proportional Integral Derivative (PID)
3/// controller.
4pub mod pid;
5#[allow(dead_code)]
6mod rate_limiting;
7/// Contains a simple proportional output implementation of the [`SurbBalancerController`] trait.
8pub mod simple;
9
10pub use controller::{BalancerStateValues, SurbBalancer, SurbBalancerConfig};
11pub(crate) use rate_limiting::MAX_WAIT_CHUNK;
12pub use rate_limiting::{RateController, RateLimitSinkExt, RateLimitStreamExt};
13
14/// Smallest possible interval for balancer sampling.
15pub const MIN_BALANCER_SAMPLING_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100);
16
17/// Allows estimating the flow of SURBs in a Session (production or consumption).
18pub trait SurbFlowEstimator {
19 /// Estimates the number of SURBs consumed.
20 ///
21 /// Value returned on each call must be equal or greater to the value returned by a previous call.
22 fn estimate_surbs_consumed(&self) -> u64;
23 /// Estimates the number of SURBs produced or received.
24 ///
25 /// Value returned on each call must be equal or greater to the value returned by a previous call.
26 fn estimate_surbs_produced(&self) -> u64;
27
28 /// Estimates the number of received SURBs the buffer evicted on overflow.
29 ///
30 /// An evicted SURB has left the buffer exactly as a consumed one has; the difference is only that
31 /// the full store dropped it on arrival instead of it being spent. Counting it keeps
32 /// [`saturating_diff`](Self::saturating_diff) equal to what the buffer actually holds. Defaults to
33 /// `0` for estimators that never observe eviction (the Entry, whose store does not overflow, and
34 /// test mocks). Value returned on each call must be equal or greater to a previous call.
35 fn estimate_surbs_evicted(&self) -> u64 {
36 0
37 }
38
39 /// Subtracts SURBs consumed and evicted from SURBs produced, saturating at zero.
40 ///
41 /// Consumed and evicted SURBs have both left the buffer, so what remains is
42 /// `produced - consumed - evicted` -- the count the buffer actually holds. Read by the Exit's
43 /// keep-alive level report and by the PIX fill-reserve gate, both of which otherwise saw a count
44 /// inflated by every SURB the full store ever evicted.
45 fn saturating_diff(&self) -> u64 {
46 self.estimate_surbs_produced()
47 .saturating_sub(self.estimate_surbs_consumed())
48 .saturating_sub(self.estimate_surbs_evicted())
49 }
50
51 /// Computes the estimated change in SURB buffer.
52 ///
53 /// This is done by computing the change in produced and consumed SURBs since the `earlier`
54 /// state and then taking their difference.
55 ///
56 /// A positive result is a surplus number of SURBs added to the buffer, a negative result is a loss of SURBs
57 /// from the buffer.
58 ///
59 /// Evictions count as SURBs leaving the buffer alongside consumption: a tick that receives some
60 /// SURBs while the full store drops others nets only what it actually retained. Without this the
61 /// accumulated level drifts above the real occupancy by every SURB the store ever evicted.
62 /// Returns `None` if `earlier` had more SURBs produced/consumed/evicted than this instance (overflow).
63 fn estimated_surb_buffer_change<E: SurbFlowEstimator>(&self, earlier: &E) -> Option<i64> {
64 match (
65 self.estimate_surbs_produced()
66 .checked_sub(earlier.estimate_surbs_produced()),
67 self.estimate_surbs_consumed()
68 .checked_sub(earlier.estimate_surbs_consumed()),
69 self.estimate_surbs_evicted()
70 .checked_sub(earlier.estimate_surbs_evicted()),
71 ) {
72 (Some(surbs_delivered_delta), Some(surbs_consumed_delta), Some(surbs_evicted_delta)) => {
73 Some(surbs_delivered_delta as i64 - surbs_consumed_delta as i64 - surbs_evicted_delta as i64)
74 }
75 _ => None,
76 }
77 }
78}
79
80/// Allows controlling the production or consumption of SURBs in a Session.
81#[cfg_attr(test, mockall::automock)]
82pub trait SurbFlowController {
83 /// Adjusts the amount of SURB production or consumption.
84 fn adjust_surb_flow(&self, surbs_per_sec: usize);
85}
86
87/// Represents the setpoint (target) and output limit of a controller.
88#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
89pub struct BalancerControllerBounds(u64, u64);
90
91impl BalancerControllerBounds {
92 /// Creates a new instance.
93 pub fn new(target: u64, output_limit: u64) -> Self {
94 Self(target, output_limit)
95 }
96
97 /// Gets the target (setpoint) of a controller.
98 #[inline]
99 pub fn target(&self) -> u64 {
100 self.0
101 }
102
103 /// Gets the output limit of a controller.
104 #[inline]
105 pub fn output_limit(&self) -> u64 {
106 self.1
107 }
108
109 /// Unpacks the controller bounds into two `u64`s (target and output limit).
110 #[inline]
111 pub fn unzip(&self) -> (u64, u64) {
112 (self.0, self.1)
113 }
114}
115
116/// Trait abstracting a controller used in the [`SurbBalancer`].
117pub trait SurbBalancerController {
118 /// Gets the current bounds of the controller.
119 fn bounds(&self) -> BalancerControllerBounds;
120 /// Updates the controller's target (setpoint) and output limit.
121 fn set_target_and_limit(&mut self, bounds: BalancerControllerBounds);
122 /// Queries the controller for the next control output based on the `current_buffer_level` of SURBs.
123 fn next_control_output(&mut self, current_buffer_level: u64) -> u64;
124 /// Discards accumulated history, leaving the bounds intact.
125 ///
126 /// Used when the buffer estimate the controller has been acting on stops meaning what it meant
127 /// -- crossing into or out of a period where the counterparty could not be observed at all.
128 /// Carrying the error accumulated under the old regime into the new one makes the controller
129 /// answer a question nobody asked.
130 fn reset(&mut self);
131}
132
133/// Implementation of [`SurbFlowEstimator`] that tracks the number of produced
134/// and consumed SURBs via two `u64`s.
135///
136/// This implementation can take "snapshots" of other `SurbFlowEstimators` (via `From` trait) by simply
137/// calling their respective methods to fill in its values.
138#[derive(Clone, Copy, Debug, Default)]
139pub struct SimpleSurbFlowEstimator {
140 /// Number of produced SURBs.
141 pub produced: u64,
142 /// Number of consumed SURBs.
143 pub consumed: u64,
144 /// Number of SURBs evicted from the buffer on overflow.
145 pub evicted: u64,
146}
147
148impl<T: SurbFlowEstimator> From<&T> for SimpleSurbFlowEstimator {
149 fn from(value: &T) -> Self {
150 Self {
151 produced: value.estimate_surbs_produced(),
152 consumed: value.estimate_surbs_consumed(),
153 evicted: value.estimate_surbs_evicted(),
154 }
155 }
156}
157
158impl SurbFlowEstimator for SimpleSurbFlowEstimator {
159 fn estimate_surbs_consumed(&self) -> u64 {
160 self.consumed
161 }
162
163 fn estimate_surbs_produced(&self) -> u64 {
164 self.produced
165 }
166
167 fn estimate_surbs_evicted(&self) -> u64 {
168 self.evicted
169 }
170}
171
172/// An implementation of `SurbFlowEstimator` that tracks the number of produced
173/// and consumed SURBs via two `AtomicU64`s.
174#[derive(Clone, Debug, Default)]
175pub struct AtomicSurbFlowEstimator {
176 /// Number of consumed SURBs.
177 pub consumed: std::sync::Arc<std::sync::atomic::AtomicU64>,
178 /// Number of produced or received SURBs.
179 pub produced: std::sync::Arc<std::sync::atomic::AtomicU64>,
180 /// Number of received SURBs the buffer evicted on overflow.
181 pub evicted: std::sync::Arc<std::sync::atomic::AtomicU64>,
182}
183
184impl SurbFlowEstimator for AtomicSurbFlowEstimator {
185 fn estimate_surbs_consumed(&self) -> u64 {
186 self.consumed.load(std::sync::atomic::Ordering::Relaxed)
187 }
188
189 fn estimate_surbs_produced(&self) -> u64 {
190 self.produced.load(std::sync::atomic::Ordering::Relaxed)
191 }
192
193 fn estimate_surbs_evicted(&self) -> u64 {
194 self.evicted.load(std::sync::atomic::Ordering::Relaxed)
195 }
196}
197
198/// Wraps a [`RateController`] as [`SurbFlowController`] with the given correction
199/// factor on time unit.
200///
201/// For example, when this is used to control the flow of keep-alive messages (carrying SURBs),
202/// the correction factor is [`KeepAliveMessage::max_surbs`](hopr_protocol_start::KeepAliveMessage::max_surbs),
203/// the number of SURBs a single keep-alive message can bear including its headers.
204///
205/// In another case, when this is used to control the egress of a Session, each outgoing packet
206/// consumes only a single SURB and therefore the correction factor is `1`.
207pub struct SurbControllerWithCorrection(pub RateController, pub u32);
208
209impl SurbFlowController for SurbControllerWithCorrection {
210 fn adjust_surb_flow(&self, surbs_per_sec: usize) {
211 self.0
212 .set_rate_per_unit(surbs_per_sec, self.1 * std::time::Duration::from_secs(1));
213 }
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[test]
221 fn test_estimated_surb_buffer_change() {
222 let estimator_1 = SimpleSurbFlowEstimator {
223 produced: 10,
224 consumed: 5,
225 evicted: 0,
226 };
227 let estimator_2 = SimpleSurbFlowEstimator {
228 produced: 15,
229 consumed: 11,
230 evicted: 0,
231 };
232 let estimator_3 = SimpleSurbFlowEstimator {
233 produced: 25,
234 consumed: 16,
235 evicted: 0,
236 };
237 assert_eq!(estimator_1.estimated_surb_buffer_change(&estimator_1), Some(0));
238 assert_eq!(estimator_2.estimated_surb_buffer_change(&estimator_1), Some(-1));
239 assert_eq!(estimator_3.estimated_surb_buffer_change(&estimator_2), Some(5));
240 assert_eq!(estimator_1.estimated_surb_buffer_change(&estimator_2), None);
241 }
242
243 /// An evicted SURB has left the buffer, so `saturating_diff` -- what the buffer actually holds --
244 /// must subtract evictions as well as consumption. Received 100, spent 20, and the full store
245 /// dropped 50 on arrival leaves 30 held, not 80.
246 #[test]
247 fn saturating_diff_subtracts_evicted_surbs() {
248 let estimator = SimpleSurbFlowEstimator {
249 produced: 100,
250 consumed: 20,
251 evicted: 50,
252 };
253 assert_eq!(estimator.saturating_diff(), 30);
254
255 let atomic = AtomicSurbFlowEstimator::default();
256 atomic.produced.store(100, std::sync::atomic::Ordering::Relaxed);
257 atomic.consumed.store(20, std::sync::atomic::Ordering::Relaxed);
258 atomic.evicted.store(50, std::sync::atomic::Ordering::Relaxed);
259 assert_eq!(atomic.saturating_diff(), 30);
260 }
261
262 /// Eviction can drive the held count to zero: after an overflow whose surplus is later drained,
263 /// consumed + evicted can meet produced, and the buffer is empty rather than "still full".
264 #[test]
265 fn saturating_diff_saturates_when_evicted_and_consumed_exceed_produced() {
266 let estimator = SimpleSurbFlowEstimator {
267 produced: 40,
268 consumed: 25,
269 evicted: 25,
270 };
271 assert_eq!(estimator.saturating_diff(), 0);
272 }
273
274 /// The per-tick buffer-change delta the balancer accumulates must also treat evictions as SURBs
275 /// leaving the buffer, or the accumulated level drifts above the real occupancy by every SURB the
276 /// store ever evicted. A tick that receives 10 while the full store evicts 8 nets +2, not +10.
277 #[test]
278 fn estimated_surb_buffer_change_subtracts_evicted_delta() {
279 let earlier = SimpleSurbFlowEstimator {
280 produced: 100,
281 consumed: 50,
282 evicted: 30,
283 };
284 let later = SimpleSurbFlowEstimator {
285 produced: 110,
286 consumed: 50,
287 evicted: 38,
288 };
289 assert_eq!(later.estimated_surb_buffer_change(&earlier), Some(2));
290
291 // A full store receiving 10 more evicts all 10: the level does not move.
292 let full_store = SimpleSurbFlowEstimator {
293 produced: 120,
294 consumed: 50,
295 evicted: 48,
296 };
297 assert_eq!(full_store.estimated_surb_buffer_change(&later), Some(0));
298 }
299}