Use the QUICKLIST encoding to dump the list payload to compatible wit…

…h Redis >= 4 (#2277)
apache · Apr 30, 2024 · 9b2ee7e · 9b2ee7e
1 parent 1834ce8
commit 9b2ee7e
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 14 deletions.
diff --git a/src/storage/rdb.cc b/src/storage/rdb.cc
@@ -730,7 +730,7 @@ Status RDB::SaveObjectType(const RedisType type) {
   } else if (type == kRedisHash) {
     robj_type = RDBTypeHash;
   } else if (type == kRedisList) {
-    robj_type = RDBTypeListQuickList2;
+    robj_type = RDBTypeListQuickList;
   } else if (type == kRedisSet) {
     robj_type = RDBTypeSet;
   } else if (type == kRedisZSet) {
@@ -892,12 +892,7 @@ Status RDB::SaveListObject(const std::vector<std::string> &elems) {
     }
 
     for (const auto &elem : elems) {
-      status = RdbSaveLen(1 /*plain container mode */);
-      if (!status.IsOK()) {
-        return {Status::RedisExecErr, status.Msg()};
-      }
-
-      status = SaveStringObject(elem);
+      auto status = rdbSaveZipListObject(elem);
       if (!status.IsOK()) {
         return {Status::RedisExecErr, status.Msg()};
       }
@@ -1005,3 +1000,35 @@ Status RDB::rdbSaveBinaryDoubleValue(double val) {
   memrev64ifbe(&val);
   return stream_->Write((const char *)(&val), sizeof(val));
 }
+
+Status RDB::rdbSaveZipListObject(const std::string &elem) {
+  // calc total ziplist size
+  uint prevlen = 0;
+  const size_t ziplist_size = zlHeaderSize + zlEndSize + elem.length() +
+                              ZipList::ZipStorePrevEntryLength(nullptr, 0, prevlen) +
+                              ZipList::ZipStoreEntryEncoding(nullptr, 0, elem.length());
+  auto zl_string = std::string(ziplist_size, '\0');
+  auto zl_ptr = reinterpret_cast<unsigned char *>(&zl_string[0]);
+
+  // set ziplist header
+  ZipList::SetZipListBytes(zl_ptr, ziplist_size, (static_cast<uint32_t>(ziplist_size)));
+  ZipList::SetZipListTailOffset(zl_ptr, ziplist_size, intrev32ifbe(zlHeaderSize));
+
+  // set ziplist entry
+  auto pos = ZipList::GetZipListEntryHead(zl_ptr, ziplist_size);
+  pos += ZipList::ZipStorePrevEntryLength(pos, ziplist_size, prevlen);
+  pos += ZipList::ZipStoreEntryEncoding(pos, ziplist_size, elem.length());
+  assert(pos + elem.length() <= zl_ptr + ziplist_size);
+  memcpy(pos, elem.c_str(), elem.length());
+
+  // set ziplist end
+  ZipList::SetZipListLength(zl_ptr, ziplist_size, 1);
+  zl_ptr[ziplist_size - 1] = zlEnd;
+
+  auto status = SaveStringObject(zl_string);
+  if (!status.IsOK()) {
+    return {Status::RedisExecErr, status.Msg()};
+  }
+
+  return Status::OK();
+}
diff --git a/src/storage/rdb.h b/src/storage/rdb.h
@@ -153,4 +153,5 @@ class RDB {
   static bool isEmptyRedisObject(const RedisObjValue &value);
   static int rdbEncodeInteger(long long value, unsigned char *enc);
   Status rdbSaveBinaryDoubleValue(double val);
+  Status rdbSaveZipListObject(const std::string &elem);
 };
diff --git a/src/storage/rdb_ziplist.cc b/src/storage/rdb_ziplist.cc
@@ -20,11 +20,9 @@
 
 #include "rdb_ziplist.h"
 
-#include "vendor/endianconv.h"
+#include <cassert>
 
-constexpr const int zlHeaderSize = 10;
-constexpr const uint8_t ZipListBigLen = 0xFE;
-constexpr const uint8_t zlEnd = 0xFF;
+#include "vendor/endianconv.h"
 
 constexpr const uint8_t ZIP_STR_MASK = 0xC0;
 constexpr const uint8_t ZIP_STR_06B = (0 << 6);
@@ -52,7 +50,7 @@ StatusOr<std::string> ZipList::Next() {
   std::string value;
   if ((encoding) < ZIP_STR_MASK) {
     // For integer type, needs to convert to uint8_t* to avoid signed extension
-    auto data = reinterpret_cast<const uint8_t*>(input_.data());
+    auto data = reinterpret_cast<const uint8_t *>(input_.data());
     if ((encoding) == ZIP_STR_06B) {
       len_bytes = 1;
       len = data[pos_] & 0x3F;
@@ -91,7 +89,7 @@ StatusOr<std::string> ZipList::Next() {
     } else if ((encoding) == ZIP_INT_24B) {
       GET_OR_RET(peekOK(3));
       int32_t i32 = 0;
-      memcpy(reinterpret_cast<uint8_t*>(&i32) + 1, input_.data() + pos_, sizeof(int32_t) - 1);
+      memcpy(reinterpret_cast<uint8_t *>(&i32) + 1, input_.data() + pos_, sizeof(int32_t) - 1);
       memrev32ifbe(&i32);
       i32 >>= 8;
       setPreEntryLen(4);  // 3byte for encoding and 1byte for the prev entry length
@@ -126,7 +124,7 @@ StatusOr<std::string> ZipList::Next() {
 StatusOr<std::vector<std::string>> ZipList::Entries() {
   GET_OR_RET(peekOK(zlHeaderSize));
   // ignore 8 bytes of total bytes and tail of zip list
-  auto zl_len = intrev16ifbe(*reinterpret_cast<const uint16_t*>(input_.data() + 8));
+  auto zl_len = intrev16ifbe(*reinterpret_cast<const uint16_t *>(input_.data() + 8));
   pos_ += zlHeaderSize;
 
   std::vector<std::string> entries;
@@ -152,3 +150,72 @@ Status ZipList::peekOK(size_t n) {
 }
 
 uint32_t ZipList::getEncodedLengthSize(uint32_t len) { return len < ZipListBigLen ? 1 : 5; }
+
+uint32_t ZipList::ZipStorePrevEntryLengthLarge(unsigned char *p, size_t zl_size, unsigned int len) {
+  uint32_t u32 = 0;
+  if (p != nullptr) {
+    p[0] = ZipListBigLen;
+    u32 = len;
+    assert(zl_size >= 1 + sizeof(uint32_t) + zlHeaderSize);
+    memcpy(p + 1, &u32, sizeof(u32));
+    memrev32ifbe(p + 1);
+  }
+  return 1 + sizeof(uint32_t);
+}
+
+uint32_t ZipList::ZipStorePrevEntryLength(unsigned char *p, size_t zl_size, unsigned int len) {
+  if (p == nullptr) {
+    return (len < ZipListBigLen) ? 1 : sizeof(uint32_t) + 1;
+  }
+  if (len < ZipListBigLen) {
+    p[0] = len;
+    return 1;
+  }
+  return ZipStorePrevEntryLengthLarge(p, zl_size, len);
+}
+
+uint32_t ZipList::ZipStoreEntryEncoding(unsigned char *p, size_t zl_size, unsigned int rawlen) {
+  unsigned char len = 1, buf[5];
+
+  /* Although encoding is given it may not be set for strings,
+   * so we determine it here using the raw length. */
+  if (rawlen <= 0x3f) {
+    if (!p) return len;
+    buf[0] = ZIP_STR_06B | rawlen;
+  } else if (rawlen <= 0x3fff) {
+    len += 1;
+    if (!p) return len;
+    buf[0] = ZIP_STR_14B | ((rawlen >> 8) & 0x3f);
+    buf[1] = rawlen & 0xff;
+  } else {
+    len += 4;
+    if (!p) return len;
+    buf[0] = ZIP_STR_32B;
+    buf[1] = (rawlen >> 24) & 0xff;
+    buf[2] = (rawlen >> 16) & 0xff;
+    buf[3] = (rawlen >> 8) & 0xff;
+    buf[4] = rawlen & 0xff;
+  }
+  assert(zl_size >= zlHeaderSize + len);
+  /* Store this length at p. */
+  memcpy(p, buf, len);
+  return len;
+}
+
+void ZipList::SetZipListBytes(unsigned char *zl, size_t zl_size, uint32_t value) {
+  assert(zl_size >= sizeof(uint32_t));
+  memcpy(zl, &value, sizeof(uint32_t));
+}
+void ZipList::SetZipListTailOffset(unsigned char *zl, size_t zl_size, uint32_t value) {
+  assert(zl_size >= sizeof(uint32_t) * 2);
+  memcpy(zl + sizeof(uint32_t), &value, sizeof(uint32_t));
+}
+void ZipList::SetZipListLength(unsigned char *zl, size_t zl_size, uint16_t value) {
+  assert(zl_size >= sizeof(uint32_t) * 2 + sizeof(uint16_t));
+  memcpy(zl + sizeof(uint32_t) * 2, &value, sizeof(uint16_t));
+}
+
+unsigned char *ZipList::GetZipListEntryHead(unsigned char *zl, size_t zl_size) {
+  assert(zl_size >= zlHeaderSize);
+  return ((zl) + zlHeaderSize);
+}
diff --git a/src/storage/rdb_ziplist.h b/src/storage/rdb_ziplist.h
@@ -25,13 +25,25 @@
 
 #include "common/status.h"
 
+constexpr const int zlHeaderSize = 10;
+constexpr const int zlEndSize = 1;
+constexpr const uint8_t ZipListBigLen = 0xFE;
+constexpr const uint8_t zlEnd = 0xFF;
+
 class ZipList {
  public:
   explicit ZipList(std::string_view input) : input_(input){};
   ~ZipList() = default;
 
   StatusOr<std::string> Next();
   StatusOr<std::vector<std::string>> Entries();
+  static uint32_t ZipStorePrevEntryLengthLarge(unsigned char *p, size_t zl_size, unsigned int len);
+  static uint32_t ZipStorePrevEntryLength(unsigned char *p, size_t zl_size, unsigned int len);
+  static uint32_t ZipStoreEntryEncoding(unsigned char *p, size_t zl_size, unsigned int rawlen);
+  static void SetZipListBytes(unsigned char *zl, size_t zl_size, uint32_t value);
+  static void SetZipListTailOffset(unsigned char *zl, size_t zl_size, uint32_t value);
+  static void SetZipListLength(unsigned char *zl, size_t zl_size, uint16_t value);
+  static unsigned char *GetZipListEntryHead(unsigned char *zl, size_t zl_size);
 
  private:
   std::string_view input_;

diff --git a/tests/gocase/unit/dump/dump_test.go b/tests/gocase/unit/dump/dump_test.go
@@ -113,10 +113,21 @@ func TestDump_List(t *testing.T) {
 	require.NoError(t, rdb.RPush(ctx, key, elements).Err())
 	serialized, err := rdb.Dump(ctx, key).Result()
 	require.NoError(t, err)
+	require.Equal(t, "\x0e\x03\x15\x15\x00\x00\x00\n\x00\x00\x00\x01\x00\x00\bkvrocks1\xff\x15\x15\x00\x00\x00\n\x00\x00\x00\x01\x00\x00\bkvrocks2\xff\x15\x15\x00\x00\x00\n\x00\x00\x00\x01\x00\x00\bkvrocks3\xff\x06\x00u\xc7\x19h\x1da\xd0\xd8", serialized)
 
 	restoredKey := fmt.Sprintf("restore_%s", key)
 	require.NoError(t, rdb.RestoreReplace(ctx, restoredKey, 0, serialized).Err())
 	require.EqualValues(t, elements, rdb.LRange(ctx, restoredKey, 0, -1).Val())
+
+	//test special case
+	elements = []string{"A", " ", "", util.RandString(0, 4000, util.Alpha)}
+	require.NoError(t, rdb.Del(ctx, key).Err())
+	require.NoError(t, rdb.RPush(ctx, key, elements).Err())
+	serialized, err = rdb.Dump(ctx, key).Result()
+	require.NoError(t, err)
+
+	require.NoError(t, rdb.RestoreReplace(ctx, restoredKey, 0, serialized).Err())
+	require.EqualValues(t, elements, rdb.LRange(ctx, restoredKey, 0, -1).Val())
 }
 
 func TestDump_Set(t *testing.T) {