Add Redis-compatible cursors for SCAN commands

kvrocks.conf

@@ -296,6 +296,12 @@ max-backup-keep-hours 24
 # Default: 16
 max-bitmap-to-string-mb 16
 
+# Whether to enable SCAN-like cursor compatible with Redis. 
+# If enabled, the cursor will be unsigned 64-bit integers.
+# If disabled, the cursor will be a string.
+# Default: no
+redis-cursor-compatible no
+
 ################################## TLS ###################################
 
 # By default, TLS/SSL is disabled, i.e. `tls-port` is set to 0.

src/commands/cmd_hash.cc

@@ -366,12 +366,13 @@ class CommandHScan : public CommandSubkeyScanBase {
     redis::Hash hash_db(svr->storage, conn->GetNamespace());
     std::vector<std::string> fields;
     std::vector<std::string> values;
-    auto s = hash_db.Scan(key_, cursor_, limit_, prefix_, &fields, &values);
+    auto key_name = svr->GetKeyNameFromCursor(cursor_, CursorType::kTypeHash);
+    auto s = hash_db.Scan(key_, key_name, limit_, prefix_, &fields, &values);
     if (!s.ok() && !s.IsNotFound()) {
       return {Status::RedisExecErr, s.ToString()};
     }
 
-    *output = GenerateOutput(fields, values);
+    *output = GenerateOutput(svr, fields, values, CursorType::kTypeHash);
     return Status::OK();
   }
 };

src/commands/cmd_server.cc

@@ -760,11 +760,11 @@ class CommandScan : public CommandScanBase {
     return Commander::Parse(args);
   }
 
-  static std::string GenerateOutput(const std::vector<std::string> &keys, std::string end_cursor) {
+  static std::string GenerateOutput(Server *svr, const std::vector<std::string> &keys, const std::string &end_cursor) {
     std::vector<std::string> list;
     if (!end_cursor.empty()) {
-      end_cursor = kCursorPrefix + end_cursor;
-      list.emplace_back(redis::BulkString(end_cursor));
+      list.emplace_back(
+          redis::BulkString(svr->GenerateCursorFromKeyName(end_cursor, CursorType::kTypeBase, kCursorPrefix)));
     } else {
       list.emplace_back(redis::BulkString("0"));
     }
@@ -776,14 +776,15 @@ class CommandScan : public CommandScanBase {
 
   Status Execute(Server *svr, Connection *conn, std::string *output) override {
     redis::Database redis_db(svr->storage, conn->GetNamespace());
+    auto key_name = svr->GetKeyNameFromCursor(cursor_, CursorType::kTypeBase);
+
     std::vector<std::string> keys;
-    std::string end_cursor;
-    auto s = redis_db.Scan(cursor_, limit_, prefix_, &keys, &end_cursor);
+    std::string end_key;
+    auto s = redis_db.Scan(key_name, limit_, prefix_, &keys, &end_key);
     if (!s.ok()) {
       return {Status::RedisExecErr, s.ToString()};
     }
-
-    *output = GenerateOutput(keys, end_cursor);
+    *output = GenerateOutput(svr, keys, end_key);
     return Status::OK();
   }
 };

src/commands/cmd_set.cc

@@ -426,12 +426,13 @@ class CommandSScan : public CommandSubkeyScanBase {
   Status Execute(Server *svr, Connection *conn, std::string *output) override {
     redis::Set set_db(svr->storage, conn->GetNamespace());
     std::vector<std::string> members;
-    auto s = set_db.Scan(key_, cursor_, limit_, prefix_, &members);
+    auto key_name = svr->GetKeyNameFromCursor(cursor_, CursorType::kTypeSet);
+    auto s = set_db.Scan(key_, key_name, limit_, prefix_, &members);
     if (!s.ok() && !s.IsNotFound()) {
       return {Status::RedisExecErr, s.ToString()};
     }
 
-    *output = CommandScanBase::GenerateOutput(members);
+    *output = CommandScanBase::GenerateOutput(svr, members, CursorType::kTypeSet);
     return Status::OK();
   }
 };

src/commands/cmd_zset.cc

@@ -1319,7 +1319,8 @@ class CommandZScan : public CommandSubkeyScanBase {
     redis::ZSet zset_db(svr->storage, conn->GetNamespace());
     std::vector<std::string> members;
     std::vector<double> scores;
-    auto s = zset_db.Scan(key_, cursor_, limit_, prefix_, &members, &scores);
+    auto key_name = svr->GetKeyNameFromCursor(cursor_, CursorType::kTypeZSet);
+    auto s = zset_db.Scan(key_, key_name, limit_, prefix_, &members, &scores);
     if (!s.ok() && !s.IsNotFound()) {
       return {Status::RedisExecErr, s.ToString()};
     }
@@ -1329,7 +1330,7 @@ class CommandZScan : public CommandSubkeyScanBase {
     for (const auto &score : scores) {
       score_strings.emplace_back(util::Float2String(score));
     }
-    *output = GenerateOutput(members, score_strings);
+    *output = GenerateOutput(svr, members, score_strings, CursorType::kTypeZSet);
     return Status::OK();
   }
 };

src/commands/scan_base.h

@@ -23,6 +23,7 @@
 #include "commander.h"
 #include "error_constants.h"
 #include "parse_util.h"
+#include "server/server.h"
 
 namespace redis {
 
@@ -63,10 +64,11 @@ class CommandScanBase : public Commander {
     }
   }
 
-  std::string GenerateOutput(const std::vector<std::string> &keys) const {
+  std::string GenerateOutput(Server *svr, const std::vector<std::string> &keys, CursorType cursor_type) const {
     std::vector<std::string> list;
     if (keys.size() == static_cast<size_t>(limit_)) {
-      list.emplace_back(redis::BulkString(keys.back()));
+      auto end_cursor = svr->GenerateCursorFromKeyName(keys.back(), cursor_type);
+      list.emplace_back(redis::BulkString(end_cursor));
     } else {
       list.emplace_back(redis::BulkString("0"));
     }
@@ -109,11 +111,13 @@ class CommandSubkeyScanBase : public CommandScanBase {
     return Commander::Parse(args);
   }
 
-  std::string GenerateOutput(const std::vector<std::string> &fields, const std::vector<std::string> &values) {
+  std::string GenerateOutput(Server *svr, const std::vector<std::string> &fields,
+                             const std::vector<std::string> &values, CursorType cursor_type) {
     std::vector<std::string> list;
     auto items_count = fields.size();
     if (items_count == static_cast<size_t>(limit_)) {
-      list.emplace_back(redis::BulkString(fields.back()));
+      auto end_cursor = svr->GenerateCursorFromKeyName(fields.back(), cursor_type);
+      list.emplace_back(redis::BulkString(end_cursor));
     } else {
       list.emplace_back(redis::BulkString("0"));
     }

src/config/config.cc

@@ -165,6 +165,7 @@ Config::Config() {
       {"unixsocketperm", true, new OctalField(&unixsocketperm, 0777, 1, INT_MAX)},
       {"log-retention-days", false, new IntField(&log_retention_days, -1, -1, INT_MAX)},
       {"persist-cluster-nodes-enabled", false, new YesNoField(&persist_cluster_nodes_enabled, true)},
+      {"redis-cursor-compatible", false, new YesNoField(&redis_cursor_compatible, false)},
 
       /* rocksdb options */
       {"rocksdb.compression", false,

src/config/config.h

@@ -144,6 +144,7 @@ struct Config {
   int pipeline_size;
   int sequence_gap;
 
+  bool redis_cursor_compatible = true;
   int log_retention_days;
   // profiling
   int profiling_sample_ratio = 0;

src/server/server.cc

@@ -28,6 +28,8 @@
 #include <sys/utsname.h>
 
 #include <atomic>
+#include <cstdint>
+#include <functional>
 #include <iomanip>
 #include <jsoncons/json.hpp>
 #include <memory>
@@ -60,6 +62,9 @@ Server::Server(engine::Storage *storage, Config *config)
     stats.commands_stats[iter.first].latency = 0;
   }
 
+  // init cursor_dict_
+  cursor_dict_ = std::make_unique<CursorDictType>();
+
 #ifdef ENABLE_OPENSSL
   // init ssl context
   if (config->tls_port) {
@@ -1753,3 +1758,51 @@ std::list<std::pair<std::string, uint32_t>> Server::GetSlaveHostAndPort() {
   slave_threads_mu_.unlock();
   return result;
 }
+
+// The numeric cursor consists of a 16-bit counter, a 16-bit time stamp, a 29-bit hash,and a 3-bit cursor type. The
+// hash is used to prevent information leakage. The time_stamp is used to prevent the generation of the same cursor in
+// the extremely short period before and after a restart.
+NumberCursor::NumberCursor(CursorType cursor_type, uint16_t counter, const std::string &key_name) {
+  auto hash = static_cast<uint32_t>(std::hash<std::string>{}(key_name));
+  auto time_stamp = static_cast<uint16_t>(util::GetTimeStamp());
+  cursor_ = static_cast<uint64_t>(counter) | static_cast<uint64_t>(time_stamp) << 16 |
+            static_cast<uint64_t>(hash) << 32 | static_cast<uint64_t>(cursor_type) << 61;
+}
+
+bool NumberCursor::IsMatch(const CursorDictElement &element, CursorType cursor_type) const {
+  return cursor_ == element.cursor.cursor_ && cursor_type == getCursorType();
+}
+
+std::string Server::GenerateCursorFromKeyName(const std::string &key_name, CursorType cursor_type, const char *prefix) {
+  if (!config_->redis_cursor_compatible) {
+    // add prefix for SCAN
+    return prefix + key_name;
+  }
+  auto counter = cursor_counter_.fetch_add(1);
+  auto number_cursor = NumberCursor(cursor_type, counter, key_name);
+  cursor_dict_->at(number_cursor.GetIndex()) = {number_cursor, key_name};
+  return number_cursor.ToString();
+}
+
+std::string Server::GetKeyNameFromCursor(const std::string &cursor, CursorType cursor_type) {
+  // When cursor is 0, cursor string is empty
+  if (cursor.empty() || !config_->redis_cursor_compatible) {
+    return cursor;
+  }
+
+  auto s = ParseInt<uint64_t>(cursor, 10);
+  // When Cursor 0 or not a Integer return empty string.
+  // Although the parameter 'cursor' is not expected to be 0, we still added a check for 0 to increase the robustness of
+  // the code.
+  if (!s.IsOK() || *s == 0) {
+    return {};
+  }
+  auto number_cursor = NumberCursor(*s);
+  // Because the index information is fully stored in the cursor, we can directly obtain the index from the cursor.
+  auto item = cursor_dict_->at(number_cursor.GetIndex());
+  if (number_cursor.IsMatch(item, cursor_type)) {
+    return item.key_name;
+  }
+
+  return {};
+}

src/server/server.h

@@ -22,6 +22,10 @@
 
 #include <inttypes.h>
 
+#include <array>
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
 #include <list>
 #include <map>
 #include <memory>
@@ -84,6 +88,39 @@ struct ChannelSubscribeNum {
   size_t subscribe_num;
 };
 
+// CURSOR_DICT_SIZE must be 2^n where n <= 16
+constexpr const size_t CURSOR_DICT_SIZE = 1024 * 16;
+static_assert((CURSOR_DICT_SIZE & (CURSOR_DICT_SIZE - 1)) == 0, "CURSOR_DICT_SIZE must be 2^n");
+static_assert(CURSOR_DICT_SIZE <= (1 << 16), "CURSOR_DICT_SIZE must be less than or equal to 2^16");
+
+enum class CursorType : uint8_t {
+  kTypeNone = 0,  // none
+  kTypeBase = 1,  // cursor for SCAN
+  kTypeHash = 2,  // cursor for HSCAN
+  kTypeSet = 3,   // cursor for SSCAN
+  kTypeZSet = 4,  // cursor for ZSCAN
+};
+struct CursorDictElement;
+
+class NumberCursor {
+ public:
+  NumberCursor() = default;
+  explicit NumberCursor(CursorType cursor_type, uint16_t counter, const std::string &key_name);
+  explicit NumberCursor(uint64_t number_cursor) : cursor_(number_cursor) {}
+  size_t GetIndex() const { return cursor_ % CURSOR_DICT_SIZE; }
+  bool IsMatch(const CursorDictElement &element, CursorType cursor_type) const;
+  std::string ToString() const { return std::to_string(cursor_); }
+
+ private:
+  CursorType getCursorType() const { return static_cast<CursorType>(cursor_ >> 61); }
+  uint64_t cursor_;
+};
+
+struct CursorDictElement {
+  NumberCursor cursor;
+  std::string key_name;
+};
+
 enum SlowLog {
   kSlowLogMaxArgc = 32,
   kSlowLogMaxString = 128,
@@ -196,6 +233,9 @@ class Server {
   void GetLatestKeyNumStats(const std::string &ns, KeyNumStats *stats);
   time_t GetLastScanTime(const std::string &ns);
 
+  std::string GenerateCursorFromKeyName(const std::string &key_name, CursorType cursor_type, const char *prefix = "");
+  std::string GetKeyNameFromCursor(const std::string &cursor, CursorType cursor_type);
+
   int DecrClientNum();
   int IncrClientNum();
   int IncrMonitorClientNum();
@@ -319,4 +359,9 @@ class Server {
   std::atomic<size_t> watched_key_size_ = 0;
   std::map<std::string, std::set<redis::Connection *>> watched_key_map_;
   std::shared_mutex watched_key_mutex_;
+
+  // SCAN ring buffer
+  std::atomic<uint16_t> cursor_counter_ = {0};
+  using CursorDictType = std::array<CursorDictElement, CURSOR_DICT_SIZE>;
+  std::unique_ptr<CursorDictType> cursor_dict_;
 };

src/storage/redis_db.cc

@@ -277,9 +277,8 @@ rocksdb::Status Database::Scan(const std::string &cursor, uint64_t limit, const
       keys->emplace_back(user_key);
       cnt++;
     }
-
     if (!storage_->IsSlotIdEncoded() || prefix.empty()) {
-      if (!keys->empty()) {
+      if (!keys->empty() && cnt >= limit) {
         end_cursor->append(user_key);
       }
       break;

tests/gocase/integration/cli/cli_test.go

@@ -340,6 +340,14 @@ func TestRedisCli(t *testing.T) {
 		require.Len(t, strings.Split(r, "\n"), 10)
 	})
 
+	t.Run("Use reids-cli --bigkeys", func(t *testing.T) {
+		runCli(t, srv, nil, "--bigkeys").Success()
+	})
+
+	t.Run("Use reids-cli --memkeys", func(t *testing.T) {
+		runCli(t, srv, nil, "--memkeys").Success()
+	})
+
 	formatArgs := func(args ...string) string {
 		cmd := fmt.Sprintf("*%d\r\n", len(args))
 		for _, arg := range args {

tests/gocase/unit/scan/scan_test.go

@@ -52,13 +52,26 @@ func TestScanEmptyKey(t *testing.T) {
 	require.Equal(t, []string{"", "fab", "fiz", "foobar"}, keys)
 }
 
-func TestScan(t *testing.T) {
+func TestScanWithNumberCursor(t *testing.T) {
 	srv := util.StartServer(t, map[string]string{})
 	defer srv.Close()
+	ctx := context.Background()
+	rdb := srv.NewClient()
+	defer func() { require.NoError(t, rdb.Close()) }()
+	ScanTest(t, rdb, ctx)
+}
 
+func TestScanWithStringCursor(t *testing.T) {
+	srv := util.StartServer(t, map[string]string{})
+	defer srv.Close()
 	ctx := context.Background()
 	rdb := srv.NewClient()
 	defer func() { require.NoError(t, rdb.Close()) }()
+	require.NoError(t, rdb.ConfigSet(ctx, "redis-cursor-compatible", "no").Err())
+	ScanTest(t, rdb, ctx)
+}
+
+func ScanTest(t *testing.T, rdb *redis.Client, ctx context.Context) {
 
 	t.Run("SCAN Basic", func(t *testing.T) {
 		require.NoError(t, rdb.FlushDB(ctx).Err())