Skip to main content

h2/proto/streams/
counts.rs

1use super::*;
2
3#[derive(Debug)]
4struct Budget {
5    spent: usize,
6    max: usize,
7}
8
9#[derive(Debug)]
10pub(super) struct BudgetExhausted;
11
12impl Budget {
13    fn new(max: usize) -> Self {
14        Budget { spent: 0, max }
15    }
16
17    fn consume(&mut self, amount: usize) -> Result<(), BudgetExhausted> {
18        let spent = self.spent.checked_add(amount).ok_or(BudgetExhausted)?;
19        if spent > self.max {
20            return Err(BudgetExhausted);
21        }
22        self.spent = spent;
23        Ok(())
24    }
25
26    fn replenish(&mut self, amount: usize) {
27        self.spent = self.spent.saturating_sub(amount);
28    }
29
30    fn set_max(&mut self, max: usize) {
31        self.max = max;
32    }
33
34    #[cfg(test)]
35    fn available(&self) -> usize {
36        self.max.saturating_sub(self.spent)
37    }
38}
39
40#[derive(Debug)]
41pub(super) struct Counts {
42    /// Acting as a client or server. This allows us to track which values to
43    /// inc / dec.
44    peer: peer::Dyn,
45
46    /// Maximum number of locally initiated streams
47    max_send_streams: usize,
48
49    /// Current number of remote initiated streams
50    num_send_streams: usize,
51
52    /// Maximum number of remote initiated streams
53    max_recv_streams: usize,
54
55    /// Current number of locally initiated streams
56    num_recv_streams: usize,
57
58    /// Maximum number of pending locally reset streams
59    max_local_reset_streams: usize,
60
61    /// Current number of pending locally reset streams
62    num_local_reset_streams: usize,
63
64    /// Max number of "pending accept" streams that were remotely reset
65    max_remote_reset_streams: usize,
66
67    /// Current number of "pending accept" streams that were remotely reset
68    num_remote_reset_streams: usize,
69
70    /// Maximum number of locally reset streams due to protocol error across
71    /// the lifetime of the connection.
72    ///
73    /// When this gets exceeded, we issue GOAWAYs.
74    max_local_error_reset_streams: Option<usize>,
75
76    /// Total number of locally reset streams due to protocol error across the
77    /// lifetime of the connection.
78    num_local_error_reset_streams: usize,
79
80    /// connection-level budget for DATA framing overhead.
81    data_frame_budget: Budget,
82
83    /// Whether the DATA frame budget follows the target connection window.
84    data_frame_budget_is_auto: bool,
85
86    /// Number of empty, non-final DATA frames received over the lifetime of
87    /// the connection.
88    num_recv_empty_data_frames: usize,
89}
90
91impl Counts {
92    /// Create a new `Counts` using the provided configuration values.
93    pub fn new(peer: peer::Dyn, config: &Config) -> Self {
94        Counts {
95            peer,
96            max_send_streams: config.initial_max_send_streams,
97            num_send_streams: 0,
98            max_recv_streams: config.remote_max_initiated.unwrap_or(usize::MAX),
99            num_recv_streams: 0,
100            max_local_reset_streams: config.local_reset_max,
101            num_local_reset_streams: 0,
102            max_remote_reset_streams: config.remote_reset_max,
103            num_remote_reset_streams: 0,
104            max_local_error_reset_streams: config.local_max_error_reset_streams,
105            num_local_error_reset_streams: 0,
106            data_frame_budget: Budget::new(config.data_frame_budget.resolve(None)),
107            data_frame_budget_is_auto: config.data_frame_budget.is_auto(),
108            num_recv_empty_data_frames: 0,
109        }
110    }
111
112    pub fn set_target_connection_window_size(&mut self, size: WindowSize) {
113        if self.data_frame_budget_is_auto {
114            self.data_frame_budget
115                .set_max(DataFrameBudget::Auto.resolve(Some(size)));
116        }
117    }
118
119    /// Records the framing overhead of a DATA frame.
120    pub fn record_data_frame(&mut self, payload_len: usize) -> Result<(), BudgetExhausted> {
121        if payload_len == 0 {
122            self.num_recv_empty_data_frames = self
123                .num_recv_empty_data_frames
124                .checked_add(1)
125                .ok_or(BudgetExhausted)?;
126            if self.num_recv_empty_data_frames > MAX_RECV_EMPTY_DATA_FRAMES {
127                return Err(BudgetExhausted);
128            }
129            Ok(())
130        } else if payload_len < DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD {
131            self.data_frame_budget
132                .consume(DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD - payload_len)
133        } else {
134            self.data_frame_budget
135                .replenish(payload_len - DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD);
136            Ok(())
137        }
138    }
139
140    /// Releases the framing overhead of a DATA frame that is no longer
141    /// buffered internally.
142    pub fn release_data_frame(&mut self, payload_len: usize) {
143        if payload_len != 0 && payload_len < DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD {
144            self.data_frame_budget
145                .replenish(DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD - payload_len);
146        }
147    }
148
149    /// Returns true when the next opened stream will reach capacity of outbound streams
150    ///
151    /// The number of client send streams is incremented in prioritize; send_request has to guess if
152    /// it should wait before allowing another request to be sent.
153    pub fn next_send_stream_will_reach_capacity(&self) -> bool {
154        self.max_send_streams <= (self.num_send_streams + 1)
155    }
156
157    /// Returns the current peer
158    pub fn peer(&self) -> peer::Dyn {
159        self.peer
160    }
161
162    pub fn has_streams(&self) -> bool {
163        self.num_send_streams != 0 || self.num_recv_streams != 0
164    }
165
166    /// Returns true if we can issue another local reset due to protocol error.
167    pub fn can_inc_num_local_error_resets(&self) -> bool {
168        if let Some(max) = self.max_local_error_reset_streams {
169            max > self.num_local_error_reset_streams
170        } else {
171            true
172        }
173    }
174
175    pub fn inc_num_local_error_resets(&mut self) {
176        assert!(self.can_inc_num_local_error_resets());
177
178        // Increment the number of remote initiated streams
179        self.num_local_error_reset_streams += 1;
180    }
181
182    pub(crate) fn max_local_error_resets(&self) -> Option<usize> {
183        self.max_local_error_reset_streams
184    }
185
186    /// Returns true if the receive stream concurrency can be incremented
187    pub fn can_inc_num_recv_streams(&self) -> bool {
188        self.max_recv_streams > self.num_recv_streams
189    }
190
191    /// Increments the number of concurrent receive streams.
192    ///
193    /// # Panics
194    ///
195    /// Panics on failure as this should have been validated before hand.
196    pub fn inc_num_recv_streams(&mut self, stream: &mut store::Ptr) {
197        assert!(self.can_inc_num_recv_streams());
198        assert!(!stream.is_counted);
199
200        // Increment the number of remote initiated streams
201        self.num_recv_streams += 1;
202        stream.is_counted = true;
203    }
204
205    /// Returns true if the send stream concurrency can be incremented
206    pub fn can_inc_num_send_streams(&self) -> bool {
207        self.max_send_streams > self.num_send_streams
208    }
209
210    /// Increments the number of concurrent send streams.
211    ///
212    /// # Panics
213    ///
214    /// Panics on failure as this should have been validated before hand.
215    pub fn inc_num_send_streams(&mut self, stream: &mut store::Ptr) {
216        assert!(self.can_inc_num_send_streams());
217        assert!(!stream.is_counted);
218
219        // Increment the number of remote initiated streams
220        self.num_send_streams += 1;
221        stream.is_counted = true;
222    }
223
224    /// Returns true if the number of pending reset streams can be incremented.
225    pub fn can_inc_num_reset_streams(&self) -> bool {
226        self.max_local_reset_streams > self.num_local_reset_streams
227    }
228
229    /// Increments the number of pending reset streams.
230    ///
231    /// # Panics
232    ///
233    /// Panics on failure as this should have been validated before hand.
234    pub fn inc_num_reset_streams(&mut self) {
235        assert!(self.can_inc_num_reset_streams());
236
237        self.num_local_reset_streams += 1;
238    }
239
240    pub(crate) fn max_remote_reset_streams(&self) -> usize {
241        self.max_remote_reset_streams
242    }
243
244    /// Returns true if the number of pending REMOTE reset streams can be
245    /// incremented.
246    pub(crate) fn can_inc_num_remote_reset_streams(&self) -> bool {
247        self.max_remote_reset_streams > self.num_remote_reset_streams
248    }
249
250    /// Increments the number of pending REMOTE reset streams.
251    ///
252    /// # Panics
253    ///
254    /// Panics on failure as this should have been validated before hand.
255    pub(crate) fn inc_num_remote_reset_streams(&mut self) {
256        assert!(self.can_inc_num_remote_reset_streams());
257
258        self.num_remote_reset_streams += 1;
259    }
260
261    pub(crate) fn dec_num_remote_reset_streams(&mut self) {
262        assert!(self.num_remote_reset_streams > 0);
263
264        self.num_remote_reset_streams -= 1;
265    }
266
267    pub fn apply_remote_settings(&mut self, settings: &frame::Settings, is_initial: bool) {
268        match settings.max_concurrent_streams() {
269            Some(val) => self.max_send_streams = val as usize,
270            None if is_initial => self.max_send_streams = usize::MAX,
271            None => {}
272        }
273    }
274
275    /// Run a block of code that could potentially transition a stream's state.
276    ///
277    /// If the stream state transitions to closed, this function will perform
278    /// all necessary cleanup.
279    ///
280    /// TODO: Is this function still needed?
281    pub fn transition<F, U>(&mut self, mut stream: store::Ptr, f: F) -> U
282    where
283        F: FnOnce(&mut Self, &mut store::Ptr) -> U,
284    {
285        // TODO: Does this need to be computed before performing the action?
286        let is_pending_reset = stream.is_pending_reset_expiration();
287
288        // Run the action
289        let ret = f(self, &mut stream);
290
291        self.transition_after(stream, is_pending_reset);
292
293        ret
294    }
295
296    // TODO: move this to macro?
297    pub fn transition_after(&mut self, mut stream: store::Ptr, is_reset_counted: bool) {
298        tracing::trace!(
299            "transition_after; stream={:?}; state={:?}; is_closed={:?}; \
300             pending_send_empty={:?}; buffered_send_data={}; \
301             num_recv={}; num_send={}",
302            stream.id,
303            stream.state,
304            stream.is_closed(),
305            stream.pending_send.is_empty(),
306            stream.buffered_send_data,
307            self.num_recv_streams,
308            self.num_send_streams
309        );
310
311        if stream.is_closed() {
312            if !stream.is_pending_reset_expiration() {
313                stream.unlink();
314                if is_reset_counted {
315                    self.dec_num_reset_streams();
316                }
317            }
318
319            if !stream.state.is_scheduled_reset() && stream.is_counted {
320                tracing::trace!("dec_num_streams; stream={:?}", stream.id);
321                // Decrement the number of active streams.
322                self.dec_num_streams(&mut stream);
323            }
324        }
325
326        // Release the stream if it requires releasing
327        if stream.is_released() {
328            stream.remove();
329        }
330    }
331
332    /// Returns the maximum number of streams that can be initiated by this
333    /// peer.
334    pub(crate) fn max_send_streams(&self) -> usize {
335        self.max_send_streams
336    }
337
338    /// Returns the maximum number of streams that can be initiated by the
339    /// remote peer.
340    pub(crate) fn max_recv_streams(&self) -> usize {
341        self.max_recv_streams
342    }
343
344    fn dec_num_streams(&mut self, stream: &mut store::Ptr) {
345        assert!(stream.is_counted);
346
347        if self.peer.is_local_init(stream.id) {
348            assert!(self.num_send_streams > 0);
349            self.num_send_streams -= 1;
350            stream.is_counted = false;
351        } else {
352            assert!(self.num_recv_streams > 0);
353            self.num_recv_streams -= 1;
354            stream.is_counted = false;
355        }
356    }
357
358    fn dec_num_reset_streams(&mut self) {
359        assert!(self.num_local_reset_streams > 0);
360        self.num_local_reset_streams -= 1;
361    }
362}
363
364impl Drop for Counts {
365    fn drop(&mut self) {
366        use std::thread;
367
368        if !thread::panicking() {
369            debug_assert!(!self.has_streams());
370        }
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377    use crate::frame::DEFAULT_INITIAL_WINDOW_SIZE;
378
379    fn counts_config() -> Config {
380        Config {
381            initial_max_send_streams: 0,
382            local_max_buffer_size: 0,
383            local_next_stream_id: 2.into(),
384            local_push_enabled: false,
385            extended_connect_protocol_enabled: false,
386            local_reset_duration: Duration::ZERO,
387            local_reset_max: 0,
388            remote_reset_max: 0,
389            remote_init_window_sz: DEFAULT_INITIAL_WINDOW_SIZE,
390            remote_max_initiated: None,
391            local_max_error_reset_streams: None,
392            data_frame_budget: DataFrameBudget::Configured(DEFAULT_DATA_FRAME_BUDGET),
393        }
394    }
395
396    fn counts() -> Counts {
397        Counts::new(peer::Dyn::Server, &counts_config())
398    }
399
400    #[test]
401    fn budget_is_bounded() {
402        let mut budget = Budget::new(10);
403
404        budget.consume(4).unwrap();
405        budget.replenish(20);
406        assert_eq!(budget.available(), 10);
407    }
408
409    #[test]
410    fn budget_reports_exhaustion_without_underflowing() {
411        let mut budget = Budget::new(10);
412
413        budget.consume(10).unwrap();
414        assert!(budget.consume(1).is_err());
415        assert_eq!(budget.available(), 0);
416    }
417
418    #[test]
419    fn resizing_budget_preserves_outstanding_charges() {
420        let mut budget = Budget::new(100);
421        budget.consume(60).unwrap();
422
423        budget.set_max(200);
424        assert_eq!(budget.available(), 140);
425
426        budget.set_max(40);
427        assert_eq!(budget.available(), 0);
428        assert!(budget.consume(1).is_err());
429
430        budget.replenish(30);
431        assert_eq!(budget.available(), 10);
432    }
433
434    #[test]
435    fn auto_budget_tracks_target_window_and_configured_budget_does_not() {
436        let mut auto = Counts::new(
437            peer::Dyn::Server,
438            &Config {
439                data_frame_budget: DataFrameBudget::Auto,
440                ..counts_config()
441            },
442        );
443        auto.set_target_connection_window_size(1024 * 1024);
444        assert_eq!(auto.data_frame_budget.max, 512 * 1024);
445        for _ in 0..1_000 {
446            auto.record_data_frame(171).unwrap();
447        }
448
449        let mut configured = Counts::new(
450            peer::Dyn::Server,
451            &Config {
452                data_frame_budget: DataFrameBudget::Configured(123),
453                ..counts_config()
454            },
455        );
456        configured.set_target_connection_window_size(1024 * 1024);
457        assert_eq!(configured.data_frame_budget.max, 123);
458    }
459
460    #[test]
461    fn good_sized_data_frames_do_not_exhaust_budget() {
462        let mut counts = counts();
463
464        for _ in 0..1_000_000 {
465            counts
466                .record_data_frame(DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD)
467                .unwrap();
468        }
469    }
470
471    #[test]
472    fn consumed_small_data_frames_do_not_exhaust_budget() {
473        let mut counts = counts();
474
475        for _ in 0..1_000_000 {
476            counts.record_data_frame(1).unwrap();
477            counts.release_data_frame(1);
478        }
479    }
480
481    #[test]
482    fn empty_data_frames_do_not_consume_data_frame_budget() {
483        let mut counts = counts();
484        counts.data_frame_budget = Budget::new(0);
485
486        for _ in 0..MAX_RECV_EMPTY_DATA_FRAMES {
487            counts.record_data_frame(0).unwrap();
488        }
489
490        // Empty frames have their own limit, while a non-empty small frame
491        // still consumes the independently configured DATA frame budget.
492        assert!(counts.record_data_frame(0).is_err());
493        assert!(counts.record_data_frame(1).is_err());
494    }
495
496    #[test]
497    fn large_data_frames_do_not_replenish_empty_data_frame_limit() {
498        let mut counts = counts();
499
500        for _ in 0..MAX_RECV_EMPTY_DATA_FRAMES {
501            counts.record_data_frame(0).unwrap();
502            counts
503                .record_data_frame(DEFAULT_DATA_FRAME_OVERHEAD_THRESHOLD * 2)
504                .unwrap();
505        }
506        assert!(counts.record_data_frame(0).is_err());
507    }
508}