diff --git a/Cargo.toml b/Cargo.toml index bc219f15..bb164968 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,6 +65,10 @@ name = "info" crate-type = ["cdylib"] required-features = [] +[[example]] +name = "scan_keys" +crate-type = ["cdylib"] + [dependencies] bitflags = "1.2" libc = "0.2" diff --git a/examples/scan_keys.rs b/examples/scan_keys.rs new file mode 100644 index 00000000..406b42c2 --- /dev/null +++ b/examples/scan_keys.rs @@ -0,0 +1,24 @@ +#[macro_use] +extern crate redis_module; + +use redis_module::{Context, KeysCursor, RedisResult, RedisString, RedisValue}; + +fn scan_keys(ctx: &Context, _args: Vec) -> RedisResult { + let cursor = KeysCursor::new(); + let mut res = Vec::new(); + while cursor.scan(ctx, &|_ctx, key_name, _key| { + res.push(RedisValue::BulkRedisString(key_name)); + }) {} + Ok(RedisValue::Array(res)) +} + +////////////////////////////////////////////////////// + +redis_module! { + name: "scan", + version: 1, + data_types: [], + commands: [ + ["scan_keys", scan_keys, "readonly", 0, 0, 0], + ], +} diff --git a/src/context/keys_cursor.rs b/src/context/keys_cursor.rs new file mode 100644 index 00000000..4ddf4d8f --- /dev/null +++ b/src/context/keys_cursor.rs @@ -0,0 +1,68 @@ +use crate::context::Context; +use crate::key::RedisKey; +use crate::raw; +use crate::redismodule::RedisString; +use std::ffi::c_void; + +pub struct KeysCursor { + inner_cursor: *mut raw::RedisModuleScanCursor, +} + +extern "C" fn scan_callback)>( + ctx: *mut raw::RedisModuleCtx, + key_name: *mut raw::RedisModuleString, + key: *mut raw::RedisModuleKey, + private_data: *mut ::std::os::raw::c_void, +) { + let context = Context::new(ctx); + let key_name = RedisString::new(ctx, key_name); + let redis_key = if key.is_null() { + None + } else { + Some(RedisKey::from_raw_parts(ctx, key)) + }; + let callback = unsafe { &mut *(private_data.cast::()) }; + callback(&context, key_name, redis_key.as_ref()); + + // we are not the owner of the key, so we must take the underline *mut raw::RedisModuleKey so it will not be freed. + redis_key.map(|v| v.take()); +} + +impl KeysCursor { + pub fn new() -> KeysCursor { + let inner_cursor = unsafe { raw::RedisModule_ScanCursorCreate.unwrap()() }; + KeysCursor { inner_cursor } + } + + pub fn scan)>( + &self, + ctx: &Context, + callback: &F, + ) -> bool { + let res = unsafe { + raw::RedisModule_Scan.unwrap()( + ctx.ctx, + self.inner_cursor, + Some(scan_callback::), + callback as *const F as *mut c_void, + ) + }; + res != 0 + } + + pub fn restart(&self) { + unsafe { raw::RedisModule_ScanCursorRestart.unwrap()(self.inner_cursor) }; + } +} + +impl Default for KeysCursor { + fn default() -> Self { + Self::new() + } +} + +impl Drop for KeysCursor { + fn drop(&mut self) { + unsafe { raw::RedisModule_ScanCursorDestroy.unwrap()(self.inner_cursor) }; + } +} diff --git a/src/context/mod.rs b/src/context/mod.rs index 05270425..1c86709f 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -22,6 +22,8 @@ pub mod blocked; pub mod info; +pub mod keys_cursor; + /// `Context` is a structure that's designed to give us a high-level interface to /// the Redis module API by abstracting away the raw C FFI calls. pub struct Context { diff --git a/src/key.rs b/src/key.rs index ee20e8d8..ef4e7945 100644 --- a/src/key.rs +++ b/src/key.rs @@ -37,11 +37,24 @@ pub struct RedisKey { } impl RedisKey { + pub(crate) fn take(mut self) -> *mut raw::RedisModuleKey { + let res = self.key_inner; + self.key_inner = std::ptr::null_mut(); + res + } + pub fn open(ctx: *mut raw::RedisModuleCtx, key: &RedisString) -> Self { let key_inner = raw::open_key(ctx, key.inner, to_raw_mode(KeyMode::Read)); Self { ctx, key_inner } } + pub(crate) fn from_raw_parts( + ctx: *mut raw::RedisModuleCtx, + key_inner: *mut raw::RedisModuleKey, + ) -> Self { + Self { ctx, key_inner } + } + /// # Panics /// /// Will panic if `RedisModule_ModuleTypeGetValue` is missing in redismodule.h @@ -128,7 +141,9 @@ impl RedisKey { impl Drop for RedisKey { // Frees resources appropriately as a RedisKey goes out of scope. fn drop(&mut self) { - raw::close_key(self.key_inner); + if !self.key_inner.is_null() { + raw::close_key(self.key_inner); + } } } diff --git a/src/lib.rs b/src/lib.rs index 75893f9d..ec0907db 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,7 @@ pub use crate::context::thread_safe::{DetachedFromClient, ThreadSafeContext}; #[cfg(feature = "experimental-api")] pub use crate::raw::NotifyEvent; +pub use crate::context::keys_cursor::KeysCursor; pub use crate::context::Context; pub use crate::raw::*; pub use crate::redismodule::*; diff --git a/tests/integration.rs b/tests/integration.rs index 8847625d..374715e0 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -168,3 +168,29 @@ fn test_string() -> Result<()> { Ok(()) } + +#[test] +fn test_scan() -> Result<()> { + let port: u16 = 6486; + let _guards = vec![start_redis_server_with_module("scan_keys", port) + .with_context(|| "failed to start redis server")?]; + let mut con = + get_redis_connection(port).with_context(|| "failed to connect to redis server")?; + + redis::cmd("set") + .arg(&["x", "1"]) + .query(&mut con) + .with_context(|| "failed to run string.set")?; + + redis::cmd("set") + .arg(&["y", "1"]) + .query(&mut con) + .with_context(|| "failed to run string.set")?; + + let mut res: Vec = redis::cmd("scan_keys").query(&mut con)?; + res.sort(); + + assert_eq!(&res, &["x", "y"]); + + Ok(()) +}