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 peer: peer::Dyn,
45
46 max_send_streams: usize,
48
49 num_send_streams: usize,
51
52 max_recv_streams: usize,
54
55 num_recv_streams: usize,
57
58 max_local_reset_streams: usize,
60
61 num_local_reset_streams: usize,
63
64 max_remote_reset_streams: usize,
66
67 num_remote_reset_streams: usize,
69
70 max_local_error_reset_streams: Option<usize>,
75
76 num_local_error_reset_streams: usize,
79
80 data_frame_budget: Budget,
82
83 data_frame_budget_is_auto: bool,
85
86 num_recv_empty_data_frames: usize,
89}
90
91impl Counts {
92 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 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 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 pub fn next_send_stream_will_reach_capacity(&self) -> bool {
154 self.max_send_streams <= (self.num_send_streams + 1)
155 }
156
157 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 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 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 pub fn can_inc_num_recv_streams(&self) -> bool {
188 self.max_recv_streams > self.num_recv_streams
189 }
190
191 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 self.num_recv_streams += 1;
202 stream.is_counted = true;
203 }
204
205 pub fn can_inc_num_send_streams(&self) -> bool {
207 self.max_send_streams > self.num_send_streams
208 }
209
210 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 self.num_send_streams += 1;
221 stream.is_counted = true;
222 }
223
224 pub fn can_inc_num_reset_streams(&self) -> bool {
226 self.max_local_reset_streams > self.num_local_reset_streams
227 }
228
229 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 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 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 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 let is_pending_reset = stream.is_pending_reset_expiration();
287
288 let ret = f(self, &mut stream);
290
291 self.transition_after(stream, is_pending_reset);
292
293 ret
294 }
295
296 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 self.dec_num_streams(&mut stream);
323 }
324 }
325
326 if stream.is_released() {
328 stream.remove();
329 }
330 }
331
332 pub(crate) fn max_send_streams(&self) -> usize {
335 self.max_send_streams
336 }
337
338 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 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}