mirrors/dragonfly
1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2025-05-11 10:25:47 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

bug(rdb loader): When reading from zstd uncompressed buf skip ensure … (#525)

Browse source 
* bug(rdb loader): When reading from zstd uncompressed buf skip ensure read flow

Signed-off-by: adi_holden <adi@dragonflydb.io>
This commit is contained in:
adiholden 2022-12-04 11:14:24 +02:00 • committed by GitHub
parent 74d1839f97
commit c7974a4e80
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 38 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

21
src/server/rdb_load.cc
View file

@ -41,8 +41,10 @@ ABSL_DECLARE_FLAG(bool, use_set2);
#define SET_OR_RETURN(expr, dest) \
do { \
auto exp_val = (expr); \
if (!exp_val) \
if (!exp_val) { \
VLOG(1) << "Error while calling " #expr; \
return exp_val.error(); \
} \
dest = exp_val.value(); \
} while (0)
@ -230,6 +232,7 @@ io::Result<base::IoBuf*> ZstdDecompressImpl::Decompress(std::string_view str) {
LOG(ERROR) << "Invalid ZSTD compressed string";
return Unexpected(errc::invalid_encoding);
}
uncompressed_mem_buf_.Reserve(uncomp_size + 1);
// Uncompress string to membuf
@ -1705,6 +1708,19 @@ error_code RdbLoader::Load(io::Source* src) {
return kOk;
}
std::error_code RdbLoaderBase::EnsureRead(size_t min_sz) {
// In the flow of reading compressed data, we store the uncompressed data to in uncompressed
// buffer. When parsing entries we call ensure read with 9 bytes to read the length of key/value.
// If the key/value is very small (less than 9 bytes) the remainded data in uncompressed buffer
// might contain less than 9 bytes. We need to make sure that we dont read from sink to the
// uncompressed buffer and therefor in this flow we return here.
if (mem_buf_ != &origin_mem_buf_)
return std::error_code{};
if (mem_buf_->InputLen() >= min_sz)
return std::error_code{};
return EnsureReadInternal(min_sz);
}
error_code RdbLoaderBase::EnsureReadInternal(size_t min_sz) {
DCHECK_LT(mem_buf_->InputLen(), min_sz);
@ -1792,7 +1808,8 @@ error_code RdbLoaderBase::HandleCompressedBlob() {
}
error_code RdbLoaderBase::HandleCompressedBlobFinish() {
// TODO validate that all uncompressed data was fetched
CHECK_NE(&origin_mem_buf_, mem_buf_);
CHECK_EQ(mem_buf_->InputLen(), size_t(0));
mem_buf_ = &origin_mem_buf_;
return kOk;
}

7
src/server/rdb_load.h
View file

@ -132,12 +132,7 @@ class RdbLoaderBase {
static size_t StrLen(const RdbVariant& tset);
std::error_code EnsureRead(size_t min_sz) {
if (mem_buf_->InputLen() >= min_sz)
return std::error_code{};
return EnsureReadInternal(min_sz);
}
std::error_code EnsureRead(size_t min_sz);
std::error_code EnsureReadInternal(size_t min_sz);

13
src/server/rdb_test.cc
View file

@ -159,6 +159,19 @@ TEST_F(RdbTest, ComressionModeSaveDragonflyAndReload) {
}
}
TEST_F(RdbTest, RdbLoaderOnReadCompressedDataShouldNotEnterEnsureReadFlow) {
SetFlag(&FLAGS_compression_mode, 2);
for (int i = 0; i < 1000; ++i) {
Run({"set", StrCat(i), "1"});
}
RespExpr resp = Run({"save", "df"});
ASSERT_EQ(resp, "OK");
auto save_info = service_->server_family().GetLastSaveInfo();
resp = Run({"debug", "load", save_info->file_name});
ASSERT_EQ(resp, "OK");
}
TEST_F(RdbTest, Reload) {
absl::FlagSaver fs;