hopr_transport_session/balancer/
pid.rs

1use std::str::FromStr;
2
3use anyhow::anyhow;
4use pid::Pid;
5
6use crate::{
7    balancer::{BalancerControllerBounds, SurbBalancerController},
8    errors,
9    errors::SessionManagerError,
10};
11
12/// Carries finite Proportional, Integral and Derivative controller gains for a PID controller.
13#[derive(Clone, Copy, Debug, PartialEq)]
14pub struct PidControllerGains(f64, f64, f64);
15
16impl PidControllerGains {
17    /// Creates PID controller gains, returns an error if the gains are not finite.
18    pub fn new(p: f64, i: f64, d: f64) -> errors::Result<Self> {
19        if p.is_finite() && i.is_finite() && d.is_finite() {
20            Ok(Self(p, i, d))
21        } else {
22            Err(SessionManagerError::other(anyhow!("gains must be finite")).into())
23        }
24    }
25
26    /// Uses PID controller gains from the env variables or uses the defaults if not set.
27    pub fn from_env_or_default() -> Self {
28        let default = Self::default();
29        Self(
30            std::env::var("HOPR_BALANCER_PID_P_GAIN")
31                .ok()
32                .and_then(|v| f64::from_str(&v).ok())
33                .unwrap_or(default.0),
34            std::env::var("HOPR_BALANCER_PID_I_GAIN")
35                .ok()
36                .and_then(|v| f64::from_str(&v).ok())
37                .unwrap_or(default.1),
38            std::env::var("HOPR_BALANCER_PID_D_GAIN")
39                .ok()
40                .and_then(|v| f64::from_str(&v).ok())
41                .unwrap_or(default.2),
42        )
43    }
44
45    /// P gain.
46    #[inline]
47    pub fn p(&self) -> f64 {
48        self.0
49    }
50
51    /// I gain.
52    #[inline]
53    pub fn i(&self) -> f64 {
54        self.1
55    }
56
57    /// D gain.
58    #[inline]
59    pub fn d(&self) -> f64 {
60        self.2
61    }
62}
63
64// Safe to implement Eq, because the floats are finite
65impl Eq for PidControllerGains {}
66
67// Default coefficients for the PID controller
68// This might be tweaked in the future.
69const DEFAULT_P_GAIN: f64 = 0.6;
70const DEFAULT_I_GAIN: f64 = 0.7;
71const DEFAULT_D_GAIN: f64 = 0.2;
72
73impl Default for PidControllerGains {
74    fn default() -> Self {
75        Self(DEFAULT_P_GAIN, DEFAULT_I_GAIN, DEFAULT_D_GAIN)
76    }
77}
78
79impl TryFrom<(f64, f64, f64)> for PidControllerGains {
80    type Error = errors::TransportSessionError;
81
82    fn try_from(value: (f64, f64, f64)) -> Result<Self, Self::Error> {
83        Self::new(value.0, value.1, value.2)
84    }
85}
86
87/// Implementation of [`SurbBalancerController`] using a PID controller.
88#[derive(Clone, Copy, Debug)]
89pub struct PidBalancerController(Pid<f64>);
90
91impl PidBalancerController {
92    /// Creates new instance given the `setpoint`, `output_limit` and PID gains (P,I and D).
93    pub fn new(setpoint: u64, output_limit: u64, gains: PidControllerGains) -> Self {
94        let mut pid = Pid::new(setpoint as f64, output_limit as f64);
95        pid.p(gains.p(), output_limit as f64);
96        pid.i(gains.i(), output_limit as f64);
97        pid.d(gains.d(), output_limit as f64);
98        Self(pid)
99    }
100
101    /// Creates new instance with setpoint and output limit set to 0.
102    ///
103    /// Needs to be [reconfigured](SurbBalancerController::set_target_and_limit) in order to function
104    /// correctly.
105    pub fn from_gains(gains: PidControllerGains) -> Self {
106        Self::new(0, 0, gains)
107    }
108}
109
110impl Default for PidBalancerController {
111    /// The default instance does nothing unless [reconfigured](SurbBalancerController::set_target_and_limit).
112    fn default() -> Self {
113        Self::new(0, 0, PidControllerGains::default())
114    }
115}
116
117impl SurbBalancerController for PidBalancerController {
118    fn bounds(&self) -> BalancerControllerBounds {
119        BalancerControllerBounds::new(self.0.setpoint as u64, self.0.output_limit as u64)
120    }
121
122    fn set_target_and_limit(&mut self, bounds: BalancerControllerBounds) {
123        let mut pid = Pid::new(bounds.target() as f64, bounds.output_limit() as f64);
124        pid.p(self.0.kp, bounds.output_limit() as f64);
125        pid.i(self.0.ki, bounds.output_limit() as f64);
126        pid.d(self.0.kd, bounds.output_limit() as f64);
127        self.0 = pid;
128    }
129
130    fn next_control_output(&mut self, current_buffer_level: u64) -> u64 {
131        self.0.next_control_output(current_buffer_level as f64).output.max(0.0) as u64
132    }
133
134    fn reset(&mut self) {
135        // Only the integral carries history across calls; P and D are derived from the current
136        // sample, so clearing the accumulated error is the whole of "start again".
137        self.0.reset_integral_term();
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    /// A starved buffer must command production. This is the controller's entire purpose, and a
146    /// cluster run showed it returning 0 for an error of -9803 -- maximally starved.
147    #[test]
148    fn pid_should_command_production_when_the_buffer_is_empty() {
149        let mut c = PidBalancerController::from_gains(PidControllerGains::default());
150        c.set_target_and_limit(BalancerControllerBounds::new(7_000, 5_000));
151
152        let out = c.next_control_output(0);
153        assert!(out > 0, "empty buffer against a 7000 target must produce, got {out}");
154    }
155
156    /// The output limit is what every PID term is clamped to. A zero limit silently turns the
157    /// controller into a no-op that reports a huge error and commands nothing -- which is
158    /// indistinguishable, from the outside, from a controller that thinks the buffer is full.
159    #[test]
160    fn pid_with_a_zero_output_limit_commands_nothing_however_starved() {
161        let mut c = PidBalancerController::from_gains(PidControllerGains::default());
162        c.set_target_and_limit(BalancerControllerBounds::new(7_000, 0));
163
164        assert_eq!(0, c.next_control_output(0), "a zero limit clamps every term to zero");
165    }
166
167    /// Reproduces the exact operating point observed on a cluster: target 9803, empty buffer,
168    /// default limit. Production logged `output=0` there, and 15 minutes per observation is too
169    /// slow a loop to debug in -- so pin it here at 80ms instead.
170    #[test]
171    fn pid_should_produce_at_the_operating_point_seen_in_production() {
172        let mut c = PidBalancerController::from_gains(PidControllerGains::default());
173        c.set_target_and_limit(BalancerControllerBounds::new(
174            9_803,
175            crate::SurbBalancerConfig::default().max_surbs_per_sec,
176        ));
177
178        let out = c.next_control_output(0);
179        assert!(out > 0, "target 9803 with an empty buffer produced {out}");
180    }
181
182    /// The very first output after configuration, which is the one that matters on a cold start or
183    /// straight after the estimate is discarded.
184    #[test]
185    fn pid_first_output_after_reconfiguration_should_not_be_zero() {
186        let mut c = PidBalancerController::default();
187        c.set_target_and_limit(BalancerControllerBounds::new(7_000, 5_000));
188
189        assert!(
190            c.next_control_output(0) > 0,
191            "first output after reconfiguration was zero"
192        );
193    }
194
195    /// The exact live bounds, read off a cluster run: target 9803, output limit 1960, empty buffer.
196    /// Production logs `output=0` here while the same controller with a 5000 limit produces.
197    #[test]
198    fn pid_should_produce_at_the_live_bounds() {
199        let mut c = PidBalancerController::from_gains(PidControllerGains::default());
200        c.set_target_and_limit(BalancerControllerBounds::new(9_803, 1_960));
201
202        let out = c.next_control_output(0);
203        assert!(out > 0, "target 9803 / limit 1960 with an empty buffer produced {out}");
204    }
205
206    /// A buffer that was healthy and then drained must be refilled. This is the return-path failure
207    /// exactly: the session runs at target, the return path breaks, the buffer empties -- and
208    /// production commands nothing despite an error of -9803.
209    #[test]
210    fn pid_should_recover_after_a_healthy_buffer_drains() {
211        let mut c = PidBalancerController::from_gains(PidControllerGains::default());
212        c.set_target_and_limit(BalancerControllerBounds::new(9_803, 1_960));
213
214        // Healthy: sitting at target, so the controller has nothing to do for a while.
215        for _ in 0..100 {
216            c.next_control_output(9_803);
217        }
218
219        // The return path breaks and the buffer drains.
220        let out = c.next_control_output(0);
221        assert!(out > 0, "an emptied buffer after a healthy period produced {out}");
222    }
223
224    // --- PidControllerGains tests ---
225
226    #[test]
227    fn gains_default_values_are_stable() {
228        let gains = PidControllerGains::default();
229        insta::assert_yaml_snapshot!((gains.p(), gains.i(), gains.d()));
230    }
231
232    #[test]
233    fn gains_finite_values_are_accepted() -> anyhow::Result<()> {
234        let gains = PidControllerGains::new(1.0, 2.0, 3.0)?;
235        insta::assert_yaml_snapshot!((gains.p(), gains.i(), gains.d()));
236        Ok(())
237    }
238
239    #[test]
240    fn gains_infinity_is_rejected() {
241        assert!(PidControllerGains::new(f64::INFINITY, 0.0, 0.0).is_err());
242        assert!(PidControllerGains::new(0.0, f64::NEG_INFINITY, 0.0).is_err());
243        assert!(PidControllerGains::new(0.0, 0.0, f64::NAN).is_err());
244    }
245
246    #[test]
247    fn gains_try_from_tuple() -> anyhow::Result<()> {
248        let gains = PidControllerGains::try_from((0.5, 0.3, 0.1))?;
249        insta::assert_yaml_snapshot!((gains.p(), gains.i(), gains.d()));
250        Ok(())
251    }
252
253    #[test]
254    fn gains_try_from_tuple_with_nan_fails() {
255        assert!(PidControllerGains::try_from((f64::NAN, 0.0, 0.0)).is_err());
256    }
257
258    #[test]
259    fn gains_eq_works() -> anyhow::Result<()> {
260        let a = PidControllerGains::new(1.0, 2.0, 3.0)?;
261        let b = PidControllerGains::new(1.0, 2.0, 3.0)?;
262        assert_eq!(a, b);
263        Ok(())
264    }
265
266    // --- PidBalancerController tests ---
267
268    #[test]
269    fn controller_default_has_zero_bounds() {
270        let ctrl = PidBalancerController::default();
271        assert_eq!(ctrl.bounds().unzip(), (0, 0));
272    }
273
274    #[test]
275    fn controller_new_stores_bounds() {
276        let gains = PidControllerGains::default();
277        let ctrl = PidBalancerController::new(100, 50, gains);
278        assert_eq!(ctrl.bounds().unzip(), (100, 50));
279    }
280
281    #[test]
282    fn controller_set_target_and_limit_updates_bounds() {
283        let mut ctrl = PidBalancerController::default();
284        ctrl.set_target_and_limit(BalancerControllerBounds::new(200, 100));
285        assert_eq!(ctrl.bounds().unzip(), (200, 100));
286    }
287
288    #[test]
289    fn controller_step_response_snapshot() {
290        // Apply a step input (setpoint=100, buffer=0) and observe outputs over N steps
291        let gains = PidControllerGains::default();
292        let mut ctrl = PidBalancerController::new(100, 200, gains);
293
294        let outputs: Vec<u64> = (0..10).map(|_| ctrl.next_control_output(0)).collect();
295        insta::assert_yaml_snapshot!(outputs);
296    }
297
298    #[test]
299    fn controller_at_setpoint_outputs_zero_or_near_zero() {
300        let gains = PidControllerGains::default();
301        let mut ctrl = PidBalancerController::new(100, 200, gains);
302
303        // When buffer level equals setpoint, output should converge toward 0
304        let output = ctrl.next_control_output(100);
305        // First call at setpoint: P=0, I=0, D=0 → output=0
306        assert_eq!(output, 0);
307    }
308
309    #[test]
310    fn controller_above_setpoint_clamps_to_zero() {
311        let gains = PidControllerGains::default();
312        let mut ctrl = PidBalancerController::new(100, 200, gains);
313
314        // Buffer well above setpoint — PID error is negative, output clamped to 0
315        let output = ctrl.next_control_output(200);
316        assert_eq!(output, 0);
317    }
318
319    #[test]
320    fn controller_convergence_from_empty_buffer() {
321        // Simulate filling a buffer from 0 toward setpoint=100
322        let gains = PidControllerGains::default();
323        let mut ctrl = PidBalancerController::new(100, 200, gains);
324
325        let mut buffer: f64 = 0.0;
326        let mut history = Vec::new();
327
328        for _ in 0..20 {
329            let output = ctrl.next_control_output(buffer as u64);
330            buffer += output as f64;
331            buffer = buffer.min(200.0); // clamp to limit
332            history.push(buffer as u64);
333        }
334
335        insta::assert_yaml_snapshot!(history);
336    }
337
338    #[test]
339    fn controller_from_gains_uses_defaults_for_bounds() {
340        let gains = PidControllerGains::default();
341        let ctrl = PidBalancerController::from_gains(gains);
342        assert_eq!(ctrl.bounds().unzip(), (0, 0));
343    }
344}
hopr_transport_session/balancer/pid.rs

hopr_transport_session/balancer/
pid.rs
