From 8e79223e4c86c4be898d3fe891753201b96d0fb9 Mon Sep 17 00:00:00 2001 From: Roman Gershman Date: Fri, 18 Oct 2024 14:05:47 +0300 Subject: [PATCH 1/2] chore: remove ToUpper calls in main_service Also, test for IsPaused() first to avoid doing more checks for common-case. Signed-off-by: Roman Gershman --- src/server/command_registry.cc | 28 ++++++++++++++++- src/server/command_registry.h | 3 ++ src/server/hset_family.cc | 4 +-- src/server/main_service.cc | 57 +++++++--------------------------- src/server/main_service.h | 1 - src/server/server_state.cc | 4 --- src/server/server_state.h | 4 ++- 7 files changed, 47 insertions(+), 54 deletions(-) diff --git a/src/server/command_registry.cc b/src/server/command_registry.cc index aad459e08d80..06df8bf145f6 100644 --- a/src/server/command_registry.cc +++ b/src/server/command_registry.cc @@ -163,7 +163,7 @@ void CommandRegistry::StartFamily() { } std::string_view CommandRegistry::RenamedOrOriginal(std::string_view orig) const { - if (cmd_rename_map_.contains(orig)) { + if (!cmd_rename_map_.empty() && cmd_rename_map_.contains(orig)) { return cmd_rename_map_.find(orig)->second; } return orig; @@ -173,6 +173,32 @@ CommandRegistry::FamiliesVec CommandRegistry::GetFamilies() { return std::move(family_of_commands_); } +std::pair CommandRegistry::FindExtended(string_view cmd, + CmdArgList tail_args) const { + if (cmd == RenamedOrOriginal("ACL"sv)) { + if (tail_args.empty()) { + return {Find(cmd), {}}; + } + + auto second_cmd = absl::AsciiStrToUpper(ArgS(tail_args, 0)); + string full_cmd = absl::StrCat(cmd, " ", second_cmd); + + return {Find(full_cmd), tail_args.subspan(1)}; + } + + const CommandId* res = Find(cmd); + if (!res) + return {nullptr, {}}; + + // A workaround for XGROUP HELP that does not fit our static taxonomy of commands. + if (tail_args.size() == 1 && res->name() == "XGROUP") { + if (absl::EqualsIgnoreCase(ArgS(tail_args, 0), "HELP")) { + res = Find("_XGROUP_HELP"); + } + } + return {res, tail_args}; +} + namespace CO { const char* OptName(CO::CommandOpt fl) { diff --git a/src/server/command_registry.h b/src/server/command_registry.h index 82ce102a785d..20d596a33b89 100644 --- a/src/server/command_registry.h +++ b/src/server/command_registry.h @@ -190,6 +190,9 @@ class CommandRegistry { using FamiliesVec = std::vector>; FamiliesVec GetFamilies(); + std::pair FindExtended(std::string_view cmd, + facade::CmdArgList tail_args) const; + private: absl::flat_hash_map cmd_map_; absl::flat_hash_map cmd_rename_map_; diff --git a/src/server/hset_family.cc b/src/server/hset_family.cc index 65abbb946141..8c55fe8eff3c 100644 --- a/src/server/hset_family.cc +++ b/src/server/hset_family.cc @@ -1125,8 +1125,8 @@ void HSetFamily::HRandField(CmdArgList args, ConnectionContext* cntx) { } if (args.size() == 3) { - ToUpper(&args[2]); - if (ArgS(args, 2) != "WITHVALUES") + string arg = absl::AsciiStrToUpper(ArgS(args, 2)); + if (arg != "WITHVALUES") return cntx->SendError(kSyntaxErr); else with_values = true; diff --git a/src/server/main_service.cc b/src/server/main_service.cc index 0ede303cab93..16df17b83a2a 100644 --- a/src/server/main_service.cc +++ b/src/server/main_service.cc @@ -1241,13 +1241,13 @@ std::optional Service::VerifyCommandState(const CommandId* cid, CmdA void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx) { absl::Cleanup clear_last_error( [cntx]() { std::ignore = cntx->reply_builder()->ConsumeLastError(); }); - CHECK(!args.empty()); + DCHECK(!args.empty()); DCHECK_NE(0u, shard_set->size()) << "Init was not called"; ServerState& etl = *ServerState::tlocal(); - ToUpper(&args[0]); - const auto [cid, args_no_cmd] = FindCmd(args); + string cmd = absl::AsciiStrToUpper(ArgS(args, 0)); + const auto [cid, args_no_cmd] = registry_.FindExtended(cmd, args.subspan(1)); if (cid == nullptr) { return cntx->SendError(ReportUnknownCmd(ArgS(args, 0))); @@ -1264,7 +1264,7 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx) } // Don't interrupt running multi commands or admin connections. - if (!dispatching_in_multi && cntx->conn() && !cntx->conn()->IsPrivileged()) { + if (etl.IsPaused() && !dispatching_in_multi && cntx->conn() && !cntx->conn()->IsPrivileged()) { bool is_write = cid->IsWriteOnly(); is_write |= cid->name() == "PUBLISH" || cid->name() == "EVAL" || cid->name() == "EVALSHA"; is_write |= cid->name() == "EXEC" && dfly_cntx->conn_state.exec_info.is_write; @@ -1535,8 +1535,8 @@ size_t Service::DispatchManyCommands(absl::Span args_list, return 0; for (auto args : args_list) { - ToUpper(&args[0]); - const auto [cid, tail_args] = FindCmd(args); + string cmd = absl::AsciiStrToUpper(ArgS(args, 0)); + const auto [cid, tail_args] = registry_.FindExtended(cmd, args.subspan(1)); // MULTI...EXEC commands need to be collected into a single context, so squashing is not // possible @@ -1850,10 +1850,10 @@ void Service::CallFromScript(ConnectionContext* cntx, Interpreter::CallArgs& ca) if (ca.async) { auto& info = cntx->conn_state.script_info; - ToUpper(&ca.args[0]); + string cmd = absl::AsciiStrToUpper(ArgS(ca.args, 0)); // Full command verification happens during squashed execution - if (auto* cid = registry_.Find(ArgS(ca.args, 0)); cid != nullptr) { + if (auto* cid = registry_.Find(cmd); cid != nullptr) { auto replies = ca.error_abort ? ReplyMode::ONLY_ERR : ReplyMode::NONE; info->async_cmds.emplace_back(std::move(*ca.buffer), cid, ca.args.subspan(1), replies); info->async_cmds_heap_mem += info->async_cmds.back().UsedMemory(); @@ -1995,35 +1995,6 @@ optional StartMultiEval(DbIndex dbid, CmdArgList keys, ScriptMgr::ScriptPa return false; } -static std::string FullAclCommandFromArgs(CmdArgList args, std::string_view name) { - ToUpper(&args[1]); - auto res = absl::StrCat(name, " ", ArgS(args, 1)); - return res; -} - -std::pair Service::FindCmd(CmdArgList args) const { - const std::string_view command = facade::ToSV(args[0]); - std::string_view acl = "ACL"; - if (command == registry_.RenamedOrOriginal(acl)) { - if (args.size() == 1) { - return {registry_.Find(ArgS(args, 0)), args}; - } - return {registry_.Find(FullAclCommandFromArgs(args, command)), args.subspan(2)}; - } - - const CommandId* res = registry_.Find(ArgS(args, 0)); - if (!res) - return {nullptr, args}; - - // A workaround for XGROUP HELP that does not fit our static taxonomy of commands. - if (args.size() == 2 && res->name() == "XGROUP") { - if (absl::EqualsIgnoreCase(ArgS(args, 1), "HELP")) { - res = registry_.Find("_XGROUP_HELP"); - } - } - return {res, args.subspan(1)}; -} - static bool CanRunSingleShardMulti(optional sid, const ScriptMgr::ScriptParams& params, const Transaction& tx) { if (!sid.has_value()) { @@ -2426,8 +2397,7 @@ void Service::PUnsubscribe(CmdArgList args, ConnectionContext* cntx) { // Not a real implementation. Serves as a decorator to accept some function commands // for testing. void Service::Function(CmdArgList args, ConnectionContext* cntx) { - ToUpper(&args[0]); - string_view sub_cmd = ArgS(args, 0); + string sub_cmd = absl::AsciiStrToUpper(ArgS(args, 0)); if (sub_cmd == "FLUSH") { return cntx->SendOk(); @@ -2475,8 +2445,7 @@ void Service::Pubsub(CmdArgList args, ConnectionContext* cntx) { return; } - ToUpper(&args[0]); - string_view subcmd = ArgS(args, 0); + string subcmd = absl::AsciiStrToUpper(ArgS(args, 0)); if (subcmd == "HELP") { string_view help_arr[] = { @@ -2552,8 +2521,7 @@ void Service::Command(CmdArgList args, ConnectionContext* cntx) { return; } - ToUpper(&args[0]); - string_view subcmd = ArgS(args, 0); + string subcmd = absl::AsciiStrToUpper(ArgS(args, 0)); // COUNT if (subcmd == "COUNT") { @@ -2562,8 +2530,7 @@ void Service::Command(CmdArgList args, ConnectionContext* cntx) { // INFO [cmd] if (subcmd == "INFO" && args.size() == 2) { - ToUpper(&args[1]); - string_view cmd = ArgS(args, 1); + string cmd = absl::AsciiStrToUpper(ArgS(args, 1)); if (const auto* cid = registry_.Find(cmd); cid) { rb->StartArray(1); diff --git a/src/server/main_service.h b/src/server/main_service.h index 9c829199b4c0..5195a1d0e195 100644 --- a/src/server/main_service.h +++ b/src/server/main_service.h @@ -64,7 +64,6 @@ class Service : public facade::ServiceInterface { facade::ConnectionContext* CreateContext(util::FiberSocketBase* peer, facade::Connection* owner) final; - std::pair FindCmd(CmdArgList args) const; const CommandId* FindCmd(std::string_view) const; CommandRegistry* mutable_registry() { diff --git a/src/server/server_state.cc b/src/server/server_state.cc index d4dab12d2a40..a906c2de2ffd 100644 --- a/src/server/server_state.cc +++ b/src/server/server_state.cc @@ -155,10 +155,6 @@ void ServerState::AwaitPauseState(bool is_write) { }); } -bool ServerState::IsPaused() const { - return (client_pauses_[0] + client_pauses_[1]) > 0; -} - void ServerState::DecommitMemory(uint8_t flags) { if (flags & kDataHeap) { mi_heap_collect(data_heap(), true); diff --git a/src/server/server_state.h b/src/server/server_state.h index f4b0cbac624f..f8be766a5816 100644 --- a/src/server/server_state.h +++ b/src/server/server_state.h @@ -272,7 +272,9 @@ class ServerState { // public struct - to allow initialization. // @is_write controls whether the command is a write command or not. void AwaitPauseState(bool is_write); - bool IsPaused() const; + bool IsPaused() const { + return (client_pauses_[0] + client_pauses_[1]) > 0; + } SlowLogShard& GetSlowLog() { return slow_log_shard_; From 6fec6905257cd489fa3aa831bb0ad91c9b5f22e6 Mon Sep 17 00:00:00 2001 From: Roman Gershman Date: Fri, 18 Oct 2024 16:29:52 +0300 Subject: [PATCH 2/2] chore: fixes --- src/server/main_service.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server/main_service.cc b/src/server/main_service.cc index 16df17b83a2a..ac99fe683e7e 100644 --- a/src/server/main_service.cc +++ b/src/server/main_service.cc @@ -1250,7 +1250,7 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx) const auto [cid, args_no_cmd] = registry_.FindExtended(cmd, args.subspan(1)); if (cid == nullptr) { - return cntx->SendError(ReportUnknownCmd(ArgS(args, 0))); + return cntx->SendError(ReportUnknownCmd(cmd)); } ConnectionContext* dfly_cntx = static_cast(cntx);