Skip to content

Commit

Permalink
Introduce the buffer-pool feature
Browse files Browse the repository at this point in the history
  • Loading branch information
blackbeam committed Dec 29, 2021
1 parent 1edd17e commit d9efbdc
Show file tree
Hide file tree
Showing 3 changed files with 144 additions and 0 deletions.
25 changes: 25 additions & 0 deletions src/buffer_pool/disabled.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#![cfg(not(feature = "buffer-pool"))]

use std::ops::Deref;

#[derive(Debug)]
#[repr(transparent)]
pub struct Buffer(Vec<u8>);

impl AsMut<Vec<u8>> for Buffer {
fn as_mut(&mut self) -> &mut Vec<u8> {
&mut self.0
}
}

impl Deref for Buffer {
type Target = [u8];

fn deref(&self) -> &Self::Target {
self.0.deref()
}
}

pub const fn get_buffer() -> Buffer {
Buffer(Vec::new())
}
103 changes: 103 additions & 0 deletions src/buffer_pool/enabled.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
#![cfg(feature = "buffer-pool")]

use crossbeam::queue::ArrayQueue;
use once_cell::sync::Lazy;

use std::{mem::replace, ops::Deref, sync::Arc};

const DEFAULT_MYSQL_BUFFER_POOL_CAP: usize = 128;
const DEFAULT_MYSQL_BUFFER_SIZE_CAP: usize = 4 * 1024 * 1024;

static BUFFER_POOL: Lazy<Arc<BufferPool>> = Lazy::new(|| Default::default());

#[inline(always)]
pub fn get_buffer() -> Buffer {
BUFFER_POOL.get()
}

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

impl Inner {
fn get(self: &Arc<Self>) -> Buffer {
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) }

Buffer(buf, Some(self.clone()))
}

fn put(&self, mut buf: Vec<u8>) {
buf.shrink_to(self.buffer_cap);
let _ = self.pool.push(buf);
}
}

/// Smart pointer to a buffer pool.
#[derive(Debug, Clone)]
pub struct BufferPool(Option<Arc<Inner>>);

impl BufferPool {
pub fn new() -> Self {
let pool_cap = std::env::var("RUST_MYSQL_BUFFER_POOL_CAP")
.ok()
.and_then(|x| x.parse().ok())
.unwrap_or(DEFAULT_MYSQL_BUFFER_POOL_CAP);

let buffer_cap = std::env::var("RUST_MYSQL_BUFFER_SIZE_CAP")
.ok()
.and_then(|x| x.parse().ok())
.unwrap_or(DEFAULT_MYSQL_BUFFER_SIZE_CAP);

Self((pool_cap > 0).then(|| {
Arc::new(Inner {
buffer_cap,
pool: ArrayQueue::new(pool_cap),
})
}))
}

pub fn get(self: &Arc<Self>) -> Buffer {
match self.0 {
Some(ref inner) => inner.get(),
None => Buffer(Vec::new(), None),
}
}
}

impl Default for BufferPool {
fn default() -> Self {
Self::new()
}
}

#[derive(Debug)]
pub struct Buffer(Vec<u8>, Option<Arc<Inner>>);

impl AsMut<Vec<u8>> for Buffer {
fn as_mut(&mut self) -> &mut Vec<u8> {
&mut self.0
}
}

impl Deref for Buffer {
type Target = [u8];

fn deref(&self) -> &Self::Target {
self.0.deref()
}
}

impl Drop for Buffer {
fn drop(&mut self) {
if let Some(ref inner) = self.1 {
inner.put(replace(&mut self.0, vec![]));
}
}
}
16 changes: 16 additions & 0 deletions src/buffer_pool/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
// Copyright (c) 2021 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

mod disabled;
mod enabled;

#[cfg(feature = "buffer-pool")]
pub use enabled::{get_buffer, Buffer};

#[cfg(not(feature = "buffer-pool"))]
pub use disabled::{get_buffer, Buffer};

0 comments on commit d9efbdc

Please sign in to comment.