remem/
lib.rs

1//! # Thread-safe object reuse
2//!
3//! In many programs the allocator can be a bottleneck, especially when
4//! allocating larger structures in a short amount of time. The allocator will
5//! have to spend time finding free memory to fit new allocations, and
6//! defragmenting memory to free up new space for new ones.
7//!
8//! However often we'll be allocating in a loop, and right after we drop an
9//! object we'll want a new object of the same size. This is where `remem` comes
10//! in: it allows you to reuse memory in a thread-safe way.
11//!
12//! This is useful when writing networked services, performing file reads, or
13//! anything else that might allocate a lot. Internally it's implemented using a
14//! "Treiber stack" which is a really fast algorithm that makes `remem` safe to
15//! use between threads!
16//!
17//! # Example
18//!
19//! ```rust
20//! use remem::Pool;
21//! use std::thread;
22//!
23//! let p = Pool::new(|| vec![0usize; 1024]);
24//!
25//! // Create a new handle onto the pool and send it to a new thread.
26//! let p2 = p.clone();
27//! let t = thread::spawn(move || {
28//!     // Get a new vec from the pool and push two values into it.
29//!     let mut v = p2.get();
30//!     v.push(1);
31//!     v.push(2);
32//! });
33//!
34//! // Meanwhile we can still access the original handle from the main
35//! // thread and use it to get new vecs from.
36//! let mut v = p.get();
37//! v.push(1);
38//! v.push(2);
39//!
40//! // Wait for the other thread to complete
41//! t.join().unwrap();
42//!
43//! // When the vec is dropped, it's returned to the pool and is ready to be
44//! // used again from a next call to `p.get()`.
45//! drop(v);
46//! ```
47use std::ops::{Deref, DerefMut};
48use std::sync::Arc;
49use treiber_stack::TreiberStack as Stack;
50
51mod treiber_stack;
52
53struct Internal<T> {
54    stack: Stack<T>,
55    create: Box<dyn Fn() -> T + Send + Sync>,
56    clear: Box<dyn Fn(&mut T) + Send + Sync>,
57}
58
59impl<T> Internal<T> {
60    fn new<C, D>(create: C, clear: D) -> Self
61    where
62        C: Fn() -> T + Send + Sync + 'static,
63        D: Fn(&mut T) -> () + Send + Sync + 'static,
64    {
65        Internal {
66            stack: Stack::new(),
67            create: Box::new(create),
68            clear: Box::new(clear),
69        }
70    }
71}
72
73/// A pool of reusable memory.
74pub struct Pool<T> {
75    internal: Arc<Internal<T>>,
76}
77
78impl<T> Pool<T> {
79    /// Create a new Pool from an initializer function.
80    pub fn new<C>(create: C) -> Pool<T>
81    where
82        C: Fn() -> T + Send + Sync + 'static,
83    {
84        Pool {
85            internal: Arc::new(Internal::new(create, |_| {})),
86        }
87    }
88
89    /// Create a new Pool from an initializer function and a clear function.
90    ///
91    /// The clear function will be called whenever the `ItemGuard` is dropped,
92    /// and provides an opportunity to clear values and remove potentially
93    /// sensitive information from the items before returning it to the memory
94    /// pool.
95    ///
96    /// Note that `drop` in Rust is not guaranteed to run, so this function
97    /// is not a replacement for proper security measures.
98    pub fn with_clear<C, D>(create: C, clear: D) -> Pool<T>
99    where
100        C: Fn() -> T + Send + Sync + 'static,
101        D: Fn(&mut T) -> () + Send + Sync + 'static,
102    {
103        Pool {
104            internal: Arc::new(Internal::new(create, clear)),
105        }
106    }
107
108    /// Get an item from the pool.
109    pub fn get<'a>(&'a self) -> ItemGuard<'a, T> {
110        let pool = &self.internal;
111        let item = pool.stack.pop();
112        ItemGuard {
113            item: Some(item.unwrap_or_else(|| (*self.internal.create)())),
114            pool: self,
115        }
116    }
117
118    /// Store an item back inside the pool.
119    fn push(&self, mut item: T) {
120        (*self.internal.clear)(&mut item);
121        self.internal.stack.push(item);
122    }
123}
124
125impl<T> Clone for Pool<T> {
126    fn clone(&self) -> Self {
127        Pool {
128            internal: self.internal.clone(),
129        }
130    }
131}
132
133/// RAII structure used to reintroduce an item into the pool when dropped.
134pub struct ItemGuard<'a, T> {
135    item: Option<T>,
136    pool: &'a Pool<T>,
137}
138
139impl<'a, T> Drop for ItemGuard<'a, T> {
140    fn drop(&mut self) {
141        self.pool.push(self.item.take().unwrap())
142    }
143}
144
145impl<'a, T> Deref for ItemGuard<'a, T> {
146    type Target = T;
147
148    fn deref(&self) -> &Self::Target {
149        self.item.as_ref().unwrap()
150    }
151}
152
153impl<'a, T> DerefMut for ItemGuard<'a, T> {
154    fn deref_mut(&mut self) -> &mut Self::Target {
155        self.item.as_mut().unwrap()
156    }
157}
remem/lib.rs

remem/
lib.rs
