Add COUNT support to SSCAN

src/Dispatcher.cc

@@ -346,12 +346,26 @@ RedisEncodedResponse RedisDispatcher::dispatch(RedisRequest &request, LogIndex c
       return Formatter::integer(count);
     }
     case RedisCommand::SSCAN: {
-      if(request.size() != 3) return errArgs(request, commit);
-      if(request[2] != "0") return Formatter::err("invalid cursor");
-      std::vector<std::string> members;
-      rocksdb::Status st = store.smembers(request[1], members);
+      if(request.size() < 3) return errArgs(request, commit);
+
+      ScanCommandArguments args = parseScanCommand(request.begin()+2, request.end());
+      if(!args.error.empty()) {
+        return Formatter::err(args.error);
+      }
+
+      // No support for MATCH here, maybe add later
+      if(!args.match.empty()) {
+        return Formatter::err("syntax error");
+      }
+
+      std::string newcursor;
+      std::vector<std::string> vec;
+      rocksdb::Status st = store.sscan(request[1], args.cursor, args.count, newcursor, vec);
       if(!st.ok()) return Formatter::fromStatus(st);
-      return Formatter::scan("0", members);
+
+      if(newcursor == "") newcursor = "0";
+      else newcursor = "next:" + newcursor;
+      return Formatter::scan(newcursor, vec);
     }
     case RedisCommand::LLEN: {
       if(request.size() != 2) return errArgs(request, commit);

src/StateMachine.cc

@@ -429,6 +429,32 @@ rocksdb::Status StateMachine::hscan(const std::string &key, const std::string &c
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status StateMachine::sscan(const std::string &key, const std::string &cursor, size_t count, std::string &newCursor, std::vector<std::string> &res) {
+  Snapshot snapshot(db);
+  if(!assertKeyType(snapshot, key, KeyType::kSet)) return wrong_type();
+
+  FieldLocator locator(KeyType::kSet, key, cursor);
+  res.clear();
+
+  newCursor = "";
+  IteratorPtr iter(db->NewIterator(snapshot.opts()));
+  for(iter->Seek(locator.toSlice()); iter->Valid(); iter->Next()) {
+    std::string tmp = iter->key().ToString();
+
+    if(!StringUtils::startswith(tmp, locator.getPrefix())) break;
+
+    std::string fieldname = std::string(tmp.begin()+locator.getPrefixSize(), tmp.end());
+    if(res.size() >= count) {
+      newCursor = fieldname;
+      break;
+    }
+
+    res.push_back(fieldname);
+  }
+
+  return rocksdb::Status::OK();
+}
+
 rocksdb::Status StateMachine::hvals(const std::string &key, std::vector<std::string> &vals) {
   Snapshot snapshot(db);
   if(!assertKeyType(snapshot, key, KeyType::kHash)) return wrong_type();

src/StateMachine.hh

@@ -96,6 +96,7 @@ public:
   rocksdb::Status srem(const std::string &key, const VecIterator &start, const VecIterator &end, int64_t &removed, LogIndex index = 0);
   rocksdb::Status smembers(const std::string &key, std::vector<std::string> &members);
   rocksdb::Status scard(const std::string &key, size_t &count);
+  rocksdb::Status sscan(const std::string &key, const std::string &cursor, size_t count, std::string &newCursor, std::vector<std::string> &res);
   rocksdb::Status set(const std::string& key, const std::string& value, LogIndex index = 0);
   rocksdb::Status get(const std::string &key, std::string &value);
   rocksdb::Status del(const VecIterator &start, const VecIterator &end, int64_t &removed, LogIndex index = 0);

test/e2e.cc

@@ -725,3 +725,30 @@ TEST_F(Raft_e2e, hincrbymulti) {
   ASSERT_REPLY(tunnel(leaderID)->exec("hget", "h2", "h3"), "24");
   ASSERT_REPLY(tunnel(leaderID)->exec("hget", "h4", "h8"), "13");
 }
+
+TEST_F(Raft_e2e, sscan) {
+  spinup(0); spinup(1); spinup(2);
+  RETRY_ASSERT_TRUE(checkStateConsensus(0, 1, 2));
+
+  int leaderID = getLeaderID();
+
+  redisReplyPtr reply = tunnel(leaderID)->exec("sscan", "myset", "0", "asdf", "123").get();
+  ASSERT_ERR(reply, "ERR syntax error");
+
+  ASSERT_REPLY(tunnel(leaderID)->exec("sadd", "myset", "a", "b", "c", "d", "e", "f", "g"), 7);
+
+  reply = tunnel(leaderID)->exec("sscan", "myset", "0", "COUNT", "3").get();
+  ASSERT_REPLY(reply, std::make_pair("next:d", make_vec("a", "b", "c")));
+
+  reply = tunnel(leaderID)->exec("sscan", "myset", "next:d", "COUNT", "2").get();
+  ASSERT_REPLY(reply, std::make_pair("next:f", make_vec("d", "e")));
+
+  reply = tunnel(leaderID)->exec("sscan", "myset", "next:f", "COUNT", "2").get();
+  ASSERT_REPLY(reply, std::make_pair("0", make_vec("f", "g")));
+
+  reply = tunnel(leaderID)->exec("sscan", "myset", "next:zz").get();
+  ASSERT_REPLY(reply, std::make_pair("0", make_vec()));
+
+  reply = tunnel(leaderID)->exec("sscan", "not-existing", "next:zz").get();
+  ASSERT_REPLY(reply, std::make_pair("0", make_vec()));
+}