From 257ff2cb95068d41dc431a21cbab256579511947 Mon Sep 17 00:00:00 2001
From: Tomer Filiba <tomer@sweet.security>
Date: Thu, 22 Aug 2024 15:03:03 +0300
Subject: [PATCH] Add config option to mlock the mmap'ed headers of the shard

---
 Cargo.lock   | 5 +++--
 Cargo.toml   | 1 +
 src/lib.rs   | 2 ++
 src/shard.rs | 4 ++++
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1472fa5..bc13fee 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -77,6 +77,7 @@ dependencies = [
  "bytemuck",
  "databuf",
  "fslock",
+ "libc",
  "memmap",
  "parking_lot",
  "rand",
@@ -157,9 +158,9 @@ dependencies = [
 
 [[package]]
 name = "libc"
-version = "0.2.155"
+version = "0.2.158"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c"
+checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439"
 
 [[package]]
 name = "lock_api"
diff --git a/Cargo.toml b/Cargo.toml
index c02e3d4..a39cd52 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,7 @@ parking_lot = "0.12.3"
 uuid = { version = "1.10.0" }
 rand = "0.8.5"
 fslock = "0.2.1"
+libc = "0.2.158"
 
 [features]
 whitebox_testing = []
diff --git a/src/lib.rs b/src/lib.rs
index 499ecbe..6149b30 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -108,6 +108,7 @@ pub struct Config {
     pub max_concurrent_list_ops: u32, // number of keyed locks for concurrent list ops
     pub truncate_up: bool, // whether or not to truncate up shard files to their max size (spare files)
     pub clear_on_unsupported_version: bool, // whether or not to clear the DB if the version is unsupported
+    pub mlock_headers: bool, // whether or not to mlock the shard headers to RAM (POSIX only)
 }
 
 impl Default for Config {
@@ -121,6 +122,7 @@ impl Default for Config {
             max_concurrent_list_ops: 64,
             truncate_up: true,
             clear_on_unsupported_version: false,
+            mlock_headers: false,
         }
     }
 }
diff --git a/src/shard.rs b/src/shard.rs
index 18f5ba4..365e7f6 100644
--- a/src/shard.rs
+++ b/src/shard.rs
@@ -272,6 +272,10 @@ impl Shard {
 
         let mut mmap = unsafe { MmapOptions::new().len(HEADER_SIZE as usize).map_mut(&file) }?;
 
+        if cfg!(target_family = "unix") {
+            unsafe { libc::mlock(mmap.as_ptr() as *const _, mmap.len()) };
+        }
+
         let header = unsafe { &mut *(mmap.as_mut_ptr() as *mut ShardHeader) };
         header.metadata.magic = SHARD_FILE_MAGIC;
         header.metadata.version = SHARD_FILE_VERSION;