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
//! # Thread-safe object reuse
//!
//! In many programs the allocator can be a bottleneck, especially when
//! allocating larger structures in a short amount of time. The allocator will
//! have to spend time finding free memory to fit new allocations, and
//! defragmenting memory to free up new space for new ones.
//!
//! However often we'll be allocating in a loop, and right after we drop an
//! object we'll want a new object of the same size. This is where `remem` comes
//! in: it allows you to reuse memory in a thread-safe way.
//!
//! This is useful when writing networked services, performing file reads, or
//! anything else that might allocate a lot. Internally it's implemented using a
//! "Treiber stack" which is a really fast algorithm that makes `remem` safe to
//! use between threads!
//!
//! # Example
//!
//! ```rust
//! use remem::Pool;
//! use std::thread;
//!
//! let p = Pool::new(|| vec![0usize; 1024]);
//!
//! // Create a new handle onto the pool and send it to a new thread.
//! let p2 = p.clone();
//! let t = thread::spawn(move || {
//!     // Get a new vec from the pool and push two values into it.
//!     let mut v = p2.get();
//!     v.push(1);
//!     v.push(2);
//! });
//!
//! // Meanwhile we can still access the original handle from the main
//! // thread and use it to get new vecs from.
//! let mut v = p.get();
//! v.push(1);
//! v.push(2);
//!
//! // Wait for the other thread to complete
//! t.join().unwrap();
//!
//! // When the vec is dropped, it's returned to the pool and is ready to be
//! // used again from a next call to `p.get()`.
//! drop(v);
//! ```
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
use treiber_stack::TreiberStack as Stack;

mod treiber_stack;

struct Internal<T> {
    stack: Stack<T>,
    create: Box<dyn Fn() -> T + Send + Sync>,
    clear: Box<dyn Fn(&mut T) + Send + Sync>,
}

impl<T> Internal<T> {
    fn new<C, D>(create: C, clear: D) -> Self
    where
        C: Fn() -> T + Send + Sync + 'static,
        D: Fn(&mut T) -> () + Send + Sync + 'static,
    {
        Internal {
            stack: Stack::new(),
            create: Box::new(create),
            clear: Box::new(clear),
        }
    }
}

/// A pool of reusable memory.
pub struct Pool<T> {
    internal: Arc<Internal<T>>,
}

impl<T> Pool<T> {
    /// Create a new Pool from an initializer function.
    pub fn new<C>(create: C) -> Pool<T>
    where
        C: Fn() -> T + Send + Sync + 'static,
    {
        Pool {
            internal: Arc::new(Internal::new(create, |_| {})),
        }
    }

    /// Create a new Pool from an initializer function and a clear function.
    ///
    /// The clear function will be called whenever the `ItemGuard` is dropped,
    /// and provides an opportunity to clear values and remove potentially
    /// sensitive information from the items before returning it to the memory
    /// pool.
    ///
    /// Note that `drop` in Rust is not guaranteed to run, so this function
    /// is not a replacement for proper security measures.
    pub fn with_clear<C, D>(create: C, clear: D) -> Pool<T>
    where
        C: Fn() -> T + Send + Sync + 'static,
        D: Fn(&mut T) -> () + Send + Sync + 'static,
    {
        Pool {
            internal: Arc::new(Internal::new(create, clear)),
        }
    }

    /// Get an item from the pool.
    pub fn get<'a>(&'a self) -> ItemGuard<'a, T> {
        let pool = &self.internal;
        let item = pool.stack.pop();
        ItemGuard {
            item: Some(item.unwrap_or_else(|| (*self.internal.create)())),
            pool: self,
        }
    }

    /// Store an item back inside the pool.
    fn push(&self, mut item: T) {
        (*self.internal.clear)(&mut item);
        self.internal.stack.push(item);
    }
}

impl<T> Clone for Pool<T> {
    fn clone(&self) -> Self {
        Pool {
            internal: self.internal.clone(),
        }
    }
}

/// RAII structure used to reintroduce an item into the pool when dropped.
pub struct ItemGuard<'a, T> {
    item: Option<T>,
    pool: &'a Pool<T>,
}

impl<'a, T> Drop for ItemGuard<'a, T> {
    fn drop(&mut self) {
        self.pool.push(self.item.take().unwrap())
    }
}

impl<'a, T> Deref for ItemGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.item.as_ref().unwrap()
    }
}

impl<'a, T> DerefMut for ItemGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.item.as_mut().unwrap()
    }
}