feat(server): Add support for command aliasing (#4932)

Add support for command aliasing using command_alias flag

Signed-off-by: Abhijat Malviya <abhijat@dragonflydb.io>

Author: abhijat

src/server/acl/acl_family.cc

@@ -591,7 +591,7 @@ void AclFamily::DryRun(CmdArgList args, const CommandContext& cmd_cntx) {
 
   string command = absl::AsciiStrToUpper(ArgS(args, 1));
   auto* cid = cmd_registry_->Find(command);
-  if (!cid) {
+  if (!cid || cid->IsAlias()) {
     auto error = absl::StrCat("Command '", command, "' not found");
     rb->SendError(error);
     return;
@@ -1062,7 +1062,7 @@ std::pair<AclFamily::OptCommand, bool> AclFamily::MaybeParseAclCommand(
     std::string_view command) const {
   if (absl::StartsWith(command, "+")) {
     auto res = cmd_registry_->Find(command.substr(1));
-    if (!res) {
+    if (!res || res->IsAlias()) {
       return {};
     }
     std::pair<size_t, uint64_t> cmd{res->GetFamily(), res->GetBitIndex()};
@@ -1071,7 +1071,7 @@ std::pair<AclFamily::OptCommand, bool> AclFamily::MaybeParseAclCommand(
 
   if (absl::StartsWith(command, "-")) {
     auto res = cmd_registry_->Find(command.substr(1));
-    if (!res) {
+    if (!res || res->IsAlias()) {
       return {};
     }
     std::pair<size_t, uint64_t> cmd{res->GetFamily(), res->GetBitIndex()};

src/server/acl/acl_family_test.cc

@@ -19,6 +19,7 @@
 using namespace testing;
 
 ABSL_DECLARE_FLAG(std::vector<std::string>, rename_command);
+ABSL_DECLARE_FLAG(std::vector<std::string>, command_alias);
 
 namespace dfly {
 
@@ -29,6 +30,7 @@ class AclFamilyTest : public BaseFamilyTest {
 class AclFamilyTestRename : public BaseFamilyTest {
   void SetUp() override {
     absl::SetFlag(&FLAGS_rename_command, {"ACL=ROCKS"});
+    absl::SetFlag(&FLAGS_command_alias, {"___SET=SET"});
     ResetService();
   }
 };
@@ -538,4 +540,22 @@ TEST_F(AclFamilyTest, TestPubSub) {
   EXPECT_THAT(vec[9], "resetchannels &foo");
 }
 
+TEST_F(AclFamilyTest, TestAlias) {
+  auto resp = Run({"ACL", "SETUSER", "luke", "+___SET"});
+  EXPECT_THAT(resp, ErrArg("ERR Unrecognized parameter +___SET"));
+
+  resp = Run({"ACL", "SETUSER", "leia", "-___SET"});
+  EXPECT_THAT(resp, ErrArg("ERR Unrecognized parameter -___SET"));
+
+  resp = Run({"ACL", "SETUSER", "anakin", "+SET"});
+  EXPECT_EQ(resp, "OK");
+
+  resp = Run({"ACL", "SETUSER", "jarjar", "allcommands"});
+  EXPECT_EQ(resp, "OK");
+
+  resp = Run({"ACL", "DRYRUN", "jarjar", "___SET"});
+  EXPECT_THAT(resp, ErrArg("ERR Command '___SET' not found"));
+  EXPECT_EQ(Run({"ACL", "DRYRUN", "jarjar", "SET"}), "OK");
+}
+
 }  // namespace dfly

src/server/acl/validator.cc

@@ -66,6 +66,10 @@ bool ValidateCommand(const std::vector<uint64_t>& acl_commands, const CommandId&
     return true;
   }
 
+  if (id.IsAlias()) {
+    return false;
+  }
+
   std::pair<bool, AclLog::Reason> auth_res;
 
   if (id.IsPubSub() || id.IsShardedPSub()) {

src/server/command_registry.cc

@@ -22,14 +22,16 @@ using namespace std;
 ABSL_FLAG(vector<string>, rename_command, {},
           "Change the name of commands, format is: <cmd1_name>=<cmd1_new_name>, "
           "<cmd2_name>=<cmd2_new_name>");
-ABSL_FLAG(vector<string>, command_alias, {},
-          "Add an alias for given commands, format is: <alias>=<original>, "
-          "<alias>=<original>");
 ABSL_FLAG(vector<string>, restricted_commands, {},
           "Commands restricted to connections on the admin port");
 
 ABSL_FLAG(vector<string>, oom_deny_commands, {},
           "Additinal commands that will be marked as denyoom");
+
+ABSL_FLAG(vector<string>, command_alias, {},
+          "Add an alias for given command(s), format is: <alias>=<original>, <alias>=<original>. "
+          "Aliases must be set identically on replicas, if applicable");
+
 namespace dfly {
 
 using namespace facade;
@@ -75,16 +77,17 @@ uint32_t ImplicitAclCategories(uint32_t mask) {
   return out;
 }
 
-absl::flat_hash_map<std::string, std::string> ParseCmdlineArgMap(
-    const absl::Flag<std::vector<std::string>>& flag, const bool allow_duplicates = false) {
+using CmdLineMapping = absl::flat_hash_map<std::string, std::string>;
+
+CmdLineMapping ParseCmdlineArgMap(const absl::Flag<std::vector<std::string>>& flag) {
   const auto& mappings = absl::GetFlag(flag);
-  absl::flat_hash_map<std::string, std::string> parsed_mappings;
+  CmdLineMapping parsed_mappings;
   parsed_mappings.reserve(mappings.size());
 
   for (const std::string& mapping : mappings) {
-    std::vector<std::string_view> kv = absl::StrSplit(mapping, '=');
+    absl::InlinedVector<std::string_view, 2> kv = absl::StrSplit(mapping, '=');
     if (kv.size() != 2) {
-      LOG(ERROR) << "Malformed command " << mapping << " for " << flag.Name()
+      LOG(ERROR) << "Malformed command '" << mapping << "' for " << flag.Name()
                  << ", expected key=value";
       exit(1);
     }
@@ -97,15 +100,27 @@ absl::flat_hash_map<std::string, std::string> ParseCmdlineArgMap(
       exit(1);
     }
 
-    const bool inserted = parsed_mappings.emplace(std::move(key), std::move(value)).second;
-    if (!allow_duplicates && !inserted) {
+    if (!parsed_mappings.emplace(std::move(key), std::move(value)).second) {
       LOG(ERROR) << "Duplicate insert to " << flag.Name() << " not allowed";
       exit(1);
     }
   }
   return parsed_mappings;
 }
 
+CmdLineMapping OriginalToAliasMap() {
+  CmdLineMapping original_to_alias;
+  CmdLineMapping alias_to_original = ParseCmdlineArgMap(FLAGS_command_alias);
+  original_to_alias.reserve(alias_to_original.size());
+  std::for_each(std::make_move_iterator(alias_to_original.begin()),
+                std::make_move_iterator(alias_to_original.end()),
+                [&original_to_alias](auto&& pair) {
+                  original_to_alias.emplace(std::move(pair.second), std::move(pair.first));
+                });
+
+  return original_to_alias;
+}
+
 }  // namespace
 
 CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first_key,
@@ -115,6 +130,17 @@ CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first
   implicit_acl_ = !acl_categories.has_value();
 }
 
+CommandId CommandId::Clone(const std::string_view name) const {
+  CommandId cloned =
+      CommandId{name.data(), opt_mask_, arity_, first_key_, last_key_, acl_categories_};
+  cloned.handler_ = handler_;
+  cloned.opt_mask_ = opt_mask_ | CO::HIDDEN;
+  cloned.acl_categories_ = acl_categories_;
+  cloned.implicit_acl_ = implicit_acl_;
+  cloned.is_alias_ = true;
+  return cloned;
+}
+
 bool CommandId::IsTransactional() const {
   if (first_key_ > 0 || (opt_mask_ & CO::GLOBAL_TRANS) || (opt_mask_ & CO::NO_KEY_TRANSACTIONAL))
     return true;
@@ -130,16 +156,15 @@ bool CommandId::IsMultiTransactional() const {
   return CO::IsTransKind(name()) || CO::IsEvalKind(name());
 }
 
-uint64_t CommandId::Invoke(CmdArgList args, const CommandContext& cmd_cntx,
-                           std::string_view orig_cmd_name) const {
+uint64_t CommandId::Invoke(CmdArgList args, const CommandContext& cmd_cntx) const {
   int64_t before = absl::GetCurrentTimeNanos();
   handler_(args, cmd_cntx);
   int64_t after = absl::GetCurrentTimeNanos();
 
   ServerState* ss = ServerState::tlocal();  // Might have migrated thread, read after invocation
   int64_t execution_time_usec = (after - before) / 1000;
 
-  auto& ent = command_stats_[ss->thread_index()][orig_cmd_name];
+  auto& ent = command_stats_[ss->thread_index()];
 
   ++ent.first;
   ent.second += execution_time_usec;
@@ -169,7 +194,6 @@ optional<facade::ErrorReply> CommandId::Validate(CmdArgList tail_args) const {
 
 CommandRegistry::CommandRegistry() {
   cmd_rename_map_ = ParseCmdlineArgMap(FLAGS_rename_command);
-  cmd_aliases_ = ParseCmdlineArgMap(FLAGS_command_alias, true);
 
   for (string name : GetFlag(FLAGS_restricted_commands)) {
     restricted_cmds_.emplace(AsciiStrToUpper(name));
@@ -181,9 +205,20 @@ CommandRegistry::CommandRegistry() {
 }
 
 void CommandRegistry::Init(unsigned int thread_count) {
+  const CmdLineMapping original_to_alias = OriginalToAliasMap();
+  absl::flat_hash_map<std::string, CommandId> alias_to_command_id;
+  alias_to_command_id.reserve(original_to_alias.size());
   for (auto& [_, cmd] : cmd_map_) {
     cmd.Init(thread_count);
+    if (auto it = original_to_alias.find(cmd.name()); it != original_to_alias.end()) {
+      auto alias_cmd = cmd.Clone(it->second);
+      alias_cmd.Init(thread_count);
+      alias_to_command_id.insert({it->second, std::move(alias_cmd)});
+    }
   }
+  std::copy(std::make_move_iterator(alias_to_command_id.begin()),
+            std::make_move_iterator(alias_to_command_id.end()),
+            std::inserter(cmd_map_, cmd_map_.end()));
 }
 
 CommandRegistry& CommandRegistry::operator<<(CommandId cmd) {
@@ -212,7 +247,7 @@ CommandRegistry& CommandRegistry::operator<<(CommandId cmd) {
 
   if (!is_sub_command || absl::StartsWith(cmd.name(), "ACL")) {
     cmd.SetBitIndex(1ULL << bit_index_);
-    family_of_commands_.back().push_back(std::string(k));
+    family_of_commands_.back().emplace_back(k);
     ++bit_index_;
   } else {
     DCHECK(absl::StartsWith(k, family_of_commands_.back().back()));
@@ -266,10 +301,6 @@ std::pair<const CommandId*, ArgSlice> CommandRegistry::FindExtended(string_view
   return {res, tail_args};
 }
 
-bool CommandRegistry::IsAlias(std::string_view cmd) const {
-  return cmd_aliases_.contains(cmd);
-}
-
 namespace CO {
 
 const char* OptName(CO::CommandOpt fl) {

src/server/command_registry.h

@@ -71,9 +71,8 @@ static_assert(!IsEvalKind(""));
 
 };  // namespace CO
 
-// Per thread vector of command stats. Each entry is:
-// command invocation string -> {cmd_calls, cmd_latency_agg in usec}.
-using CmdCallStats = absl::flat_hash_map<std::string, std::pair<uint64_t, uint64_t>>;
+// Per thread vector of command stats. Each entry is {cmd_calls, cmd_latency_agg in usec}.
+using CmdCallStats = std::pair<uint64_t, uint64_t>;
 
 struct CommandContext {
   CommandContext(Transaction* _tx, facade::SinkReplyBuilder* _rb, ConnectionContext* cntx)
@@ -94,6 +93,8 @@ class CommandId : public facade::CommandId {
 
   CommandId(CommandId&&) = default;
 
+  [[nodiscard]] CommandId Clone(std::string_view name) const;
+
   void Init(unsigned thread_count) {
     command_stats_ = std::make_unique<CmdCallStats[]>(thread_count);
   }
@@ -103,10 +104,8 @@ class CommandId : public facade::CommandId {
   using ArgValidator = fu2::function_base<true, true, fu2::capacity_default, false, false,
                                           std::optional<facade::ErrorReply>(CmdArgList) const>;
 
-  // Invokes the command handler. Returns the invoke time in usec. The invoked_by parameter is set
-  // to the string passed in by user, if available. If not set, defaults to command name.
-  uint64_t Invoke(CmdArgList args, const CommandContext& cmd_cntx,
-                  std::string_view orig_cmd_name) const;
+  // Returns the invoke time in usec.
+  uint64_t Invoke(CmdArgList args, const CommandContext& cmd_cntx) const;
 
   // Returns error if validation failed, otherwise nullopt
   std::optional<facade::ErrorReply> Validate(CmdArgList tail_args) const;
@@ -144,7 +143,7 @@ class CommandId : public facade::CommandId {
   }
 
   void ResetStats(unsigned thread_index) {
-    command_stats_[thread_index].clear();
+    command_stats_[thread_index] = {0, 0};
   }
 
   CmdCallStats GetStats(unsigned thread_index) const {
@@ -156,11 +155,16 @@ class CommandId : public facade::CommandId {
       acl_categories_ |= mask;
   }
 
+  bool IsAlias() const {
+    return is_alias_;
+  }
+
  private:
   bool implicit_acl_;
   std::unique_ptr<CmdCallStats[]> command_stats_;
   Handler3 handler_;
   ArgValidator validator_;
+  bool is_alias_{false};
 };
 
 class CommandRegistry {
@@ -172,16 +176,8 @@ class CommandRegistry {
   CommandRegistry& operator<<(CommandId cmd);
 
   const CommandId* Find(std::string_view cmd) const {
-    if (const auto it = cmd_map_.find(cmd); it != cmd_map_.end()) {
-      return &it->second;
-    }
-
-    if (const auto it = cmd_aliases_.find(cmd); it != cmd_aliases_.end()) {
-      if (const auto alias_lookup = cmd_map_.find(it->second); alias_lookup != cmd_map_.end()) {
-        return &alias_lookup->second;
-      }
-    }
-    return nullptr;
+    auto it = cmd_map_.find(cmd);
+    return it == cmd_map_.end() ? nullptr : &it->second;
   }
 
   CommandId* Find(std::string_view cmd) {
@@ -203,17 +199,13 @@ class CommandRegistry {
     }
   }
 
-  void MergeCallStats(
-      unsigned thread_index,
-      std::function<void(std::string_view, const CmdCallStats::mapped_type&)> cb) const {
-    for (const auto& [_, cmd_id] : cmd_map_) {
-      for (const auto& [cmd_name, call_stats] : cmd_id.GetStats(thread_index)) {
-        if (call_stats.first == 0) {
-          continue;
-        }
-
-        cb(cmd_name, call_stats);
-      }
+  void MergeCallStats(unsigned thread_index,
+                      std::function<void(std::string_view, const CmdCallStats&)> cb) const {
+    for (const auto& k_v : cmd_map_) {
+      auto src = k_v.second.GetStats(thread_index);
+      if (src.first == 0)
+        continue;
+      cb(k_v.second.name(), src);
     }
   }
 
@@ -227,16 +219,9 @@ class CommandRegistry {
   std::pair<const CommandId*, facade::ArgSlice> FindExtended(std::string_view cmd,
                                                              facade::ArgSlice tail_args) const;
 
-  bool IsAlias(std::string_view cmd) const;
-
  private:
   absl::flat_hash_map<std::string, CommandId> cmd_map_;
   absl::flat_hash_map<std::string, std::string> cmd_rename_map_;
-  // Stores a mapping from alias to original command. During the find operation, the first lookup is
-  // done in the cmd_map_, then in the alias map. This results in two lookups but only for commands
-  // which are not in original map, ie either typos or aliases. While it would be faster, we cannot
-  // store iterators into cmd_map_ here as they may be invalidated on rehashing.
-  absl::flat_hash_map<std::string, std::string> cmd_aliases_;
   absl::flat_hash_set<std::string> restricted_cmds_;
   absl::flat_hash_set<std::string> oomdeny_cmds_;