prometheus-client-mmap 0.20.3-x86_64-linux-musl → 0.22.0-x86_64-linux-musl

data/ext/fast_mmaped_file_rs/src/mmap.rs

@@ -1,9 +1,28 @@
+use magnus::exception::*;
+use magnus::prelude::*;
+use magnus::rb_sys::{AsRawValue, FromRawValue};
 use magnus::typed_data::Obj;
 use magnus::value::Fixnum;
-use magnus::{Integer, RArray, RClass, RHash, RString, Value};
+use magnus::{block::Proc, eval, scan_args, Error, Integer, RArray, RClass, RHash, RString, Value};
+use rb_sys::rb_str_new_static;
+use std::fs::File;
+use std::io::{prelude::*, SeekFrom};
+use std::mem;
+use std::path::Path;
+use std::ptr::NonNull;
+use std::sync::RwLock;
 
+use crate::err;
+use crate::error::MmapError;
 use crate::file_entry::FileEntry;
 use crate::map::EntryMap;
+use crate::raw_entry::RawEntry;
+use crate::util::{self, CheckedOps};
+use crate::Result;
+use crate::HEADER_SIZE;
+use inner::InnerMmap;
+
+mod inner;
 
 /// A Rust struct wrapped in a Ruby object, providing access to a memory-mapped
 /// file used to store, update, and read out Prometheus metrics.
@@ -35,9 +54,21 @@ use crate::map::EntryMap;
 /// |  K2 Value   |
 /// +-+-+-+-+-+-+-+
 /// ```
+//
+// The API imposed by `magnus` requires all methods to use shared borrows.
+// This means we can't store any mutable state in the top-level struct,
+// and must store the interior data behind a `RwLock`, which adds run-time
+// checks that mutable operations have no concurrent read or writes.
+//
+// We are further limited by the need to support subclassing in Ruby, which
+// requires us to define an allocation function for the class, the
+// `magnus::class::define_alloc_func()` function. This needs a support the
+// `Default` trait, so a `File` cannot directly help by the object being
+// constructed. Having the `RwLock` hold an `Option` of the interior object
+// resolves this.
 #[derive(Debug, Default)]
 #[magnus::wrap(class = "FastMmapedFileRs", free_immediately, size)]
-pub struct MmapedFile;
+pub struct MmapedFile(RwLock<Option<InnerMmap>>);
 
 impl MmapedFile {
     /// call-seq:
@@ -50,14 +81,43 @@ impl MmapedFile {
     ///
     ///     Creates a mapping that's shared with all other processes
     ///     mapping the same area of the file.
-    pub fn new(_klass: RClass, _args: &[Value]) -> magnus::error::Result<Obj<Self>> {
-        Ok(Obj::wrap(Self))
+    pub fn new(klass: RClass, args: &[Value]) -> magnus::error::Result<Obj<Self>> {
+        let args = scan_args::scan_args::<(RString,), (), (), (), (), ()>(args)?;
+        let path = args.required.0;
+
+        let lock = MmapedFile(RwLock::new(None));
+        let obj = Obj::wrap_as(lock, klass);
+
+        let _: Value = obj.funcall("initialize", (path,))?;
+
+        Ok(obj)
     }
 
     /// Initialize a new `FastMmapedFileRs` object. This must be defined in
     /// order for inheritance to work.
-    pub fn initialize(_rb_self: Obj<Self>, _fname: String) -> magnus::error::Result<()> {
-        unimplemented!();
+    pub fn initialize(rb_self: Obj<Self>, fname: String) -> magnus::error::Result<()> {
+        let file = File::options()
+            .read(true)
+            .write(true)
+            .open(&fname)
+            .map_err(|_| err!(arg_error(), "Can't open {}", fname))?;
+
+        let inner = InnerMmap::new(fname.into(), file)?;
+        rb_self.insert_inner(inner)?;
+
+        let weak_klass = RClass::from_value(eval("ObjectSpace::WeakMap")?)
+            .ok_or_else(|| err!(no_method_error(), "unable to create WeakMap"))?;
+        let weak_obj_tracker = weak_klass.new_instance(())?;
+
+        // We will need to iterate over strings backed by the mmapped file, but
+        // don't want to prevent the GC from reaping them when the Ruby code
+        // has finished with them. `ObjectSpace::WeakMap` allows us to track
+        // them without extending their lifetime.
+        //
+        // https://ruby-doc.org/core-3.0.0/ObjectSpace/WeakMap.html
+        rb_self.ivar_set("@weak_obj_tracker", weak_obj_tracker)?;
+
+        Ok(())
     }
 
     /// Read the list of files provided from Ruby and convert them to a Prometheus
@@ -89,8 +149,23 @@ impl MmapedFile {
     ///   self[start, length]
     ///
     /// return a substring of <em>lenght</em> characters from <em>start</em>
-    pub fn slice(_rb_self: Obj<Self>, _args: &[Value]) -> magnus::error::Result<RString> {
-        unimplemented!();
+    pub fn slice(rb_self: Obj<Self>, args: &[Value]) -> magnus::error::Result<RString> {
+        // The C implemenation would trigger a GC cycle via `rb_gc_force_recycle`
+        // if the `MM_PROTECT` flag is set, but in practice this is never used.
+        // We omit this logic, particularly because `rb_gc_force_recycle` is a
+        // no-op as of Ruby 3.1.
+
+        let str = (*rb_self).str(rb_self)?;
+        let res = str.funcall("[]", args);
+
+        // The C implementation does this, perhaps to validate that the len we
+        // provided is actually being used.
+        (*rb_self).inner_mut(|inner| {
+            inner.set_len(str.len());
+            Ok(())
+        })?;
+
+        res
     }
 
     /// Document-method: msync
@@ -100,8 +175,23 @@ impl MmapedFile {
     /// call-seq: msync
     ///
     /// flush the file
-    pub fn sync(&self, _args: &[Value]) -> magnus::error::Result<()> {
-        unimplemented!();
+    pub fn sync(&self, args: &[Value]) -> magnus::error::Result<()> {
+        use nix::sys::mman::MsFlags;
+
+        let mut ms_async = false;
+        let args = scan_args::scan_args::<(), (Option<i32>,), (), (), (), ()>(args)?;
+
+        if let Some(flag) = args.optional.0 {
+            let flag = MsFlags::from_bits(flag).unwrap_or(MsFlags::empty());
+            ms_async = flag.contains(MsFlags::MS_ASYNC);
+        }
+
+        // The `memmap2` crate does not support the `MS_INVALIDATE` flag. We ignore that
+        // flag if passed in, checking only for `MS_ASYNC`. In practice no arguments are ever
+        // passed to this function, but we do this to maintain compatibility with the
+        // C implementation.
+        self.inner_mut(|inner| inner.flush(ms_async))
+            .map_err(|e| e.into())
     }
 
     /// Document-method: munmap
@@ -110,42 +200,514 @@ impl MmapedFile {
     /// call-seq: munmap
     ///
     /// terminate the association
-    pub fn munmap(_rb_self: Obj<Self>) -> magnus::error::Result<()> {
-        unimplemented!();
+    pub fn munmap(rb_self: Obj<Self>) -> magnus::error::Result<()> {
+        let rs_self = &*rb_self;
+
+        rs_self.inner_mut(|inner| {
+            // truncate file to actual used size
+            inner.truncate_file()?;
+
+            // We are about to release the backing mmap for Ruby's String
+            // objects. If Ruby attempts to read from them the program will
+            // segfault. We update the length of all Strings to zero so Ruby
+            // does not attempt to access the now invalid address between now
+            // and when GC eventually reaps the objects.
+            //
+            // See the following for more detail:
+            // https://gitlab.com/gitlab-org/ruby/gems/prometheus-client-mmap/-/issues/39
+            // https://gitlab.com/gitlab-org/ruby/gems/prometheus-client-mmap/-/issues/41
+            // https://gitlab.com/gitlab-org/ruby/gems/prometheus-client-mmap/-/merge_requests/80
+            inner.set_len(0);
+            Ok(())
+        })?;
+
+        // Update each String object to be zero-length.
+        rs_self.update_weak_map(rb_self)?;
+
+        // Remove the `InnerMmap` from the `RwLock`. This will drop
+        // end of this function, unmapping and closing the file.
+        let _ = rs_self.take_inner()?;
+        Ok(())
     }
 
     /// Fetch the `used` header from the `.db` file, the length
     /// in bytes of the data written to the file.
     pub fn load_used(&self) -> magnus::error::Result<Integer> {
-        unimplemented!();
+        let used = self.inner(|inner| inner.load_used())?;
+
+        Ok(Integer::from_u64(used as u64))
     }
 
     /// Update the `used` header for the `.db` file, the length
     /// in bytes of the data written to the file.
-    pub fn save_used(_rb_self: Obj<Self>, _used: Fixnum) -> magnus::error::Result<Fixnum> {
-        unimplemented!();
+    pub fn save_used(rb_self: Obj<Self>, used: Fixnum) -> magnus::error::Result<Fixnum> {
+        let rs_self = &*rb_self;
+        let used_uint = used.to_u32()?;
+
+        // If the underlying mmap is smaller than the header, then resize to fit.
+        // The file has already been expanded to page size when first opened, so
+        // even if the map is less than HEADER_SIZE, we're not at risk of a
+        // SIGBUS.
+        if rs_self.capacity() < HEADER_SIZE {
+            rs_self.expand_to_fit(rb_self, HEADER_SIZE)?;
+        }
+
+        rs_self.inner_mut(|inner| inner.save_used(used_uint))?;
+
+        Ok(used)
     }
 
     /// Fetch the value associated with a key from the mmap.
     /// If no entry is present, initialize with the default
     /// value provided.
     pub fn fetch_entry(
-        _rb_self: Obj<Self>,
-        _positions: RHash,
-        _key: RString,
-        _default_value: f64,
+        rb_self: Obj<Self>,
+        positions: RHash,
+        key: RString,
+        default_value: f64,
     ) -> magnus::error::Result<f64> {
-        unimplemented!();
+        let rs_self = &*rb_self;
+        let position: Option<Fixnum> = positions.lookup(key)?;
+
+        if let Some(pos) = position {
+            let pos = pos.to_usize()?;
+            return rs_self
+                .inner(|inner| inner.load_value(pos))
+                .map_err(|e| e.into());
+        }
+
+        rs_self.check_expand(rb_self, key.len())?;
+
+        let value_offset: usize = rs_self.inner_mut(|inner| {
+            // SAFETY: We must not call any Ruby code for the lifetime of this borrow.
+            unsafe { inner.initialize_entry(key.as_slice(), default_value) }
+        })?;
+
+        // CAST: no-op on 64-bit, widening on 32-bit.
+        positions.aset(key, Integer::from_u64(value_offset as u64))?;
+
+        rs_self.load_value(value_offset)
     }
 
     /// Update the value of an existing entry, if present. Otherwise create a new entry
     /// for the key.
     pub fn upsert_entry(
-        _rb_self: Obj<Self>,
-        _positions: RHash,
-        _key: RString,
-        _value: f64,
+        rb_self: Obj<Self>,
+        positions: RHash,
+        key: RString,
+        value: f64,
     ) -> magnus::error::Result<f64> {
-        unimplemented!();
+        let rs_self = &*rb_self;
+        let position: Option<Fixnum> = positions.lookup(key)?;
+
+        if let Some(pos) = position {
+            let pos = pos.to_usize()?;
+            return rs_self
+                .inner_mut(|inner| {
+                    inner.save_value(pos, value)?;
+
+                    // TODO just return `value` here instead of loading it?
+                    // This is how the C implementation did it, but I don't
+                    // see what the extra load gains us.
+                    inner.load_value(pos)
+                })
+                .map_err(|e| e.into());
+        }
+
+        rs_self.check_expand(rb_self, key.len())?;
+
+        let value_offset: usize = rs_self.inner_mut(|inner| {
+            // SAFETY: We must not call any Ruby code for the lifetime of this borrow.
+            unsafe { inner.initialize_entry(key.as_slice(), value) }
+        })?;
+
+        // CAST: no-op on 64-bit, widening on 32-bit.
+        positions.aset(key, Integer::from_u64(value_offset as u64))?;
+
+        rs_self.load_value(value_offset)
+    }
+
+    /// Creates a Ruby String containing the section of the mmapped file that
+    /// has been written to.
+    fn str(&self, rb_self: Obj<Self>) -> magnus::error::Result<RString> {
+        let val_id = (*rb_self).inner(|inner| {
+            let ptr = inner.as_ptr();
+            let len = inner.len();
+
+            // SAFETY: This is safe so long as the data provided to Ruby meets its
+            // requirements. When unmapping the file this will no longer be the
+            // case, see the comment on `munmap` how we handle this.
+            Ok(unsafe { rb_str_new_static(ptr as _, len as _) })
+        })?;
+
+        let val = unsafe { Value::from_raw(val_id) };
+
+        // UNWRAP: We created this value as a string above.
+        let str = RString::from_value(val).unwrap();
+
+        let tracker: Value = rb_self.ivar_get("@weak_obj_tracker")?;
+
+        // Use the string's Id as the key in the `WeakMap`.
+        let key = str.as_raw();
+        let _: Value = tracker.funcall("[]=", (key, str))?;
+
+        Ok(str)
+    }
+
+    /// If we reallocate, any live Ruby strings provided by the `str()` method
+    /// will be will be invalidated. We need to iterate over them using and
+    /// update their heap pointers to the newly allocated memory region.
+    fn update_weak_map(&self, rb_self: Obj<Self>) -> magnus::error::Result<()> {
+        let tracker: Value = rb_self.ivar_get("@weak_obj_tracker")?;
+
+        let (new_ptr, new_len) = (*rb_self).inner(|inner| {
+            // Pointers are not `Send`, convert it to usize to allow capture in closure.
+            let new_ptr = inner.as_ptr() as usize;
+            let new_len = util::cast_chk::<_, i64>(inner.len(), "mmap len")?;
+
+            Ok((new_ptr, new_len))
+        })?;
+
+        // Allocate a block with a closure that updates the string details.
+        let block = Proc::from_fn(move |args, _block| {
+            let val = args
+                .get(0)
+                .ok_or_else(|| err!(arg_error(), "no argument received from `each_value`"))?;
+            let str = RString::from_value(*val)
+                .ok_or_else(|| err!(arg_error(), "weakmap value was not a string"))?;
+
+            // SAFETY: We're messing with Ruby's internals here, YOLO.
+            unsafe {
+                // Convert the magnus wrapper type to a raw string exposed by `rb_sys`,
+                // which provides access to its internals.
+                let mut raw_str = Self::rb_string_internal(str);
+
+                // Update the string to point to the new mmapped file.
+                // We're matching the behavior of Ruby's `str_new_static` function.
+                // See https://github.com/ruby/ruby/blob/e51014f9c05aa65cbf203442d37fef7c12390015/string.c#L1030-L1053
+                raw_str.as_mut().as_.heap.ptr = new_ptr as _;
+                raw_str.as_mut().as_.heap.len = new_len;
+                raw_str.as_mut().as_.heap.aux.capa = new_len;
+            }
+            Ok(())
+        });
+
+        // Execute the block.
+        let _: Value = tracker.funcall_with_block("each_value", (), block)?;
+        Ok(())
+    }
+
+    /// Check that the mmap is large enough to contain the value to be added,
+    /// and expand it to fit if necessary.
+    fn check_expand(&self, rb_self: Obj<Self>, key_len: usize) -> magnus::error::Result<()> {
+        // CAST: no-op on 32-bit, widening on 64-bit.
+        let used = self.inner(|inner| inner.load_used())? as usize;
+        let entry_len = RawEntry::calc_total_len(key_len)?;
+
+        // We need the mmapped region to contain at least one byte beyond the
+        // written data to create a NUL- terminated C string. Validate that
+        // new length does not exactly match the length of the mmap or exceed
+        // it.
+        while self.capacity() <= used.add_chk(entry_len)? {
+            self.expand_to_fit(rb_self, self.capacity().mul_chk(2)?)?;
+        }
+
+        Ok(())
+    }
+
+    /// Expand the underlying file until it is long enough to fit `target_cap`.
+    /// This will remove the existing mmap, expand the file, then update any
+    /// strings held by the `WeakMap` to point to the newly mmapped address.
+    fn expand_to_fit(&self, rb_self: Obj<Self>, target_cap: usize) -> magnus::error::Result<()> {
+        if target_cap < self.capacity() {
+            return Err(err!(arg_error(), "Can't reduce the size of mmap"));
+        }
+
+        let mut new_cap = self.capacity();
+        while new_cap < target_cap {
+            new_cap = new_cap.mul_chk(2)?;
+        }
+
+        if new_cap != self.capacity() {
+            // Drop the old mmap.
+            let (mut file, path) = self.take_inner()?.munmap();
+
+            self.expand_file(&mut file, &path, target_cap)?;
+
+            // Re-mmap the expanded file.
+            let new_inner = InnerMmap::reestablish(path, file, target_cap)?;
+
+            self.insert_inner(new_inner)?;
+
+            return self.update_weak_map(rb_self);
+        }
+
+        Ok(())
+    }
+
+    /// Use lseek(2) to seek past the end of the file and write a NUL byte. This
+    /// creates a file hole that expands the size of the file without consuming
+    /// disk space until it is actually written to.
+    fn expand_file(&self, file: &mut File, path: &Path, len: usize) -> Result<()> {
+        if len == 0 {
+            return Err(MmapError::overflowed(0, -1, "adding"));
+        }
+
+        // CAST: no-op on 64-bit, widening on 32-bit.
+        let len = len as u64;
+
+        match file.seek(SeekFrom::Start(len - 1)) {
+            Ok(_) => {}
+            Err(_) => {
+                return Err(MmapError::WithErrno(format!("Can't lseek {}", len - 1)));
+            }
+        }
+
+        match file.write(&[0x0]) {
+            Ok(1) => {}
+            _ => {
+                return Err(MmapError::WithErrno(format!(
+                    "Can't extend {}",
+                    path.display()
+                )));
+            }
+        }
+
+        Ok(())
+    }
+
+    /// The total capacity of the underlying mmap.
+    #[inline]
+    fn capacity(&self) -> usize {
+        // UNWRAP: This is actually infallible, but we need to
+        // wrap it in a `Result` for use with `inner()`.
+        self.inner(|inner| Ok(inner.capacity())).unwrap()
+    }
+
+    fn load_value(&self, position: usize) -> magnus::error::Result<f64> {
+        self.inner(|inner| inner.load_value(position))
+            .map_err(|e| e.into())
+    }
+
+    /// Takes a closure with immutable access to InnerMmap. Will fail if the inner
+    /// object has a mutable borrow or has been dropped.
+    fn inner<F, T>(&self, func: F) -> Result<T>
+    where
+        F: FnOnce(&InnerMmap) -> Result<T>,
+    {
+        let inner_opt = self.0.try_read().map_err(|_| MmapError::ConcurrentAccess)?;
+
+        let inner = inner_opt.as_ref().ok_or(MmapError::UnmappedFile)?;
+
+        func(inner)
+    }
+
+    /// Takes a closure with mutable access to InnerMmap. Will fail if the inner
+    /// object has an existing mutable borrow, or has been dropped.
+    fn inner_mut<F, T>(&self, func: F) -> Result<T>
+    where
+        F: FnOnce(&mut InnerMmap) -> Result<T>,
+    {
+        let mut inner_opt = self
+            .0
+            .try_write()
+            .map_err(|_| MmapError::ConcurrentAccess)?;
+
+        let inner = inner_opt.as_mut().ok_or(MmapError::UnmappedFile)?;
+
+        func(inner)
+    }
+
+    /// Take ownership of the `InnerMmap` from the `RwLock`.
+    /// Will fail if a mutable borrow is already held or the inner
+    /// object has been dropped.
+    fn take_inner(&self) -> Result<InnerMmap> {
+        let mut inner_opt = self
+            .0
+            .try_write()
+            .map_err(|_| MmapError::ConcurrentAccess)?;
+        match (*inner_opt).take() {
+            Some(i) => Ok(i),
+            None => Err(MmapError::UnmappedFile),
+        }
+    }
+
+    /// Move `new_inner` into the `RwLock`.
+    /// Will return an error if a mutable borrow is already held.
+    fn insert_inner(&self, new_inner: InnerMmap) -> Result<()> {
+        let mut inner_opt = self
+            .0
+            .try_write()
+            .map_err(|_| MmapError::ConcurrentAccess)?;
+        (*inner_opt).replace(new_inner);
+
+        Ok(())
+    }
+
+    /// Convert `magnus::RString` into the raw binding used by `rb_sys::RString`.
+    /// We need this to manually change the pointer and length values for strings
+    /// when moving the mmap to a new file.
+    unsafe fn rb_string_internal(rb_str: RString) -> NonNull<rb_sys::RString> {
+        mem::transmute::<RString, NonNull<rb_sys::RString>>(rb_str)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use magnus::error::Error;
+    use magnus::eval;
+    use magnus::Range;
+    use nix::unistd::{sysconf, SysconfVar};
+    use std::mem::size_of;
+
+    use super::*;
+    use crate::raw_entry::RawEntry;
+    use crate::testhelper::TestFile;
+
+    /// Create a wrapped MmapedFile object.
+    fn create_obj() -> Obj<MmapedFile> {
+        let TestFile {
+            file: _file,
+            path,
+            dir: _dir,
+        } = TestFile::new(&[0u8; 8]);
+
+        let path_str = path.display().to_string();
+        let rpath = RString::new(&path_str);
+
+        eval!("FastMmapedFileRs.new(path)", path = rpath).unwrap()
+    }
+
+    /// Add three entries to the mmap. Expected length is 56, 3x 16-byte
+    /// entries with 8-byte header.
+    fn populate_entries(rb_self: &Obj<MmapedFile>) -> RHash {
+        let positions = RHash::from_value(eval("{}").unwrap()).unwrap();
+
+        MmapedFile::upsert_entry(*rb_self, positions, RString::new("a"), 0.0).unwrap();
+        MmapedFile::upsert_entry(*rb_self, positions, RString::new("b"), 1.0).unwrap();
+        MmapedFile::upsert_entry(*rb_self, positions, RString::new("c"), 2.0).unwrap();
+
+        positions
+    }
+
+    #[test]
+    fn test_new() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let TestFile {
+            file,
+            path,
+            dir: _dir,
+        } = TestFile::new(&[0u8; 8]);
+
+        let path_str = path.display().to_string();
+        let rpath = RString::new(&path_str);
+
+        // Object created successfully
+        let result: std::result::Result<Obj<MmapedFile>, Error> =
+            eval!("FastMmapedFileRs.new(path)", path = rpath);
+        assert!(result.is_ok());
+
+        // Weak map added
+        let obj = result.unwrap();
+        let weak_tracker: Value = obj.ivar_get("@weak_obj_tracker").unwrap();
+        assert_eq!("ObjectSpace::WeakMap", weak_tracker.class().inspect());
+
+        // File expanded to page size
+        let page_size = sysconf(SysconfVar::PAGE_SIZE).unwrap().unwrap() as u64;
+        let stat = file.metadata().unwrap();
+        assert_eq!(page_size, stat.len());
+
+        // Used set to header size
+        assert_eq!(
+            HEADER_SIZE as u64,
+            obj.load_used().unwrap().to_u64().unwrap()
+        );
+    }
+
+    #[test]
+    fn test_slice() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let obj = create_obj();
+        let _ = populate_entries(&obj);
+
+        // Validate header updated with new length
+        let header_range = Range::new(0, HEADER_SIZE, true).unwrap().as_value();
+        let header_slice = MmapedFile::slice(obj, &[header_range]).unwrap();
+        assert_eq!([56, 0, 0, 0, 0, 0, 0, 0], unsafe {
+            header_slice.as_slice()
+        });
+
+        let value_range = Range::new(HEADER_SIZE, 24, true).unwrap().as_value();
+        let value_slice = MmapedFile::slice(obj, &[value_range]).unwrap();
+
+        // Validate string length
+        assert_eq!(1u32.to_ne_bytes(), unsafe { &value_slice.as_slice()[0..4] });
+
+        // Validate string and padding
+        assert_eq!("a   ", unsafe {
+            String::from_utf8_lossy(&value_slice.as_slice()[4..8])
+        });
+
+        // Validate value
+        assert_eq!(0.0f64.to_ne_bytes(), unsafe {
+            &value_slice.as_slice()[8..16]
+        });
+    }
+
+    #[test]
+    fn test_slice_resize() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let obj = create_obj();
+        let _ = populate_entries(&obj);
+
+        let rs_self = &*obj;
+
+        let value_range = Range::new(HEADER_SIZE, 24, false).unwrap().as_value();
+        let value_slice = MmapedFile::slice(obj, &[value_range]).unwrap();
+
+        rs_self.expand_to_fit(obj, rs_self.capacity() * 2).unwrap();
+
+        // If we haven't updated the pointer to the newly remapped file this will segfault.
+        let _: Value = eval!("puts str", str = value_slice).unwrap();
+    }
+
+    #[test]
+    fn test_dont_fill_mmap() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let obj = create_obj();
+        let positions = populate_entries(&obj);
+
+        let rs_self = &*obj;
+
+        rs_self.expand_to_fit(obj, 1024).unwrap();
+
+        let current_used = rs_self.inner(|inner| inner.load_used()).unwrap() as usize;
+        let current_cap = rs_self.inner(|inner| Ok(inner.len())).unwrap();
+
+        // Create a new entry that exactly fills the capacity of the mmap.
+        let val_len =
+            current_cap - current_used - HEADER_SIZE - size_of::<f64>() - size_of::<u32>();
+        assert_eq!(
+            current_cap,
+            RawEntry::calc_total_len(val_len).unwrap() + current_used
+        );
+
+        let str = String::from_utf8(vec![b'A'; val_len]).unwrap();
+        MmapedFile::upsert_entry(obj, positions, RString::new(&str), 1.0).unwrap();
+
+        // Validate that we have expanded the mmap, ensuring a trailing NUL.
+        assert!(rs_self.capacity() > current_cap);
     }
 }