Skip to content

Commit 53d6b64

Browse files
romangeromange
authored
chore: factor out rdb_load utilities into separate files (#4315)
* chore: factor out rdb_load utilities into separate files rdb_load.cc is huge and contains many auxillary classes. This PR moves DecompressImpl and ErrorRdb code into detail/ It also fixes minor bugs around error conditions with de-compression: a. Do not check-fail on invalid opcode and return error_code instead. b. Print correctly LZ4 errors. Signed-off-by: Roman Gershman <roman@dragonflydb.io> * chore: fixes --------- Signed-off-by: Roman Gershman <roman@dragonflydb.io>
1 parent 027eff2 commit 53d6b64

File tree

7 files changed

+278
-203
lines changed

7 files changed

+278
-203
lines changed

‎src/server/CMakeLists.txt

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,12 @@ endif()
4545

4646
add_library(dragonfly_lib bloom_family.cc
4747
config_registry.cc conn_context.cc debugcmd.cc dflycmd.cc engine_shard.cc
48-
engine_shard_set.cc family_utils.cc
48+
engine_shard_set.cc error.cc family_utils.cc
4949
generic_family.cc hset_family.cc http_api.cc json_family.cc
5050
list_family.cc main_service.cc memory_cmd.cc rdb_load.cc rdb_save.cc replica.cc
5151
protocol_client.cc
5252
snapshot.cc script_mgr.cc server_family.cc
53+
detail/decompress.cc
5354
detail/save_stages_controller.cc
5455
detail/snapshot_storage.cc
5556
set_family.cc stream_family.cc string_family.cc

‎src/server/detail/decompress.cc

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
// Copyright 2024, DragonflyDB authors. All rights reserved.
2+
// See LICENSE for licensing terms.
3+
//
4+
5+
#include "server/detail/decompress.h"
6+
7+
#include <lz4frame.h>
8+
#include <zstd.h>
9+
10+
#include "base/logging.h"
11+
#include "server/error.h"
12+
#include "server/rdb_extensions.h"
13+
14+
namespace dfly {
15+
16+
namespace detail {
17+
18+
using io::IoBuf;
19+
using rdb::errc;
20+
using namespace std;
21+
22+
inline auto Unexpected(errc ev) {
23+
return nonstd::make_unexpected(RdbError(ev));
24+
}
25+
26+
class ZstdDecompress : public DecompressImpl {
27+
public:
28+
ZstdDecompress() {
29+
dctx_ = ZSTD_createDCtx();
30+
}
31+
~ZstdDecompress() {
32+
ZSTD_freeDCtx(dctx_);
33+
}
34+
35+
io::Result<io::IoBuf*> Decompress(std::string_view str);
36+
37+
private:
38+
ZSTD_DCtx* dctx_;
39+
};
40+
41+
io::Result<io::IoBuf*> ZstdDecompress::Decompress(std::string_view str) {
42+
// Prepare membuf memory to uncompressed string.
43+
auto uncomp_size = ZSTD_getFrameContentSize(str.data(), str.size());
44+
if (uncomp_size == ZSTD_CONTENTSIZE_UNKNOWN) {
45+
LOG(ERROR) << "Zstd compression missing frame content size";
46+
return Unexpected(errc::invalid_encoding);
47+
}
48+
if (uncomp_size == ZSTD_CONTENTSIZE_ERROR) {
49+
LOG(ERROR) << "Invalid ZSTD compressed string";
50+
return Unexpected(errc::invalid_encoding);
51+
}
52+
53+
uncompressed_mem_buf_.Reserve(uncomp_size + 1);
54+
55+
// Uncompress string to membuf
56+
IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
57+
if (dest.size() < uncomp_size) {
58+
return Unexpected(errc::out_of_memory);
59+
}
60+
size_t const d_size =
61+
ZSTD_decompressDCtx(dctx_, dest.data(), dest.size(), str.data(), str.size());
62+
if (d_size == 0 || d_size != uncomp_size) {
63+
LOG(ERROR) << "Invalid ZSTD compressed string";
64+
return Unexpected(errc::rdb_file_corrupted);
65+
}
66+
uncompressed_mem_buf_.CommitWrite(d_size);
67+
68+
// Add opcode of compressed blob end to membuf.
69+
dest = uncompressed_mem_buf_.AppendBuffer();
70+
if (dest.size() < 1) {
71+
return Unexpected(errc::out_of_memory);
72+
}
73+
dest[0] = RDB_OPCODE_COMPRESSED_BLOB_END;
74+
uncompressed_mem_buf_.CommitWrite(1);
75+
76+
return &uncompressed_mem_buf_;
77+
}
78+
79+
class Lz4Decompress : public DecompressImpl {
80+
public:
81+
Lz4Decompress() {
82+
auto result = LZ4F_createDecompressionContext(&dctx_, LZ4F_VERSION);
83+
CHECK(!LZ4F_isError(result));
84+
}
85+
~Lz4Decompress() {
86+
auto result = LZ4F_freeDecompressionContext(dctx_);
87+
CHECK(!LZ4F_isError(result));
88+
}
89+
90+
io::Result<base::IoBuf*> Decompress(std::string_view str);
91+
92+
private:
93+
LZ4F_dctx* dctx_;
94+
};
95+
96+
io::Result<base::IoBuf*> Lz4Decompress::Decompress(std::string_view data) {
97+
LZ4F_frameInfo_t frame_info;
98+
size_t frame_size = data.size();
99+
100+
// Get content size from frame data
101+
size_t consumed = frame_size; // The nb of bytes consumed from data will be written into consumed
102+
size_t res = LZ4F_getFrameInfo(dctx_, &frame_info, data.data(), &consumed);
103+
if (LZ4F_isError(res)) {
104+
LOG(ERROR) << "LZ4F_getFrameInfo failed with error " << LZ4F_getErrorName(res);
105+
return Unexpected(errc::rdb_file_corrupted);
106+
;
107+
}
108+
109+
if (frame_info.contentSize == 0) {
110+
LOG(ERROR) << "Missing frame content size";
111+
return Unexpected(errc::rdb_file_corrupted);
112+
}
113+
114+
// reserve place for uncompressed data and end opcode
115+
size_t reserve = frame_info.contentSize + 1;
116+
uncompressed_mem_buf_.Reserve(reserve);
117+
IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
118+
if (dest.size() < reserve) {
119+
return Unexpected(errc::out_of_memory);
120+
}
121+
122+
// Uncompress data to membuf
123+
string_view src = data.substr(consumed);
124+
size_t src_size = src.size();
125+
126+
size_t ret = 1;
127+
while (ret != 0) {
128+
IoBuf::Bytes dest = uncompressed_mem_buf_.AppendBuffer();
129+
size_t dest_capacity = dest.size();
130+
131+
// It will read up to src_size bytes from src,
132+
// and decompress data into dest, of capacity dest_capacity
133+
// The nb of bytes consumed from src will be written into src_size
134+
// The nb of bytes decompressed into dest will be written into dest_capacity
135+
ret = LZ4F_decompress(dctx_, dest.data(), &dest_capacity, src.data(), &src_size, nullptr);
136+
if (LZ4F_isError(ret)) {
137+
LOG(ERROR) << "LZ4F_decompress failed with error " << LZ4F_getErrorName(ret);
138+
return Unexpected(errc::rdb_file_corrupted);
139+
}
140+
consumed += src_size;
141+
142+
uncompressed_mem_buf_.CommitWrite(dest_capacity);
143+
src = src.substr(src_size);
144+
src_size = src.size();
145+
}
146+
if (consumed != frame_size) {
147+
return Unexpected(errc::rdb_file_corrupted);
148+
}
149+
if (uncompressed_mem_buf_.InputLen() != frame_info.contentSize) {
150+
return Unexpected(errc::rdb_file_corrupted);
151+
}
152+
153+
// Add opcode of compressed blob end to membuf.
154+
dest = uncompressed_mem_buf_.AppendBuffer();
155+
if (dest.size() < 1) {
156+
return Unexpected(errc::out_of_memory);
157+
}
158+
dest[0] = RDB_OPCODE_COMPRESSED_BLOB_END;
159+
uncompressed_mem_buf_.CommitWrite(1);
160+
161+
return &uncompressed_mem_buf_;
162+
}
163+
164+
unique_ptr<DecompressImpl> DecompressImpl::CreateLZ4() {
165+
return make_unique<Lz4Decompress>();
166+
}
167+
168+
unique_ptr<DecompressImpl> DecompressImpl::CreateZstd() {
169+
return make_unique<ZstdDecompress>();
170+
}
171+
172+
} // namespace detail
173+
} // namespace dfly

‎src/server/detail/decompress.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// Copyright 2024, DragonflyDB authors. All rights reserved.
2+
// See LICENSE for licensing terms.
3+
//
4+
#pragma once
5+
6+
#include <memory>
7+
8+
#include "io/io.h"
9+
#include "io/io_buf.h"
10+
11+
namespace dfly {
12+
13+
namespace detail {
14+
15+
class DecompressImpl {
16+
public:
17+
static std::unique_ptr<DecompressImpl> CreateLZ4();
18+
static std::unique_ptr<DecompressImpl> CreateZstd();
19+
20+
DecompressImpl() : uncompressed_mem_buf_{1U << 14} {
21+
}
22+
virtual ~DecompressImpl() {
23+
}
24+
25+
virtual io::Result<io::IoBuf*> Decompress(std::string_view str) = 0;
26+
27+
protected:
28+
io::IoBuf uncompressed_mem_buf_;
29+
};
30+
31+
} // namespace detail
32+
} // namespace dfly

‎src/server/error.cc

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
// Copyright 2024, DragonflyDB authors. All rights reserved.
2+
// See LICENSE for licensing terms.
3+
//
4+
5+
#include "server/error.h"
6+
7+
#include <absl/strings/str_cat.h>
8+
9+
using namespace std;
10+
11+
namespace dfly {
12+
namespace rdb {
13+
14+
class error_category : public std::error_category {
15+
public:
16+
const char* name() const noexcept final {
17+
return "dragonfly.rdbload";
18+
}
19+
20+
string message(int ev) const final;
21+
22+
error_condition default_error_condition(int ev) const noexcept final;
23+
24+
bool equivalent(int ev, const error_condition& condition) const noexcept final {
25+
return condition.value() == ev && &condition.category() == this;
26+
}
27+
28+
bool equivalent(const error_code& error, int ev) const noexcept final {
29+
return error.value() == ev && &error.category() == this;
30+
}
31+
};
32+
33+
string error_category::message(int ev) const {
34+
switch (ev) {
35+
case errc::wrong_signature:
36+
return "Wrong signature while trying to load from rdb file";
37+
case errc::out_of_memory:
38+
return "Out of memory, or used memory is too high";
39+
default:
40+
return absl::StrCat("Internal error when loading RDB file ", ev);
41+
break;
42+
}
43+
}
44+
45+
error_condition error_category::default_error_condition(int ev) const noexcept {
46+
return error_condition{ev, *this};
47+
}
48+
49+
static error_category rdb_category;
50+
51+
} // namespace rdb
52+
53+
error_code RdbError(rdb::errc ev) {
54+
return error_code{static_cast<int>(ev), rdb::rdb_category};
55+
}
56+
57+
} // namespace dfly

‎src/server/error.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,6 @@ enum errc {
7979

8080
} // namespace rdb
8181

82+
std::error_code RdbError(rdb::errc ev);
83+
8284
} // namespace dfly

0 commit comments

Comments
 (0)