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

use std::cell;
use std::rc::Rc;

struct PoolState {
    bufs: Vec<Vec<u8>>,
    packet_size: usize,
}

pub struct HeapPool {
    state: Rc<cell::RefCell<PoolState>>,
}
impl HeapPool {
    pub fn new(packet_count: usize, packet_size: usize) -> HeapPool {
        HeapPool {
            state: Rc::new(cell::RefCell::new(PoolState {
                bufs: Self::mk_bufs(packet_count, packet_size),
                packet_size,
            })),
        }
    }

    fn mk_bufs(packet_count: usize, packet_size: usize) -> Vec<Vec<u8>> {
        (0..packet_count)
            .map(|_| Vec::with_capacity(packet_size))
            .collect()
    }
}
impl crate::BufferPool for HeapPool {
    type P = HeapPacket;

    fn allocate(&self) -> Option<Self::P> {
        // TODO: maybe this should allocate N at once, for efficiency?

        let mut state = self.state.borrow_mut();
        let mut buf = state.bufs.pop()?;
        // TODO: try to cook up a scheme to avoid zeroing-out the buffer
        //       e.g. https://github.com/tbu-/buffer ?
        buf.clear();
        buf.resize(state.packet_size, 0);
        Some(HeapPacket {
            state: self.state.clone(),
            buf: Some(buf),
        })
    }
}
impl Clone for HeapPool {
    fn clone(&self) -> Self {
        HeapPool {
            state: self.state.clone(),
        }
    }
}

pub struct HeapPacket {
    state: Rc<cell::RefCell<PoolState>>,
    // the Vec is wrapped in an Option simply so that we can extract it and pass it by value back
    // to the Heap pool in the Drop impl.  buf will be 'Some' at all other times.
    buf: Option<Vec<u8>>,
}
impl crate::Packet for HeapPacket {
    fn payload(&self) -> &[u8] {
        &self.buf.as_ref().unwrap()[..]
    }
    fn payload_mut(&mut self) -> &mut [u8] {
        &mut self.buf.as_mut().unwrap()[..]
    }
    fn truncate(&mut self, size: usize) {
        assert_ne!(size, 0);
        assert!(self.buf.as_ref().unwrap().len() >= size);
        self.buf.as_mut().unwrap().truncate(size)
    }
}
impl Drop for HeapPacket {
    fn drop(&mut self) {
        self.state.borrow_mut().bufs.push(self.buf.take().unwrap())
    }
}

#[cfg(test)]
mod test {
    use crate::heap_pool::HeapPool;
    use crate::BufferPool;
    use crate::Packet;

    #[test]
    fn it_works() {
        let pool = HeapPool::new(2, 1500);
        {
            let _one = pool.allocate().unwrap();
            let mut two = pool.allocate().unwrap();
            assert!(pool.allocate().is_none());
            {
                let data = two.payload_mut();
                data[0] = 123;
            }
            assert_eq!(123, two.payload()[0]);
        }
        // Now all the above pool usage is done, the allocations should have returned to the pool
        let _one = pool.allocate().unwrap();
        let _two = pool.allocate().unwrap();
        assert!(pool.allocate().is_none());
    }
}