changeset 51208:8c4e8d06432e

rust-mixed-index: move the mmap keepalive into a function The same code will be used for keeping the new index mmap around.
author Raphaël Gomès <rgomes@octobus.net>
date Thu, 29 Jun 2023 16:09:57 +0200
parents 8ade5e6cdb61
children 1ef4a36a934d
files rust/hg-cpython/src/revlog.rs
diffstat 1 files changed, 41 insertions(+), 20 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-cpython/src/revlog.rs	Thu Jun 29 15:00:46 2023 +0200
+++ b/rust/hg-cpython/src/revlog.rs	Thu Jun 29 16:09:57 2023 +0200
@@ -314,6 +314,44 @@
 
 });
 
+/// Take a (potentially) mmap'ed buffer, and return the underlying Python
+/// buffer along with the Rust slice into said buffer. We need to keep the
+/// Python buffer around, otherwise we'd get a dangling pointer once the buffer
+/// is freed from Python's side.
+///
+/// # Safety
+///
+/// The caller must make sure that the buffer is kept around for at least as
+/// long as the slice.
+#[deny(unsafe_op_in_unsafe_fn)]
+unsafe fn mmap_keeparound(
+    py: Python,
+    data: PyObject,
+) -> PyResult<(
+    PyBuffer,
+    Box<dyn std::ops::Deref<Target = [u8]> + Send + 'static>,
+)> {
+    let buf = PyBuffer::get(py, &data)?;
+    let len = buf.item_count();
+
+    // Build a slice from the mmap'ed buffer data
+    let cbuf = buf.buf_ptr();
+    let bytes = if std::mem::size_of::<u8>() == buf.item_size()
+        && buf.is_c_contiguous()
+        && u8::is_compatible_format(buf.format())
+    {
+        unsafe { std::slice::from_raw_parts(cbuf as *const u8, len) }
+    } else {
+        return Err(PyErr::new::<ValueError, _>(
+            py,
+            "Nodemap data buffer has an invalid memory representation"
+                .to_string(),
+        ));
+    };
+
+    Ok((buf, Box::new(bytes)))
+}
+
 impl MixedIndex {
     fn new(py: Python, cindex: PyObject) -> PyResult<MixedIndex> {
         Self::create_instance(
@@ -427,29 +465,12 @@
         docket: PyObject,
         nm_data: PyObject,
     ) -> PyResult<PyObject> {
-        let buf = PyBuffer::get(py, &nm_data)?;
+        // Safety: we keep the buffer around inside the class as `nodemap_mmap`
+        let (buf, bytes) = unsafe { mmap_keeparound(py, nm_data)? };
         let len = buf.item_count();
-
-        // Build a slice from the mmap'ed buffer data
-        let cbuf = buf.buf_ptr();
-        let bytes = if std::mem::size_of::<u8>() == buf.item_size()
-            && buf.is_c_contiguous()
-            && u8::is_compatible_format(buf.format())
-        {
-            unsafe { std::slice::from_raw_parts(cbuf as *const u8, len) }
-        } else {
-            return Err(PyErr::new::<ValueError, _>(
-                py,
-                "Nodemap data buffer has an invalid memory representation"
-                    .to_string(),
-            ));
-        };
-
-        // Keep a reference to the mmap'ed buffer, otherwise we get a dangling
-        // pointer.
         self.nodemap_mmap(py).borrow_mut().replace(buf);
 
-        let mut nt = NodeTree::load_bytes(Box::new(bytes), len);
+        let mut nt = NodeTree::load_bytes(bytes, len);
 
         let data_tip = docket
             .getattr(py, "tip_rev")?