feat(streams): implement rdb load for streams

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/redis/object.c

@@ -369,10 +369,11 @@ void freeModuleObject(robj *o) {
     zfree(mv);
 }
 
+#endif
+
 void freeStreamObject(robj *o) {
   freeStream(o->ptr);
 }
-#endif
 
 void incrRefCount(robj *o) {
     if (o->refcount < OBJ_FIRST_SPECIAL_REFCOUNT) {
@@ -399,8 +400,7 @@ void decrRefCount(robj *o) {
             // freeModuleObject(o);
         break;
         case OBJ_STREAM:
-            serverPanic("Unsupported OBJ_STREAM type");
-            //freeStreamObject(o);
+            freeStreamObject(o);
         break;
         default: serverPanic("Unknown object type"); break;
         }

src/redis/t_stream.c

@@ -2515,15 +2515,12 @@ streamCG *streamLookupCG(stream *s, sds groupname) {
     return (cg == raxNotFound) ? NULL : cg;
 }
 
-#if ROMAN_ENABLE
 /* Create a consumer with the specified name in the group 'cg' and return.
  * If the consumer exists, return NULL. As a side effect, when the consumer
  * is successfully created, the key space will be notified and dirty++ unless
  * the SCC_NO_NOTIFY or SCC_NO_DIRTIFY flags is specified. */
 streamConsumer *streamCreateConsumer(streamCG *cg, sds name, robj *key, int dbid, int flags) {
     if (cg == NULL) return NULL;
-    int notify = !(flags & SCC_NO_NOTIFY);
-    int dirty = !(flags & SCC_NO_DIRTIFY);
     streamConsumer *consumer = zmalloc(sizeof(*consumer));
     int success = raxTryInsert(cg->consumers,(unsigned char*)name,
                                sdslen(name),consumer,NULL);
@@ -2534,11 +2531,9 @@ streamConsumer *streamCreateConsumer(streamCG *cg, sds name, robj *key, int dbid
     consumer->name = sdsdup(name);
     consumer->pel = raxNew();
     consumer->seen_time = mstime();
-    if (dirty) server.dirty++;
-    if (notify) notifyKeyspaceEvent(NOTIFY_STREAM,"xgroup-createconsumer",key,dbid);
+
     return consumer;
 }
-#endif 
 
 /* Lookup the consumer with the specified name in the group 'cg'. Its last
  * seen time is updated unless the SLC_NO_REFRESH flag is specified. */
@@ -3943,6 +3938,8 @@ NULL
     }
 }
 
+#endif
+
 /* Validate the integrity stream listpack entries structure. Both in term of a
  * valid listpack, but also that the structure of the entries matches a valid
  * stream. return 1 if valid 0 if not valid. */
@@ -4035,4 +4032,3 @@ int streamValidateListpackIntegrity(unsigned char *lp, size_t size, int deep) {
     return 1;
 }
 
-#endif

src/server/CMakeLists.txt

@@ -22,7 +22,8 @@ cxx_test(list_family_test dfly_test_lib LABELS DFLY)
 cxx_test(set_family_test dfly_test_lib LABELS DFLY)
 cxx_test(stream_family_test dfly_test_lib LABELS DFLY)
 cxx_test(string_family_test dfly_test_lib LABELS DFLY)
-cxx_test(rdb_test dfly_test_lib DATA testdata/empty.rdb testdata/redis6_small.rdb LABELS DFLY)
+cxx_test(rdb_test dfly_test_lib DATA testdata/empty.rdb testdata/redis6_small.rdb
+         testdata/redis6_stream.rdb LABELS DFLY)
 cxx_test(zset_family_test dfly_test_lib LABELS DFLY)
 cxx_test(blocking_controller_test dragonfly_lib LABELS DFLY)
 

src/server/rdb_load.cc

@@ -10,6 +10,7 @@ extern "C" {
 #include "redis/listpack.h"
 #include "redis/lzfP.h" /* LZF compression library */
 #include "redis/rdb.h"
+#include "redis/stream.h"
 #include "redis/util.h"
 #include "redis/ziplist.h"
 #include "redis/zmalloc.h"
@@ -834,6 +835,9 @@ io::Result<robj*> RdbLoader::ReadObj(int rdbtype) {
     case RDB_TYPE_LIST_QUICKLIST:
       res_obj = ReadListQuicklist(rdbtype);
       break;
+    case RDB_TYPE_STREAM_LISTPACKS:
+      res_obj = ReadStreams();
+      break;
     default:
       LOG(ERROR) << "Unsupported rdb type " << rdbtype;
       return Unexpected(errc::invalid_encoding);
@@ -1264,6 +1268,179 @@ io::Result<robj*> RdbLoader::ReadListQuicklist(int rdbtype) {
   return res;
 }
 
+::io::Result<robj*> RdbLoader::ReadStreams() {
+  uint64_t listpacks;
+  SET_OR_UNEXPECT(LoadLen(nullptr), listpacks);
+
+  robj* o = createStreamObject();
+  stream* s = (stream*)o->ptr;
+
+  auto cleanup = absl::Cleanup([&] { decrRefCount(o); });
+
+  while (listpacks--) {
+    /* Get the master ID, the one we'll use as key of the radix tree
+     * node: the entries inside the listpack itself are delta-encoded
+     * relatively to this ID. */
+    sds nodekey;
+    SET_OR_UNEXPECT(ReadKey(), nodekey);
+    auto cleanup2 = absl::Cleanup([&] { sdsfree(nodekey); });
+
+    if (sdslen(nodekey) != sizeof(streamID)) {
+      LOG(ERROR) << "Stream node key entry is not the size of a stream ID";
+
+      return Unexpected(errc::rdb_file_corrupted);
+    }
+
+    /* Load the listpack. */
+    OpaqueBuf fetch;
+    SET_OR_UNEXPECT(FetchGenericString(RDB_LOAD_PLAIN), fetch);
+    if (fetch.second == 0 || fetch.first == nullptr) {
+      LOG(ERROR) << "Stream listpacks loading failed";
+      return Unexpected(errc::rdb_file_corrupted);
+    }
+    DCHECK(fetch.first);
+
+    uint8_t* lp = (uint8_t*)fetch.first;
+
+    if (!streamValidateListpackIntegrity(lp, fetch.second, 0)) {
+      zfree(lp);
+      LOG(ERROR) << "Stream listpack integrity check failed.";
+      return Unexpected(errc::rdb_file_corrupted);
+    }
+
+    unsigned char* first = lpFirst(lp);
+    if (first == NULL) {
+      /* Serialized listpacks should never be empty, since on
+       * deletion we should remove the radix tree key if the
+       * resulting listpack is empty. */
+      LOG(ERROR) << "Empty listpack inside stream";
+      zfree(lp);
+      return Unexpected(errc::rdb_file_corrupted);
+    }
+
+    /* Insert the key in the radix tree. */
+    int retval = raxTryInsert(s->rax_tree, (unsigned char*)nodekey, sizeof(streamID), lp, NULL);
+    if (!retval) {
+      LOG(ERROR) << "Listpack re-added with existing key";
+      return Unexpected(errc::duplicate_key);
+    }
+  }
+
+  /* Load total number of items inside the stream. */
+  SET_OR_UNEXPECT(LoadLen(nullptr), s->length);
+
+  /* Load the last entry ID. */
+  SET_OR_UNEXPECT(LoadLen(nullptr), s->last_id.ms);
+  SET_OR_UNEXPECT(LoadLen(nullptr), s->last_id.seq);
+
+  /* Consumer groups loading */
+  uint64_t cgroups_count;
+  SET_OR_UNEXPECT(LoadLen(nullptr), cgroups_count);
+
+  while (cgroups_count--) {
+    /* Get the consumer group name and ID. We can then create the
+     * consumer group ASAP and populate its structure as
+     * we read more data. */
+    streamID cg_id;
+
+    sds cgname;
+    SET_OR_UNEXPECT(ReadKey(), cgname);
+
+    auto cleanup2 = absl::Cleanup([&] { sdsfree(cgname); });
+    SET_OR_UNEXPECT(LoadLen(nullptr), cg_id.ms);
+    SET_OR_UNEXPECT(LoadLen(nullptr), cg_id.seq);
+
+    streamCG* cgroup = streamCreateCG(s, cgname, sdslen(cgname), &cg_id, 0);
+    if (cgroup == NULL) {
+      LOG(ERROR) << "Duplicated consumer group name " << cgname;
+      return Unexpected(errc::duplicate_key);
+    }
+    std::move(cleanup2).Invoke();
+    // no need to free cgroup because it's attached to s.
+
+    /* Load the global PEL for this consumer group, however we'll
+     * not yet populate the NACK structures with the message
+     * owner, since consumers for this group and their messages will
+     * be read as a next step. So for now leave them not resolved
+     * and later populate it. */
+    uint64_t pel_size;
+    SET_OR_UNEXPECT(LoadLen(nullptr), pel_size);
+
+    while (pel_size--) {
+      uint8_t rawid[sizeof(streamID)];
+      error_code ec = FetchBuf(sizeof(rawid), rawid);
+      if (ec) {
+        LOG(ERROR) << "Stream PEL ID loading failed.";
+        return make_unexpected(ec);
+      }
+
+      streamNACK* nack = streamCreateNACK(NULL);
+      auto cleanup2 = absl::Cleanup([&] { streamFreeNACK(nack); });
+
+      SET_OR_UNEXPECT(FetchInt<int64_t>(), nack->delivery_time);
+      SET_OR_UNEXPECT(LoadLen(nullptr), nack->delivery_count);
+
+      if (!raxTryInsert(cgroup->pel, rawid, sizeof(rawid), nack, NULL)) {
+        LOG(ERROR) << "Duplicated global PEL entry loading stream consumer group";
+        return Unexpected(errc::duplicate_key);
+      }
+      std::move(cleanup2).Cancel();
+    }
+
+    /* Now that we loaded our global PEL, we need to load the
+     * consumers and their local PELs. */
+    uint64_t consumers_num;
+    SET_OR_UNEXPECT(LoadLen(nullptr), consumers_num);
+
+    while (consumers_num--) {
+      sds cname;
+      SET_OR_UNEXPECT(ReadKey(), cname);
+
+      streamConsumer* consumer =
+          streamCreateConsumer(cgroup, cname, NULL, 0, SCC_NO_NOTIFY | SCC_NO_DIRTIFY);
+      sdsfree(cname);
+      if (!consumer) {
+        LOG(ERROR) << "Duplicate stream consumer detected.";
+        return Unexpected(errc::duplicate_key);
+      }
+      // no need to free consumer because it's attached to cgroup.
+
+      SET_OR_UNEXPECT(FetchInt<int64_t>(), consumer->seen_time);
+
+      /* Load the PEL about entries owned by this specific
+       * consumer. */
+      SET_OR_UNEXPECT(LoadLen(nullptr), pel_size);
+      while (pel_size--) {
+        unsigned char rawid[sizeof(streamID)];
+        error_code ec = FetchBuf(sizeof(rawid), rawid);
+        if (ec) {
+          LOG(ERROR) << "Stream PEL ID loading failed.";
+          return make_unexpected(ec);
+        }
+        streamNACK* nack = (streamNACK*)raxFind(cgroup->pel, rawid, sizeof(rawid));
+        if (nack == raxNotFound) {
+          LOG(ERROR) << "Consumer entry not found in group global PEL";
+          return Unexpected(errc::rdb_file_corrupted);
+        }
+
+        /* Set the NACK consumer, that was left to NULL when
+         * loading the global PEL. Then set the same shared
+         * NACK structure also in the consumer-specific PEL. */
+        nack->consumer = consumer;
+        if (!raxTryInsert(consumer->pel, rawid, sizeof(rawid), nack, NULL)) {
+          LOG(ERROR) << "Duplicated consumer PEL entry loading a stream consumer group";
+          streamFreeNACK(nack);
+          return Unexpected(errc::duplicate_key);
+        }
+      }
+    }  // while (consumers_num)
+  }    // while (cgroup_num)
+
+  std::move(cleanup).Cancel();
+
+  return o;
+}
+
 void RdbLoader::ResizeDb(size_t key_num, size_t expire_num) {
   DCHECK_LT(key_num, 1U << 31);
   DCHECK_LT(expire_num, 1U << 31);

src/server/rdb_load.h

@@ -72,6 +72,7 @@ class RdbLoader {
   ::io::Result<robj*> ReadZSet(int rdbtype);
   ::io::Result<robj*> ReadZSetZL();
   ::io::Result<robj*> ReadListQuicklist(int rdbtype);
+  ::io::Result<robj*> ReadStreams();
 
   std::error_code EnsureRead(size_t min_sz) {
     if (mem_buf_.InputLen() >= min_sz)

src/server/rdb_test.cc

@@ -115,6 +115,17 @@ TEST_F(RdbTest, LoadSmall6) {
   EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(1), IntArg(1)));
 }
 
+TEST_F(RdbTest, LoadStream) {
+  io::FileSource fs = GetSource("redis6_stream.rdb");
+  RdbLoader loader(service_->script_mgr());
+
+  // must run in proactor thread in order to avoid polluting the serverstate
+  // in the main, testing thread.
+  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+
+  ASSERT_FALSE(ec) << ec.message();
+}
+
 TEST_F(RdbTest, Reload) {
   absl::FlagSaver fs;