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
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use super::Credits;
use crate::stream::send::{error::Error, flow};
use s2n_quic_core::{ensure, varint::VarInt};
use std::sync::{Condvar, Mutex};
pub struct State {
state: Mutex<Inner>,
notify: Condvar,
}
impl State {
#[inline]
pub fn new(initial_flow_offset: VarInt) -> Self {
Self {
state: Mutex::new(Inner {
stream_offset: VarInt::ZERO,
flow_offset: initial_flow_offset,
is_finished: false,
}),
notify: Condvar::new(),
}
}
}
struct Inner {
/// Monotonic offset which tracks where the application is currently writing
stream_offset: VarInt,
/// Monotonic offset which indicates the maximum offset the application can write to
flow_offset: VarInt,
/// Indicates that the stream has been finalized
is_finished: bool,
}
impl State {
/// Called by the background worker to release flow credits
///
/// Callers MUST ensure the provided offset is monotonic.
#[inline]
pub fn release(&self, flow_offset: VarInt) -> Result<(), Error> {
let mut guard = self.state.lock().map_err(|_| Error::FatalError)?;
// only notify subscribers if we actually increment the offset
debug_assert!(
guard.flow_offset < flow_offset,
"flow offsets MUST be monotonic"
);
ensure!(guard.flow_offset < flow_offset, Ok(()));
guard.flow_offset = flow_offset;
drop(guard);
self.notify.notify_all();
Ok(())
}
/// Called by the application to acquire flow credits
#[inline]
pub fn acquire(&self, mut request: flow::Request) -> Result<Credits, Error> {
let mut guard = self.state.lock().map_err(|_| Error::FatalError)?;
loop {
ensure!(!guard.is_finished, Err(Error::FinalSizeChanged));
// TODO check for an error
let current_offset = guard.stream_offset;
let flow_offset = guard.flow_offset;
debug_assert!(
current_offset <= flow_offset,
"current_offset={current_offset} should be <= flow_offset={flow_offset}"
);
let Some(flow_credits) = flow_offset
.as_u64()
.checked_sub(current_offset.as_u64())
.filter(|v| {
// if we're finishing the stream and don't have any buffered data, then we
// don't need any flow control
if request.len == 0 && request.is_fin {
true
} else {
*v > 0
}
})
else {
guard = self.notify.wait(guard).map_err(|_| Error::FatalError)?;
continue;
};
// clamp the request to the flow credits we have
request.clamp(flow_credits);
// update the stream offset with the given request
guard.stream_offset = current_offset
.checked_add_usize(request.len)
.ok_or(Error::PayloadTooLarge)?;
// update the finished status
guard.is_finished |= request.is_fin;
// drop the lock before notifying all of the waiting Condvar
drop(guard);
// notify the other handles when we finish
if request.is_fin {
self.notify.notify_all();
}
// the offset was correctly updated so return our acquired credits
let credits = request.response(current_offset);
return Ok(credits);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::stream::send::path;
use std::{
sync::atomic::{AtomicU64, Ordering},
thread,
};
#[test]
fn concurrent_flow() {
let mut initial_offset = VarInt::from_u8(255);
let expected_len = VarInt::from_u16(u16::MAX);
let state = State::new(initial_offset);
let path_info = path::Info {
max_datagram_size: 1500,
send_quantum: 10,
ecn: Default::default(),
next_expected_control_packet: Default::default(),
};
let total = AtomicU64::new(0);
let workers = 5;
let worker_counts = Vec::from_iter((0..workers).map(|_| AtomicU64::new(0)));
thread::scope(|s| {
let total = &total;
let path_info = &path_info;
let state = &state;
for (idx, count) in worker_counts.iter().enumerate() {
s.spawn(move || {
thread::sleep(core::time::Duration::from_millis(10));
let mut buffer_len = 1;
let mut is_fin = false;
let max_segments = 10;
let max_header_len = 50;
let mut max_offset = VarInt::ZERO;
loop {
let mut request = flow::Request {
len: buffer_len,
initial_len: buffer_len,
is_fin,
};
request.clamp(path_info.max_flow_credits(max_header_len, max_segments));
let Ok(credits) = state.acquire(request) else {
break;
};
eprintln!(
"thread={idx} offset={}..{}",
credits.offset,
credits.offset + credits.len
);
buffer_len += 1;
buffer_len = buffer_len.min(
expected_len
.as_u64()
.saturating_sub(credits.offset.as_u64())
.saturating_sub(credits.len as u64)
as usize,
);
assert!(max_offset <= credits.offset);
max_offset = credits.offset;
if buffer_len == 0 {
is_fin = true;
}
total.fetch_add(credits.len as _, Ordering::Relaxed);
count.fetch_add(credits.len as _, Ordering::Relaxed);
}
});
}
s.spawn(|| {
let mut credits = 10;
while initial_offset < expected_len {
thread::sleep(core::time::Duration::from_millis(1));
initial_offset = (initial_offset + credits).min(expected_len);
credits += 1;
let _ = state.release(initial_offset);
}
});
});
assert_eq!(total.load(Ordering::Relaxed), expected_len.as_u64());
let mut at_least_one_write = true;
for (idx, count) in worker_counts.into_iter().enumerate() {
let count = count.load(Ordering::Relaxed);
eprintln!("thread={idx}, count={}", count);
if count == 0 {
at_least_one_write = false;
}
}
let _ = at_least_one_write;
// TODO the Mutex mechanism doesn't fairly distribute between workers so don't make this
// assertion until we can do something more reliable
/*
assert!(
at_least_one_write,
"all workers need to write at least one byte"
);
*/
}
}