chore: remove locks from acl validation in the hot path (#1765)

* remove locking the registry when Validating users
* deep copies the acl categories in the connection context
* streamline updates to the acl categories via the proactor pool

Author: kostasrim

src/server/acl/acl_family.cc

@@ -3,6 +3,8 @@
 
 #include "server/acl/acl_family.h"
 
+#include <glog/logging.h>
+
 #include <cctype>
 #include <optional>
 #include <variant>
@@ -12,6 +14,7 @@
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "core/overloaded.h"
+#include "facade/dragonfly_connection.h"
 #include "facade/facade_types.h"
 #include "server/acl/acl_commands_def.h"
 #include "server/command_registry.h"
@@ -157,12 +160,31 @@ std::variant<User::UpdateRequest, ErrorReply> ParseAclSetUser(CmdArgList args) {
 
 }  // namespace
 
+void AclFamily::StreamUpdatesToAllProactorConnections(std::string_view user, uint32_t update_cat) {
+  auto update_cb = [user, update_cat]([[maybe_unused]] size_t id, util::Connection* conn) {
+    DCHECK(conn);
+    auto connection = static_cast<facade::Connection*>(conn);
+    auto ctx = static_cast<ConnectionContext*>(connection->cntx());
+    if (ctx && user == ctx->authed_username) {
+      ctx->acl_categories = update_cat;
+    }
+  };
+
+  if (main_listener_) {
+    main_listener_->TraverseConnections(update_cb);
+  }
+}
+
 void AclFamily::SetUser(CmdArgList args, ConnectionContext* cntx) {
   std::string_view username = facade::ToSV(args[0]);
   auto req = ParseAclSetUser(args.subspan(1));
   auto error_case = [cntx](ErrorReply&& error) { (*cntx)->SendError(error); };
-  auto update_case = [username, cntx](User::UpdateRequest&& req) {
-    ServerState::tlocal()->user_registry->MaybeAddAndUpdate(username, std::move(req));
+  auto update_case = [username, cntx, this](User::UpdateRequest&& req) {
+    auto& registry = ServerState::tlocal()->user_registry;
+    auto user_with_lock = registry->MaybeAddAndUpdateWithLock(username, std::move(req));
+    if (user_with_lock.exists) {
+      StreamUpdatesToAllProactorConnections(username, user_with_lock.user.AclCategory());
+    }
     (*cntx)->SendOk();
   };
 
@@ -171,8 +193,15 @@ void AclFamily::SetUser(CmdArgList args, ConnectionContext* cntx) {
 
 using CI = dfly::CommandId;
 
-#define HFUNC(x) SetHandler(&AclFamily::x)
+using MemberFunc = void (AclFamily::*)(CmdArgList args, ConnectionContext* cntx);
+
+inline CommandId::Handler HandlerFunc(AclFamily* acl, MemberFunc f) {
+  return [=](CmdArgList args, ConnectionContext* cntx) { return (acl->*f)(args, cntx); };
+}
+
+#define HFUNC(x) SetHandler(HandlerFunc(this, &AclFamily::x))
 
+constexpr uint32_t kAcl = acl::CONNECTION;
 constexpr uint32_t kList = acl::ADMIN | acl::SLOW | acl::DANGEROUS;
 constexpr uint32_t kSetUser = acl::ADMIN | acl::SLOW | acl::DANGEROUS;
 
@@ -184,7 +213,7 @@ constexpr uint32_t kSetUser = acl::ADMIN | acl::SLOW | acl::DANGEROUS;
 // easy to handle that case explicitly in `DispatchCommand`.
 
 void AclFamily::Register(dfly::CommandRegistry* registry) {
-  *registry << CI{"ACL", CO::NOSCRIPT | CO::LOADING, 0, 0, 0, 0, acl::kList}.HFUNC(Acl);
+  *registry << CI{"ACL", CO::NOSCRIPT | CO::LOADING, 0, 0, 0, 0, acl::kAcl}.HFUNC(Acl);
   *registry << CI{"ACL LIST", CO::ADMIN | CO::NOSCRIPT | CO::LOADING, 1, 0, 0, 0, acl::kList}.HFUNC(
       List);
   *registry << CI{"ACL SETUSER", CO::ADMIN | CO::NOSCRIPT | CO::LOADING, -2, 0, 0, 0, acl::kSetUser}
@@ -193,4 +222,8 @@ void AclFamily::Register(dfly::CommandRegistry* registry) {
 
 #undef HFUNC
 
+void AclFamily::Init(facade::Listener* main_listener) {
+  main_listener_ = main_listener;
+}
+
 }  // namespace dfly::acl

src/server/acl/acl_family.h

@@ -4,22 +4,38 @@
 
 #pragma once
 
+#include <cstdint>
+#include <string_view>
+#include <vector>
+
+#include "facade/dragonfly_listener.h"
+#include "helio/util/proactor_pool.h"
 #include "server/common.h"
 
 namespace dfly {
+
 class ConnectionContext;
 class CommandRegistry;
 
 namespace acl {
 
-class AclFamily {
+class AclFamily final {
  public:
-  static void Register(CommandRegistry* registry);
+  AclFamily() = default;
+
+  void Register(CommandRegistry* registry);
+  void Init(facade::Listener* listener);
 
  private:
-  static void Acl(CmdArgList args, ConnectionContext* cntx);
-  static void List(CmdArgList args, ConnectionContext* cntx);
-  static void SetUser(CmdArgList args, ConnectionContext* cntx);
+  void Acl(CmdArgList args, ConnectionContext* cntx);
+  void List(CmdArgList args, ConnectionContext* cntx);
+  void SetUser(CmdArgList args, ConnectionContext* cntx);
+
+  // Helper function that updates all open connections and their
+  // respective ACL fields on all the available proactor threads
+  void StreamUpdatesToAllProactorConnections(std::string_view user, uint32_t update_cat);
+
+  facade::Listener* main_listener_{nullptr};
 };
 
 }  // namespace acl

src/server/acl/user_registry.cc

@@ -55,14 +55,28 @@ bool UserRegistry::AuthUser(std::string_view username, std::string_view password
   return user->second.IsActive() && user->second.HasPassword(password);
 }
 
-UserRegistry::RegistryViewWithLock::RegistryViewWithLock(std::shared_lock<util::SharedMutex> mu,
+UserRegistry::RegistryViewWithLock::RegistryViewWithLock(std::shared_lock<util::SharedMutex> lk,
                                                          const RegistryType& registry)
-    : registry(registry), registry_mu_(std::move(mu)) {
+    : registry(registry), registry_lk_(std::move(lk)) {
 }
 
 UserRegistry::RegistryViewWithLock UserRegistry::GetRegistryWithLock() const {
   std::shared_lock<util::SharedMutex> lock(mu_);
   return {std::move(lock), registry_};
 }
 
+UserRegistry::UserViewWithLock::UserViewWithLock(std::shared_lock<util::SharedMutex> lk,
+                                                 const User& user, bool exists)
+    : user(user), exists(exists), registry_lk_(std::move(lk)) {
+}
+
+UserRegistry::UserViewWithLock UserRegistry::MaybeAddAndUpdateWithLock(std::string_view username,
+                                                                       User::UpdateRequest req) {
+  std::shared_lock<util::SharedMutex> lock(mu_);
+  const bool exists = registry_.contains(username);
+  auto& user = registry_[username];
+  user.Update(std::move(req));
+  return {std::move(lock), user, exists};
+}
+
 }  // namespace dfly::acl

src/server/acl/user_registry.h

@@ -30,7 +30,6 @@ class UserRegistry {
   // If the user with name `username` does not exist, it's added in the store with
   // the exact fields found in req
   // If the user exists, the bitfields are updated with a `logical and` operation
-  // TODO change return time to communicate back results to acl commands
   void MaybeAddAndUpdate(std::string_view username, User::UpdateRequest req);
 
   // Acquires a write lock on mu_
@@ -55,16 +54,29 @@ class UserRegistry {
   // Helper class for accessing the registry with a ReadLock outside the scope of UserRegistry
   class RegistryViewWithLock {
    public:
-    RegistryViewWithLock(std::shared_lock<util::SharedMutex> mu, const RegistryType& registry);
+    RegistryViewWithLock(std::shared_lock<util::SharedMutex> lk, const RegistryType& registry);
     const RegistryType& registry;
 
    private:
-    std::shared_lock<util::SharedMutex> registry_mu_;
+    std::shared_lock<util::SharedMutex> registry_lk_;
   };
 
   // Helper function used for printing users via ACL LIST
   RegistryViewWithLock GetRegistryWithLock() const;
 
+  // Helper class for accessing a user with a ReadLock outside the scope of UserRegistry
+  class UserViewWithLock {
+   public:
+    UserViewWithLock(std::shared_lock<util::SharedMutex> lk, const User& user, bool exists);
+    const User& user;
+    const bool exists;
+
+   private:
+    std::shared_lock<util::SharedMutex> registry_lk_;
+  };
+
+  UserViewWithLock MaybeAddAndUpdateWithLock(std::string_view username, User::UpdateRequest req);
+
  private:
   RegistryType registry_;
   // TODO add abseil mutex attributes

src/server/acl/validator.cc

@@ -4,16 +4,12 @@
 
 #include "server/acl/validator.h"
 
-#include "server/server_state.h"
-
 namespace dfly::acl {
 
 [[nodiscard]] bool IsUserAllowedToInvokeCommand(const ConnectionContext& cntx,
                                                 const facade::CommandId& id) {
-  auto& registry = *ServerState::tlocal()->user_registry;
-  auto credentials = registry.GetCredentials(cntx.authed_username);
   auto command_credentials = id.acl_categories();
-  return (credentials.acl_categories & command_credentials) != 0;
+  return (cntx.acl_categories & command_credentials) != 0;
 }
 
 }  // namespace dfly::acl

src/server/acl/validator.h

@@ -4,8 +4,6 @@
 
 #pragma once
 
-#include <string_view>
-
 #include "facade/command_id.h"
 #include "server/conn_context.h"
 

src/server/conn_context.h

@@ -7,6 +7,7 @@
 #include <absl/container/fixed_array.h>
 #include <absl/container/flat_hash_set.h>
 
+#include "acl/acl_commands_def.h"
 #include "core/fibers.h"
 #include "facade/conn_context.h"
 #include "facade/reply_capture.h"
@@ -197,6 +198,7 @@ class ConnectionContext : public facade::ConnectionContext {
   FlowInfo* replication_flow;
 
   std::string authed_username{"default"};
+  uint32_t acl_categories{acl::ALL};
 
  private:
   void EnableMonitoring(bool enable) {

src/server/dfly_main.cc

@@ -330,6 +330,9 @@ bool RunEngine(ProactorPool* pool, AcceptServer* acceptor) {
   Listener* main_listener = nullptr;
 
   std::vector<facade::Listener*> listeners;
+  // If we ever add a new listener, plz don't change this,
+  // we depend on tcp listener to be at the front since we later
+  // need to pass it to the AclFamily::Init
   if (!tcp_disabled) {
     main_listener = new Listener{Protocol::REDIS, &service};
     listeners.push_back(main_listener);

src/server/main_service.cc

@@ -664,7 +664,12 @@ void Service::Init(util::AcceptServer* acceptor, std::vector<facade::Listener*>
   }
 
   shard_set->Init(shard_num, !opts.disable_time_update);
-
+  const auto tcp_disabled = GetFlag(FLAGS_port) == 0u;
+  // We assume that listeners.front() is the main_listener
+  // see dfly_main RunEngine
+  if (!tcp_disabled && !listeners.empty()) {
+    acl_family_.Init(listeners.front());
+  }
   request_latency_usec.Init(&pp_);
   StringFamily::Init(&pp_);
   GenericFamily::Init(&pp_);
@@ -2132,7 +2137,7 @@ void Service::RegisterCommands() {
   BitOpsFamily::Register(&registry_);
   HllFamily::Register(&registry_);
   SearchFamily::Register(&registry_);
-  acl::AclFamily::Register(&registry_);
+  acl_family_.Register(&registry_);
 
   server_family_.Register(&registry_);
   cluster_family_.Register(&registry_);

src/server/main_service.h

@@ -9,6 +9,7 @@
 #include "base/varz_value.h"
 #include "core/interpreter.h"
 #include "facade/service_interface.h"
+#include "server/acl/acl_family.h"
 #include "server/acl/user_registry.h"
 #include "server/cluster/cluster_family.h"
 #include "server/command_registry.h"
@@ -165,6 +166,7 @@ class Service : public facade::ServiceInterface {
   util::ProactorPool& pp_;
 
   acl::UserRegistry user_registry_;
+  acl::AclFamily acl_family_;
   ServerFamily server_family_;
   ClusterFamily cluster_family_;
   CommandRegistry registry_;