Add support for custom allocator

rust-lang · Nov 22, 2024 · 21f1511 · 21f1511
1 parent a2e05da
commit 21f1511
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 0 deletions.
diff --git a/libgit2-sys/lib.rs b/libgit2-sys/lib.rs
@@ -925,6 +925,13 @@ pub struct git_writestream {
     pub free: Option<extern "C" fn(*mut git_writestream)>,
 }
 
+#[repr(C)]
+pub struct git_allocator {
+    pub gmalloc: *const extern "C" fn(size_t, *const c_char, c_int) -> *mut c_void,
+    pub grealloc: *const extern "C" fn(*mut c_void, size_t, *const c_char, c_int) -> *mut c_void,
+    pub gfree: *const extern "C" fn(*mut c_void),
+}
+
 git_enum! {
     pub enum git_attr_value_t {
         GIT_ATTR_VALUE_UNSPECIFIED = 0,

diff --git a/src/opts.rs b/src/opts.rs
@@ -33,6 +33,29 @@ where
     Ok(())
 }
 
+/// Set the memory allocator to a different memory allocator.
+/// 
+/// This allocator will then be used to make all memory allocations for
+/// libgit2 operations. If the given `allocator` is None, then the
+/// system default will be restored.
+pub unsafe fn set_allocator(
+    gmalloc: *const extern "C" fn(libc::size_t, *const core::ffi::c_char, core::ffi::c_int) -> *mut core::ffi::c_void,
+    grealloc: *const extern "C" fn(*mut core::ffi::c_void, libc::size_t, *const core::ffi::c_char, core::ffi::c_int) -> *mut core::ffi::c_void,
+    gfree: *const extern "C" fn(*mut core::ffi::c_void),
+) -> Result<(), Error> {
+    crate::init();
+    let allocator = raw::git_allocator{
+        gmalloc: gmalloc,
+        gfree: gfree,
+        grealloc: grealloc,
+    };
+    try_call!(raw::git_libgit2_opts(
+        raw::GIT_OPT_SET_ALLOCATOR as libc::c_int,
+        &allocator as *const raw::git_allocator
+    ));
+    Ok(())
+}
+
 /// Reset the search path for a given level of config data to the default
 /// (generally based on environment variables).
 ///
@@ -416,6 +439,8 @@ pub unsafe fn set_server_timeout_in_milliseconds(timeout: libc::c_int) -> Result
 
 #[cfg(test)]
 mod test {
+    use crate::test::repo_init;
+
     use super::*;
 
     #[test]
@@ -447,6 +472,40 @@ mod test {
         }
     }
 
+    static mut ALLOC_CALLED: bool = false;
+    static mut FREE_CALLED: bool = false;
+
+    #[test]
+    fn custom_allocator() {
+        unsafe {
+            extern "C" fn gmalloc(size: libc::size_t, _: *const core::ffi::c_char, _: core::ffi::c_int) -> *mut core::ffi::c_void {
+                unsafe {
+                    ALLOC_CALLED = true;
+                    libc::malloc(size)
+                }
+            }
+            extern "C" fn grealloc(ptr: *mut core::ffi::c_void, size: libc::size_t, _: *const core::ffi::c_char, _: core::ffi::c_int) -> *mut core::ffi::c_void {
+                unsafe {
+                    ALLOC_CALLED = true;
+                    libc::realloc(ptr, size)
+                }
+            }
+            extern "C" fn gfree(ptr: *mut core::ffi::c_void) {
+                unsafe {
+                    FREE_CALLED = true;
+                    libc::free(ptr)
+                }
+            }
+            assert!(set_allocator(
+                gmalloc as *const extern "C" fn(libc::size_t, *const core::ffi::c_char, core::ffi::c_int) -> *mut core::ffi::c_void,
+                grealloc as *const extern "C" fn(*mut core::ffi::c_void, libc::size_t, *const core::ffi::c_char, core::ffi::c_int) -> *mut core::ffi::c_void,
+                gfree as *const extern "C" fn(*mut core::ffi::c_void),
+            ).is_ok());
+            repo_init();
+            assert!(ALLOC_CALLED); 
+            assert!(FREE_CALLED);}
+    }
+
     #[test]
     fn server_connect_timeout() {
         unsafe {