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
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
//! Pooled buffers with a total budget, wiped on release.
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::limits::SizeClass;
/// A byte buffer from a pool: a fixed capacity (its size class) and a filled length.
///
/// # Examples
///
/// ```
/// use net_io::{BufferPool, SizeClass};
/// let mut pool = BufferPool::standalone(1 << 20);
/// let mut buf = pool.take(SizeClass::SMALL).unwrap();
/// buf.spare()[..5].copy_from_slice(b"hello");
/// buf.set_len(5);
/// assert_eq!(buf.as_slice(), b"hello");
/// pool.put(buf);
/// ```
pub struct Buffer {
data: Box<[u8]>,
len: usize,
class: SizeClass,
}
impl Buffer {
/// A buffer not drawn from any pool's free list (the caller already charged the budget).
pub(crate) fn detached(class: SizeClass) -> Buffer {
Buffer::new(class)
}
/// An empty stand-in that owns no memory, left where a buffer travels with a
/// driver operation.
#[cfg(feature = "reactor")]
pub(crate) fn placeholder() -> Buffer {
Buffer {
data: Box::default(),
len: 0,
class: SizeClass::SMALL,
}
}
fn new(class: SizeClass) -> Buffer {
Buffer {
data: vec![0u8; class.bytes() as usize].into_boxed_slice(),
len: 0,
class,
}
}
/// The filled part.
#[must_use]
pub fn as_slice(&self) -> &[u8] {
&self.data[..self.len]
}
/// The filled part, mutable.
pub fn as_mut_slice(&mut self) -> &mut [u8] {
&mut self.data[..self.len]
}
/// The unfilled tail.
pub fn spare(&mut self) -> &mut [u8] {
&mut self.data[self.len..]
}
/// The filled length.
#[must_use]
pub fn len(&self) -> usize {
self.len
}
/// Whether nothing is filled.
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
/// Sets the filled length; saturates at the capacity.
pub fn set_len(&mut self, n: usize) {
self.len = n.min(self.data.len());
}
/// The capacity in bytes.
#[must_use]
pub fn capacity(&self) -> usize {
self.data.len()
}
/// The size class.
#[must_use]
pub fn class(&self) -> SizeClass {
self.class
}
/// Empties the buffer (does not wipe; the pool wipes on release).
pub fn clear(&mut self) {
self.len = 0;
}
/// Overwrites the filled part with zeroes and empties the buffer.
pub(crate) fn wipe(&mut self) {
let used = &mut self.data[..self.len];
used.fill(0);
// The buffer stays reachable and is read again by the next owner, so the
// stores cannot be removed as dead; black_box makes that explicit.
std::hint::black_box(&*used);
self.len = 0;
}
/// Moves the bytes `from..` to the front.
pub(crate) fn compact(&mut self, from: usize) {
let from = from.min(self.len);
self.data.copy_within(from..self.len, 0);
self.len -= from;
}
}
impl fmt::Debug for Buffer {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Buffer")
.field("len", &self.len)
.field("capacity", &self.data.len())
.finish()
}
}
/// Counts of one pool.
///
/// ```
/// let pool = net_io::BufferPool::standalone(1 << 20);
/// assert_eq!(pool.stats().allocated, 0);
/// ```
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct PoolStats {
/// Buffers this pool allocated.
pub allocated: u64,
/// Buffers in its free lists.
pub free: u64,
/// Bytes left in the shared budget.
pub budget_left: u64,
/// Takes refused because the budget was spent.
pub budget_refusals: u64,
}
/// A worker's buffer pool: free lists per size class, drawing new buffers from a
/// budget shared by every worker of a server.
pub struct BufferPool {
lists: Vec<(SizeClass, Vec<Buffer>)>,
budget: Arc<AtomicU64>,
wipe: bool,
allocated: u64,
refusals: u64,
}
impl BufferPool {
pub(crate) fn new(budget: Arc<AtomicU64>, wipe: bool) -> BufferPool {
BufferPool {
lists: Vec::with_capacity(4),
budget,
wipe,
allocated: 0,
refusals: 0,
}
}
/// A pool with its own budget of `bytes`, for use outside a server.
///
/// ```
/// let mut pool = net_io::BufferPool::standalone(64 * 1024);
/// assert!(pool.take(net_io::SizeClass::LARGE).is_some());
/// assert!(pool.take(net_io::SizeClass::LARGE).is_none()); // budget spent
/// ```
#[must_use]
pub fn standalone(bytes: u64) -> BufferPool {
BufferPool::new(Arc::new(AtomicU64::new(bytes)), true)
}
/// A buffer of `class`, or `None` when the free list is empty and the budget is spent.
pub fn take(&mut self, class: SizeClass) -> Option<Buffer> {
if let Some(buf) = self.list(class).pop() {
return Some(buf);
}
let bytes = u64::from(class.bytes());
let mut left = self.budget.load(Ordering::Relaxed);
loop {
if left < bytes {
self.refusals += 1;
return None;
}
match self.budget.compare_exchange_weak(
left,
left - bytes,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(actual) => left = actual,
}
}
self.allocated += 1;
Some(Buffer::new(class))
}
/// Returns a buffer; its used part is wiped first unless wiping is off.
///
/// A buffer from another worker's pool is accepted (it is just memory).
pub fn put(&mut self, mut buf: Buffer) {
if self.wipe {
buf.wipe();
} else {
buf.clear();
}
let class = buf.class;
self.list(class).push(buf);
}
/// Counts of this pool.
#[must_use]
pub fn stats(&self) -> PoolStats {
PoolStats {
allocated: self.allocated,
free: self.lists.iter().map(|(_, l)| l.len() as u64).sum(),
budget_left: self.budget.load(Ordering::Relaxed),
budget_refusals: self.refusals,
}
}
fn list(&mut self, class: SizeClass) -> &mut Vec<Buffer> {
let pos = match self.lists.iter().position(|(c, _)| *c == class) {
Some(p) => p,
None => {
self.lists.push((class, Vec::new()));
self.lists.len() - 1
}
};
&mut self.lists[pos].1
}
}
impl fmt::Debug for BufferPool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BufferPool")
.field("stats", &self.stats())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pool_reuses_and_wipes() {
let mut pool = BufferPool::standalone(1 << 20);
let mut b = pool.take(SizeClass::SMALL).unwrap();
b.spare()[..6].copy_from_slice(b"secret");
b.set_len(6);
pool.put(b);
let mut b = pool.take(SizeClass::SMALL).unwrap();
assert_eq!(pool.stats().allocated, 1);
assert!(b.is_empty());
b.set_len(6);
assert_eq!(b.as_slice(), &[0u8; 6]);
}
#[test]
fn test_pool_budget_bounds_allocation() {
let mut pool = BufferPool::standalone(u64::from(SizeClass::SMALL.bytes()) * 2);
let a = pool.take(SizeClass::SMALL).unwrap();
let _b = pool.take(SizeClass::SMALL).unwrap();
assert!(pool.take(SizeClass::SMALL).is_none());
assert_eq!(pool.stats().budget_refusals, 1);
pool.put(a);
assert!(pool.take(SizeClass::SMALL).is_some());
}
#[test]
fn test_set_len_saturates_and_compact_moves_tail() {
let mut pool = BufferPool::standalone(1 << 20);
let mut b = pool.take(SizeClass::SMALL).unwrap();
b.set_len(usize::MAX);
assert_eq!(b.len(), b.capacity());
b.set_len(0);
b.spare()[..4].copy_from_slice(b"abcd");
b.set_len(4);
b.compact(2);
assert_eq!(b.as_slice(), b"cd");
b.compact(99);
assert!(b.is_empty());
}
#[test]
fn test_pool_without_wipe_keeps_bytes_but_empties() {
let mut pool = BufferPool::new(Arc::new(AtomicU64::new(1 << 20)), false);
let mut b = pool.take(SizeClass::SMALL).unwrap();
b.spare()[0] = 9;
b.set_len(1);
pool.put(b);
let b = pool.take(SizeClass::SMALL).unwrap();
assert!(b.is_empty());
}
}