Skip to content

Commit

Permalink
Ignore max-db-size limit when deleting data or writing aux informatio…
Browse files Browse the repository at this point in the history 
…ns (#2047)
  • Loading branch information
@caipengbo
caipengbo authored Jan 30, 2024
1 parent 9d618e0 commit fd33e55
Show file tree
Hide file tree
Showing 15 changed files with 80 additions and 38 deletions.
2 changes: 1 addition & 1 deletion src/commands/cmd_hash.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -439,7 +439,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandHGet>("hget", 3, "read-only", 1, 1, 1
MakeCmdAttr<CommandHIncrByFloat>("hincrbyfloat", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandHMSet>("hset", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandHSetNX>("hsetnx", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandHDel>("hdel", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandHDel>("hdel", -3, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandHStrlen>("hstrlen", 3, "read-only", 1, 1, 1),
MakeCmdAttr<CommandHExists>("hexists", 3, "read-only", 1, 1, 1),
MakeCmdAttr<CommandHLen>("hlen", 2, "read-only", 1, 1, 1),
Expand Down
6 changes: 3 additions & 3 deletions src/commands/cmd_json.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -600,15 +600,15 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", 4, "write", 1, 1
MakeCmdAttr<CommandJsonType>("json.type", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonArrAppend>("json.arrappend", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonArrInsert>("json.arrinsert", -5, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonArrTrim>("json.arrtrim", 5, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonClear>("json.clear", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonArrTrim>("json.arrtrim", 5, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandJsonClear>("json.clear", -2, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandJsonToggle>("json.toggle", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonArrLen>("json.arrlen", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonMerge>("json.merge", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonObjkeys>("json.objkeys", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonArrPop>("json.arrpop", -2, "write", 1, 1, 1),
MakeCmdAttr<CommanderJsonArrIndex>("json.arrindex", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonDel>("json.del", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonDel>("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<CommandJsonDel>("json.forget", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonNumIncrBy>("json.numincrby", 4, "write", 1, 1, 1),
Expand Down
4 changes: 2 additions & 2 deletions src/commands/cmd_key.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -350,8 +350,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandTTL>("ttl", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandPExpireAt>("pexpireat", 3, "write", 1, 1, 1),
MakeCmdAttr<CommandExpireTime>("expiretime", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandPExpireTime>("pexpiretime", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandDel>("del", -2, "write", 1, -1, 1),
MakeCmdAttr<CommandDel>("unlink", -2, "write", 1, -1, 1),
MakeCmdAttr<CommandDel>("del", -2, "write no-dbsize-check", 1, -1, 1),
MakeCmdAttr<CommandDel>("unlink", -2, "write no-dbsize-check", 1, -1, 1),
MakeCmdAttr<CommandRename>("rename", 3, "write", 1, 2, 1),
MakeCmdAttr<CommandRenameNX>("renamenx", 3, "write", 1, 2, 1), )

Expand Down
4 changes: 2 additions & 2 deletions src/commands/cmd_list.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -861,9 +861,9 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandBLPop>("blpop", -3, "write no-script"
MakeCmdAttr<CommandLPush>("lpush", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandLPushX>("lpushx", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandLRange>("lrange", 4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandLRem>("lrem", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandLRem>("lrem", 4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandLSet>("lset", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandLTrim>("ltrim", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandLTrim>("ltrim", 4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandLMPop>("lmpop", -4, "write", CommandLMPop::keyRangeGen),
MakeCmdAttr<CommandRPop>("rpop", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandRPopLPUSH>("rpoplpush", 3, "write", 1, 2, 1),
Expand Down
20 changes: 16 additions & 4 deletions src/commands/cmd_server.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<int32_t>(args[2], {0, 1}, 10);
if (!val) {
return {Status::RedisParseErr, "invalid debug dbsize-limit value"};
}

dbsize_limit_ = static_cast<bool>(val);
return Status::OK();
}
return {Status::RedisInvalidCmd, "Syntax error, DEBUG SLEEP <seconds>|PROTOCOL <type>"};
return {Status::RedisInvalidCmd, "Syntax error, DEBUG SLEEP <seconds>|PROTOCOL <type>|DBSIZE-LIMIT <0|1>"};
}

Status Execute(Server *srv, Connection *conn, std::string *output) override {
Expand Down Expand Up @@ -638,8 +646,11 @@ class CommandDebug : public Commander {
"Wrong protocol type name. Please use one of the following: "
"string|integer|double|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();
}
Expand All @@ -648,6 +659,7 @@ class CommandDebug : public Commander {
std::string subcommand_;
std::string protocol_type_;
uint64_t microsecond_ = 0;
bool dbsize_limit_ = false;
};

class CommandCommand : public Commander {
Expand Down Expand Up @@ -1318,8 +1330,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loadin
MakeCmdAttr<CommandConfig>("config", -2, "read-only", 0, 0, 0, GenerateConfigFlag),
MakeCmdAttr<CommandNamespace>("namespace", -3, "read-only exclusive", 0, 0, 0),
MakeCmdAttr<CommandKeys>("keys", 2, "read-only", 0, 0, 0),
MakeCmdAttr<CommandFlushDB>("flushdb", 1, "write", 0, 0, 0),
MakeCmdAttr<CommandFlushAll>("flushall", 1, "write", 0, 0, 0),
MakeCmdAttr<CommandFlushDB>("flushdb", 1, "write no-dbsize-check", 0, 0, 0),
MakeCmdAttr<CommandFlushAll>("flushall", 1, "write no-dbsize-check", 0, 0, 0),
MakeCmdAttr<CommandDBSize>("dbsize", -1, "read-only", 0, 0, 0),
MakeCmdAttr<CommandSlowlog>("slowlog", -2, "read-only", 0, 0, 0),
MakeCmdAttr<CommandPerfLog>("perflog", -2, "read-only", 0, 0, 0),
Expand Down
2 changes: 1 addition & 1 deletion src/commands/cmd_set.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -438,7 +438,7 @@ class CommandSScan : public CommandSubkeyScanBase {
};

REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandSAdd>("sadd", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandSRem>("srem", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandSRem>("srem", -3, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandSCard>("scard", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandSMembers>("smembers", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandSIsMember>("sismember", 3, "read-only", 1, 1, 1),
Expand Down
2 changes: 1 addition & 1 deletion src/commands/cmd_sortedint.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ class CommandSortedintRevRangeByValue : public CommandSortedintRangeByValue {
};

REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandSortedintAdd>("siadd", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandSortedintRem>("sirem", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandSortedintRem>("sirem", -3, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandSortedintCard>("sicard", 2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandSortedintExists>("siexists", -3, "read-only", 1, 1, 1),
MakeCmdAttr<CommandSortedintRange>("sirange", -4, "read-only", 1, 1, 1),
Expand Down
4 changes: 2 additions & 2 deletions src/commands/cmd_stream.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1190,14 +1190,14 @@ class CommandXSetId : public Commander {
};

REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandXAdd>("xadd", -5, "write", 1, 1, 1),
MakeCmdAttr<CommandXDel>("xdel", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandXDel>("xdel", -3, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandXGroup>("xgroup", -4, "write", 2, 2, 1),
MakeCmdAttr<CommandXLen>("xlen", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandXInfo>("xinfo", -2, "read-only", 0, 0, 0),
MakeCmdAttr<CommandXRange>("xrange", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandXRevRange>("xrevrange", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandXRead>("xread", -4, "read-only", 0, 0, 0),
MakeCmdAttr<CommandXTrim>("xtrim", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandXTrim>("xtrim", -4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandXSetId>("xsetid", -3, "write", 1, 1, 1))

} // namespace redis
2 changes: 1 addition & 1 deletion src/commands/cmd_string.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -626,7 +626,7 @@ REDIS_REGISTER_COMMANDS(
MakeCmdAttr<CommandGetSet>("getset", 3, "write", 1, 1, 1),
MakeCmdAttr<CommandGetRange>("getrange", 4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandSubStr>("substr", 4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandGetDel>("getdel", 2, "write", 1, 1, 1),
MakeCmdAttr<CommandGetDel>("getdel", 2, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandSetRange>("setrange", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandMGet>("mget", -2, "read-only", 1, -1, 1),
MakeCmdAttr<CommandAppend>("append", 3, "write", 1, 1, 1), MakeCmdAttr<CommandSet>("set", -3, "write", 1, 1, 1),
Expand Down
8 changes: 4 additions & 4 deletions src/commands/cmd_zset.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1537,10 +1537,10 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandZAdd>("zadd", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandZRevRangeByLex>("zrevrangebylex", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZRangeByScore>("zrangebyscore", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZRank>("zrank", -3, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZRem>("zrem", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByRank>("zremrangebyrank", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByScore>("zremrangebyscore", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByLex>("zremrangebylex", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandZRem>("zrem", -3, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByRank>("zremrangebyrank", 4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByScore>("zremrangebyscore", 4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandZRemRangeByLex>("zremrangebylex", 4, "write no-dbsize-check", 1, 1, 1),
MakeCmdAttr<CommandZRevRangeByScore>("zrevrangebyscore", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZRevRank>("zrevrank", -3, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZScore>("zscore", 3, "read-only", 1, 1, 1),
Expand Down
27 changes: 15 additions & 12 deletions src/commands/commander.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions src/server/redis_connection.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -501,6 +501,11 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *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(
Expand Down
4 changes: 0 additions & 4 deletions src/storage/storage.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -639,10 +639,6 @@ rocksdb::Status Storage::Write(const rocksdb::WriteOptions &options, rocksdb::Wr
}

rocksdb::Status Storage::writeToDB(const rocksdb::WriteOptions &options, rocksdb::WriteBatch *updates) {
if (db_size_limit_reached_) {
return rocksdb::Status::SpaceLimit();
}

// Put replication id logdata at the end of write batch
if (replid_.length() == kReplIdLength) {
updates->PutLogData(ServerLogData(kReplIdLog, replid_).Encode());
Expand Down
4 changes: 3 additions & 1 deletion src/storage/storage.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,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<std::shared_mutex> ReadLockGuard();
Expand Down Expand Up @@ -254,7 +256,7 @@ class Storage {
Config *config_ = nullptr;
std::vector<rocksdb::ColumnFamilyHandle *> cf_handles_;
LockManager lock_mgr_;
bool db_size_limit_reached_ = false;
std::atomic<bool> db_size_limit_reached_{false};

DBStats db_stats_;

Expand Down
24 changes: 24 additions & 0 deletions tests/gocase/unit/debug/debug_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,3 +119,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())
})
}

0 comments on commit fd33e55

Please sign in to comment.