From ebb7bf6bbd0d3443240b6120a1991b922b58e81a Mon Sep 17 00:00:00 2001 From: Myth Date: Wed, 24 Jan 2024 10:30:26 +0800 Subject: [PATCH 1/2] Ignore max-db-size limit when deleting data or writing aux information --- src/storage/redis_db.cc | 2 +- src/storage/storage.cc | 18 ++++++++++-------- src/storage/storage.h | 6 ++++-- src/types/redis_hash.cc | 2 +- src/types/redis_json.cc | 2 +- src/types/redis_list.cc | 4 ++-- src/types/redis_set.cc | 2 +- src/types/redis_sortedint.cc | 2 +- src/types/redis_stream.cc | 4 ++-- src/types/redis_zset.cc | 2 +- 10 files changed, 24 insertions(+), 20 deletions(-) diff --git a/src/storage/redis_db.cc b/src/storage/redis_db.cc index 63e8ef571cc..150a776577b 100644 --- a/src/storage/redis_db.cc +++ b/src/storage/redis_db.cc @@ -200,7 +200,7 @@ rocksdb::Status Database::MDel(const std::vector &keys, uint64_t *deleted if (*deleted_cnt == 0) return rocksdb::Status::OK(); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Database::Exists(const std::vector &keys, int *ret) { diff --git a/src/storage/storage.cc b/src/storage/storage.cc index a600813fae8..110f82d8ba9 100644 --- a/src/storage/storage.cc +++ b/src/storage/storage.cc @@ -601,16 +601,18 @@ void Storage::MultiGet(const rocksdb::ReadOptions &options, rocksdb::ColumnFamil } } -rocksdb::Status Storage::Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates) { +rocksdb::Status Storage::Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, + bool ignore_max_db_size) { if (is_txn_mode_) { // The batch won't be flushed until the transaction was committed or rollback return rocksdb::Status::OK(); } - return writeToDB(options, updates); + return writeToDB(options, updates, ignore_max_db_size); } -rocksdb::Status Storage::writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates) { - if (db_size_limit_reached_) { +rocksdb::Status Storage::writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, + bool ignore_max_db_size) { + if (!ignore_max_db_size && db_size_limit_reached_) { return rocksdb::Status::SpaceLimit(); } @@ -626,7 +628,7 @@ rocksdb::Status Storage::Delete(const rocksdb::WriteOptions &options, rocksdb::C const rocksdb::Slice &key) { auto batch = GetWriteBatchBase(); batch->Delete(cf_handle, key); - return Write(options, batch->GetWriteBatch()); + return Write(options, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Storage::DeleteRange(const std::string &first_key, const std::string &last_key) { @@ -642,7 +644,7 @@ rocksdb::Status Storage::DeleteRange(const std::string &first_key, const std::st return s; } - return Write(write_opts_, batch->GetWriteBatch()); + return Write(write_opts_, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Storage::FlushScripts(const rocksdb::WriteOptions &options, rocksdb::ColumnFamilyHandle *cf_handle) { @@ -657,7 +659,7 @@ rocksdb::Status Storage::FlushScripts(const rocksdb::WriteOptions &options, rock return s; } - return Write(options, batch->GetWriteBatch()); + return Write(options, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } Status Storage::ReplicaApplyWriteBatch(std::string &&raw_batch) { @@ -816,7 +818,7 @@ Status Storage::WriteToPropagateCF(const std::string &key, const std::string &va auto batch = GetWriteBatchBase(); auto cf = GetCFHandle(kPropagateColumnFamilyName); batch->Put(cf, key, value); - auto s = Write(write_opts_, batch->GetWriteBatch()); + auto s = Write(write_opts_, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); if (!s.ok()) { return {Status::NotOK, s.ToString()}; } diff --git a/src/storage/storage.h b/src/storage/storage.h index 7e37d82f7ea..f457618ad53 100644 --- a/src/storage/storage.h +++ b/src/storage/storage.h @@ -147,7 +147,8 @@ class Storage { rocksdb::Iterator *NewIterator(const rocksdb::ReadOptions &options, rocksdb::ColumnFamilyHandle *column_family); rocksdb::Iterator *NewIterator(const rocksdb::ReadOptions &options); - [[nodiscard]] rocksdb::Status Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates); + [[nodiscard]] rocksdb::Status Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, + bool ignore_max_db_size = false); const rocksdb::WriteOptions &DefaultWriteOptions() { return write_opts_; } rocksdb::ReadOptions DefaultScanOptions() const; rocksdb::ReadOptions DefaultMultiGetOptions() const; @@ -263,7 +264,8 @@ class Storage { rocksdb::WriteOptions write_opts_ = rocksdb::WriteOptions(); - rocksdb::Status writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates); + rocksdb::Status writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, + bool ignore_max_db_size = false); void recordKeyspaceStat(const rocksdb::ColumnFamilyHandle *column_family, const rocksdb::Status &s); }; diff --git a/src/types/redis_hash.cc b/src/types/redis_hash.cc index d7a1ad8b7bc..ff0fc195280 100644 --- a/src/types/redis_hash.cc +++ b/src/types/redis_hash.cc @@ -227,7 +227,7 @@ rocksdb::Status Hash::Delete(const Slice &user_key, const std::vector &fi std::string bytes; metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Hash::MSet(const Slice &user_key, const std::vector &field_values, bool nx, diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc index fb3add5feb0..47840ba5082 100644 --- a/src/types/redis_json.cc +++ b/src/types/redis_json.cc @@ -95,7 +95,7 @@ rocksdb::Status Json::del(const Slice &ns_key) { batch->Delete(metadata_cf_handle_, ns_key); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Json::Info(const std::string &user_key, JsonStorageFormat *storage_format) { diff --git a/src/types/redis_list.cc b/src/types/redis_list.cc index ed67007e905..590fafe68b4 100644 --- a/src/types/redis_list.cc +++ b/src/types/redis_list.cc @@ -250,7 +250,7 @@ rocksdb::Status List::Rem(const Slice &user_key, int count, const Slice &elem, u } *removed_cnt = to_delete_indexes.size(); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status List::Insert(const Slice &user_key, const Slice &pivot, const Slice &elem, bool before, int *new_size) { @@ -658,6 +658,6 @@ rocksdb::Status List::Trim(const Slice &user_key, int start, int stop) { std::string bytes; metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } } // namespace redis diff --git a/src/types/redis_set.cc b/src/types/redis_set.cc index 98677e27560..603957fe237 100644 --- a/src/types/redis_set.cc +++ b/src/types/redis_set.cc @@ -122,7 +122,7 @@ rocksdb::Status Set::Remove(const Slice &user_key, const std::vector &mem batch->Delete(metadata_cf_handle_, ns_key); } } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Set::Card(const Slice &user_key, uint64_t *size) { diff --git a/src/types/redis_sortedint.cc b/src/types/redis_sortedint.cc index 6670eb7474f..aa883e5dc1d 100644 --- a/src/types/redis_sortedint.cc +++ b/src/types/redis_sortedint.cc @@ -93,7 +93,7 @@ rocksdb::Status Sortedint::Remove(const Slice &user_key, const std::vectorPut(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status Sortedint::Card(const Slice &user_key, uint64_t *size) { diff --git a/src/types/redis_stream.cc b/src/types/redis_stream.cc index 48fe7928c45..84296e8e741 100644 --- a/src/types/redis_stream.cc +++ b/src/types/redis_stream.cc @@ -550,7 +550,7 @@ rocksdb::Status Stream::DeleteEntries(const Slice &stream_name, const std::vecto batch->Put(metadata_cf_handle_, ns_key, bytes); } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } // If `options` is StreamLenOptions{} the function just returns the number of entries in the stream. @@ -950,7 +950,7 @@ rocksdb::Status Stream::Trim(const Slice &stream_name, const StreamTrimOptions & metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } return rocksdb::Status::OK(); diff --git a/src/types/redis_zset.cc b/src/types/redis_zset.cc index 328182b5518..447c31a6691 100644 --- a/src/types/redis_zset.cc +++ b/src/types/redis_zset.cc @@ -544,7 +544,7 @@ rocksdb::Status ZSet::Remove(const Slice &user_key, const std::vector &me metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); } rocksdb::Status ZSet::Rank(const Slice &user_key, const Slice &member, bool reversed, int *member_rank, From 6417414d93069c599ea6cde9a489583bd3972d48 Mon Sep 17 00:00:00 2001 From: Myth Date: Fri, 26 Jan 2024 11:31:31 +0800 Subject: [PATCH 2/2] Add no-dbsize-check command flag --- src/commands/cmd_hash.cc | 2 +- src/commands/cmd_json.cc | 6 +++--- src/commands/cmd_key.cc | 4 ++-- src/commands/cmd_list.cc | 4 ++-- src/commands/cmd_server.cc | 20 ++++++++++++++++---- src/commands/cmd_set.cc | 2 +- src/commands/cmd_sortedint.cc | 2 +- src/commands/cmd_stream.cc | 4 ++-- src/commands/cmd_string.cc | 2 +- src/commands/cmd_zset.cc | 8 ++++---- src/commands/commander.h | 27 +++++++++++++++------------ src/server/redis_connection.cc | 5 +++++ src/storage/redis_db.cc | 2 +- src/storage/storage.cc | 20 +++++++------------- src/storage/storage.h | 10 +++++----- src/types/redis_hash.cc | 2 +- src/types/redis_json.cc | 2 +- src/types/redis_list.cc | 4 ++-- src/types/redis_set.cc | 2 +- src/types/redis_sortedint.cc | 2 +- src/types/redis_stream.cc | 4 ++-- src/types/redis_zset.cc | 2 +- tests/gocase/unit/debug/debug_test.go | 24 ++++++++++++++++++++++++ 23 files changed, 99 insertions(+), 61 deletions(-) diff --git a/src/commands/cmd_hash.cc b/src/commands/cmd_hash.cc index c62aabd572c..342ec4e0247 100644 --- a/src/commands/cmd_hash.cc +++ b/src/commands/cmd_hash.cc @@ -432,7 +432,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("hget", 3, "read-only", 1, 1, 1 MakeCmdAttr("hincrbyfloat", 4, "write", 1, 1, 1), MakeCmdAttr("hset", -4, "write", 1, 1, 1), MakeCmdAttr("hsetnx", -4, "write", 1, 1, 1), - MakeCmdAttr("hdel", -3, "write", 1, 1, 1), + MakeCmdAttr("hdel", -3, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("hstrlen", 3, "read-only", 1, 1, 1), MakeCmdAttr("hexists", 3, "read-only", 1, 1, 1), MakeCmdAttr("hlen", 2, "read-only", 1, 1, 1), diff --git a/src/commands/cmd_json.cc b/src/commands/cmd_json.cc index 8cd49c51e3f..54a28271eae 100644 --- a/src/commands/cmd_json.cc +++ b/src/commands/cmd_json.cc @@ -600,15 +600,15 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("json.set", 4, "write", 1, 1 MakeCmdAttr("json.type", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.arrappend", -4, "write", 1, 1, 1), MakeCmdAttr("json.arrinsert", -5, "write", 1, 1, 1), - MakeCmdAttr("json.arrtrim", 5, "write", 1, 1, 1), - MakeCmdAttr("json.clear", -2, "write", 1, 1, 1), + MakeCmdAttr("json.arrtrim", 5, "write no-dbsize-check", 1, 1, 1), + MakeCmdAttr("json.clear", -2, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("json.toggle", -2, "write", 1, 1, 1), MakeCmdAttr("json.arrlen", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.merge", 4, "write", 1, 1, 1), MakeCmdAttr("json.objkeys", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.arrpop", -2, "write", 1, 1, 1), MakeCmdAttr("json.arrindex", -4, "read-only", 1, 1, 1), - MakeCmdAttr("json.del", -2, "write", 1, 1, 1), + MakeCmdAttr("json.del", -2, "write no-dbsize-check", 1, 1, 1), // JSON.FORGET is an alias for JSON.DEL, refer: https://redis.io/commands/json.forget/ MakeCmdAttr("json.forget", -2, "write", 1, 1, 1), MakeCmdAttr("json.numincrby", 4, "write", 1, 1, 1), diff --git a/src/commands/cmd_key.cc b/src/commands/cmd_key.cc index f94f87fecf5..2eacdd1ef74 100644 --- a/src/commands/cmd_key.cc +++ b/src/commands/cmd_key.cc @@ -350,8 +350,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("ttl", 2, "read-only", 1, 1, 1), MakeCmdAttr("pexpireat", 3, "write", 1, 1, 1), MakeCmdAttr("expiretime", 2, "read-only", 1, 1, 1), MakeCmdAttr("pexpiretime", 2, "read-only", 1, 1, 1), - MakeCmdAttr("del", -2, "write", 1, -1, 1), - MakeCmdAttr("unlink", -2, "write", 1, -1, 1), + MakeCmdAttr("del", -2, "write no-dbsize-check", 1, -1, 1), + MakeCmdAttr("unlink", -2, "write no-dbsize-check", 1, -1, 1), MakeCmdAttr("rename", 3, "write", 1, 2, 1), MakeCmdAttr("renamenx", 3, "write", 1, 2, 1), ) diff --git a/src/commands/cmd_list.cc b/src/commands/cmd_list.cc index 726d2a70889..f354d64cc4b 100644 --- a/src/commands/cmd_list.cc +++ b/src/commands/cmd_list.cc @@ -861,9 +861,9 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("blpop", -3, "write no-script" MakeCmdAttr("lpush", -3, "write", 1, 1, 1), MakeCmdAttr("lpushx", -3, "write", 1, 1, 1), MakeCmdAttr("lrange", 4, "read-only", 1, 1, 1), - MakeCmdAttr("lrem", 4, "write", 1, 1, 1), + MakeCmdAttr("lrem", 4, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("lset", 4, "write", 1, 1, 1), - MakeCmdAttr("ltrim", 4, "write", 1, 1, 1), + MakeCmdAttr("ltrim", 4, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("lmpop", -4, "write", CommandLMPop::keyRangeGen), MakeCmdAttr("rpop", -2, "write", 1, 1, 1), MakeCmdAttr("rpoplpush", 3, "write", 1, 2, 1), diff --git a/src/commands/cmd_server.cc b/src/commands/cmd_server.cc index 45a6aeaf9bc..850730c0844 100644 --- a/src/commands/cmd_server.cc +++ b/src/commands/cmd_server.cc @@ -594,8 +594,16 @@ class CommandDebug : public Commander { } else if (subcommand_ == "protocol" && args.size() == 3) { protocol_type_ = util::ToLower(args[2]); return Status::OK(); + } else if (subcommand_ == "dbsize-limit" && args.size() == 3) { + auto val = ParseInt(args[2], {0, 1}, 10); + if (!val) { + return {Status::RedisParseErr, "invalid debug dbsize-limit value"}; + } + + dbsize_limit_ = static_cast(val); + return Status::OK(); } - return {Status::RedisInvalidCmd, "Syntax error, DEBUG SLEEP |PROTOCOL "}; + return {Status::RedisInvalidCmd, "Syntax error, DEBUG SLEEP |PROTOCOL |DBSIZE-LIMIT <0|1>"}; } Status Execute(Server *srv, Connection *conn, std::string *output) override { @@ -636,8 +644,11 @@ class CommandDebug : public Commander { "Wrong protocol type name. Please use one of the following: " "string|integer|array|set|bignum|true|false|null"); } + } else if (subcommand_ == "dbsize-limit") { + srv->storage->SetDBSizeLimit(dbsize_limit_); + *output = redis::SimpleString("OK"); } else { - return {Status::RedisInvalidCmd, "Unknown subcommand, should be DEBUG or PROTOCOL"}; + return {Status::RedisInvalidCmd, "Unknown subcommand, should be DEBUG, PROTOCOL or DBSIZE-LIMIT"}; } return Status::OK(); } @@ -646,6 +657,7 @@ class CommandDebug : public Commander { std::string subcommand_; std::string protocol_type_; uint64_t microsecond_ = 0; + bool dbsize_limit_ = false; }; class CommandCommand : public Commander { @@ -1316,8 +1328,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("auth", 2, "read-only ok-loadin MakeCmdAttr("config", -2, "read-only", 0, 0, 0, GenerateConfigFlag), MakeCmdAttr("namespace", -3, "read-only exclusive", 0, 0, 0), MakeCmdAttr("keys", 2, "read-only", 0, 0, 0), - MakeCmdAttr("flushdb", 1, "write", 0, 0, 0), - MakeCmdAttr("flushall", 1, "write", 0, 0, 0), + MakeCmdAttr("flushdb", 1, "write no-dbsize-check", 0, 0, 0), + MakeCmdAttr("flushall", 1, "write no-dbsize-check", 0, 0, 0), MakeCmdAttr("dbsize", -1, "read-only", 0, 0, 0), MakeCmdAttr("slowlog", -2, "read-only", 0, 0, 0), MakeCmdAttr("perflog", -2, "read-only", 0, 0, 0), diff --git a/src/commands/cmd_set.cc b/src/commands/cmd_set.cc index ced252234b2..213a6768279 100644 --- a/src/commands/cmd_set.cc +++ b/src/commands/cmd_set.cc @@ -438,7 +438,7 @@ class CommandSScan : public CommandSubkeyScanBase { }; REDIS_REGISTER_COMMANDS(MakeCmdAttr("sadd", -3, "write", 1, 1, 1), - MakeCmdAttr("srem", -3, "write", 1, 1, 1), + MakeCmdAttr("srem", -3, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("scard", 2, "read-only", 1, 1, 1), MakeCmdAttr("smembers", 2, "read-only", 1, 1, 1), MakeCmdAttr("sismember", 3, "read-only", 1, 1, 1), diff --git a/src/commands/cmd_sortedint.cc b/src/commands/cmd_sortedint.cc index b668a0a69e6..a97e357f154 100644 --- a/src/commands/cmd_sortedint.cc +++ b/src/commands/cmd_sortedint.cc @@ -250,7 +250,7 @@ class CommandSortedintRevRangeByValue : public CommandSortedintRangeByValue { }; REDIS_REGISTER_COMMANDS(MakeCmdAttr("siadd", -3, "write", 1, 1, 1), - MakeCmdAttr("sirem", -3, "write", 1, 1, 1), + MakeCmdAttr("sirem", -3, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("sicard", 2, "read-only", 1, 1, 1), MakeCmdAttr("siexists", -3, "read-only", 1, 1, 1), MakeCmdAttr("sirange", -4, "read-only", 1, 1, 1), diff --git a/src/commands/cmd_stream.cc b/src/commands/cmd_stream.cc index 7ba408859c1..76e2146b7b0 100644 --- a/src/commands/cmd_stream.cc +++ b/src/commands/cmd_stream.cc @@ -1190,14 +1190,14 @@ class CommandXSetId : public Commander { }; REDIS_REGISTER_COMMANDS(MakeCmdAttr("xadd", -5, "write", 1, 1, 1), - MakeCmdAttr("xdel", -3, "write", 1, 1, 1), + MakeCmdAttr("xdel", -3, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("xgroup", -4, "write", 2, 2, 1), MakeCmdAttr("xlen", -2, "read-only", 1, 1, 1), MakeCmdAttr("xinfo", -2, "read-only", 0, 0, 0), MakeCmdAttr("xrange", -4, "read-only", 1, 1, 1), MakeCmdAttr("xrevrange", -2, "read-only", 1, 1, 1), MakeCmdAttr("xread", -4, "read-only", 0, 0, 0), - MakeCmdAttr("xtrim", -4, "write", 1, 1, 1), + MakeCmdAttr("xtrim", -4, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("xsetid", -3, "write", 1, 1, 1)) } // namespace redis diff --git a/src/commands/cmd_string.cc b/src/commands/cmd_string.cc index debf5e3fb7b..3f7b5090b8c 100644 --- a/src/commands/cmd_string.cc +++ b/src/commands/cmd_string.cc @@ -626,7 +626,7 @@ REDIS_REGISTER_COMMANDS( MakeCmdAttr("getset", 3, "write", 1, 1, 1), MakeCmdAttr("getrange", 4, "read-only", 1, 1, 1), MakeCmdAttr("substr", 4, "read-only", 1, 1, 1), - MakeCmdAttr("getdel", 2, "write", 1, 1, 1), + MakeCmdAttr("getdel", 2, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("setrange", 4, "write", 1, 1, 1), MakeCmdAttr("mget", -2, "read-only", 1, -1, 1), MakeCmdAttr("append", 3, "write", 1, 1, 1), MakeCmdAttr("set", -3, "write", 1, 1, 1), diff --git a/src/commands/cmd_zset.cc b/src/commands/cmd_zset.cc index 7d80829df5c..b69e080a9f6 100644 --- a/src/commands/cmd_zset.cc +++ b/src/commands/cmd_zset.cc @@ -1534,10 +1534,10 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("zadd", -4, "write", 1, 1, 1), MakeCmdAttr("zrevrangebylex", -4, "read-only", 1, 1, 1), MakeCmdAttr("zrangebyscore", -4, "read-only", 1, 1, 1), MakeCmdAttr("zrank", -3, "read-only", 1, 1, 1), - MakeCmdAttr("zrem", -3, "write", 1, 1, 1), - MakeCmdAttr("zremrangebyrank", 4, "write", 1, 1, 1), - MakeCmdAttr("zremrangebyscore", 4, "write", 1, 1, 1), - MakeCmdAttr("zremrangebylex", 4, "write", 1, 1, 1), + MakeCmdAttr("zrem", -3, "write no-dbsize-check", 1, 1, 1), + MakeCmdAttr("zremrangebyrank", 4, "write no-dbsize-check", 1, 1, 1), + MakeCmdAttr("zremrangebyscore", 4, "write no-dbsize-check", 1, 1, 1), + MakeCmdAttr("zremrangebylex", 4, "write no-dbsize-check", 1, 1, 1), MakeCmdAttr("zrevrangebyscore", -4, "read-only", 1, 1, 1), MakeCmdAttr("zrevrank", -3, "read-only", 1, 1, 1), MakeCmdAttr("zscore", 3, "read-only", 1, 1, 1), diff --git a/src/commands/commander.h b/src/commands/commander.h index 3330337c746..85cdde7651d 100644 --- a/src/commands/commander.h +++ b/src/commands/commander.h @@ -52,18 +52,19 @@ class Connection; struct CommandAttributes; enum CommandFlags : uint64_t { - kCmdWrite = 1ULL << 0, // "write" flag - kCmdReadOnly = 1ULL << 1, // "read-only" flag - kCmdReplication = 1ULL << 2, // "replication" flag - kCmdPubSub = 1ULL << 3, // "pub-sub" flag - kCmdScript = 1ULL << 4, // "script" flag - kCmdLoading = 1ULL << 5, // "ok-loading" flag - kCmdMulti = 1ULL << 6, // "multi" flag - kCmdExclusive = 1ULL << 7, // "exclusive" flag - kCmdNoMulti = 1ULL << 8, // "no-multi" flag - kCmdNoScript = 1ULL << 9, // "no-script" flag - kCmdROScript = 1ULL << 10, // "ro-script" flag for read-only script commands - kCmdCluster = 1ULL << 11, // "cluster" flag + kCmdWrite = 1ULL << 0, // "write" flag + kCmdReadOnly = 1ULL << 1, // "read-only" flag + kCmdReplication = 1ULL << 2, // "replication" flag + kCmdPubSub = 1ULL << 3, // "pub-sub" flag + kCmdScript = 1ULL << 4, // "script" flag + kCmdLoading = 1ULL << 5, // "ok-loading" flag + kCmdMulti = 1ULL << 6, // "multi" flag + kCmdExclusive = 1ULL << 7, // "exclusive" flag + kCmdNoMulti = 1ULL << 8, // "no-multi" flag + kCmdNoScript = 1ULL << 9, // "no-script" flag + kCmdROScript = 1ULL << 10, // "ro-script" flag for read-only script commands + kCmdCluster = 1ULL << 11, // "cluster" flag + kCmdNoDBSizeCheck = 1ULL << 12, // "no-dbsize-check" flag }; class Commander { @@ -178,6 +179,8 @@ inline uint64_t ParseCommandFlags(const std::string &description, const std::str flags |= kCmdROScript; else if (flag == "cluster") flags |= kCmdCluster; + else if (flag == "no-dbsize-check") + flags |= kCmdNoDBSizeCheck; else { std::cout << fmt::format("Encountered non-existent flag '{}' in command {} in command attribute parsing", flag, cmd_name) diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc index 99d579fe136..aa830ac8791 100644 --- a/src/server/redis_connection.cc +++ b/src/server/redis_connection.cc @@ -501,6 +501,11 @@ void Connection::ExecuteCommands(std::deque *to_process_cmds) { continue; } + if ((cmd_flags & kCmdWrite) && !(cmd_flags & kCmdNoDBSizeCheck) && srv_->storage->ReachedDBSizeLimit()) { + Reply(redis::Error("ERR write command not allowed when reached max-db-size.")); + continue; + } + if (!config->slave_serve_stale_data && srv_->IsSlave() && cmd_name != "info" && cmd_name != "slaveof" && srv_->GetReplicationState() != kReplConnected) { Reply( diff --git a/src/storage/redis_db.cc b/src/storage/redis_db.cc index 150a776577b..63e8ef571cc 100644 --- a/src/storage/redis_db.cc +++ b/src/storage/redis_db.cc @@ -200,7 +200,7 @@ rocksdb::Status Database::MDel(const std::vector &keys, uint64_t *deleted if (*deleted_cnt == 0) return rocksdb::Status::OK(); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status Database::Exists(const std::vector &keys, int *ret) { diff --git a/src/storage/storage.cc b/src/storage/storage.cc index 110f82d8ba9..e178d59592d 100644 --- a/src/storage/storage.cc +++ b/src/storage/storage.cc @@ -601,21 +601,15 @@ void Storage::MultiGet(const rocksdb::ReadOptions &options, rocksdb::ColumnFamil } } -rocksdb::Status Storage::Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, - bool ignore_max_db_size) { +rocksdb::Status Storage::Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates) { if (is_txn_mode_) { // The batch won't be flushed until the transaction was committed or rollback return rocksdb::Status::OK(); } - return writeToDB(options, updates, ignore_max_db_size); + return writeToDB(options, updates); } -rocksdb::Status Storage::writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, - bool ignore_max_db_size) { - if (!ignore_max_db_size && db_size_limit_reached_) { - return rocksdb::Status::SpaceLimit(); - } - +rocksdb::Status Storage::writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates) { // Put replication id logdata at the end of write batch if (replid_.length() == kReplIdLength) { updates->PutLogData(ServerLogData(kReplIdLog, replid_).Encode()); @@ -628,7 +622,7 @@ rocksdb::Status Storage::Delete(const rocksdb::WriteOptions &options, rocksdb::C const rocksdb::Slice &key) { auto batch = GetWriteBatchBase(); batch->Delete(cf_handle, key); - return Write(options, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return Write(options, batch->GetWriteBatch()); } rocksdb::Status Storage::DeleteRange(const std::string &first_key, const std::string &last_key) { @@ -644,7 +638,7 @@ rocksdb::Status Storage::DeleteRange(const std::string &first_key, const std::st return s; } - return Write(write_opts_, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return Write(write_opts_, batch->GetWriteBatch()); } rocksdb::Status Storage::FlushScripts(const rocksdb::WriteOptions &options, rocksdb::ColumnFamilyHandle *cf_handle) { @@ -659,7 +653,7 @@ rocksdb::Status Storage::FlushScripts(const rocksdb::WriteOptions &options, rock return s; } - return Write(options, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return Write(options, batch->GetWriteBatch()); } Status Storage::ReplicaApplyWriteBatch(std::string &&raw_batch) { @@ -818,7 +812,7 @@ Status Storage::WriteToPropagateCF(const std::string &key, const std::string &va auto batch = GetWriteBatchBase(); auto cf = GetCFHandle(kPropagateColumnFamilyName); batch->Put(cf, key, value); - auto s = Write(write_opts_, batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + auto s = Write(write_opts_, batch->GetWriteBatch()); if (!s.ok()) { return {Status::NotOK, s.ToString()}; } diff --git a/src/storage/storage.h b/src/storage/storage.h index f457618ad53..fff10086bbe 100644 --- a/src/storage/storage.h +++ b/src/storage/storage.h @@ -147,8 +147,7 @@ class Storage { rocksdb::Iterator *NewIterator(const rocksdb::ReadOptions &options, rocksdb::ColumnFamilyHandle *column_family); rocksdb::Iterator *NewIterator(const rocksdb::ReadOptions &options); - [[nodiscard]] rocksdb::Status Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, - bool ignore_max_db_size = false); + [[nodiscard]] rocksdb::Status Write(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates); const rocksdb::WriteOptions &DefaultWriteOptions() { return write_opts_; } rocksdb::ReadOptions DefaultScanOptions() const; rocksdb::ReadOptions DefaultMultiGetOptions() const; @@ -172,6 +171,8 @@ class Storage { void PurgeOldBackups(uint32_t num_backups_to_keep, uint32_t backup_max_keep_hours); uint64_t GetTotalSize(const std::string &ns = kDefaultNamespace); void CheckDBSizeLimit(); + bool ReachedDBSizeLimit() { return db_size_limit_reached_; } + void SetDBSizeLimit(bool limit) { db_size_limit_reached_ = limit; } void SetIORateLimit(int64_t max_io_mb); std::shared_lock ReadLockGuard(); @@ -244,7 +245,7 @@ class Storage { Config *config_ = nullptr; std::vector cf_handles_; LockManager lock_mgr_; - bool db_size_limit_reached_ = false; + std::atomic db_size_limit_reached_{false}; DBStats db_stats_; @@ -264,8 +265,7 @@ class Storage { rocksdb::WriteOptions write_opts_ = rocksdb::WriteOptions(); - rocksdb::Status writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates, - bool ignore_max_db_size = false); + rocksdb::Status writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates); void recordKeyspaceStat(const rocksdb::ColumnFamilyHandle *column_family, const rocksdb::Status &s); }; diff --git a/src/types/redis_hash.cc b/src/types/redis_hash.cc index ff0fc195280..d7a1ad8b7bc 100644 --- a/src/types/redis_hash.cc +++ b/src/types/redis_hash.cc @@ -227,7 +227,7 @@ rocksdb::Status Hash::Delete(const Slice &user_key, const std::vector &fi std::string bytes; metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status Hash::MSet(const Slice &user_key, const std::vector &field_values, bool nx, diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc index 47840ba5082..fb3add5feb0 100644 --- a/src/types/redis_json.cc +++ b/src/types/redis_json.cc @@ -95,7 +95,7 @@ rocksdb::Status Json::del(const Slice &ns_key) { batch->Delete(metadata_cf_handle_, ns_key); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status Json::Info(const std::string &user_key, JsonStorageFormat *storage_format) { diff --git a/src/types/redis_list.cc b/src/types/redis_list.cc index 590fafe68b4..ed67007e905 100644 --- a/src/types/redis_list.cc +++ b/src/types/redis_list.cc @@ -250,7 +250,7 @@ rocksdb::Status List::Rem(const Slice &user_key, int count, const Slice &elem, u } *removed_cnt = to_delete_indexes.size(); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status List::Insert(const Slice &user_key, const Slice &pivot, const Slice &elem, bool before, int *new_size) { @@ -658,6 +658,6 @@ rocksdb::Status List::Trim(const Slice &user_key, int start, int stop) { std::string bytes; metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } } // namespace redis diff --git a/src/types/redis_set.cc b/src/types/redis_set.cc index 603957fe237..98677e27560 100644 --- a/src/types/redis_set.cc +++ b/src/types/redis_set.cc @@ -122,7 +122,7 @@ rocksdb::Status Set::Remove(const Slice &user_key, const std::vector &mem batch->Delete(metadata_cf_handle_, ns_key); } } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status Set::Card(const Slice &user_key, uint64_t *size) { diff --git a/src/types/redis_sortedint.cc b/src/types/redis_sortedint.cc index aa883e5dc1d..6670eb7474f 100644 --- a/src/types/redis_sortedint.cc +++ b/src/types/redis_sortedint.cc @@ -93,7 +93,7 @@ rocksdb::Status Sortedint::Remove(const Slice &user_key, const std::vectorPut(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status Sortedint::Card(const Slice &user_key, uint64_t *size) { diff --git a/src/types/redis_stream.cc b/src/types/redis_stream.cc index 84296e8e741..48fe7928c45 100644 --- a/src/types/redis_stream.cc +++ b/src/types/redis_stream.cc @@ -550,7 +550,7 @@ rocksdb::Status Stream::DeleteEntries(const Slice &stream_name, const std::vecto batch->Put(metadata_cf_handle_, ns_key, bytes); } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } // If `options` is StreamLenOptions{} the function just returns the number of entries in the stream. @@ -950,7 +950,7 @@ rocksdb::Status Stream::Trim(const Slice &stream_name, const StreamTrimOptions & metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } return rocksdb::Status::OK(); diff --git a/src/types/redis_zset.cc b/src/types/redis_zset.cc index 447c31a6691..328182b5518 100644 --- a/src/types/redis_zset.cc +++ b/src/types/redis_zset.cc @@ -544,7 +544,7 @@ rocksdb::Status ZSet::Remove(const Slice &user_key, const std::vector &me metadata.Encode(&bytes); batch->Put(metadata_cf_handle_, ns_key, bytes); } - return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch(), /*ignore_max_db_size*/ true); + return storage_->Write(storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } rocksdb::Status ZSet::Rank(const Slice &user_key, const Slice &member, bool reversed, int *member_rank, diff --git a/tests/gocase/unit/debug/debug_test.go b/tests/gocase/unit/debug/debug_test.go index 1aae088f0da..5a418abc675 100644 --- a/tests/gocase/unit/debug/debug_test.go +++ b/tests/gocase/unit/debug/debug_test.go @@ -117,3 +117,27 @@ func TestDebugProtocolV3(t *testing.T) { require.EqualValues(t, false, val) }) } + +func TestDebugDBSizeLimit(t *testing.T) { + srv := util.StartServer(t, map[string]string{}) + defer srv.Close() + + ctx := context.Background() + rdb := srv.NewClient() + defer func() { require.NoError(t, rdb.Close()) }() + + t.Run("debug ignore dbsize check", func(t *testing.T) { + r := rdb.Do(ctx, "SET", "k1", "v1") + require.NoError(t, r.Err()) + + r = rdb.Do(ctx, "DEBUG", "DBSIZE-LIMIT", "1") + require.NoError(t, r.Err()) + + r = rdb.Do(ctx, "SET", "k2", "v2") + require.Error(t, r.Err()) + util.ErrorRegexp(t, r.Err(), "ERR.*not allowed.*") + + r = rdb.Do(ctx, "DEL", "k1") + require.NoError(t, r.Err()) + }) +}