feat(logging): Move the logger remote_command handler to command_manager (#2087)

acelyc111 · web-flow · commit af3677e3f4f1 · 2024-08-07T14:06:06.000+08:00
diff --git a/src/utils/command_manager.cpp b/src/utils/command_manager.cpp
@@ -29,13 +29,17 @@
 #include <fmt/format.h>
 // IWYU pragma: no_include <ext/alloc_traits.h>
 #include <stdlib.h>
+#include <algorithm>
 #include <chrono>
 #include <limits>
 #include <sstream> // IWYU pragma: keep
 #include <thread>
+#include <type_traits>
 #include <unordered_set>
 #include <utility>
 
+#include "gutil/map_util.h"
+
 namespace dsn {
 
 std::unique_ptr<command_deregister>
@@ -44,19 +48,19 @@ command_manager::register_command(const std::vector<std::string> &commands,
                                   const std::string &args,
                                   command_handler handler)
 {
-    auto *c = new command_instance();
-    c->commands = commands;
-    c->help = help;
-    c->args = args;
-    c->handler = std::move(handler);
+    auto ch = std::make_shared<commands_handler>();
+    ch->commands = commands;
+    ch->help = help;
+    ch->args = args;
+    ch->handler = std::move(handler);
 
     utils::auto_write_lock l(_lock);
     for (const auto &cmd : commands) {
         CHECK(!cmd.empty(), "should not register empty command");
-        CHECK(_handlers.emplace(cmd, c).second, "command '{}' already registered", cmd);
+        gutil::InsertOrDie(&_handler_by_cmd, cmd, ch);
     }
 
-    return std::make_unique<command_deregister>(reinterpret_cast<uintptr_t>(c));
+    return std::make_unique<command_deregister>(reinterpret_cast<uintptr_t>(ch.get()));
 }
 
 std::unique_ptr<command_deregister> command_manager::register_bool_command(
@@ -94,35 +98,39 @@ command_manager::register_multiple_commands(const std::vector<std::string> &comm
                             handler);
 }
 
-void command_manager::deregister_command(uintptr_t handle)
+void command_manager::deregister_command(uintptr_t cmd_id)
 {
-    auto c = reinterpret_cast<command_instance *>(handle);
-    CHECK_NOTNULL(c, "cannot deregister a null handle");
+    const auto ch = reinterpret_cast<commands_handler *>(cmd_id);
+    CHECK_NOTNULL(ch, "cannot deregister a null command id");
     utils::auto_write_lock l(_lock);
-    for (const std::string &cmd : c->commands) {
-        _handlers.erase(cmd);
+    for (const auto &cmd : ch->commands) {
+        _handler_by_cmd.erase(cmd);
     }
 }
 
+void command_manager::add_global_cmd(std::unique_ptr<command_deregister> cmd)
+{
+    utils::auto_write_lock l(_lock);
+    _cmds.push_back(std::move(cmd));
+}
+
 bool command_manager::run_command(const std::string &cmd,
                                   const std::vector<std::string> &args,
                                   /*out*/ std::string &output)
 {
-    command_instance *h = nullptr;
+    std::shared_ptr<commands_handler> ch;
     {
         utils::auto_read_lock l(_lock);
-        auto it = _handlers.find(cmd);
-        if (it != _handlers.end())
-            h = it->second;
+        ch = gutil::FindPtrOrNull(_handler_by_cmd, cmd);
     }
 
-    if (h == nullptr) {
-        output = std::string("unknown command '") + cmd + "'";
+    if (!ch) {
+        output = fmt::format("unknown command '{}'", cmd);
         return false;
-    } else {
-        output = h->handler(args);
-        return true;
     }
+
+    output = ch->handler(args);
+    return true;
 }
 
 std::string command_manager::set_bool(bool &value,
@@ -159,36 +167,34 @@ std::string command_manager::set_bool(bool &value,
 
 command_manager::command_manager()
 {
-    _cmds.emplace_back(
-        register_multiple_commands({"help", "h", "H", "Help"},
-                                   "Display help information",
-                                   "[command]",
-                                   [this](const std::vector<std::string> &args) {
-                                       std::stringstream ss;
-                                       if (args.empty()) {
-                                           std::unordered_set<command_instance *> cmds;
-                                           utils::auto_read_lock l(_lock);
-                                           for (const auto &c : this->_handlers) {
-                                               // Multiple commands with the same handler are print
-                                               // only once.
-                                               if (cmds.insert(c.second.get()).second) {
-                                                   ss << c.second->help << std::endl;
-                                               }
-                                           }
-                                       } else {
-                                           utils::auto_read_lock l(_lock);
-                                           auto it = _handlers.find(args[0]);
-                                           if (it == _handlers.end()) {
-                                               ss << "cannot find command '" << args[0] << "'";
-                                           } else {
-                                               ss.width(6);
-                                               ss << std::left << it->second->help << std::endl
-                                                  << it->second->args << std::endl;
-                                           }
-                                       }
-
-                                       return ss.str();
-                                   }));
+    _cmds.emplace_back(register_multiple_commands(
+        {"help", "h", "H", "Help"},
+        "Display help information",
+        "[command]",
+        [this](const std::vector<std::string> &args) {
+            std::stringstream ss;
+            if (args.empty()) {
+                std::unordered_set<commands_handler *> chs;
+                utils::auto_read_lock l(_lock);
+                for (const auto &[_, ch] : _handler_by_cmd) {
+                    // Multiple commands with the same handler are print only once.
+                    if (gutil::InsertIfNotPresent(&chs, ch.get())) {
+                        ss << ch->help << std::endl;
+                    }
+                }
+            } else {
+                utils::auto_read_lock l(_lock);
+                const auto ch = gutil::FindPtrOrNull(_handler_by_cmd, args[0]);
+                if (!ch) {
+                    ss << "cannot find command '" << args[0] << "'";
+                } else {
+                    ss.width(6);
+                    ss << std::left << ch->help << std::endl << ch->args << std::endl;
+                }
+            }
+
+            return ss.str();
+        }));
 
     _cmds.emplace_back(register_multiple_commands(
         {"repeat", "r", "R", "Repeat"},
@@ -241,10 +247,10 @@ command_manager::command_manager()
 command_manager::~command_manager()
 {
     _cmds.clear();
-    CHECK(_handlers.empty(),
+    CHECK(_handler_by_cmd.empty(),
           "All commands must be deregistered before command_manager is destroyed, however '{}' is "
           "still registered",
-          _handlers.begin()->first);
+          _handler_by_cmd.begin()->first);
 }
 
 } // namespace dsn
diff --git a/src/utils/command_manager.h b/src/utils/command_manager.h
@@ -28,17 +28,16 @@
 
 #include <fmt/core.h>
 #include <fmt/format.h>
+#include <nlohmann/json.hpp>
+#include <nlohmann/json_fwd.hpp>
 // IWYU pragma: no_include <ext/alloc_traits.h>
 #include <stdint.h>
 #include <functional>
 #include <map>
 #include <memory>
-#include <nlohmann/json.hpp>
-#include <nlohmann/json_fwd.hpp>
 #include <string>
 #include <vector>
 
-#include "utils/autoref_ptr.h"
 #include "utils/fmt_logging.h"
 #include "utils/ports.h"
 #include "utils/singleton.h"
@@ -47,7 +46,6 @@
 #include "utils/synchronize.h"
 
 namespace dsn {
-
 class command_deregister;
 
 class command_manager : public ::dsn::utils::singleton<command_manager>
@@ -85,22 +83,25 @@ class command_manager : public ::dsn::utils::singleton<command_manager>
             });
     }
 
-    // Register a single 'command' with the 'help' description, its arguments are describe in
+    // Register a single 'command' with the 'help' description, its arguments are described in
     // 'args'.
     std::unique_ptr<command_deregister>
     register_single_command(const std::string &command,
                             const std::string &help,
                             const std::string &args,
                             command_handler handler) WARN_UNUSED_RESULT;
 
-    // Register multiple 'commands' with the 'help' description, their arguments are describe in
+    // Register multiple 'commands' with the 'help' description, their arguments are described in
     // 'args'.
     std::unique_ptr<command_deregister>
     register_multiple_commands(const std::vector<std::string> &commands,
                                const std::string &help,
                                const std::string &args,
                                command_handler handler) WARN_UNUSED_RESULT;
 
+    // Register a global command which is not associated with any objects.
+    void add_global_cmd(std::unique_ptr<command_deregister> cmd);
+
     bool run_command(const std::string &cmd,
                      const std::vector<std::string> &args,
                      /*out*/ std::string &output);
@@ -112,7 +113,7 @@ class command_manager : public ::dsn::utils::singleton<command_manager>
     command_manager();
     ~command_manager();
 
-    struct command_instance : public ref_counter
+    struct commands_handler
     {
         std::vector<std::string> commands;
         std::string help;
@@ -126,7 +127,7 @@ class command_manager : public ::dsn::utils::singleton<command_manager>
                      const std::string &args,
                      command_handler handler) WARN_UNUSED_RESULT;
 
-    void deregister_command(uintptr_t handle);
+    void deregister_command(uintptr_t cmd_id);
 
     static std::string
     set_bool(bool &value, const std::string &name, const std::vector<std::string> &args);
@@ -177,9 +178,8 @@ class command_manager : public ::dsn::utils::singleton<command_manager>
         return msg.dump(2);
     }
 
-    typedef ref_ptr<command_instance> command_instance_ptr;
     utils::rw_lock_nr _lock;
-    std::map<std::string, command_instance_ptr> _handlers;
+    std::map<std::string, std::shared_ptr<commands_handler>> _handler_by_cmd;
 
     std::vector<std::unique_ptr<command_deregister>> _cmds;
 };
diff --git a/src/utils/logging_provider.h b/src/utils/logging_provider.h
@@ -64,16 +64,13 @@ class logging_provider
 
     virtual void flush() = 0;
 
-    void deregister_commands() { _cmds.clear(); }
-
 protected:
     static std::unique_ptr<logging_provider> _logger;
 
     static logging_provider *create_default_instance();
 
     logging_provider(log_level_t stderr_start_level) : _stderr_start_level(stderr_start_level) {}
 
-    std::vector<std::unique_ptr<command_deregister>> _cmds;
     const log_level_t _stderr_start_level;
 };
 
diff --git a/src/utils/simple_logger.cpp b/src/utils/simple_logger.cpp
@@ -36,7 +36,7 @@
 #include <cstdint>
 #include <ctime>
 #include <functional>
-#include <memory>
+#include <mutex>
 #include <type_traits>
 #include <utility>
 #include <vector>
@@ -234,37 +234,39 @@ simple_logger::simple_logger(const char *log_dir, const char *role_name)
 
     create_log_file();
 
-    // TODO(yingchun): simple_logger is destroyed after command_manager, so will cause crash like
-    //  "assertion expression: [_handlers.empty()] All commands must be deregistered before
-    //  command_manager is destroyed, however 'flush-log' is still registered".
-    //  We need to fix it.
-    _cmds.emplace_back(::dsn::command_manager::instance().register_single_command(
-        "flush-log",
-        "Flush log to stderr or file",
-        "",
-        [this](const std::vector<std::string> &args) {
-            this->flush();
-            return "Flush done.";
-        }));
-
-    _cmds.emplace_back(::dsn::command_manager::instance().register_single_command(
-        "reset-log-start-level",
-        "Reset the log start level",
-        "[DEBUG | INFO | WARNING | ERROR | FATAL]",
-        [](const std::vector<std::string> &args) {
-            log_level_t start_level;
-            if (args.size() == 0) {
-                start_level = enum_from_string(FLAGS_logging_start_level, LOG_LEVEL_INVALID);
-            } else {
-                std::string level_str = "LOG_LEVEL_" + args[0];
-                start_level = enum_from_string(level_str.c_str(), LOG_LEVEL_INVALID);
-                if (start_level == LOG_LEVEL_INVALID) {
-                    return "ERROR: invalid level '" + args[0] + "'";
-                }
-            }
-            set_log_start_level(start_level);
-            return std::string("OK, current level is ") + enum_to_string(start_level);
-        }));
+    static std::once_flag flag;
+    std::call_once(flag, [&]() {
+        ::dsn::command_manager::instance().add_global_cmd(
+            ::dsn::command_manager::instance().register_single_command(
+                "flush-log",
+                "Flush log to stderr or file",
+                "",
+                [this](const std::vector<std::string> &args) {
+                    this->flush();
+                    return "Flush done.";
+                }));
+
+        ::dsn::command_manager::instance().add_global_cmd(
+            ::dsn::command_manager::instance().register_single_command(
+                "reset-log-start-level",
+                "Reset the log start level",
+                "[DEBUG | INFO | WARNING | ERROR | FATAL]",
+                [](const std::vector<std::string> &args) {
+                    log_level_t start_level;
+                    if (args.size() == 0) {
+                        start_level =
+                            enum_from_string(FLAGS_logging_start_level, LOG_LEVEL_INVALID);
+                    } else {
+                        std::string level_str = "LOG_LEVEL_" + args[0];
+                        start_level = enum_from_string(level_str.c_str(), LOG_LEVEL_INVALID);
+                        if (start_level == LOG_LEVEL_INVALID) {
+                            return "ERROR: invalid level '" + args[0] + "'";
+                        }
+                    }
+                    set_log_start_level(start_level);
+                    return std::string("OK, current level is ") + enum_to_string(start_level);
+                }));
+    });
 }
 
 void simple_logger::create_log_file()
diff --git a/src/utils/test/logger.cpp b/src/utils/test/logger.cpp
