Add new SCRIPT STATS subcommand

Signed-off-by: artikell <[email protected]>
valkey-io · Dec 10, 2024 · b103f0f · b103f0f
1 parent 4f61034
commit b103f0f
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 0 deletions.
diff --git a/src/commands.def b/src/commands.def
@@ -5458,6 +5458,23 @@ struct COMMAND_ARG SCRIPT_SHOW_Args[] = {
 {MAKE_ARG("sha1",ARG_TYPE_STRING,-1,NULL,NULL,NULL,CMD_ARG_NONE,0,NULL)},
 };
 
+/********** SCRIPT STATS ********************/
+
+#ifndef SKIP_CMD_HISTORY_TABLE
+/* SCRIPT STATS history */
+#define SCRIPT_STATS_History NULL
+#endif
+
+#ifndef SKIP_CMD_TIPS_TABLE
+/* SCRIPT STATS tips */
+#define SCRIPT_STATS_Tips NULL
+#endif
+
+#ifndef SKIP_CMD_KEY_SPECS_TABLE
+/* SCRIPT STATS key specs */
+#define SCRIPT_STATS_Keyspecs NULL
+#endif
+
 /* SCRIPT command table */
 struct COMMAND_STRUCT SCRIPT_Subcommands[] = {
 {MAKE_CMD("debug","Sets the debug mode of server-side Lua scripts.","O(1)","3.2.0",CMD_DOC_NONE,NULL,NULL,"scripting",COMMAND_GROUP_SCRIPTING,SCRIPT_DEBUG_History,0,SCRIPT_DEBUG_Tips,0,scriptCommand,3,CMD_NOSCRIPT,ACL_CATEGORY_SCRIPTING,SCRIPT_DEBUG_Keyspecs,0,NULL,1),.args=SCRIPT_DEBUG_Args},
@@ -5467,6 +5484,7 @@ struct COMMAND_STRUCT SCRIPT_Subcommands[] = {
 {MAKE_CMD("kill","Terminates a server-side Lua script during execution.","O(1)","2.6.0",CMD_DOC_NONE,NULL,NULL,"scripting",COMMAND_GROUP_SCRIPTING,SCRIPT_KILL_History,0,SCRIPT_KILL_Tips,2,scriptCommand,2,CMD_NOSCRIPT|CMD_ALLOW_BUSY,ACL_CATEGORY_SCRIPTING,SCRIPT_KILL_Keyspecs,0,NULL,0)},
 {MAKE_CMD("load","Loads a server-side Lua script to the script cache.","O(N) with N being the length in bytes of the script body.","2.6.0",CMD_DOC_NONE,NULL,NULL,"scripting",COMMAND_GROUP_SCRIPTING,SCRIPT_LOAD_History,0,SCRIPT_LOAD_Tips,2,scriptCommand,3,CMD_NOSCRIPT|CMD_STALE,ACL_CATEGORY_SCRIPTING,SCRIPT_LOAD_Keyspecs,0,NULL,1),.args=SCRIPT_LOAD_Args},
 {MAKE_CMD("show","Show server-side Lua script in the script cache.","O(1).","8.0.0",CMD_DOC_NONE,NULL,NULL,"scripting",COMMAND_GROUP_SCRIPTING,SCRIPT_SHOW_History,0,SCRIPT_SHOW_Tips,0,scriptCommand,3,CMD_NOSCRIPT,ACL_CATEGORY_SCRIPTING,SCRIPT_SHOW_Keyspecs,0,NULL,1),.args=SCRIPT_SHOW_Args},
+{MAKE_CMD("stats","Returns information about a script during execution.","O(1).","8.0.0",CMD_DOC_NONE,NULL,NULL,"scripting",COMMAND_GROUP_SCRIPTING,SCRIPT_STATS_History,0,SCRIPT_STATS_Tips,0,scriptCommand,2,CMD_NOSCRIPT|CMD_ALLOW_BUSY,ACL_CATEGORY_SCRIPTING,SCRIPT_STATS_Keyspecs,0,NULL,0)},
 {0}
 };
 

diff --git a/src/commands/script-stats.json b/src/commands/script-stats.json
@@ -0,0 +1,55 @@
+{
+    "STATS": {
+        "summary": "Returns information about a script during execution.",
+        "complexity": "O(1).",
+        "group": "scripting",
+        "since": "8.0.0",
+        "arity": 2,
+        "container": "SCRIPT",
+        "function": "scriptCommand",
+        "command_flags": [
+            "NOSCRIPT",
+            "ALLOW_BUSY"
+        ],
+        "acl_categories": [
+            "SCRIPTING"
+        ],
+        "reply_schema": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+                "running_script": {
+                    "description": "Information about the running script.",
+                    "oneOf": [
+                        {
+                            "description": "If there's no in-flight script.",
+                            "type": "null"
+                        },
+                        {
+                            "description": "A map with the information about the running script.",
+                            "type": "object",
+                            "additionalProperties": false,
+                            "properties": {
+                                "name": {
+                                    "description": "The name of the script.",
+                                    "type": "string"
+                                },
+                                "command": {
+                                    "description": "The command and arguments used for invoking the script.",
+                                    "type": "array",
+                                    "items": {
+                                        "type": "string"
+                                    }
+                                },
+                                "duration_ms": {
+                                    "description": "The script's runtime duration in milliseconds.",
+                                    "type": "integer"
+                                }
+                            }
+                        }
+                    ]
+                }
+            }
+        }
+    }
+}
diff --git a/src/eval.c b/src/eval.c
@@ -688,6 +688,8 @@ void scriptCommand(client *c) {
             "    Load a script into the scripts cache without executing it.",
             "SHOW <sha1>",
             "    Show a script from the scripts cache.",
+            "STATS",
+            "    Returns information about a script during execution.",
             NULL,
         };
         addReplyHelp(c, help);
@@ -750,6 +752,25 @@ void scriptCommand(client *c) {
         } else {
             addReplyErrorObject(c, shared.noscripterr);
         }
+    } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr, "stats")) {
+        addReplyMapLen(c, 1);
+
+        addReplyBulkCString(c, "running_script");
+        if (!scriptIsRunning()) {
+            addReplyNull(c);
+        } else {
+            addReplyMapLen(c, 3);
+            addReplyBulkCString(c, "name");
+            addReplyBulkCString(c, scriptCurrFunction());
+            addReplyBulkCString(c, "command");
+            client *script_client = scriptGetCaller();
+            addReplyArrayLen(c, script_client->argc);
+            for (int i = 0; i < script_client->argc; ++i) {
+                addReplyBulkCBuffer(c, script_client->argv[i]->ptr, sdslen(script_client->argv[i]->ptr));
+            }
+            addReplyBulkCString(c, "duration_ms");
+            addReplyLongLong(c, scriptRunDuration());
+        }
     } else {
         addReplySubcommandSyntaxError(c);
     }

diff --git a/tests/unit/scripting.tcl b/tests/unit/scripting.tcl
@@ -1161,10 +1161,13 @@ start_server {tags {"scripting"}} {
     test {Timedout read-only scripts can be killed by SCRIPT KILL} {
         set rd [valkey_deferring_client]
         r config set lua-time-limit 10
+        assert_match {running_script {}} [r SCRIPT STATS]
+
         run_script_on_connection $rd {while true do end} 0
         after 200
         catch {r ping} e
         assert_match {BUSY*} $e
+        assert_match {running_script {name * command * duration_ms *}} [r SCRIPT STATS]
         kill_script
         after 200 ; # Give some time to Lua to call the hook again...
         assert_equal [r ping] "PONG"