Merge branch 'cleanup-redis-allocator-usage'

Author: gavrie

.circleci/config.yml

@@ -56,4 +56,4 @@ jobs:
 
       - run:
           name: Run all tests
-          command: cargo test --features experimental-api --all
+          command: cargo test --features test,experimental-api --all-targets

Cargo.toml

@@ -36,3 +36,7 @@ cc = "1.0"
 default = []
 experimental-api = []
 
+# Workaround to allow cfg(feature = "test") in dependencies:
+# https://github.com/rust-lang/rust/issues/59168#issuecomment-472653680
+# This requires running the tests with `--features test`
+test = []

build.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env sh
+cargo build --features experimental-api --all --all-targets
+

src/alloc.rs

@@ -1,34 +1,17 @@
-use std::alloc::{GlobalAlloc, Layout, System};
+use std::alloc::{GlobalAlloc, Layout};
 use std::os::raw::c_void;
-use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
 
 use crate::raw;
 
 pub struct RedisAlloc;
 
-static USE_REDIS_ALLOC: AtomicBool = AtomicBool::new(false);
-
 unsafe impl GlobalAlloc for RedisAlloc {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let use_redis = USE_REDIS_ALLOC.load(SeqCst);
-        if use_redis {
-            // complete the requested size to be aligned with the requested layout.align()
-            let size = (layout.size() + layout.align() - 1) & (!(layout.align() - 1));
-            return raw::RedisModule_Alloc.unwrap()(size) as *mut u8;
-        }
-        System.alloc(layout)
+        let size = (layout.size() + layout.align() - 1) & (!(layout.align() - 1));
+        return raw::RedisModule_Alloc.unwrap()(size) as *mut u8;
     }
 
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        let use_redis = USE_REDIS_ALLOC.load(SeqCst);
-        if use_redis {
-            return raw::RedisModule_Free.unwrap()(ptr as *mut c_void);
-        }
-        System.dealloc(ptr, layout);
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        return raw::RedisModule_Free.unwrap()(ptr as *mut c_void);
     }
 }
-
-pub fn use_redis_alloc() {
-    USE_REDIS_ALLOC.store(true, SeqCst);
-    eprintln!("Now using Redis allocator");
-}

src/context.rs

@@ -57,8 +57,8 @@ impl Context {
     }
 
     pub fn is_keys_position_request(&self) -> bool {
-        // HACK: Used for tests, where the context is null
-        if self.ctx.is_null() {
+        // We want this to be available in tests where we don't have an actual Redis to call
+        if cfg!(feature = "test") {
             return false;
         }
 

src/lib.rs

@@ -28,6 +28,10 @@ mod key;
 pub use crate::context::Context;
 pub use crate::redismodule::*;
 
+/// Ideally this would be `#[cfg(not(test))]`, but that doesn't work:
+/// https://github.com/rust-lang/rust/issues/59168#issuecomment-472653680
+/// The workaround is to use the `test` feature instead.
+#[cfg(not(feature = "test"))]
 #[global_allocator]
 static ALLOC: crate::alloc::RedisAlloc = crate::alloc::RedisAlloc;
 

src/macros.rs

@@ -73,7 +73,7 @@ macro_rules! redis_module {
               ]),* $(,)*
         ] $(,)*
     ) => {
-        use std::os::raw::c_int;
+        use std::os::raw::{c_int, c_char};
         use std::ffi::CString;
         use std::slice;
 
@@ -87,25 +87,27 @@ macro_rules! redis_module {
             _argv: *mut *mut raw::RedisModuleString,
             _argc: c_int,
         ) -> c_int {
-            {
-                // We use an explicit block here to make sure all memory allocated before we
-                // switch to the Redis allocator will be out of scope and thus deallocated.
-                let module_name = CString::new($module_name).unwrap();
-                let module_version = $module_version as c_int;
 
-                if unsafe { raw::Export_RedisModule_Init(
-                    ctx,
-                    module_name.as_ptr(),
-                    module_version,
-                    raw::REDISMODULE_APIVER_1 as c_int,
-                ) } == raw::Status::Err as c_int { return raw::Status::Err as c_int; }
+            // We use a statically sized buffer to avoid allocating.
+            // This is needed since we use a custom allocator that relies on the Redis allocator,
+            // which isn't yet ready at this point.
+            let mut name_buffer = [0; 64];
+            unsafe {
+                std::ptr::copy(
+                    $module_name.as_ptr(),
+                    name_buffer.as_mut_ptr(),
+                    $module_name.len(),
+                );
             }
 
-            if true {
-                redis_module::alloc::use_redis_alloc();
-            } else {
-                eprintln!("*** NOT USING Redis allocator ***");
-            }
+            let module_version = $module_version as c_int;
+
+            if unsafe { raw::Export_RedisModule_Init(
+                ctx,
+                name_buffer.as_ptr() as *const c_char,
+                module_version,
+                raw::REDISMODULE_APIVER_1 as c_int,
+            ) } == raw::Status::Err as c_int { return raw::Status::Err as c_int; }
 
             $(
                 if $init_func(ctx) == raw::Status::Err as c_int {

test.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env sh
+cargo test --all-targets --features test,experimental-api