[pbs-devel] [PATCH proxmox-backup 1/3] tools: add mmap_buffer module

[Stefan Reiter <s.reiter@proxmox.com>]

Can allocate a large buffer of fixed-size chunks via a memory mapped
file, to allow the kernel to easily swap these pages back to disk on
memory pressure.

Designed to be used with an LRU cache or similar.

Signed-off-by: Stefan Reiter <s.reiter@proxmox.com>
---
 src/tools.rs             |   1 +
 src/tools/mmap_buffer.rs | 113 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)
 create mode 100644 src/tools/mmap_buffer.rs

diff --git a/src/tools.rs b/src/tools.rs
index 08af55e5..8f22e6c8 100644
--- a/src/tools.rs
+++ b/src/tools.rs
@@ -34,6 +34,7 @@ pub mod json;
 pub mod logrotate;
 pub mod loopdev;
 pub mod lru_cache;
+pub mod mmap_buffer;
 pub mod nom;
 pub mod runtime;
 pub mod serde_filter;
diff --git a/src/tools/mmap_buffer.rs b/src/tools/mmap_buffer.rs
new file mode 100644
index 00000000..764303c2
--- /dev/null
+++ b/src/tools/mmap_buffer.rs
@@ -0,0 +1,113 @@
+//! Provides a buffer allocation API using a memory mapped file.
+use anyhow::{bail, Error};
+use nix::sys::mman;
+use nix::unistd::{ftruncate, unlink};
+use std::sync::{Arc, Mutex};
+
+use proxmox::tools::fs::{make_tmp_file, CreateOptions};
+use proxmox::tools::mmap::Mmap;
+
+pub struct MmapBufferEntry {
+    pub buf: &'static mut [u8],
+    entry_idx: usize,
+    parent: Arc<MmapBuffer>,
+}
+
+pub struct MmapBuffer {
+    entry_size: usize,
+    entry_count: usize,
+    bitmap: Mutex<Vec<u64>>,
+    mmap: Mmap<u8>,
+}
+
+impl MmapBuffer {
+    /// Create a new mmap buffer containing _count [u8] slices of length _size. This requires root
+    /// to allocate the temp file securely.
+    pub fn new(entry_size: usize, entry_count: usize) -> Result<Arc<Self>, Error> {
+        let bytes = entry_size * entry_count;
+        let (fd, path) = make_tmp_file(
+            "/var/cache/proxmox-backup/mmap-buffer",
+            CreateOptions::default(),
+        )?;
+        unlink(&path)?;
+        ftruncate(fd.0, bytes as i64)?;
+        let mmap = unsafe {
+            Mmap::map_fd(
+                fd.0,
+                0,
+                bytes,
+                mman::ProtFlags::PROT_READ | mman::ProtFlags::PROT_WRITE,
+                mman::MapFlags::MAP_SHARED,
+            )?
+        };
+        Ok(Arc::new(Self {
+            entry_size,
+            entry_count,
+            bitmap: Mutex::new(vec![0; (entry_count+63)/64]),
+            mmap,
+        }))
+    }
+
+    #[inline]
+    fn bitmap_idx(lin: usize) -> (usize, u64) {
+        let idx = lin / 64;
+        let mask = 1 << (lin & 0x3F);
+        (idx, mask)
+    }
+
+    /// Returns one of the slices available in the buffer. The slice will be returned to this
+    /// instance on drop. If entry_count slices are already handed out, this function will fail.
+    pub fn allocate(self: &Arc<Self>) -> Result<MmapBufferEntry, Error> {
+        let bitmap = &mut (*self.bitmap.lock().unwrap());
+        for i in 0..self.entry_count as usize {
+            let (idx, mask) = Self::bitmap_idx(i);
+            if (bitmap[idx] & mask) == 0 {
+                bitmap[idx] |= mask;
+                let offset = self.entry_size * i;
+                let end = offset + self.entry_size;
+                // Safety:
+                // * mut: self.bitmap prevents multiple borrows on one region
+                // * 'static: the MmapBufferEntry contains a clone of the Arc pointing to ourselves
+                let data: &'static mut [u8] = unsafe {
+                    std::slice::from_raw_parts_mut(
+                        self.mmap[offset..end].as_ptr() as *mut u8,
+                        end - offset,
+                    )
+                };
+                let me = Arc::clone(&self);
+                return Ok(MmapBufferEntry {
+                    buf: data,
+                    entry_idx: i,
+                    parent: me,
+                });
+            }
+        }
+
+        bail!("mmap-buffer: cannot allocate more entries in buffer");
+    }
+
+    fn deallocate(self: &Arc<Self>, offset: usize) {
+        let bitmap = &mut (*self.bitmap.lock().unwrap());
+        let (idx, mask) = Self::bitmap_idx(offset);
+        bitmap[idx] &= !mask;
+    }
+}
+
+impl Drop for MmapBufferEntry {
+    fn drop(&mut self) {
+        self.parent.deallocate(self.entry_idx);
+    }
+}
+
+impl std::ops::Deref for MmapBufferEntry {
+    type Target = [u8];
+    fn deref(&self) -> &Self::Target {
+        self.buf
+    }
+}
+
+impl std::ops::DerefMut for MmapBufferEntry {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.buf
+    }
+}
-- 
2.20.1