1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
//! Independent, bounded capacity for setup, application and control traffic.
use std::sync::{
Arc,
atomic::{AtomicU64, AtomicUsize, Ordering},
};
/// Connection class. Classify authenticated peers before admitting application work.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Pool {
/// Connections whose handshake has not finished.
Handshake,
/// Ordinary application sessions.
Application,
/// Lifecycle control sessions; never dispatch application work on this budget.
Control,
}
/// Independent limits. A zero limit disables its pool.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Limits {
/// Maximum pending handshakes.
pub handshakes: usize,
/// Maximum application sessions.
pub application: usize,
/// Maximum control sessions, even when application capacity is exhausted.
pub control: usize,
}
impl Default for Limits {
fn default() -> Self {
Self {
handshakes: 32,
application: 256,
control: 32,
}
}
}
/// One pool's observation; counters are independent atomic samples.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PoolSnapshot {
/// Configured capacity.
pub limit: usize,
/// Permits currently held.
pub active: usize,
/// Failed admissions since construction.
pub rejected: u64,
}
struct Capacity {
limit: usize,
active: AtomicUsize,
rejected: AtomicU64,
}
/// Shared capacity, with no waiting queue or background task.
pub struct Admission {
pools: [Capacity; 3],
}
impl Admission {
/// Construct separate budgets. Keep one instance per service generation.
pub fn new(limits: Limits) -> Self {
Self {
pools: [limits.handshakes, limits.application, limits.control].map(|limit| Capacity {
limit,
active: AtomicUsize::new(0),
rejected: AtomicU64::new(0),
}),
}
}
/// Acquire without waiting. Hold the permit for the whole session.
pub fn try_acquire(self: &Arc<Self>, pool: Pool) -> Option<Permit> {
let capacity = &self.pools[pool as usize];
if capacity
.active
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |active| {
(active < capacity.limit).then(|| active + 1)
})
.is_err()
{
capacity.rejected.fetch_add(1, Ordering::Relaxed);
return None;
}
Some(Permit {
admission: Arc::clone(self),
pool,
})
}
/// Read occupancy without waiting for a connection or writer lock.
pub fn snapshot(&self, pool: Pool) -> PoolSnapshot {
let capacity = &self.pools[pool as usize];
PoolSnapshot {
limit: capacity.limit,
active: capacity.active.load(Ordering::Relaxed),
rejected: capacity.rejected.load(Ordering::Relaxed),
}
}
}
/// Releases one admission slot on drop, including cancellation and errors.
#[must_use = "hold the permit until the session ends"]
pub struct Permit {
admission: Arc<Admission>,
pool: Pool,
}
impl Drop for Permit {
fn drop(&mut self) {
self.admission.pools[self.pool as usize]
.active
.fetch_sub(1, Ordering::Relaxed);
}
}