expose rocksdb perf/iostat context with ANALYZE command

apache · Dec 21, 2023 · b71c53d · b71c53d
1 parent a25c054
commit b71c53d
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 1 deletion.
diff --git a/src/commands/cmd_server.cc b/src/commands/cmd_server.cc
@@ -18,6 +18,9 @@
  *
  */
 
+#include <rocksdb/iostats_context.h>
+#include <rocksdb/perf_context.h>
+
 #include "command_parser.h"
 #include "commander.h"
 #include "commands/scan_base.h"
@@ -1123,6 +1126,75 @@ class CommandRdb : public Commander {
   uint32_t db_index_ = 0;
 };
 
+class CommandAnalyze : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) {
+    if (args.size() <= 1) return {Status::RedisExecErr, errInvalidSyntax};
+    for (int i = 1; i < args.size(); ++i) {
+      true_args_.push_back(args[i]);
+    }
+    return Status::OK();
+  }
+  Status Execute(Server *srv, Connection *conn, std::string *output) {
+    auto commands = redis::CommandTable::Get();
+    auto cmd_iter = commands->find(util::ToLower(true_args_[0]));
+    if (cmd_iter == commands->end()) {
+      // unsupported redis command
+      return {Status::RedisExecErr, errInvalidSyntax};
+    }
+    auto redis_cmd = cmd_iter->second;
+    auto cmd = redis_cmd->factory();
+    cmd->SetAttributes(redis_cmd);
+    cmd->SetArgs(true_args_);
+
+    auto s = cmd->Parse(true_args_);
+    if (!s.IsOK()) {
+      return s;
+    }
+
+    auto prev_perf_level = rocksdb::GetPerfLevel();
+    rocksdb::SetPerfLevel(rocksdb::PerfLevel::kEnableTimeExceptForMutex);
+    rocksdb::get_perf_context()->Reset();
+    rocksdb::get_iostats_context()->Reset();
+
+    std::string tmp_output;
+    s = cmd->Execute(srv, conn, &tmp_output);
+    if (!s.IsOK()) {
+      return s;
+    }
+
+    if (tmp_output[0] == '-') {
+      *output = tmp_output;
+      return s;
+    }
+
+    std::string perf_context = rocksdb::get_perf_context()->ToString(true);
+    std::string iostats_context = rocksdb::get_iostats_context()->ToString(true);
+    rocksdb::get_perf_context()->Reset();
+    rocksdb::get_iostats_context()->Reset();
+    rocksdb::SetPerfLevel(prev_perf_level);
+    if (perf_context.empty()) {
+      *output = tmp_output;
+      return Status::OK();  // request without db operation
+    }
+
+    if (tmp_output.at(0) == '*') {
+      // if multi bulkstring, append directly
+      tmp_output.append(redis::BulkString(perf_context));
+      tmp_output.append(redis::BulkString(iostats_context));
+      *output = tmp_output;
+    } else {
+      tmp_output = redis::ExtractData(tmp_output);
+      *output = redis::MultiBulkString(
+          {std::move(tmp_output), "(perf_context): " + perf_context, "(iostats_context): " + iostats_context});
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::vector<std::string> true_args_;
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loading", 0, 0, 0),
                         MakeCmdAttr<CommandPing>("ping", -1, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandSelect>("select", 2, "read-only", 0, 0, 0),
@@ -1157,6 +1229,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandAuth>("auth", 2, "read-only ok-loadin
                         MakeCmdAttr<CommandFlushBackup>("flushbackup", 1, "read-only no-script", 0, 0, 0),
                         MakeCmdAttr<CommandSlaveOf>("slaveof", 3, "read-only exclusive no-script", 0, 0, 0),
                         MakeCmdAttr<CommandStats>("stats", 1, "read-only", 0, 0, 0),
-                        MakeCmdAttr<CommandRdb>("rdb", -3, "write exclusive", 0, 0, 0), )
+                        MakeCmdAttr<CommandRdb>("rdb", -3, "write exclusive", 0, 0, 0),
+                        MakeCmdAttr<CommandAnalyze>("analyze", -1, "", 0, 0, 0), )
 
 }  // namespace redis
diff --git a/src/server/redis_reply.cc b/src/server/redis_reply.cc
@@ -34,6 +34,35 @@ std::string BulkString(const std::string &data) { return "$" + std::to_string(da
 
 std::string NilString() { return "$-1" CRLF; }
 
+std::string ExtractData(const std::string &reply) {
+  char code = reply[0];
+  switch (code) {
+    case '+': {
+      size_t pos = reply.find(CRLF);
+      return reply.substr(1, pos - 1);
+    }
+    case '-': {
+      size_t pos = reply.find(CRLF);
+      return reply.substr(1, pos - 1);
+    }
+    case '$': {
+      // nilstring
+      if (reply[1] == '-') {
+        return "nil";
+      } else {
+        // bulkstring
+        size_t pos = reply.find(CRLF);
+        if (pos != std::string::npos) {
+          size_t len = std::stoi(reply.substr(1, pos - 1));
+          return reply.substr(pos + 2, len);
+        }
+      }
+    }
+    default:
+      return reply;
+  }
+}
+
 std::string MultiBulkString(const std::vector<std::string> &values, bool output_nil_for_empty_string) {
   std::string result = "*" + std::to_string(values.size()) + CRLF;
   for (const auto &value : values) {

diff --git a/src/server/redis_reply.h b/src/server/redis_reply.h
@@ -53,4 +53,5 @@ std::string MultiBulkString(const std::vector<std::string> &values, bool output_
 std::string MultiBulkString(const std::vector<std::string> &values, const std::vector<rocksdb::Status> &statuses);
 std::string Command2RESP(const std::vector<std::string> &cmd_args);
 
+std::string ExtractData(const std::string &reply);
 }  // namespace redis