Skip to content

Commit

Permalink
feat(cmd): make key range fields private in CommandAttributes (#2621)
Browse files Browse the repository at this point in the history
  • Loading branch information
@PragmaTwice
PragmaTwice authored Oct 24, 2024
1 parent c525d7a commit 5fccd57
Show file tree
Hide file tree
Showing 6 changed files with 90 additions and 96 deletions.
12 changes: 6 additions & 6 deletions src/cluster/cluster.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -832,16 +832,16 @@ bool Cluster::IsWriteForbiddenSlot(int slot) const {

Status Cluster::CanExecByMySelf(const redis::CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
redis::Connection *conn, lua::ScriptRunCtx *script_run_ctx) {
std::vector<int> keys_indexes;
std::vector<int> key_indexes;

// No keys
if (auto s = redis::CommandTable::GetKeysFromCommand(attributes, cmd_tokens, &keys_indexes); !s.IsOK())
return Status::OK();
auto s = redis::CommandTable::GetKeysFromCommand(attributes, cmd_tokens);
if (!s) return Status::OK();
key_indexes = *s;

if (keys_indexes.empty()) return Status::OK();
if (key_indexes.empty()) return Status::OK();

int slot = -1;
for (auto i : keys_indexes) {
for (auto i : key_indexes) {
if (i >= static_cast<int>(cmd_tokens.size())) break;

int cur_slot = GetSlotIdFromKey(cmd_tokens[i]);
Expand Down
12 changes: 5 additions & 7 deletions src/commands/cmd_server.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -676,18 +676,16 @@ class CommandCommand : public Commander {
return {Status::RedisUnknownCmd, "Invalid command specified"};
}

std::vector<int> keys_indexes;
auto s = CommandTable::GetKeysFromCommand(
cmd_iter->second, std::vector<std::string>(args_.begin() + 2, args_.end()), &keys_indexes);
if (!s.IsOK()) return s;
auto key_indexes = GET_OR_RET(CommandTable::GetKeysFromCommand(
cmd_iter->second, std::vector<std::string>(args_.begin() + 2, args_.end())));

if (keys_indexes.size() == 0) {
if (key_indexes.size() == 0) {
return {Status::RedisExecErr, "Invalid arguments specified for command"};
}

std::vector<std::string> keys;
keys.reserve(keys_indexes.size());
for (const auto &key_index : keys_indexes) {
keys.reserve(key_indexes.size());
for (const auto &key_index : key_indexes) {
keys.emplace_back(args_[key_index + 2]);
}
*output = conn->MultiBulkString(keys);
Expand Down
52 changes: 22 additions & 30 deletions src/commands/commander.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,10 @@ std::string CommandTable::GetCommandInfo(const CommandAttributes *command_attrib
command_flags.append(redis::MultiLen(1));
command_flags.append(redis::BulkString(command_attributes->InitialFlags() & kCmdWrite ? "write" : "readonly"));
command.append(command_flags);
command.append(redis::Integer(command_attributes->key_range.first_key));
command.append(redis::Integer(command_attributes->key_range.last_key));
command.append(redis::Integer(command_attributes->key_range.key_step));
auto key_range = command_attributes->InitialKeyRange().ValueOr({0, 0, 0});
command.append(redis::Integer(key_range.first_key));
command.append(redis::Integer(key_range.last_key));
command.append(redis::Integer(key_range.key_step));
return command;
}

Expand All @@ -79,40 +80,31 @@ void CommandTable::GetCommandsInfo(std::string *info, const std::vector<std::str
}
}

void DumpKeyRange(std::vector<int> &keys_index, int argc, const CommandKeyRange &range) {
auto last = range.last_key;
if (last < 0) last = argc + last;

for (int i = range.first_key; i <= last; i += range.key_step) {
keys_index.emplace_back(i);
}
}

Status CommandTable::GetKeysFromCommand(const CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
std::vector<int> *keys_index) {
if (attributes->key_range.first_key == 0) {
return {Status::NotOK, "The command has no key arguments"};
}

StatusOr<std::vector<int>> CommandTable::GetKeysFromCommand(const CommandAttributes *attributes,
const std::vector<std::string> &cmd_tokens) {
int argc = static_cast<int>(cmd_tokens.size());

if ((attributes->arity > 0 && attributes->arity != argc) || argc < -attributes->arity) {
if (!attributes->CheckArity(argc)) {
return {Status::NotOK, "Invalid number of arguments specified for command"};
}

if (attributes->key_range.first_key > 0) {
DumpKeyRange(*keys_index, argc, attributes->key_range);
} else if (attributes->key_range.first_key == -1) {
DumpKeyRange(*keys_index, argc, attributes->key_range_gen(cmd_tokens));
} else if (attributes->key_range.first_key == -2) {
for (const auto &range : attributes->key_range_vec_gen(cmd_tokens)) {
DumpKeyRange(*keys_index, argc, range);
}
} else {
return {Status::NotOK, "The key range specification is invalid"};
Status status;
std::vector<int> key_indexes;

attributes->ForEachKeyRange(
[&](const std::vector<std::string> &, CommandKeyRange key_range) {
key_range.ForEachKeyIndex([&](int i) { key_indexes.push_back(i); }, cmd_tokens.size());
},
cmd_tokens,
[&](const auto &) {
status = {Status::NotOK, "The command has no key arguments"};
});

if (!status) {
return status;
}

return Status::OK();
return key_indexes;
}

bool CommandTable::IsExists(const std::string &name) {
Expand Down
83 changes: 53 additions & 30 deletions src/commands/commander.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,9 +141,18 @@ struct CommandKeyRange {
template <typename F>
void ForEachKey(F &&f, const std::vector<std::string> &args) const {
for (size_t i = first_key; last_key > 0 ? i <= size_t(last_key) : i <= args.size() + last_key; i += key_step) {
if (i >= args.size()) continue;
std::forward<F>(f)(args[i]);
}
}

template <typename F>
void ForEachKeyIndex(F &&f, size_t arg_size) const {
for (size_t i = first_key; last_key > 0 ? i <= size_t(last_key) : i <= arg_size + last_key; i += key_step) {
if (i >= arg_size) continue;
std::forward<F>(f)(i);
}
}
};

using CommandKeyRangeGen = std::function<CommandKeyRange(const std::vector<std::string> &)>;
Expand All @@ -161,20 +170,20 @@ struct CommandAttributes {
: name(std::move(name)),
arity(arity),
category(category),
key_range{0, 0, 0},
factory(std::move(factory)),
flags_(flags),
flag_gen_(std::move(flag_gen)) {}
flag_gen_(std::move(flag_gen)),
key_range_{0, 0, 0} {}

CommandAttributes(std::string name, int arity, CommandCategory category, uint64_t flags, AdditionalFlagGen flag_gen,
CommandKeyRange key_range, CommanderFactory factory)
: name(std::move(name)),
arity(arity),
category(category),
key_range(key_range),
factory(std::move(factory)),
flags_(flags),
flag_gen_(std::move(flag_gen)) {
flag_gen_(std::move(flag_gen)),
key_range_(key_range) {
if (key_range.first_key <= 0 || key_range.key_step <= 0 ||
(key_range.last_key >= 0 && key_range.last_key < key_range.first_key)) {
std::cout << fmt::format("Encountered invalid key range in command {}", this->name) << std::endl;
Expand All @@ -187,22 +196,22 @@ struct CommandAttributes {
: name(std::move(name)),
arity(arity),
category(category),
key_range{-1, 0, 0},
key_range_gen(std::move(key_range)),
factory(std::move(factory)),
flags_(flags),
flag_gen_(std::move(flag_gen)) {}
flag_gen_(std::move(flag_gen)),
key_range_{-1, 0, 0},
key_range_gen_(std::move(key_range)) {}

CommandAttributes(std::string name, int arity, CommandCategory category, uint64_t flags, AdditionalFlagGen flag_gen,
CommandKeyRangeVecGen key_range, CommanderFactory factory)
: name(std::move(name)),
arity(arity),
category(category),
key_range{-2, 0, 0},
key_range_vec_gen(std::move(key_range)),
factory(std::move(factory)),
flags_(flags),
flag_gen_(std::move(flag_gen)) {}
flag_gen_(std::move(flag_gen)),
key_range_{-2, 0, 0},
key_range_vec_gen_(std::move(key_range)) {}

// command name
std::string name;
Expand All @@ -215,16 +224,6 @@ struct CommandAttributes {
// category of this command, e.g. key, string, hash
CommandCategory category;

// static determined key range
// if key_range.first_key == 0, there's no key in this command args
CommandKeyRange key_range;

// if key_range.first_key == -1, key_range_gen is used instead
CommandKeyRangeGen key_range_gen;

// if key_range.first_key == -2, key_range_vec_gen is used instead
CommandKeyRangeVecGen key_range_vec_gen;

// commander object generator
CommanderFactory factory;

Expand All @@ -240,33 +239,57 @@ struct CommandAttributes {
return !((arity > 0 && cmd_size != arity) || (arity < 0 && cmd_size < -arity));
}

template <typename F>
void ForEachKeyRange(F &&f, const std::vector<std::string> &args) const {
if (key_range.first_key > 0) {
std::forward<F>(f)(args, key_range);
} else if (key_range.first_key == -1) {
redis::CommandKeyRange range = key_range_gen(args);
StatusOr<CommandKeyRange> InitialKeyRange() const {
if (key_range_.first_key >= 0) return key_range_;
return {Status::NotOK, "key range is unavailable without command arguments"};
}

template <typename F, typename G>
void ForEachKeyRange(F &&f, const std::vector<std::string> &args, G &&g) const {
if (key_range_.first_key > 0) {
std::forward<F>(f)(args, key_range_);
} else if (key_range_.first_key == -1) {
redis::CommandKeyRange range = key_range_gen_(args);

if (range.first_key > 0) {
std::forward<F>(f)(args, range);
}
} else if (key_range.first_key == -2) {
std::vector<redis::CommandKeyRange> vec_range = key_range_vec_gen(args);
} else if (key_range_.first_key == -2) {
std::vector<redis::CommandKeyRange> vec_range = key_range_vec_gen_(args);

for (const auto &range : vec_range) {
if (range.first_key > 0) {
std::forward<F>(f)(args, range);
}
}
} else if (key_range_.first_key == 0) {
// otherwise, if there's no key inside the command arguments
// e.g. FLUSHALL, with "write" flag but no key specified
std::forward<G>(g)(args);
}
}

template <typename F>
void ForEachKeyRange(F &&f, const std::vector<std::string> &args) const {
ForEachKeyRange(std::forward<F>(f), args, [](const auto &) {});
}

private:
// bitmap of enum CommandFlags
uint64_t flags_;

// additional flags regarding to dynamic command arguments
AdditionalFlagGen flag_gen_;

// static determined key range
// if key_range.first_key == 0, there's no key in this command args
CommandKeyRange key_range_;

// if key_range.first_key == -1, key_range_gen is used instead
CommandKeyRangeGen key_range_gen_;

// if key_range.first_key == -2, key_range_vec_gen is used instead
CommandKeyRangeVecGen key_range_vec_gen_;
};

using CommandMap = std::map<std::string, const CommandAttributes *>;
Expand Down Expand Up @@ -374,8 +397,8 @@ struct CommandTable {
static void GetAllCommandsInfo(std::string *info);
static void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names);
static std::string GetCommandInfo(const CommandAttributes *command_attributes);
static Status GetKeysFromCommand(const CommandAttributes *attributes, const std::vector<std::string> &cmd_tokens,
std::vector<int> *keys_indexes);
static StatusOr<std::vector<int>> GetKeysFromCommand(const CommandAttributes *attributes,
const std::vector<std::string> &cmd_tokens);

static size_t Size();
static bool IsExists(const std::string &name);
Expand Down
3 changes: 1 addition & 2 deletions src/common/string_util.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,7 @@ std::string EscapeString(std::string_view s);
std::string StringNext(std::string s);

template <typename T, typename F>
std::string StringJoin(
const T &con, F &&f = [](const auto &v) -> decltype(auto) { return v; }, std::string_view sep = ", ") {
std::string StringJoin(const T &con, F &&f, std::string_view sep = ", ") {
std::string res;
bool is_first = true;
for (const auto &v : con) {
Expand Down
24 changes: 3 additions & 21 deletions src/server/server.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1984,27 +1984,9 @@ void Server::updateAllWatchedKeys() {

void Server::UpdateWatchedKeysFromArgs(const std::vector<std::string> &args, const redis::CommandAttributes &attr) {
if ((attr.GenerateFlags(args) & redis::kCmdWrite) && watched_key_size_ > 0) {
if (attr.key_range.first_key > 0) {
updateWatchedKeysFromRange(args, attr.key_range);
} else if (attr.key_range.first_key == -1) {
redis::CommandKeyRange range = attr.key_range_gen(args);

if (range.first_key > 0) {
updateWatchedKeysFromRange(args, range);
}
} else if (attr.key_range.first_key == -2) {
std::vector<redis::CommandKeyRange> vec_range = attr.key_range_vec_gen(args);

for (const auto &range : vec_range) {
if (range.first_key > 0) {
updateWatchedKeysFromRange(args, range);
}
}

} else {
// support commands like flushdb (write flag && key range {0,0,0})
updateAllWatchedKeys();
}
attr.ForEachKeyRange([this](const std::vector<std::string> &args,
redis::CommandKeyRange range) { updateWatchedKeysFromRange(args, range); },
args, [this](const std::vector<std::string> &) { updateAllWatchedKeys(); });
}
}

Expand Down

0 comments on commit 5fccd57

Please sign in to comment.