use flume::{Receiver, Sender};
#[derive(Debug, Clone)]
pub struct BufferPool {
tx: Sender<Vec<u8>>,
rx: Receiver<Vec<u8>>,
buffer_size: usize,
}
impl BufferPool {
pub fn new(capacity: usize, buffer_size: usize) -> Self {
let (tx, rx) = flume::bounded(capacity);
for _ in 0..capacity {
let _ = tx.send(vec![0u8; buffer_size]);
}
Self {
tx,
rx,
buffer_size,
}
}
pub async fn lease(&self) -> Vec<u8> {
match self.rx.recv_async().await {
Ok(buf) => buf,
Err(_) => vec![0u8; self.buffer_size],
}
}
pub fn release(&self, mut buf: Vec<u8>) {
if buf.capacity() < self.buffer_size {
buf = vec![0u8; self.buffer_size];
} else {
buf.resize(self.buffer_size, 0);
}
let _ = self.tx.send(buf);
}
}