(kvrocks) branch unstable updated: feat(script): allow blocking commands in scripting (#2740)

[twice]

This is an automated email from the ASF dual-hosted git repository.

twice pushed a commit to branch unstable
in repository https://gitbox.apache.org/repos/asf/kvrocks.git


The following commit(s) were added to refs/heads/unstable by this push:
     new 81691c20 feat(script): allow blocking commands in scripting (#2740)
81691c20 is described below

commit 81691c20dcfbdc887cc52c54f97f93cd4e5aea7c
Author: Twice <twice.m...@gmail.com>
AuthorDate: Tue Jan 28 14:56:30 2025 +0800

    feat(script): allow blocking commands in scripting (#2740)
---
 src/commands/blocking_commander.h             | 2 +-
 src/server/redis_connection.cc                | 6 ++++++
 src/server/redis_connection.h                 | 2 ++
 src/storage/scripting.cc                      | 8 +-------
 tests/gocase/unit/scripting/scripting_test.go | 8 +++++++-
 5 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/commands/blocking_commander.h 
b/src/commands/blocking_commander.h
index 27496819..2571e456 100644
--- a/src/commands/blocking_commander.h
+++ b/src/commands/blocking_commander.h
@@ -52,7 +52,7 @@ class BlockingCommander : public Commander,
   // to start the blocking process
   // usually put to the end of the Execute method
   Status StartBlocking(int64_t timeout, std::string *output) {
-    if (conn_->IsInExec()) {
+    if (conn_->IsInExec() || conn_->IsInScript()) {
       *output = NoopReply(conn_);
       return Status::OK();  // no blocking in multi-exec
     }
diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc
index 3f8e0dd2..a4926dfb 100644
--- a/src/server/redis_connection.cc
+++ b/src/server/redis_connection.cc
@@ -496,6 +496,12 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> 
*to_process_cmds) {
       continue;
     }
 
+    ScopeExit in_script_exit{[this] { in_script_ = false; }, false};
+    if (attributes->category == CommandCategory::Script || 
attributes->category == CommandCategory::Function) {
+      in_script_ = true;
+      in_script_exit.Enable();
+    }
+
     SetLastCmd(cmd_name);
     {
       std::optional<MultiLockGuard> guard;
diff --git a/src/server/redis_connection.h b/src/server/redis_connection.h
index e7d277f6..e8b44d94 100644
--- a/src/server/redis_connection.h
+++ b/src/server/redis_connection.h
@@ -171,6 +171,7 @@ class Connection : public EvbufCallbackBase<Connection> {
   // Multi exec
   void SetInExec() { in_exec_ = true; }
   bool IsInExec() const { return in_exec_; }
+  bool IsInScript() const { return in_script_; }
   bool IsMultiError() const { return multi_error_; }
   void ResetMultiExec();
   std::deque<redis::CommandTokens> *GetMultiExecCommands() { return 
&multi_cmds_; }
@@ -210,6 +211,7 @@ class Connection : public EvbufCallbackBase<Connection> {
   bool multi_error_ = false;
   std::atomic<bool> is_running_ = false;
   std::deque<redis::CommandTokens> multi_cmds_;
+  bool in_script_ = false;
 
   bool importing_ = false;
   RESP protocol_version_ = RESP::v2;
diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc
index b2c0eebc..6a93ad9f 100644
--- a/src/storage/scripting.cc
+++ b/src/storage/scripting.cc
@@ -756,7 +756,7 @@ int RedisGenericCommand(lua_State *lua, int raise_error) {
 
   auto attributes = cmd->GetAttributes();
   if (!attributes->CheckArity(argc)) {
-    PushError(lua, "Wrong number of args calling Redis command From Lua 
script");
+    PushError(lua, "Wrong number of args while calling Redis command from Lua 
script");
     return raise_error ? RaiseError(lua) : 1;
   }
 
@@ -772,12 +772,6 @@ int RedisGenericCommand(lua_State *lua, int raise_error) {
     return raise_error ? RaiseError(lua) : 1;
   }
 
-  // TODO: fix blocking commands to make them work in scripting
-  if (cmd_flags & redis::kCmdBlocking) {
-    PushError(lua, "This Redis command is not allowed from scripts");
-    return raise_error ? RaiseError(lua) : 1;
-  }
-
   std::string cmd_name = attributes->name;
 
   auto *conn = script_run_ctx->conn;
diff --git a/tests/gocase/unit/scripting/scripting_test.go 
b/tests/gocase/unit/scripting/scripting_test.go
index cd04351a..c3753ecc 100644
--- a/tests/gocase/unit/scripting/scripting_test.go
+++ b/tests/gocase/unit/scripting/scripting_test.go
@@ -198,10 +198,16 @@ return {type(foo),foo == false}
        })
 
        t.Run("EVAL - Scripts can't run certain commands", func(t *testing.T) {
-               r := rdb.Eval(ctx, `return redis.pcall('blpop','x',0)`, 
[]string{})
+               r := rdb.Eval(ctx, `return redis.pcall('shutdown')`, []string{})
                require.ErrorContains(t, r.Err(), "not allowed")
        })
 
+       t.Run("EVAL - Scripts can run blocking commands and get immediate 
result", func(t *testing.T) {
+               r := rdb.Eval(ctx, `return redis.pcall('blpop', KEYS[1], 0)`, 
[]string{"key_for_blpop_script"})
+               require.Equal(t, r.Val(), nil)
+               require.ErrorContains(t, r.Err(), "nil")
+       })
+
        t.Run("EVAL - Scripts can run certain commands", func(t *testing.T) {
                r := rdb.Eval(ctx, `redis.pcall('randomkey'); return 
redis.pcall('set','x','ciao')`, []string{})
                require.NoError(t, r.Err())