From 46957db593bafe2fe3f70f4138c2eda69c2f5be2 Mon Sep 17 00:00:00 2001
From: chrisxu333 <chrisxu333@gmail.com>
Date: Mon, 18 Dec 2023 07:54:35 -0700
Subject: [PATCH] expose rocksdb perf/iostat context with ANALYZE command

---
 src/commands/cmd_server.cc | 71 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/src/commands/cmd_server.cc b/src/commands/cmd_server.cc
index b47b63b79df..c06f13f42cd 100644
--- 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,71 @@ class CommandRdb : public Commander {
   uint32_t db_index_ = 0;
 };
 
+class CommandAnalyze : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    if (args.size() <= 1) return {Status::RedisExecErr, errInvalidSyntax};
+    for (int i = 1; i < args.size(); ++i) {
+      command_args_.push_back(args[i]);
+    }
+    return Status::OK();
+  }
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    auto commands = redis::CommandTable::Get();
+    auto cmd_iter = commands->find(util::ToLower(command_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(command_args_);
+
+    int arity = cmd->GetAttributes()->arity;
+    if ((arity > 0 && command_args_.size() != arity) || (arity < 0 && command_args_.size() < -arity)) {
+      *output = redis::Error("ERR wrong number of arguments");
+      return {Status::RedisExecErr, errWrongNumOfArguments};
+    }
+
+    auto s = cmd->Parse(command_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 command_output;
+    s = cmd->Execute(srv, conn, &command_output);
+    if (!s.IsOK()) {
+      return s;
+    }
+
+    if (command_output[0] == '-') {
+      *output = command_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);
+
+    *output = redis::MultiLen(3);  // command output + perf context + iostats context
+    *output += command_output;
+    *output += redis::BulkString(perf_context);
+    *output += redis::BulkString(iostats_context);
+    return Status::OK();
+  }
+
+ private:
+  std::vector<std::string> command_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 +1225,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