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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
use tokio::sync::watch;
use crate::Error;
#[derive(Debug, Clone, Copy)]
struct CreditState {
used: u64,
max: u64,
/// Bytes freed/consumed since the last window update (recv-side only).
released: u64,
/// Set to true when the credit is closed (session teardown).
closed: bool,
}
/// Tracks used/max credit for flow control.
///
/// Works for both send and recv flow control:
/// - **Send**: `try_claim`/`claim` to reserve credit, `release` for rollback, `increase_max` on peer's MAX_DATA.
/// - **Recv**: `receive` to validate incoming data, `consume` to track app consumption and trigger window updates.
///
/// Clone is cheap (watch::Sender is internally Arc'd).
#[derive(Clone, Debug)]
pub struct Credit {
inner: watch::Sender<CreditState>,
}
impl Credit {
/// Create with initial max (used starts at 0).
pub fn new(max: u64) -> Self {
Self {
inner: watch::Sender::new(CreditState {
used: 0,
max,
released: 0,
closed: false,
}),
}
}
// --- Send-side methods ---
/// Try to claim up to `limit` units. Returns amount claimed (0 if none available).
pub fn try_claim(&self, limit: u64) -> u64 {
let mut claimed = 0;
self.inner.send_if_modified(|state| {
let available = state.max.saturating_sub(state.used);
claimed = limit.min(available);
if claimed > 0 {
state.used += claimed;
true
} else {
false
}
});
claimed
}
/// Claim up to `limit` units, waiting until credit is available.
pub async fn claim(&self, limit: u64) -> Result<u64, Error> {
loop {
if self.inner.borrow().closed {
return Err(Error::Closed);
}
let claimed = self.try_claim(limit);
if claimed > 0 {
return Ok(claimed);
}
let mut rx = self.inner.subscribe();
rx.wait_for(|state| state.used < state.max || state.closed)
.await
.map_err(|_| Error::Closed)?;
}
}
/// Claim exactly 1 unit and return the index (value of `used` before incrementing).
pub async fn claim_index(&self) -> Result<u64, Error> {
loop {
if self.inner.borrow().closed {
return Err(Error::Closed);
}
let mut index = 0;
let claimed = self.inner.send_if_modified(|state| {
if state.used < state.max {
index = state.used;
state.used += 1;
true
} else {
false
}
});
if claimed {
return Ok(index);
}
let mut rx = self.inner.subscribe();
rx.wait_for(|state| state.used < state.max || state.closed)
.await
.map_err(|_| Error::Closed)?;
}
}
/// Close the credit, causing all pending and future `claim()`/`claim_index()` calls
/// to return `Err(Error::Closed)`.
pub fn close(&self) {
self.inner.send_if_modified(|state| {
if state.closed {
false
} else {
state.closed = true;
true
}
});
}
/// Return previously claimed credit (rollback on failed send).
pub fn release(&self, amount: u64) {
self.inner.send_if_modified(|state| {
let new = state.used.saturating_sub(amount);
if new != state.used {
state.used = new;
true
} else {
false
}
});
}
/// Increase the max. Returns error if new_max < current max.
pub fn increase_max(&self, new_max: u64) -> Result<(), Error> {
let mut ok = true;
self.inner.send_if_modified(|state| {
if new_max < state.max {
ok = false;
return false;
}
if new_max == state.max {
return false;
}
state.max = new_max;
true
});
if ok {
Ok(())
} else {
Err(Error::FlowControlError)
}
}
// --- Recv-side methods ---
/// Set used to max(used, value). Returns false if value > max (flow control violation).
/// Used for stream count tracking where opening index N implies all indices 0..N.
pub fn receive_up_to(&self, value: u64) -> bool {
let mut ok = true;
self.inner.send_if_modified(|state| {
if value > state.max {
ok = false;
false
} else if value > state.used {
state.used = value;
true
} else {
false
}
});
ok
}
/// Validate and account for incoming data. Returns false if flow control is violated.
pub fn receive(&self, len: u64) -> bool {
let mut ok = true;
self.inner.send_if_modified(|state| {
if state.used + len > state.max {
ok = false;
false
} else {
state.used += len;
true
}
});
ok
}
/// Report that `len` bytes have been consumed by the application.
/// Returns `Some(new_max)` if a window update should be sent.
pub fn consume(&self, len: u64) -> Option<u64> {
let mut update = None;
self.inner.send_if_modified(|state| {
state.released += len;
// Send a window update when: used + 2*released > max
// i.e. more than half the remaining window has been consumed.
if state.used + 2 * state.released > state.max {
let new_max = state.max + state.released;
state.max = new_max;
state.released = 0;
update = Some(new_max);
true
} else {
false
}
});
update
}
}