Added keys scan API. (#265)

* Added keys scan API.

* Added comment on taken the underline RedisModuleKey.

* Apply suggestions from code review

Co-authored-by: Guy Korland <[email protected]>

Co-authored-by: Guy Korland <[email protected]>

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"

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<RedisString>) -> 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],
+    ],
+}

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<C: FnMut(&Context, RedisString, Option<&RedisKey>)>(
+    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::<C>()) };
+    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<F: FnMut(&Context, RedisString, Option<&RedisKey>)>(
+        &self,
+        ctx: &Context,
+        callback: &F,
+    ) -> bool {
+        let res = unsafe {
+            raw::RedisModule_Scan.unwrap()(
+                ctx.ctx,
+                self.inner_cursor,
+                Some(scan_callback::<F>),
+                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) };
+    }
+}

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 {

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);
+        }
     }
 }
 

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::*;

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<String> = redis::cmd("scan_keys").query(&mut con)?;
+    res.sort();
+
+    assert_eq!(&res, &["x", "y"]);
+
+    Ok(())
+}