h2/proto/streams/
counts.rs

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

h2/proto/streams/
counts.rs
