RedisJSON · MeirShpilraien · Sep 10, 2024 · Sep 3, 2024 · Sep 3, 2024 · Sep 3, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -7,7 +7,7 @@ members = [
 ]
 
 [workspace.dependencies]
-ijson = { git="https://github.com/RedisJSON/ijson", rev="e0119ac74f6c4ee918718ee122c3948b74ebeba8", default_features=false}
+ijson = { git="https://github.com/RedisJSON/ijson", rev="eede48fad51b4ace5043d3e0714f5a65481a065d", default_features=false}
 serde_json = { version="1", features = ["unbounded_depth"]}
 serde = { version = "1", features = ["derive"] }
 serde_derive = "1"

diff --git a/redis_json/Cargo.toml b/redis_json/Cargo.toml
@@ -25,11 +25,12 @@ ijson.workspace = true
 serde_json.workspace = true
 serde.workspace = true
 libc = "0.2"
-redis-module ={ version = "^2.0.7", default-features = false, features = ["min-redis-compatibility-version-7-2"] }
-redis-module-macros = "^2.0.7"
+redis-module ={ git="https://github.com/RedisLabsModules/redismodule-rs", rev="32f284bd18ee799e54f8da835c81a137a34d7f83", default-features = false, features = ["min-redis-compatibility-version-7-2"] }
+redis-module-macros = { git="https://github.com/RedisLabsModules/redismodule-rs", rev="32f284bd18ee799e54f8da835c81a137a34d7f83" }
 itertools = "0.13"
 json_path = {path="../json_path"}
 linkme = "0.3"
+lazy_static = "1"
 
 [features]
 as-library = []

diff --git a/redis_json/src/commands.rs b/redis_json/src/commands.rs
@@ -4,6 +4,7 @@
  * the Server Side Public License v1 (SSPLv1).
  */
 
+use crate::defrag::defrag_info;
 use crate::error::Error;
 use crate::formatter::ReplyFormatOptions;
 use crate::key_value::KeyValue;
@@ -1817,6 +1818,7 @@ pub fn json_debug<M: Manager>(manager: M, ctx: &Context, args: Vec<RedisString>)
                 .into())
             }
         }
+        "DEFRAG_INFO" => defrag_info(ctx),
         "HELP" => {
             let results = vec![
                 "MEMORY <key> [path] - reports memory usage",

diff --git a/redis_json/src/defrag.rs b/redis_json/src/defrag.rs
@@ -0,0 +1,106 @@
+use std::{
+    alloc::Layout,
+    os::raw::{c_int, c_void},
+};
+
+use ijson::{Defrag, DefragAllocator};
+use lazy_static::lazy_static;
+use redis_module::{
+    defrag::DefragContext, raw, redisvalue::RedisValueKey, Context, RedisGILGuard, RedisResult,
+    RedisValue,
+};
+use redis_module_macros::{defrag_end_function, defrag_start_function};
+
+use crate::redisjson::RedisJSON;
+
+#[derive(Default)]
+pub(crate) struct DefragStats {
+    defrag_started: usize,
+    defrag_ended: usize,
+    keys_defrag: usize,
+}
+
+lazy_static! {
+    pub(crate) static ref DEFRAG_STATS: RedisGILGuard<DefragStats> = RedisGILGuard::default();
+}
+
+struct DefragCtxAllocator<'dc> {
+    defrag_ctx: &'dc DefragContext,
+}
+
+impl<'dc> DefragAllocator for DefragCtxAllocator<'dc> {
+    unsafe fn realloc_ptr<T>(&mut self, ptr: *mut T, _layout: Layout) -> *mut T {
+        self.defrag_ctx.defrag_realloc(ptr)
+    }
+
+    /// Allocate memory for defrag
+    unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
+        self.defrag_ctx.defrag_alloc(layout)
+    }
+
+    /// Free memory for defrag
+    unsafe fn free<T>(&mut self, ptr: *mut T, layout: Layout) {
+        self.defrag_ctx.defrag_dealloc(ptr, layout)
+    }
+}
+
+#[defrag_start_function]
+fn defrag_start(defrag_ctx: &DefragContext) {
+    let mut defrag_stats = DEFRAG_STATS.lock(defrag_ctx);
+    defrag_stats.defrag_started += 1;
+    ijson::reinit_shared_string_cache();
+}
+
+#[defrag_end_function]
+fn defrag_end(defrag_ctx: &DefragContext) {
+    let mut defrag_stats = DEFRAG_STATS.lock(defrag_ctx);
+    defrag_stats.defrag_ended += 1;
+}
+
+#[allow(non_snake_case, unused)]
+pub 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 defrag_stats = DEFRAG_STATS.lock(&defrag_ctx);
+    defrag_stats.keys_defrag += 1;
+
+    let mut defrag_allocator = DefragCtxAllocator {
+        defrag_ctx: &defrag_ctx,
+    };
+    let value = value.cast::<*mut RedisJSON<ijson::IValue>>();
+    let new_val = defrag_allocator.realloc_ptr(*value, Layout::new::<RedisJSON<ijson::IValue>>());
+    if !new_val.is_null() {
+        *value = new_val;
+    }
+    std::ptr::write(
+        &mut (**value).data as *mut ijson::IValue,
+        std::ptr::read(*value).data.defrag(&mut defrag_allocator),
+    );
+    0
+}
+
+pub(crate) fn defrag_info(ctx: &Context) -> RedisResult {
+    let defrag_stats = DEFRAG_STATS.lock(ctx);
+    Ok(RedisValue::OrderedMap(
+        [
+            (
+                RedisValueKey::String("defrag_started".to_owned()),
+                RedisValue::Integer(defrag_stats.defrag_started as i64),
+            ),
+            (
+                RedisValueKey::String("defrag_ended".to_owned()),
+                RedisValue::Integer(defrag_stats.defrag_ended as i64),
+            ),
+            (
+                RedisValueKey::String("keys_defrag".to_owned()),
+                RedisValue::Integer(defrag_stats.keys_defrag as i64),
+            ),
+        ]
+        .into_iter()
+        .collect(),
+    ))
+}
diff --git a/redis_json/src/lib.rs b/redis_json/src/lib.rs
@@ -36,6 +36,7 @@ mod array_index;
 mod backward;
 pub mod c_api;
 pub mod commands;
+pub mod defrag;
 pub mod error;
 mod formatter;
 pub mod ivalue_manager;
@@ -73,7 +74,7 @@ pub static REDIS_JSON_TYPE: RedisType = RedisType::new(
         free_effort: None,
         unlink: None,
         copy: Some(redisjson::type_methods::copy),
-        defrag: None,
+        defrag: Some(defrag::defrag),
 
         free_effort2: None,
         unlink2: None,

diff --git a/tests/pytest/test_defrag.py b/tests/pytest/test_defrag.py
@@ -0,0 +1,83 @@
+import time
+import json
+from RLTest import Defaults
+
+Defaults.decode_responses = True
+
+def enableDefrag(env):
+    # make defrag as aggressive as possible
+    env.cmd('CONFIG', 'SET', 'hz', '100')
+    env.cmd('CONFIG', 'SET', 'active-defrag-ignore-bytes', '1')
+    env.cmd('CONFIG', 'SET', 'active-defrag-threshold-lower', '0')
+    env.cmd('CONFIG', 'SET', 'active-defrag-cycle-min', '99')
+
+    try:
+        env.cmd('CONFIG', 'SET', 'activedefrag', 'yes')
+    except Exception:
+        # If active defrag is not supported by the current Redis, simply skip the test.
+        env.skip()
+
+def defragOnObj(env, obj):
+    enableDefrag(env)
+    env.expect('JSON.SET', 'test', '$', json.dumps(obj)).ok()
+    _, _, _, _, _, keysDefrag = env.cmd('JSON.DEBUG', 'DEFRAG_INFO')
+    startTime = time.time()
+    # Wait for at least 2 defrag full cycles
+    # We verify only the 'keysDefrag' value because the other values
+    # are not promised to be updated. It depends if Redis support
+    # the start/end defrag callbacks.
+    while keysDefrag < 2:
+        time.sleep(0.1)
+        _, _, _, _, _, keysDefrag = env.cmd('JSON.DEBUG', 'DEFRAG_INFO')
+        if time.time() - startTime > 30:
+            # We will wait for up to 30 seconds and then we consider it a failure
+            env.assertTrue(False, message='Failed waiting for defrag to run')
+            return
+    # make sure json is still valid.
+    res = json.loads(env.cmd('JSON.GET', 'test', '$'))[0]
+    env.assertEqual(res, obj)
+
+def testDefragNumber(env):
+    defragOnObj(env, 1)
+
+def testDefragBigNumber(env):
+    defragOnObj(env, 100000000000000000000)
+
+def testDefragDouble(env):
+    defragOnObj(env, 1.111111111111)
+
+def testDefragNegativeNumber(env):
+    defragOnObj(env, -100000000000000000000)
+
+def testDefragNegativeDouble(env):
+    defragOnObj(env, -1.111111111111)
+
+def testDefragTrue(env):
+    defragOnObj(env, True)
+
+def testDefragFalse(env):
+    defragOnObj(env, True)
+
+def testDefragNone(env):
+    defragOnObj(env, None)
+
+def testDefragEmptyString(env):
+    defragOnObj(env, "")
+
+def testDefragString(env):
+    defragOnObj(env, "foo")
+
+def testDefragEmptyArray(env):
+    defragOnObj(env, [])
+
+def testDefragArray(env):
+    defragOnObj(env, [1, 2, 3])
+
+def testDefragEmptyObject(env):
+    defragOnObj(env, {})
+
+def testDefragObject(env):
+    defragOnObj(env, {"foo": "bar"})
+
+def testDefragComplex(env):
+    defragOnObj(env, {"foo": ["foo", 1, None, True, False, {}, {"foo": [], "bar": 1}]})