add flexible scan method usable with a while loop and encapsulate callback, handle reaming review comments

Author: DarthB

examples/scan_keys.rs

@@ -3,7 +3,8 @@
 // The example implements three commands:
 // 
 // 1. `scan_keys` - scans all keys in the database and returns their names as an array of RedisString.
-// 2. `scan_key <key>` - scans all fields and values in a hash key using a closure that stores the field/value pairs as an array of RedisString.
+// 2. `scan_key <key>` - scans all fields by using a closure and a  while loop, thus allowing an early stop. Don't use the early stop but collects all the field/value pairs as an array of RedisString.
+// 3. `scan_key_foreach <key>` - scans all fields and values in a hash key using a closure that stores the field/value pairs as an array of RedisString.
 
 use redis_module::{
     key::{KeyFlags, RedisKey}, redis_module, Context, KeysCursor, RedisError, RedisResult, RedisString, RedisValue, ScanKeyCursor
@@ -24,6 +25,27 @@ fn scan_keys(ctx: &Context, _args: Vec<RedisString>) -> RedisResult {
     Ok(RedisValue::Array(res))
 }
 
+fn scan_key(ctx: &Context, args: Vec<RedisString>) -> RedisResult {
+    // only argument is the key name
+    if args.len() != 2 {
+        return Err(RedisError::WrongArity);
+    }
+
+    let key_name = &args[1];
+    let key = ctx.open_key_with_flags(key_name, KeyFlags::NOEFFECTS | KeyFlags::NOEXPIRE | KeyFlags::ACCESS_EXPIRED );
+    let cursor  = ScanKeyCursor::new(key);
+
+    let mut res = Vec::new();
+    while cursor.scan(|_key, field, value| {
+        res.push(RedisValue::BulkRedisString(field.clone()));
+        res.push(RedisValue::BulkRedisString(value.clone()));
+    }) {
+        // here we could do something between scans if needed, like an early stop
+    }
+
+    Ok(RedisValue::Array(res))
+}
+
 /// Scans all fields and values in a hash key and returns them as an array of RedisString.
 /// The command takes one argument: the name of the hash key to scan.
 fn scan_key_foreach(ctx: &Context, args: Vec<RedisString>) -> RedisResult {
@@ -55,6 +77,7 @@ redis_module! {
     data_types: [],
     commands: [
         ["scan_keys", scan_keys, "readonly", 0, 0, 0, ""],
-        ["scan_key", scan_key_foreach, "readonly", 0, 0, 0, ""],
+        ["scan_key", scan_key, "readonly", 0, 0, 0, ""],
+        ["scan_key_foreach", scan_key_foreach, "readonly", 0, 0, 0, ""],
     ],
 }

src/context/key_cursor.rs

@@ -7,19 +7,19 @@ use crate::{key::RedisKey, raw, RedisString};
 
 /// A cursor to scan field/value pairs of a (hash) key.
 ///
-/// This is a wrapper around the [RedisModule_ScanKey](https://redis.io/docs/latest/develop/reference/modules/modules-api-ref/#redismodule_scankey) 
-/// function from the C API. It provides access via a closure given to [`ScanKeyCursor::foreach`]
-/// 
+/// It provides access via a closure given to [`ScanKeyCursor::foreach`] or if you need more control, you can use [`ScanKeyCursor::scan`] 
+/// and implement your own loop, e.g. to allow an early stop.
+///
 /// ## Example usage
-/// 
+///
 /// Here we show how to extract values to communicate them back to the Redis client. We assume that the following hash key is setup before:
-/// 
+///
 /// ```text
 /// HSET user:123 name Alice age 29 location Austin
 /// ```
-/// 
+///
 /// The following example command implementation scans all fields and values in the hash key and returns them as an array of RedisString.
-/// 
+///
 /// ```ignore
 /// fn example_scan_key_foreach(ctx: &Context) -> RedisResult {
 ///    let key = ctx.open_key_with_flags("user:123", KeyFlags::NOEFFECTS | KeyFlags::NOEXPIRE | KeyFlags::ACCESS_EXPIRED );
@@ -35,9 +35,9 @@ use crate::{key::RedisKey, raw, RedisString};
 ///    Ok(RedisValue::Array(res.take()))
 /// }
 /// ```
-/// 
+///
 /// The method will produce the following output:
-/// 
+///
 /// ```text
 /// 1) "name"
 /// 2) "Alice"
@@ -63,22 +63,31 @@ impl ScanKeyCursor {
         unsafe { raw::RedisModule_ScanCursorRestart.unwrap()(self.inner_cursor) };
     }
 
-    /// Implements a callback based foreach loop over all fields and values in the hash key, use that for optimal performance.
-    pub fn foreach<F: FnMut(&RedisKey, &RedisString, &RedisString)>(&self, f: F) {
-        // Safety: Assumption: c-side initialized the function ptr and it is is never changed,
+    pub fn scan<F: FnMut(&RedisKey, &RedisString, &RedisString)>(&self, f: F) -> bool {
+        // the following is the callback definition
+        // foreach `ScanKey` call the callback may be called multiple times
+        use pimpl::scan_callback;
+
+        // Safety: The c-side initialized the function ptr and it is is never changed,
         // i.e. after module initialization the function pointers stay valid till the end of the program.
-        let scan_key = unsafe { raw::RedisModule_ScanKey.unwrap() };
+        let res = unsafe {
+            raw::RedisModule_ScanKey.unwrap()(
+                self.key.key_inner,
+                self.inner_cursor,
+                Some(scan_callback::<F>),
+                &f as *const F as *mut c_void,
+            )
+        };
+
+        res != 0
+    }
 
-        let mut res = 1;
-        while res != 0 {
-            res = unsafe {
-                scan_key(
-                    self.key.key_inner,
-                    self.inner_cursor,
-                    Some(foreach_callback::<F>),
-                    &f as *const F as *mut c_void,
-                )
-            }
+    /// Implements a callback based foreach loop over all fields and values in the hash key, use that for optimal performance.
+    pub fn foreach<F: FnMut(&RedisKey, &RedisString, &RedisString)>(&self, mut f: F) {
+        // the following is the callback definition
+        // foreach `ScanKey` call the callback may be called multiple times
+        while self.scan(&mut f) {
+            // do nothing, the callback does the work
         }
     }
 }
@@ -89,28 +98,35 @@ impl Drop for ScanKeyCursor {
     }
 }
 
-/// The callback that is used by [`ScanKeyCursor::foreach`] as argument to `RedisModule_ScanKey`.
-///
-/// The `data` pointer is the closure given to [`ScanKeyCursor::foreach`] and the callback forwards
-/// references to the key, field and value to that closure.
-unsafe extern "C" fn foreach_callback<F: FnMut(&RedisKey, &RedisString, &RedisString)>(
-    key: *mut raw::RedisModuleKey,
-    field: *mut raw::RedisModuleString,
-    value: *mut raw::RedisModuleString,
-    data: *mut c_void,
-) {
-    let ctx = ptr::null_mut();
-    let key = RedisKey::from_raw_parts(ctx, key);
+// the module contains the private implementation details of the cursor.
+mod pimpl {
+    use super::*;
+
+    /// The callback that is used by [`ScanKeyCursor::scan`] and [`ScanKeyCursor::foreach`] as argument to `RedisModule_ScanKey`.
+    ///
+    /// The `data` pointer is the closure given to [`ScanKeyCursor::foreach`] and the callback forwards
+    /// references to the key, field and value to that closure.
+    pub(super) unsafe extern "C" fn scan_callback<
+        F: FnMut(&RedisKey, &RedisString, &RedisString),
+    >(
+        key: *mut raw::RedisModuleKey,
+        field: *mut raw::RedisModuleString,
+        value: *mut raw::RedisModuleString,
+        data: *mut c_void,
+    ) {
+        let ctx = ptr::null_mut();
+        let key = RedisKey::from_raw_parts(ctx, key);
 
-    let field = RedisString::from_redis_module_string(ctx, field);
-    let value = RedisString::from_redis_module_string(ctx, value);
+        let field = RedisString::from_redis_module_string(ctx, field);
+        let value = RedisString::from_redis_module_string(ctx, value);
 
-    let callback = unsafe { &mut *(data.cast::<F>()) };
-    callback(&key, &field, &value);
+        let callback = unsafe { &mut *(data.cast::<F>()) };
+        callback(&key, &field, &value);
 
-    // we're not the owner of field and value strings
-    field.take();
-    value.take();
+        // we're not the owner of field and value strings
+        field.take();
+        value.take();
 
-    key.take(); // we're not the owner of the key either
+        key.take(); // we're not the owner of the key either
+    }
 }

tests/integration.rs

@@ -201,7 +201,7 @@ fn test_scan() -> Result<()> {
 }
 
 #[test]
-fn test_scan_key_foreach() -> Result<()> {
+fn test_scan_key() -> Result<()> {
     let mut con = TestConnection::new("scan_keys");
     redis::cmd("hset")
         .arg(&[
@@ -210,7 +210,24 @@ fn test_scan_key_foreach() -> Result<()> {
         .query::<()>(&mut con)
         .with_context(|| "failed to hset")?;
 
-    let res: Vec<String> = redis::cmd("scan_key_fe")
+    let res: Vec<String> = redis::cmd("scan_key")
+        .arg(&["user:123"])
+        .query(&mut con)?;
+    assert_eq!(&res, &["name", "Alice", "age", "29", "location", "Austin"]);
+    Ok(())
+}
+
+#[test]
+fn test_scan_key_for_each() -> Result<()> {
+    let mut con = TestConnection::new("scan_keys");
+    redis::cmd("hset")
+        .arg(&[
+            "user:123", "name", "Alice", "age", "29", "location", "Austin",
+        ])
+        .query::<()>(&mut con)
+        .with_context(|| "failed to hset")?;
+
+    let res: Vec<String> = redis::cmd("scan_key_foreach")
         .arg(&["user:123"])
         .query(&mut con)?;
     assert_eq!(&res, &["name", "Alice", "age", "29", "location", "Austin"]);