1
use std::time::{Duration, Instant};
2
 
3
use crate::Bitrate;
4
use crate::rtp_::DataSize;
5
 
6
const SMALL_SAMPLE_THRESHOLD: DataSize = DataSize::bytes(2000);
7
const SMALL_SAMPLE_UNCERTAINTY: f64 = 25.0;
8
const UNCERTAINTY: f64 = 10.0;
9
const ESTIMATE_FLOOR: Bitrate = Bitrate::kbps(40);
10
 
11
pub struct AckedBitrateEstimator {
12
    /// The initial window to use for the first estimate.
13
    initial_window: Duration,
14
    /// The window to use for subsequent estimates after achieving the initial estimate.
15
    window: Duration,
16
    /// The estimate of the acked bitrate.
17
    estimate: Option<Bitrate>,
18
    /// The estimated variance.
19
    estimate_var: f64,
20
    /// The sum in the current window.
21
    sum: DataSize,
22
    /// The size of the current window.
23
    current_window: Duration,
24
    /// The last time the window was updated.
25
    last_update: Option<Instant>,
26
}
27
 
28
impl AckedBitrateEstimator {
29
    pub(super) fn new(initial_window: Duration, window: Duration) -> Self {
30
        Self {
31
            initial_window,
32
            window,
33
            estimate: None,
34
            estimate_var: 50.0,
35
            sum: DataSize::ZERO,
36
            current_window: Duration::ZERO,
37
            last_update: None,
38
        }
39
    }
40
 
41
    pub(super) fn update(&mut self, receive_time: Instant, packet_size: DataSize) {
42
        let window = if self.estimate.is_none() {
43
            // Use the initial, larger, window at first
44
            self.initial_window
45
        } else {
46
            self.window
47
        };
48
        let Some((sample_estimate, is_small_sample)) =
49
            self.update_window(receive_time, packet_size, window)
50
        else {
51
            // No update
52
            return;
53
        };
54
 
55
        let Some(estimate) = self.estimate else {
56
            // This is the initial estimate, use it to initialize the estimate.
57
            self.estimate = Some(sample_estimate);
58
            return;
59
        };
60
 
61
        let scale = if is_small_sample && sample_estimate < estimate {
62
            SMALL_SAMPLE_UNCERTAINTY
63
        } else {
64
            UNCERTAINTY
65
        };
66
 
67
        let sample_estimate_bps = sample_estimate.as_f64();
68
        let estimate_bps = estimate.as_f64();
69
        // Define the sample uncertainty as a function of how far away it is from the
70
        // current estimate. With low values of uncertainty_symmetry_cap_ we add more
71
        // uncertainty to increases than to decreases. For higher values we approach
72
        // symmetry.
73
        let sample_uncertainty = scale * (estimate_bps - sample_estimate_bps).abs() / estimate_bps;
74
        let sample_var = sample_uncertainty.powf(2.0);
75
 
76
        // Update a bayesian estimate of the rate, weighting it lower if the sample
77
        // uncertainty is large.
78
        // The bitrate estimate uncertainty is increased with each update to model
79
        // that the bitrate changes over time.
80
        let pred_bitrate_estimate_var = self.estimate_var + 5.0;
81
        let mut new_estimate = (sample_var * estimate_bps
82
            + pred_bitrate_estimate_var * sample_estimate_bps)
83
            / (sample_var + pred_bitrate_estimate_var);
84
 
85
        new_estimate = new_estimate.max(ESTIMATE_FLOOR.as_f64());
86
        self.estimate = Some(Bitrate::bps(new_estimate.ceil() as u64));
87
        self.estimate_var =
88
            (sample_var * pred_bitrate_estimate_var) / (sample_var + pred_bitrate_estimate_var);
89
    }
90
 
91
    pub(super) fn current_estimate(&self) -> Option<Bitrate> {
92
        self.estimate
93
    }
94
 
95
    fn update_window(
96
        &mut self,
97
        receive_time: Instant,
98
        packet_size: DataSize,
99
        window: Duration,
100
    ) -> Option<(Bitrate, bool)> {
101
        let time_moved_back = Some(receive_time) < self.last_update;
102
        if time_moved_back {
103
            // Time moved backwards, reset state
104
            self.sum = DataSize::ZERO;
105
            self.current_window = Duration::ZERO;
106
            self.last_update = Some(receive_time);
107
 
108
            return None;
109
        }
110
 
111
        if let Some(last) = self.last_update {
112
            self.current_window += receive_time - last;
113
            if receive_time - last >= window {
114
                // No update for a while, reset estimates.
115
                self.sum = DataSize::ZERO;
116
                let remaining = self.current_window.as_micros() as u64 % window.as_micros() as u64;
117
                self.current_window = Duration::from_micros(remaining);
118
            }
119
        }
120
 
121
        self.last_update = Some(receive_time);
122
 
123
        let mut estimate = None;
124
 
125
        let mut is_small = false;
126
        if self.current_window >= window {
127
            is_small = self.sum < SMALL_SAMPLE_THRESHOLD;
128
            estimate = Some(self.sum / window);
129
            self.sum = DataSize::ZERO;
130
            self.current_window -= window;
131
        }
132
 
133
        self.sum += packet_size;
134
 
135
        estimate.map(|e| (e, is_small))
136
    }
137
}
138
 
139
#[cfg(test)]
140
mod test {
141
    use super::*;
142
 
143
    #[test]
144
    fn test_no_estimate_before_initial_window_has_passed() {
145
        let now = Instant::now();
146
        let mut estimator =
147
            AckedBitrateEstimator::new(Duration::from_millis(500), Duration::from_millis(150));
148
 
149
        estimator.update(now, DataSize::bytes(950));
150
        estimator.update(now + Duration::from_millis(250), DataSize::bytes(381));
151
        estimator.update(now + Duration::from_millis(499), DataSize::bytes(1110));
152
 
153
        assert!(
154
            estimator.current_estimate().is_none(),
155
            "AckedBitrateEstiamtor should produce no estimate before the initial window is reached"
156
        );
157
 
158
        estimator.update(now + Duration::from_millis(501), DataSize::bytes(1110));
159
 
160
        let estimate = estimator.current_estimate().expect(
161
            "After the first window of time has passed AckedBitrateEstimator should produce an estimate"
162
        );
163
 
164
        assert_eq!(
165
            estimate.as_u64(),
166
            39056,
167
            "AckedBitrateEstiamtor should produce the correct bitrate"
168
        );
169
    }
170
 
171
    #[test]
172
    fn test_correct_estimate_after_initial_window() {
173
        let now = Instant::now();
174
        let mut estimator =
175
            AckedBitrateEstimator::new(Duration::from_millis(500), Duration::from_millis(150));
176
 
177
        estimator.update(now, DataSize::bytes(2500));
178
        estimator.update(now + Duration::from_millis(250), DataSize::bytes(1392));
179
        estimator.update(now + Duration::from_millis(499), DataSize::bytes(4021));
180
        estimator.update(now + Duration::from_millis(500), DataSize::bytes(0));
181
 
182
        assert!(
183
            estimator.current_estimate().is_some(),
184
            "After the first window of time has passed AckedBitrateEstimator should produce an estimate"
185
        );
186
 
187
        estimator.update(now + Duration::from_millis(550), DataSize::bytes(271));
188
        estimator.update(now + Duration::from_millis(558), DataSize::bytes(813));
189
        estimator.update(now + Duration::from_millis(648), DataSize::bytes(731));
190
        // Will not be counted, part of next window
191
        estimator.update(now + Duration::from_millis(651), DataSize::bytes(900));
192
 
193
        let estimate = estimator.current_estimate().expect(
194
            "After the first window of time has passed AckedBitrateEstimator should produce an estimate"
195
        );
196
 
197
        assert_eq!(
198
            estimate.as_u64(),
199
            108320,
200
            "AckedBitrateEstiamtor should produce the correct bitrate"
201
        );
202
    }
203
 
204
    #[test]
205
    fn test_long_gap_preserves_overflow_via_modulo_after_initial() {
206
        let now = Instant::now();
207
        let window = Duration::from_millis(150);
208
        let mut estimator = AckedBitrateEstimator::new(Duration::from_millis(500), window);
209
 
210
        // First, get past the initial 500ms window to establish an estimate
211
        // Need to send enough data over 500ms to trigger the first estimate
212
        for i in 0..50 {
213
            estimator.update(now + Duration::from_millis(i * 10), DataSize::bytes(1000));
214
        }
215
        estimator.update(now + Duration::from_millis(501), DataSize::bytes(1000));
216
        assert!(
217
            estimator.current_estimate().is_some(),
218
            "Should have initial estimate after 500ms"
219
        );
220
 
221
        // Now we're using the 150ms window
222
        // Build up current_window to 100ms
223
        estimator.update(now + Duration::from_millis(550), DataSize::bytes(1000));
224
        estimator.update(now + Duration::from_millis(650), DataSize::bytes(1000));
225
        // current_window is now 100ms
226
 
227
        // Long gap of 225ms (> 150ms window)
228
        // current_window becomes 100 + 225 = 325ms
229
        // After gap detection: current_window = 325ms % 150ms = 25ms (correct)
230
        //                  or: current_window = 150ms % 150ms = 0ms (buggy)
231
        estimator.update(now + Duration::from_millis(875), DataSize::bytes(1000));
232
 
233
        // Add 50ms
234
        estimator.update(now + Duration::from_millis(925), DataSize::bytes(1000));
235
        // Correct: current_window = 25 + 50 = 75ms
236
        // Buggy: current_window = 0 + 50 = 50ms
237
 
238
        // Add 80ms
239
        let estimate_before = estimator.current_estimate().unwrap();
240
        estimator.update(now + Duration::from_millis(1005), DataSize::bytes(1000));
241
        // Correct: current_window = 75 + 80 = 155ms (>= 150ms, triggers new sample)
242
        // Buggy: current_window = 50 + 80 = 130ms (< 150ms, no new sample)
243
        let estimate_after = estimator.current_estimate().unwrap();
244
 
245
        // With correct modulo, estimate should change (new sample produced)
246
        // With buggy modulo, estimate stays the same (no sample produced yet)
247
        assert_ne!(
248
            estimate_before.as_u64(),
249
            estimate_after.as_u64(),
250
            "Estimate should change after accumulating >= 150ms due to correct modulo overflow"
251
        );
252
    }
253
}