Skip to content

Commit

Permalink
Merge pull request blackbeam#229 from cloneable/bufferpool-optimization
Browse files Browse the repository at this point in the history
Allow specifying an initial pooled buffer size
  • Loading branch information
blackbeam authored Feb 11, 2023
2 parents 0e8d22e + 53d6adb commit ca61d55
Showing 1 changed file with 21 additions and 19 deletions.
40 changes: 21 additions & 19 deletions src/buffer_pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ use std::{mem::replace, ops::Deref, sync::Arc};

#[derive(Debug)]
pub struct BufferPool {
buffer_cap: usize,
buffer_size_cap: usize,
buffer_init_cap: usize,
pool: ArrayQueue<Vec<u8>>,
}

Expand All @@ -22,25 +23,29 @@ impl BufferPool {
.and_then(|x| x.parse().ok())
.unwrap_or(128_usize);

let buffer_cap = std::env::var("MYSQL_ASYNC_BUFFER_SIZE_CAP")
let buffer_size_cap = std::env::var("MYSQL_ASYNC_BUFFER_SIZE_CAP")
.ok()
.and_then(|x| x.parse().ok())
.unwrap_or(4 * 1024 * 1024);

let buffer_init_cap = std::env::var("MYSQL_ASYNC_BUFFER_INIT_CAP")
.ok()
.and_then(|x| x.parse().ok())
.unwrap_or(0);

Self {
pool: ArrayQueue::new(pool_cap),
buffer_cap,
buffer_size_cap,
buffer_init_cap,
}
}

pub fn get(self: &Arc<Self>) -> PooledBuf {
let mut buf = self.pool.pop().unwrap_or_default();

// SAFETY:
// 1. OK – 0 is always within capacity
// 2. OK - nothing to initialize
unsafe { buf.set_len(0) }

let buf = self
.pool
.pop()
.unwrap_or_else(|| Vec::with_capacity(self.buffer_init_cap));
debug_assert_eq!(buf.len(), 0);
PooledBuf(buf, self.clone())
}

Expand All @@ -51,15 +56,12 @@ impl BufferPool {
}

fn put(self: &Arc<Self>, mut buf: Vec<u8>) {
if buf.len() > self.buffer_cap {
// TODO: until `Vec::shrink_to` stabilization

// SAFETY:
// 1. OK – new_len <= capacity
// 2. OK - 0..new_len is initialized
unsafe { buf.set_len(self.buffer_cap) }
buf.shrink_to_fit();
}
// SAFETY:
// 1. OK – 0 is always within capacity
// 2. OK - nothing to initialize
unsafe { buf.set_len(0) }

buf.shrink_to(self.buffer_size_cap);

// ArrayQueue will make sure to drop the buffer if capacity is exceeded
let _ = self.pool.push(buf);
Expand Down

0 comments on commit ca61d55

Please sign in to comment.