chore: factor out rdb_load utilities into separate files

src/server/CMakeLists.txt

@@ -45,11 +45,12 @@ endif()
 
 add_library(dragonfly_lib bloom_family.cc
             config_registry.cc conn_context.cc debugcmd.cc dflycmd.cc engine_shard.cc
-            engine_shard_set.cc family_utils.cc
+            engine_shard_set.cc error.cc family_utils.cc
             generic_family.cc hset_family.cc http_api.cc json_family.cc
             list_family.cc main_service.cc memory_cmd.cc rdb_load.cc rdb_save.cc replica.cc
             protocol_client.cc
             snapshot.cc script_mgr.cc server_family.cc
+            detail/decompress.cc
             detail/save_stages_controller.cc
             detail/snapshot_storage.cc
             set_family.cc stream_family.cc string_family.cc

src/server/detail/decompress.cc

@@ -0,0 +1,173 @@
+// Copyright 2024, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/detail/decompress.h"
+
+#include <lz4frame.h>
+#include <zstd.h>
+
+#include "base/logging.h"
+#include "server/error.h"
+#include "server/rdb_extensions.h"
+
+namespace dfly {
+
+namespace detail {
+
+using io::IoBuf;
+using rdb::errc;
+using namespace std;
+
+inline auto Unexpected(errc ev) {
+  return nonstd::make_unexpected(RdbError(ev));
+}
+
+class ZstdDecompress : public DecompressImpl {
+ public:
+  ZstdDecompress() {
+    dctx_ = ZSTD_createDCtx();
+  }
+  ~ZstdDecompress() {
+    ZSTD_freeDCtx(dctx_);
+  }
+
+  io::Result<io::IoBuf*> Decompress(std::string_view str);
+
+ private:
+  ZSTD_DCtx* dctx_;
+};
+
+io::Result<io::IoBuf*> ZstdDecompress::Decompress(std::string_view str) {
+  // Prepare membuf memory to uncompressed string.
+  auto uncomp_size = ZSTD_getFrameContentSize(str.data(), str.size());
+  if (uncomp_size == ZSTD_CONTENTSIZE_UNKNOWN) {
+    LOG(ERROR) << "Zstd compression missing frame content size";
+    return Unexpected(errc::invalid_encoding);
+  }
+  if (uncomp_size == ZSTD_CONTENTSIZE_ERROR) {
+    LOG(ERROR) << "Invalid ZSTD compressed string";
+    return Unexpected(errc::invalid_encoding);
+  }
+
+  uncompressed_mem_buf_.Reserve(uncomp_size + 1);
+
+  // Uncompress string to membuf
+  IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
+  if (dest.size() < uncomp_size) {
+    return Unexpected(errc::out_of_memory);
+  }
+  size_t const d_size =
+      ZSTD_decompressDCtx(dctx_, dest.data(), dest.size(), str.data(), str.size());
+  if (d_size == 0 || d_size != uncomp_size) {
+    LOG(ERROR) << "Invalid ZSTD compressed string";
+    return Unexpected(errc::rdb_file_corrupted);
+  }
+  uncompressed_mem_buf_.CommitWrite(d_size);
+
+  // Add opcode of compressed blob end to membuf.
+  dest = uncompressed_mem_buf_.AppendBuffer();
+  if (dest.size() < 1) {
+    return Unexpected(errc::out_of_memory);
+  }
+  dest[0] = RDB_OPCODE_COMPRESSED_BLOB_END;
+  uncompressed_mem_buf_.CommitWrite(1);
+
+  return &uncompressed_mem_buf_;
+}
+
+class Lz4Decompress : public DecompressImpl {
+ public:
+  Lz4Decompress() {
+    auto result = LZ4F_createDecompressionContext(&dctx_, LZ4F_VERSION);
+    CHECK(!LZ4F_isError(result));
+  }
+  ~Lz4Decompress() {
+    auto result = LZ4F_freeDecompressionContext(dctx_);
+    CHECK(!LZ4F_isError(result));
+  }
+
+  io::Result<base::IoBuf*> Decompress(std::string_view str);
+
+ private:
+  LZ4F_dctx* dctx_;
+};
+
+io::Result<base::IoBuf*> Lz4Decompress::Decompress(std::string_view data) {
+  LZ4F_frameInfo_t frame_info;
+  size_t frame_size = data.size();
+
+  // Get content size from frame data
+  size_t consumed = frame_size;  // The nb of bytes consumed from data will be written into consumed
+  size_t res = LZ4F_getFrameInfo(dctx_, &frame_info, data.data(), &consumed);
+  if (LZ4F_isError(res)) {
+    LOG(ERROR) << "LZ4F_getFrameInfo failed with error " << LZ4F_getErrorName(res);
+    return Unexpected(errc::rdb_file_corrupted);
+    ;
+  }
+
+  if (frame_info.contentSize == 0) {
+    LOG(ERROR) << "Missing frame content size";
+    return Unexpected(errc::rdb_file_corrupted);
+  }
+
+  // reserve place for uncompressed data and end opcode
+  size_t reserve = frame_info.contentSize + 1;
+  uncompressed_mem_buf_.Reserve(reserve);
+  IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
+  if (dest.size() < reserve) {
+    return Unexpected(errc::out_of_memory);
+  }
+
+  // Uncompress data to membuf
+  string_view src = data.substr(consumed);
+  size_t src_size = src.size();
+
+  size_t ret = 1;
+  while (ret != 0) {
+    IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
+    size_t dest_capacity = dest.size();
+
+    // It will read up to src_size bytes from src,
+    // and decompress data into dest, of capacity dest_capacity
+    // The nb of bytes consumed from src will be written into src_size
+    // The nb of bytes decompressed into dest will be written into dest_capacity
+    ret = LZ4F_decompress(dctx_, dest.data(), &dest_capacity, src.data(), &src_size, nullptr);
+    if (LZ4F_isError(ret)) {
+      LOG(ERROR) << "LZ4F_decompress failed with error " << LZ4F_getErrorName(ret);
+      return Unexpected(errc::rdb_file_corrupted);
+    }
+    consumed += src_size;
+
+    uncompressed_mem_buf_.CommitWrite(dest_capacity);
+    src = src.substr(src_size);
+    src_size = src.size();
+  }
+  if (consumed != frame_size) {
+    return Unexpected(errc::rdb_file_corrupted);
+  }
+  if (uncompressed_mem_buf_.InputLen() != frame_info.contentSize) {
+    return Unexpected(errc::rdb_file_corrupted);
+  }
+
+  // Add opcode of compressed blob end to membuf.
+  dest = uncompressed_mem_buf_.AppendBuffer();
+  if (dest.size() < 1) {
+    return Unexpected(errc::out_of_memory);
+  }
+  dest[0] = RDB_OPCODE_COMPRESSED_BLOB_END;
+  uncompressed_mem_buf_.CommitWrite(1);
+
+  return &uncompressed_mem_buf_;
+}
+
+unique_ptr<DecompressImpl> DecompressImpl::CreateLZ4() {
+  return make_unique<Lz4Decompress>();
+}
+
+unique_ptr<DecompressImpl> DecompressImpl::CreateZstd() {
+  return make_unique<ZstdDecompress>();
+}
+
+}  // namespace detail
+}  // namespace dfly

src/server/detail/decompress.h

@@ -0,0 +1,32 @@
+// Copyright 2024, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+#pragma once
+
+#include <memory>
+
+#include "io/io.h"
+#include "io/io_buf.h"
+
+namespace dfly {
+
+namespace detail {
+
+class DecompressImpl {
+ public:
+  static std::unique_ptr<DecompressImpl> CreateLZ4();
+  static std::unique_ptr<DecompressImpl> CreateZstd();
+
+  DecompressImpl() : uncompressed_mem_buf_{1U << 14} {
+  }
+  virtual ~DecompressImpl() {
+  }
+
+  virtual io::Result<io::IoBuf*> Decompress(std::string_view str) = 0;
+
+ protected:
+  io::IoBuf uncompressed_mem_buf_;
+};
+
+}  // namespace detail
+}  // namespace dfly

src/server/error.cc

@@ -0,0 +1,57 @@
+// Copyright 2024, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/error.h"
+
+#include <absl/strings/str_cat.h>
+
+using namespace std;
+
+namespace dfly {
+namespace rdb {
+
+class error_category : public std::error_category {
+ public:
+  const char* name() const noexcept final {
+    return "dragonfly.rdbload";
+  }
+
+  string message(int ev) const final;
+
+  error_condition default_error_condition(int ev) const noexcept final;
+
+  bool equivalent(int ev, const error_condition& condition) const noexcept final {
+    return condition.value() == ev && &condition.category() == this;
+  }
+
+  bool equivalent(const error_code& error, int ev) const noexcept final {
+    return error.value() == ev && &error.category() == this;
+  }
+};
+
+string error_category::message(int ev) const {
+  switch (ev) {
+    case errc::wrong_signature:
+      return "Wrong signature while trying to load from rdb file";
+    case errc::out_of_memory:
+      return "Out of memory, or used memory is too high";
+    default:
+      return absl::StrCat("Internal error when loading RDB file ", ev);
+      break;
+  }
+}
+
+error_condition error_category::default_error_condition(int ev) const noexcept {
+  return error_condition{ev, *this};
+}
+
+static error_category rdb_category;
+
+}  // namespace rdb
+
+error_code RdbError(rdb::errc ev) {
+  return error_code{static_cast<int>(ev), rdb::rdb_category};
+}
+
+}  // namespace dfly

src/server/error.h

@@ -79,4 +79,6 @@ enum errc {
 
 }  // namespace rdb
 
+std::error_code RdbError(rdb::errc ev);
+
 }  // namespace dfly