rustpoolrefcell

How to build a pool of mutable Vecs that get reused on Drop?


I am trying to create a pool of mutable Vec objects that can be passed out to functions as needed, and reused when they're no longer needed (Since my target is WASM, I don't want to let Vecs themselves deallocate and reallocate). I have an implementation using Rc and RefCell, and I am wondering if there's a better (more efficient?) way to do this.

My current code uses Rc::strong_count to keep track of whether I've handed out a buffer and RefCell to allow mutable access to the Vec inside:

use std::{cell::RefCell, rc::Rc};

#[derive(Debug)]
struct BufferPool {
    buffers: Vec<Rc<RefCell<Vec<f64>>>>,
    buffer_size: usize,
}

impl BufferPool {
    fn new() -> Self {
        BufferPool {
            buffers: vec![],
            buffer_size: 3,
        }
    }
    fn add_buffer(&mut self) -> Rc<RefCell<Vec<f64>>> {
        self.buffers
            .push(Rc::new(RefCell::new(vec![0.; self.buffer_size])));
        Rc::clone(&self.buffers[self.buffers.len() - 1])
    }
    fn get_buffer(&mut self) -> Rc<RefCell<Vec<f64>>> {
        for buf in &self.buffers {
            // If the Rc count is 1, we haven't loaned the buffer out yet.
            if Rc::strong_count(&buf) == 1 {
                return Rc::clone(&buf);
            }
        }
        // If we made it here, there's no available buffer, so we need to create one.
        self.add_buffer()
    }
}

This code can be tested with:

#[test]
fn test_buffers() {
    let mut buffers = BufferPool::new();
    let buf_cell1 = buffers.get_buffer();
    {
        let mut buf1 = buf_cell1.borrow_mut();
        buf1[0] = 5.5;
    }
    {
        let buf_cell2 = buffers.get_buffer();
        let mut buf2 = buf_cell2.borrow_mut();
        buf2[1] = 6.6;
    }
    {
        let buf_cell3 = buffers.get_buffer();
        let mut buf3 = buf_cell3.borrow_mut();
        buf3[2] = 7.7;
    }
    dbg!(&buffers);
}

which gives the expected outputs:

 &buffers = BufferPool {
    buffers: [
        RefCell {
            value: [
                5.5,
                0.0,
                0.0,
            ],
        },
        RefCell {
            value: [
                0.0,
                6.6,
                7.7,
            ],
        },
    ],
    buffer_size: 3,
}

However, what I'm doing seems slightly inefficient since both Rc and RefCell::borrow_mut() are keeping track of whether a buffer has been "loaned out" (since RefCell has the ability to error if it's contents are double borrowed). Also, ergonomically, it's annoying that I cannot call buffers.get_buffer().borrow_mut() on one line without Rust complaining about dropped temporary values.

So, my question is: is there a better way to do this?


Solution

  • As you've noticed, providing access to objects via Rc<RefCell<T>> is functional but not very ergonomic. A better way of designing a pool is to return a wrapper that takes ownership of the value, and then puts it back in the pool when dropped. Here's a "basic" example of how to do that:

    use std::cell::RefCell;
    use std::ops::{Deref, DerefMut};
    
    #[derive(Debug)]
    struct BufferPool {
        buffers: RefCell<Vec<Vec<f32>>>,
        buffer_size: usize,
    }
    
    impl BufferPool {
        pub fn new() -> Self {
            BufferPool {
                buffers: RefCell::new(Vec::new()),
                buffer_size: 3,
            }
        }
        
        pub fn get_buffer(&self) -> BufferPoolRef {
            let mut buffers = self.buffers.borrow_mut();
            let buffer = buffers.pop().unwrap_or_else(|| vec![0.0; self.buffer_size]);
            BufferPoolRef { pool: self, buffer }
        }
        
        fn return_buffer(&self, buffer: Vec<f32>) {
            let mut buffers = self.buffers.borrow_mut();
            buffers.push(buffer);
        }
    }
    
    struct BufferPoolRef<'a> {
        pool: &'a BufferPool,
        buffer: Vec<f32>,
    }
    
    impl Deref for BufferPoolRef<'_> {
        type Target = Vec<f32>;
        
        fn deref(&self) -> &Self::Target {
            &self.buffer
        }
    }
    
    impl DerefMut for BufferPoolRef<'_> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.buffer
        }
    }
    
    impl Drop for BufferPoolRef<'_> {
        fn drop(&mut self) {
            let buffer = std::mem::take(&mut self.buffer);
            self.pool.return_buffer(buffer);
        }
    }
    
    fn main() {
        let mut buffers = BufferPool::new();
        let mut buf1 = buffers.get_buffer();
        {
            buf1[0] = 5.5;
        }
        {
            let mut buf2 = buffers.get_buffer();
            buf2[1] = 6.6;
        }
        {
            let mut buf3 = buffers.get_buffer();
            buf3[2] = 7.7;
        }
        drop(buf1);
        dbg!(&buffers);
    }
    
    [src/main.rs:71] &buffers = BufferPool {
        buffers: RefCell {
            value: [
                [
                    0.0,
                    6.6,
                    7.7,
                ],
                [
                    5.5,
                    0.0,
                    0.0,
                ],
            ],
        },
        buffer_size: 3,
    }
    

    However, instead of doing all this yourself, you could use a crate like object-pool or lifeguard. They both work on WASM and use the mechanism I've described above. Here's an implementation of BufferPool based on object-pool:

    use object_pool::{Pool, Reusable};
    
    struct BufferPool {
        pool: Pool<Vec<f32>>,
        buffer_size: usize,
    }
    
    impl BufferPool {
        pub fn new() -> Self {
            BufferPool {
                pool: Pool::new(2, || vec![0.0; 3]),
                buffer_size: 3,
            }
        }
        
        pub fn get_buffer(&self) -> Reusable<Vec<f32>> {
            self.pool.pull(|| vec![0.0; self.buffer_size])
        }
    }