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
158
159

//!Anonymous memory mapped storage.
use core::cmp;
use std::io;

/// Anonymous memory map backed storage
///
/// It provides byte storage over anonymous memory map.
///
/// When new data is written the storage may be extended
/// using following strategy:
///
/// - If new data is small enough, resize to double size.
/// - Otherwise resize to new data_size.
///
/// If needed you can use `resize` method to shrink or extend storage.
///
/// Unless otherwise stated all methods that modify content of storage doesn't perform
/// flushes, it is callee's responsibility to call it when required
pub struct Storage {
    inner: memmap::MmapMut,
    //Total size of storage
    size: usize,
    //Cursor pointing to end of storage
    end: usize,
}

impl Storage {
    /// Creates storage with specified size.
    pub fn new(size: usize) -> io::Result<Self> {
        let inner = memmap::MmapOptions::new().len(size).map_anon()?;

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

    #[inline]
    /// Returns slice with view on written data.
    pub fn as_slice(&self) -> &[u8] {
        &self.inner[..self.end]
    }

    #[inline]
    /// Returns slice with view on written data.
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        &mut self.inner[..self.end]
    }

    #[inline]
    /// Returns overall size of storage(including extra space).
    pub fn capacity(&self) -> usize {
        self.size
    }

    #[inline]
    /// Asynchronously flushes outstanding memory map modifications to disk.
    ///
    /// This method initiates flushing modified pages to durable storage, but it will not wait for
    /// the operation to complete before returning.
    pub fn flush_async(&self) -> io::Result<()> {
        self.inner.flush_async()
    }

    #[inline]
    /// Synchronously flushes outstanding memory map modifications to disk.
    pub fn flush_sync(&self) -> io::Result<()> {
        self.inner.flush()
    }

    //Increases mmap size by max(old_size*2, new_size)
    fn realloc(&mut self, new_size: usize) -> io::Result<()> {
        let new_size = cmp::max(self.size + self.size, new_size);
        let mut new_mmap = Self::new(new_size)?;
        new_mmap.copy_from_slice(self.as_slice());
        *self = new_mmap;
        Ok(())
    }

    /// Appends data to storage
    ///
    /// Note that it resizes storage if needed.
    /// Therefore error can happen.
    pub fn extend_from_slice(&mut self, data: &[u8]) -> io::Result<()> {
        if (self.size - self.end) < data.len() {
            let new_size = self.end + data.len();
            self.realloc(new_size)?;
        }
        self.inner[self.end..self.end+data.len()].copy_from_slice(data);
        self.end += data.len();
        Ok(())
    }

    /// Copies data from slice into storage.
    ///
    /// Function panics if data has greater len.
    ///
    /// **Note:** It copies data up to storage capacity.
    pub fn copy_from_slice(&mut self, data: &[u8]) {
        assert!(self.size >= data.len());
        self.end = data.len();
        self.as_mut_slice().copy_from_slice(data);
    }

    /// Resizes storage by appending or truncating.
    ///
    /// As it requires re-create mmap, the error is possible to happen.
    /// In case of error, original storage is preserved.
    ///
    /// After successfully creating new storage, previous storage content gets copied.
    pub fn resize(&mut self, new_size: usize) -> io::Result<()> {
        let mut new_mmap = Self::new(new_size)?;
        new_mmap.copy_from_slice(self.as_slice());
        *self = new_mmap;
        Ok(())
    }

    /// Clears content of storage.
    ///
    /// Note that it doesn't affect capacity in any way.
    pub fn clear(&mut self) {
        self.end = 0;
    }

    /// Creates `Vec` from content of storage.
    pub fn to_vec(&self) -> Vec<u8> {
        let mut result = Vec::with_capacity(self.end);
        result.extend_from_slice(self.as_slice());
        result
    }
}

impl Default for Storage {
    ///Creates default storage with capacity 1024
    ///
    ///Note that it panics on error
    fn default() -> Self {
        Self::new(1024).expect("To create default mmap storage")
    }
}

impl io::Write for Storage {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.extend_from_slice(buf)?;
        Ok(buf.len())
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.extend_from_slice(buf)
    }

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