diff --git a/Cargo.lock b/Cargo.lock index 71bfa8ea8..c10bb0623 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -436,7 +436,7 @@ checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" [[package]] name = "ijson" version = "0.1.3" -source = "git+https://github.com/RedisJSON/ijson?rev=e0119ac74f6c4ee918718ee122c3948b74ebeba8#e0119ac74f6c4ee918718ee122c3948b74ebeba8" +source = "git+https://github.com/RedisJSON/ijson?rev=eede48fad51b4ace5043d3e0714f5a65481a065d#eede48fad51b4ace5043d3e0714f5a65481a065d" dependencies = [ "dashmap", "hashbrown 0.13.2", @@ -820,9 +820,8 @@ dependencies = [ [[package]] name = "redis-module" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d3b95d9fe5b8681bacea6532bbc7632590ae4499cecc8e019685b515fddda71" +version = "99.99.99" +source = "git+https://github.com/RedisLabsModules/redismodule-rs?tag=v2.0.8#bce557d5ec9381d8b70eaee30c47f74bb66f458b" dependencies = [ "backtrace", "bindgen", @@ -843,9 +842,8 @@ dependencies = [ [[package]] name = "redis-module-macros" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9c92438cbeaaba0f99e2944ceeb2e34fc8fa7955576296570575df046b81197" +version = "99.99.99" +source = "git+https://github.com/RedisLabsModules/redismodule-rs?tag=v2.0.8#bce557d5ec9381d8b70eaee30c47f74bb66f458b" dependencies = [ "proc-macro2", "quote 1.0.37", @@ -856,9 +854,8 @@ dependencies = [ [[package]] name = "redis-module-macros-internals" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd1ba5724bf8bfd0c9d6f1169e45255957175b7d81944b325909024a159bc74c" +version = "99.99.99" +source = "git+https://github.com/RedisLabsModules/redismodule-rs?tag=v2.0.8#bce557d5ec9381d8b70eaee30c47f74bb66f458b" dependencies = [ "lazy_static", "proc-macro2", @@ -876,6 +873,7 @@ dependencies = [ "ijson", "itertools", "json_path", + "lazy_static", "libc", "linkme", "redis-module", diff --git a/Cargo.toml b/Cargo.toml index 6a98cf077..535926cee 100644 --- 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 index 4065b356c..16d36b360 100644 --- 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", tag="v2.0.8", default-features = false, features = ["min-redis-compatibility-version-7-2"] } +redis-module-macros = { git="https://github.com/RedisLabsModules/redismodule-rs", tag="v2.0.8" } 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 index 430520182..f18c886b6 100644 --- 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(manager: M, ctx: &Context, args: Vec) .into()) } } + "DEFRAG_INFO" => defrag_info(ctx), "HELP" => { let results = vec![ "MEMORY [path] - reports memory usage", diff --git a/redis_json/src/defrag.rs b/redis_json/src/defrag.rs new file mode 100644 index 000000000..49e6ea24c --- /dev/null +++ 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 = RedisGILGuard::default(); +} + +struct DefragCtxAllocator<'dc> { + defrag_ctx: &'dc DefragContext, +} + +impl<'dc> DefragAllocator for DefragCtxAllocator<'dc> { + unsafe fn realloc_ptr(&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(&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>(); + let new_val = defrag_allocator.realloc_ptr(*value, Layout::new::>()); + if !new_val.is_null() { + std::ptr::write(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 index c240a74c3..f5b75caf7 100644 --- 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 new file mode 100644 index 000000000..71dcd80f3 --- /dev/null +++ b/tests/pytest/test_defrag.py @@ -0,0 +1,133 @@ +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) + json_str = json.dumps(obj) + env.expect('JSON.SET', 'test', '$', json_str).ok() + for i in range(10000): + env.expect('JSON.SET', 'test%d' % i, '$', json_str).ok() + i += 1 + env.expect('JSON.SET', 'test%d' % i, '$', json_str).ok() + for i in range(10000): + env.expect('DEL', 'test%d' % i).equal(1) + i += 1 + _, _, _, _, _, 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%d' % i, '$'))[0] + env.assertEqual(res, obj) + env.assertGreater(env.cmd('info', 'Stats')['active_defrag_key_hits'], 0) + +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}]}) + +def testDefragBigJsons(env): + enableDefrag(env) + + # Disable defrag so we can actually create fragmentation + env.cmd('CONFIG', 'SET', 'activedefrag', 'no') + + env.expect('JSON.SET', 'key1', '$', "[]").ok() + env.expect('JSON.SET', 'key2', '$', "[]").ok() + + for i in range(100000): + env.cmd('JSON.ARRAPPEND', 'key1', '$', "[1.11111111111]") + env.cmd('JSON.ARRAPPEND', 'key2', '$', "[1.11111111111]") + + # Now we delete key2 which should cause fragmenation + env.expect('DEL', 'key2').equal(1) + + # wait for fragmentation for up to 30 seconds + frag = env.cmd('info', 'memory')['allocator_frag_ratio'] + startTime = time.time() + while frag < 1.4: + time.sleep(0.1) + frag = env.cmd('info', 'memory')['allocator_frag_ratio'] + 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 fragmentation, current value %s which is expected to be above 1.4.' % frag) + return + + #enable active defrag + env.cmd('CONFIG', 'SET', 'activedefrag', 'yes') + + # wait for fragmentation for go down for up to 30 seconds + frag = env.cmd('info', 'memory')['allocator_frag_ratio'] + startTime = time.time() + while frag > 1.1: + time.sleep(0.1) + frag = env.cmd('info', 'memory')['allocator_frag_ratio'] + 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 fragmentation to go down, current value %s which is expected to be bellow 1.1.' % frag) + return