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

use super::Backend;

use crate::error;

use std::cmp;
use std::io;

#[derive(Debug)]
struct Mmap {
    inner: memmap::MmapMut,
    /// End of data
    pub end: usize,
    /// Mmap total len
    pub len: usize,
}

impl Mmap {
    fn new(len: usize) -> io::Result<Self> {
        let inner = memmap::MmapOptions::new().len(len).map_anon()?;

        Ok(Self { inner, end: 0, len })
    }

    fn as_slice(&self) -> &[u8] {
        &self.inner[..self.end]
    }

    fn as_mut_slice(&mut self) -> &mut [u8] {
        &mut self.inner[..self.end]
    }

    /// Copies data to mmap and modifies data's end cursor.
    fn write(&mut self, data: &[u8]) -> error::BackendResult<()> {
        if data.len() > self.len {
            return Err(error::BackendError::Internal(format!(
                "Unexpected write beyond mmap's backend capacity."
            )));
        }
        self.end = data.len();
        self.as_mut_slice().copy_from_slice(data);
        Ok(())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }

    /// Increases mmap size by `max(old_size*2, new_size)`.
    ///
    /// Note that it doesn't copy original data
    fn resize_no_copy(&mut self, new_size: usize) -> io::Result<()> {
        let len = cmp::max(self.len + self.len, new_size);
        // Make sure we don't discard old mmap before creating new one;
        let new_mmap = Self::new(len)?;
        *self = new_mmap;
        Ok(())
    }
}

/// A backend that uses an nonymous mmap.
///
/// The `Backend` automatically creates bigger map
/// on demand using following strategy:
///
/// - If new data size allows, multiply size by 2.
/// - Otherwise new data size is used.
///
/// Note that mmap is never shrink back.
///
/// Use `Backend` methods to read and write into it.
#[derive(Debug)]
pub struct MmapStorage {
    mmap: Mmap,
}

impl MmapStorage {
    /// Creates new storage with 1024 bytes.
    pub fn new() -> error::BackendResult<Self> {
        Self::with_size(1024)
    }

    /// Creates new storage with custom size.
    pub fn with_size(len: usize) -> error::BackendResult<Self> {
        let mmap = Mmap::new(len)?;

        Ok(Self { mmap })
    }
}

impl Backend for MmapStorage {
    fn get_data(&mut self) -> error::BackendResult<Vec<u8>> {
        let mmap = self.mmap.as_slice();
        let mut buffer = Vec::with_capacity(mmap.len());
        buffer.extend_from_slice(mmap);
        Ok(buffer)
    }

    fn put_data(&mut self, data: &[u8]) -> error::BackendResult<()> {
        if self.mmap.len < data.len() {
            self.mmap.resize_no_copy(data.len())?;
        }
        self.mmap.write(data)?;
        self.mmap.flush()?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{Backend, MmapStorage};

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_mmap_storage() {
        let data = [4, 5, 1, 6, 8, 1];
        let mut storage = MmapStorage::new().expect("To crate mmap storage");

        storage.put_data(&data).expect("To put data");
        assert_eq!(storage.mmap.end, data.len());
        assert_eq!(storage.get_data().expect("To get data"), data);
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_mmap_storage_extend() {
        let data = [4, 5, 1, 6, 8, 1];
        let mut storage = MmapStorage::with_size(4).expect("To crate mmap storage");

        storage.put_data(&data).expect("To put data");
        assert_eq!(storage.mmap.end, data.len());
        assert_eq!(storage.mmap.len, 8);
        assert_eq!(storage.get_data().expect("To get data"), data);
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_mmap_storage_increase_by_new_data_size() {
        let data = [4, 5, 1, 6, 8, 1];
        let mut storage = MmapStorage::with_size(1).expect("To crate mmap storage");

        storage.put_data(&data).expect("To put data");
        assert_eq!(storage.mmap.end, data.len());
        assert_eq!(storage.mmap.len, data.len());
        assert_eq!(storage.get_data().expect("To get data"), data);
    }
}