Added support for defrag API. (#387)

* Added support for defrag API.

The PR adds support for defrag API.
The PR introduce a new struct, `DefragContext`, which provide
a safe API over `*mut raw::RedisModuleDefragCtx`.

**Notice** that we do expose an unsafe API to create `DefragContext` from `*mut raw::RedisModuleDefragCtx`.
This is because we do not have a safe API for datatype registeration. User must register an unsafe function as the defrag callback of the datatype and create the `DefragContext` from `*mut raw::RedisModuleDefragCtx`. We should consider adding a safe API for datatype creation.

In addition, the PR introduce 3 new proc macros:

* defrag_function - allows to register a function to defrag global data
* defrag_start_function - allows to register a function to defrag global data when defrag cycle starts.
* defrag_end_function - allows to register a function to defrag global data when defrag cycle ends.

Example and test for the new API were added.

* Review fixes

* Fix compilation issue

* Review fixes

Author: Meir Shpilraien (Spielrein)

examples/data_type.rs

@@ -1,12 +1,26 @@
+use lazy_static::lazy_static;
+use libc::c_int;
+use redis_module::defrag::DefragContext;
 use redis_module::native_types::RedisType;
-use redis_module::{raw, redis_module, Context, NextArg, RedisResult, RedisString};
+use redis_module::redisvalue::RedisValueKey;
+use redis_module::{
+    raw, redis_module, Context, NextArg, RedisGILGuard, RedisResult, RedisString, RedisValue,
+};
+use redis_module_macros::{defrag_end_function, defrag_function, defrag_start_function};
 use std::os::raw::c_void;
 
 #[derive(Debug)]
 struct MyType {
     data: String,
 }
 
+lazy_static! {
+    static ref NUM_KEYS_DEFRAG: RedisGILGuard<usize> = RedisGILGuard::default();
+    static ref NUM_DEFRAG_START: RedisGILGuard<usize> = RedisGILGuard::default();
+    static ref NUM_DEFRAG_END: RedisGILGuard<usize> = RedisGILGuard::default();
+    static ref NUM_DEFRAG_GLOBALS: RedisGILGuard<usize> = RedisGILGuard::default();
+}
+
 static MY_REDIS_TYPE: RedisType = RedisType::new(
     "mytype123",
     0,
@@ -30,7 +44,7 @@ static MY_REDIS_TYPE: RedisType = RedisType::new(
         free_effort: None,
         unlink: None,
         copy: None,
-        defrag: None,
+        defrag: Some(defrag),
 
         copy2: None,
         free_effort2: None,
@@ -43,6 +57,35 @@ unsafe extern "C" fn free(value: *mut c_void) {
     drop(Box::from_raw(value.cast::<MyType>()));
 }
 
+unsafe extern "C" fn defrag(
+    ctx: *mut raw::RedisModuleDefragCtx,
+    _key: *mut raw::RedisModuleString,
+    _value: *mut *mut c_void,
+) -> c_int {
+    let defrag_ctx = DefragContext::new(ctx);
+    let mut num_keys_defrag = NUM_KEYS_DEFRAG.lock(&defrag_ctx);
+    *num_keys_defrag += 1;
+    0
+}
+
+#[defrag_start_function]
+fn defrag_end(defrag_ctx: &DefragContext) {
+    let mut num_defrag_end = NUM_DEFRAG_END.lock(defrag_ctx);
+    *num_defrag_end += 1;
+}
+
+#[defrag_end_function]
+fn defrag_start(defrag_ctx: &DefragContext) {
+    let mut num_defrag_start = NUM_DEFRAG_START.lock(defrag_ctx);
+    *num_defrag_start += 1;
+}
+
+#[defrag_function]
+fn defrag_globals(defrag_ctx: &DefragContext) {
+    let mut num_defrag_globals = NUM_DEFRAG_GLOBALS.lock(defrag_ctx);
+    *num_defrag_globals += 1;
+}
+
 fn alloc_set(ctx: &Context, args: Vec<RedisString>) -> RedisResult {
     let mut args = args.into_iter().skip(1);
     let key = args.next_arg()?;
@@ -78,6 +121,35 @@ fn alloc_get(ctx: &Context, args: Vec<RedisString>) -> RedisResult {
     Ok(value)
 }
 
+fn alloc_defragstats(ctx: &Context, _args: Vec<RedisString>) -> RedisResult {
+    let num_keys_defrag = NUM_KEYS_DEFRAG.lock(ctx);
+    let num_defrag_globals = NUM_DEFRAG_GLOBALS.lock(ctx);
+    let num_defrag_start = NUM_DEFRAG_START.lock(ctx);
+    let num_defrag_end = NUM_DEFRAG_END.lock(ctx);
+    Ok(RedisValue::OrderedMap(
+        [
+            (
+                RedisValueKey::String("num_keys_defrag".to_owned()),
+                RedisValue::Integer(*num_keys_defrag as i64),
+            ),
+            (
+                RedisValueKey::String("num_defrag_globals".to_owned()),
+                RedisValue::Integer(*num_defrag_globals as i64),
+            ),
+            (
+                RedisValueKey::String("num_defrag_start".to_owned()),
+                RedisValue::Integer(*num_defrag_start as i64),
+            ),
+            (
+                RedisValueKey::String("num_defrag_end".to_owned()),
+                RedisValue::Integer(*num_defrag_end as i64),
+            ),
+        ]
+        .into_iter()
+        .collect(),
+    ))
+}
+
 //////////////////////////////////////////////////////
 
 redis_module! {
@@ -90,5 +162,6 @@ redis_module! {
     commands: [
         ["alloc.set", alloc_set, "write", 1, 1, 1],
         ["alloc.get", alloc_get, "readonly", 1, 1, 1],
+        ["alloc.defragstats", alloc_defragstats, "readonly", 0, 0, 0]
     ],
 }

examples/open_key_with_flags.rs

@@ -1,6 +1,5 @@
 use redis_module::{
-    key::KeyFlags, raw, redis_module, Context, NextArg, RedisError, RedisResult, RedisString,
-    RedisValue,
+    key::KeyFlags, redis_module, Context, NextArg, RedisError, RedisResult, RedisString, RedisValue,
 };
 use redis_module_macros::command;
 

redismodule-rs-macros/src/lib.rs

@@ -451,3 +451,72 @@ pub fn info_command_handler(_attr: TokenStream, item: TokenStream) -> TokenStrea
 pub fn info_section(item: TokenStream) -> TokenStream {
     info_section::info_section(item)
 }
+
+/// Proc macro which is set on a function that need to be called whenever server performs defrag.
+/// The function must accept a [&DefragContext]. If defrag is not supported by the Redis version
+/// the function will never be called.
+///
+/// Example:
+///
+/// ```rust,no_run,ignore
+/// #[defrag_function]
+/// fn defrag(ctx: &DefragContext) { ... }
+/// ```
+#[proc_macro_attribute]
+pub fn defrag_function(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let ast: ItemFn = match syn::parse(item) {
+        Ok(res) => res,
+        Err(e) => return e.to_compile_error().into(),
+    };
+    let gen = quote! {
+        #[linkme::distributed_slice(redis_module::defrag::DEFRAG_FUNCTIONS_LIST)]
+        #ast
+    };
+    gen.into()
+}
+
+/// Proc macro which is set on a function that need to be called whenever server start performs defrag.
+/// The function must accept a [&DefragContext]. If defrag start event is not supported by the Redis version
+/// the function will never be called.
+///
+/// Example:
+///
+/// ```rust,no_run,ignore
+/// #[defrag_start_function]
+/// fn defrag_start(ctx: &DefragContext) { ... }
+/// ```
+#[proc_macro_attribute]
+pub fn defrag_start_function(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let ast: ItemFn = match syn::parse(item) {
+        Ok(res) => res,
+        Err(e) => return e.to_compile_error().into(),
+    };
+    let gen = quote! {
+        #[linkme::distributed_slice(redis_module::defrag::DEFRAG_START_FUNCTIONS_LIST)]
+        #ast
+    };
+    gen.into()
+}
+
+/// Proc macro which is set on a function that need to be called whenever server end performs defrag.
+/// The function must accept a [&DefragContext]. If defrag end event is not supported by the Redis version
+/// the function will never be called.
+///
+/// Example:
+///
+/// ```rust,no_run,ignore
+/// #[defrag_end_function]
+/// fn defrag_end(ctx: &DefragContext) { ... }
+/// ```
+#[proc_macro_attribute]
+pub fn defrag_end_function(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let ast: ItemFn = match syn::parse(item) {
+        Ok(res) => res,
+        Err(e) => return e.to_compile_error().into(),
+    };
+    let gen = quote! {
+        #[linkme::distributed_slice(redis_module::defrag::DEFRAG_END_FUNCTIONS_LIST)]
+        #ast
+    };
+    gen.into()
+}

src/context/commands.rs

@@ -454,12 +454,12 @@ api! {[
             // the only CString pointers which are not freed are those of the key_specs, lets free them here.
             key_specs.into_iter().for_each(|v|{
                 if !v.notes.is_null() {
-                    unsafe{CString::from_raw(v.notes as *mut c_char)};
+                    drop(unsafe{CString::from_raw(v.notes as *mut c_char)});
                 }
                 if v.begin_search_type == raw::RedisModuleKeySpecBeginSearchType_REDISMODULE_KSPEC_BS_KEYWORD {
                     let keyword = unsafe{v.bs.keyword.keyword};
                     if !keyword.is_null() {
-                        unsafe{CString::from_raw(v.bs.keyword.keyword as *mut c_char)};
+                        drop(unsafe{CString::from_raw(v.bs.keyword.keyword as *mut c_char)});
                     }
                 }
             });