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
//! Buffer pooling for reducing allocation overhead
//!
//! This module provides a buffer pool for reusing memory allocations
//! across multiple HTTP requests, reducing GC pressure and improving performance.
use std::sync::Arc;
use tokio::sync::Mutex;
/// A pool of reusable buffers for HTTP response bodies
pub struct BufferPool {
buffers: Arc<Mutex<Vec<Vec<u8>>>>,
buffer_size: usize,
max_pool_size: usize,
}
impl BufferPool {
/// Create a new buffer pool
///
/// # Arguments
/// * `buffer_size` - Initial size of each buffer (e.g., 512KB for typical market data)
/// * `max_pool_size` - Maximum number of buffers to keep in the pool
pub fn new(buffer_size: usize, max_pool_size: usize) -> Self {
Self {
buffers: Arc::new(Mutex::new(Vec::with_capacity(max_pool_size))),
buffer_size,
max_pool_size,
}
}
/// Get a buffer from the pool, or create a new one if pool is empty
pub async fn get(&self) -> Vec<u8> {
let mut buffers = self.buffers.lock().await;
match buffers.pop() {
Some(mut buffer) => {
buffer.clear();
buffer
},
None => {
// Pool is empty, create a new buffer
Vec::with_capacity(self.buffer_size)
},
}
}
/// Return a buffer to the pool
pub async fn return_buffer(&self, mut buffer: Vec<u8>) {
let mut buffers = self.buffers.lock().await;
// Only return to pool if we're under the size limit
if buffers.len() < self.max_pool_size {
buffer.clear();
// Shrink if buffer grew too large
if buffer.capacity() > self.buffer_size * 2 {
buffer.shrink_to(self.buffer_size);
}
buffers.push(buffer);
}
// Otherwise, let the buffer be dropped
}
/// Get the current number of buffers in the pool
pub async fn size(&self) -> usize {
let buffers = self.buffers.lock().await;
buffers.len()
}
/// Pre-allocate buffers in the pool
pub async fn prewarm(&self, count: usize) {
let mut buffers = self.buffers.lock().await;
for _ in 0..count.min(self.max_pool_size) {
buffers.push(Vec::with_capacity(self.buffer_size));
}
}
}
impl Default for BufferPool {
fn default() -> Self {
// Default: 512KB buffers, pool of 10
Self::new(512 * 1024, 10)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_buffer_pool_get_and_return() {
let pool = BufferPool::new(1024, 5);
let buffer = pool.get().await;
assert_eq!(buffer.capacity(), 1024);
pool.return_buffer(buffer).await;
assert_eq!(pool.size().await, 1);
}
#[tokio::test]
async fn test_buffer_pool_prewarm() {
let pool = BufferPool::new(1024, 5);
pool.prewarm(3).await;
assert_eq!(pool.size().await, 3);
}
#[tokio::test]
async fn test_buffer_pool_max_size() {
let pool = BufferPool::new(1024, 2);
let buf1 = pool.get().await;
let buf2 = pool.get().await;
let buf3 = pool.get().await;
pool.return_buffer(buf1).await;
pool.return_buffer(buf2).await;
pool.return_buffer(buf3).await; // This should be dropped, not added to pool
assert_eq!(pool.size().await, 2); // Max size is 2
}
}