Add the support of the load RDB command

src/commands/cmd_server.cc

@@ -22,6 +22,7 @@
 #include "commander.h"
 #include "commands/scan_base.h"
 #include "common/io_util.h"
+#include "common/rdb_stream.h"
 #include "config/config.h"
 #include "error_constants.h"
 #include "server/redis_connection.h"
@@ -1059,8 +1060,9 @@ class CommandRestore : public Commander {
       ttl_ms_ -= now;
     }
 
-    RDB rdb(svr->storage, conn->GetNamespace(), args_[3]);
-    auto s = rdb.Restore(args_[1], ttl_ms_);
+    auto stream_ptr = std::make_shared<RdbStringStream>(args_[3]);
+    RDB rdb(svr->storage, conn->GetNamespace(), stream_ptr);
+    auto s = rdb.Restore(args_[1], args_[3], ttl_ms_);
     if (!s.IsOK()) return {Status::RedisExecErr, s.Msg()};
     *output = redis::SimpleString("OK");
     return Status::OK();
@@ -1072,6 +1074,46 @@ class CommandRestore : public Commander {
   uint64_t ttl_ms_ = 0;
 };
 
+// command format: rdb load <path> [NX]  [DB index]
+class CommandRdb : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    CommandParser parser(args, 3);
+    std::string_view ttl_flag, set_flag;
+    while (parser.Good()) {
+      if (parser.EatEqICase("NX")) {
+        is_nx_ = true;
+      } else if (parser.EatEqICase("DB")) {
+        db_index_ = GET_OR_RET(parser.TakeInt<uint32_t>());
+      } else {
+        return {Status::RedisParseErr, errInvalidSyntax};
+      }
+    }
+
+    return Status::OK();
+  }
+
+  Status Execute(Server *svr, Connection *conn, std::string *output) override {
+    rocksdb::Status db_status;
+    redis::Database redis(svr->storage, conn->GetNamespace());
+    auto type = args_[1];
+    auto path = args_[2];
+
+    auto stream_ptr = std::make_shared<RdbFileStream>(path);
+    GET_OR_RET(stream_ptr->Open());
+
+    RDB rdb(svr->storage, conn->GetNamespace(), stream_ptr);
+    GET_OR_RET(rdb.LoadRdb(db_index_, is_nx_));
+
+    *output = redis::SimpleString("OK");
+    return Status::OK();
+  }
+
+ private:
+  bool is_nx_ = false;
+  uint32_t db_index_ = 0;
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandPing>("ping", -1, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandSelect>("select", 2, "read-only", 0, 0, 0),
@@ -1105,6 +1147,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loadin
                         MakeCmdAttr<CommandLastSave>("lastsave", 1, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFlushBackup>("flushbackup", 1, "read-only no-script", 0, 0, 0),
                         MakeCmdAttr<CommandSlaveOf>("slaveof", 3, "read-only exclusive no-script", 0, 0, 0),
-                        MakeCmdAttr<CommandStats>("stats", 1, "read-only", 0, 0, 0), )
+                        MakeCmdAttr<CommandStats>("stats", 1, "read-only", 0, 0, 0),
+                        MakeCmdAttr<CommandRdb>("rdb", -3, "write exclusive", 0, 0, 0), )
 
 }  // namespace redis

src/common/rdb_stream.cc

@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include "rdb_stream.h"
+
+#include "fmt/format.h"
+#include "vendor/crc64.h"
+#include "vendor/endianconv.h"
+
+StatusOr<size_t> RdbStringStream::Read(char *buf, size_t n) {
+  if (pos_ + n > input_.size()) {
+    return {Status::NotOK, "unexpected EOF"};
+  }
+  memcpy(buf, input_.data() + pos_, n);
+  pos_ += n;
+  return n;
+}
+
+StatusOr<uint64_t> RdbStringStream::GetCheckSum() const {
+  if (input_.size() < 8) {
+    return {Status::NotOK, "invalid payload length"};
+  }
+  uint64_t crc = crc64(0, reinterpret_cast<const unsigned char *>(input_.data()), input_.size() - 8);
+  memrev64ifbe(&crc);
+  return crc;
+}
+
+Status RdbFileStream::Open() {
+  ifs_.open(file_name_, std::ifstream::in | std::ifstream::binary);
+  if (!ifs_.is_open()) {
+    return {Status::NotOK, fmt::format("failed to open rdb file: '{}': {}", file_name_, strerror(errno))};
+  }
+
+  return Status::OK();
+}
+
+StatusOr<size_t> RdbFileStream::Read(char *buf, size_t len) {
+  size_t n = 0;
+  while (len) {
+    size_t read_bytes = max_read_chunk_size_ < len ? max_read_chunk_size_ : len;
+    ifs_.read(buf, static_cast<std::streamsize>(read_bytes));
+    if (!ifs_.good()) {
+      return Status(Status::NotOK, fmt::format("read failed: {}:", strerror(errno)));
+    }
+    check_sum_ = crc64(check_sum_, (const unsigned char *)buf, read_bytes);
+    buf = (char *)buf + read_bytes;
+    len -= read_bytes;
+    total_read_bytes_ += read_bytes;
+    n += read_bytes;
+  }
+
+  return n;
+}

src/common/rdb_stream.h

@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#pragma once
+
+#include <stdint.h>
+
+#include <fstream>
+#include <string>
+
+#include "status.h"
+
+class RdbStream {
+ public:
+  RdbStream() = default;
+  virtual ~RdbStream() = default;
+
+  virtual StatusOr<size_t> Read(char *buf, size_t len) = 0;
+  virtual StatusOr<uint64_t> GetCheckSum() const = 0;
+  StatusOr<uint8_t> ReadByte() {
+    uint8_t value = 0;
+    auto s = Read(reinterpret_cast<char *>(&value), 1);
+    if (!s.IsOK()) {
+      return s;
+    }
+    return value;
+  }
+};
+
+class RdbStringStream : public RdbStream {
+ public:
+  explicit RdbStringStream(std::string_view input) : input_(input){};
+  RdbStringStream(const RdbStringStream &) = delete;
+  RdbStringStream &operator=(const RdbStringStream &) = delete;
+  ~RdbStringStream() override = default;
+
+  StatusOr<size_t> Read(char *buf, size_t len) override;
+  StatusOr<uint64_t> GetCheckSum() const override;
+
+ private:
+  std::string input_;
+  size_t pos_ = 0;
+};
+
+class RdbFileStream : public RdbStream {
+ public:
+  explicit RdbFileStream(std::string file_name, size_t chunk_size = 1024 * 1024)
+      : file_name_(std::move(file_name)), check_sum_(0), total_read_bytes_(0), max_read_chunk_size_(chunk_size){};
+  RdbFileStream(const RdbFileStream &) = delete;
+  RdbFileStream &operator=(const RdbFileStream &) = delete;
+  ~RdbFileStream() override = default;
+
+  Status Open();
+  StatusOr<size_t> Read(char *buf, size_t len) override;
+  StatusOr<uint64_t> GetCheckSum() const override { return check_sum_; }
+
+ private:
+  std::ifstream ifs_;
+  std::string file_name_;
+  uint64_t check_sum_;
+  size_t total_read_bytes_;
+  size_t max_read_chunk_size_;  // maximum single read chunk size
+};