feat(script): add SCRIPT KILL and lua-time-limit support

CMakeLists.txt

@@ -254,6 +254,10 @@ if ((CMAKE_CXX_COMPILER_ID STREQUAL "GNU") OR (CMAKE_CXX_COMPILER_ID STREQUAL "C
     endif()
 endif()
 
+if (APPLE AND ENABLE_LUAJIT)
+    set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,-no_deduplicate")
+endif()
+
 # kvrocks objects target
 file(GLOB_RECURSE KVROCKS_SRCS src/*.cc)
 list(FILTER KVROCKS_SRCS EXCLUDE REGEX src/cli/main.cc)
@@ -286,6 +290,9 @@ target_compile_definitions(kvrocks_objs PUBLIC KVROCKS_STORAGE_ENGINE=RocksDB)
 if(ENABLE_OPENSSL)
     target_compile_definitions(kvrocks_objs PUBLIC ENABLE_OPENSSL)
 endif()
+if(ENABLE_LUAJIT)
+    target_compile_definitions(kvrocks_objs PUBLIC ENABLE_LUAJIT)
+endif()
 if(ENABLE_NEW_ENCODING)
     target_compile_definitions(kvrocks_objs PUBLIC METADATA_ENCODING_VERSION=1)
 else()

cmake/luajit.cmake

@@ -47,19 +47,28 @@ FetchContent_GetProperties(luajit)
 if (NOT lua_POPULATED)
   FetchContent_Populate(luajit)
 
-  set(LUA_CFLAGS "-DLUA_ANSI -DENABLE_CJSON_GLOBAL -DREDIS_STATIC= -DLUA_USE_MKSTEMP")
+  # We use LUAJIT_CFLAGS instead of LUA_CFLAGS and we do not define LUA_ANSI.
+  # LuaJIT relies heavily on architecture-specific assembly code and compiler optimization settings,
+  # and LUA_ANSI is meant for standard Lua to restrict to ANSI C which LuaJIT does not support.
+  set(LUAJIT_CFLAGS "-DENABLE_CJSON_GLOBAL -DREDIS_STATIC= -DLUA_USE_MKSTEMP -DLUAJIT_ENABLE_CHECKHOOK")
+  set(LUAJIT_LDFLAGS "")
   if((CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang") OR
    (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang"))
-    set(LUA_CFLAGS "${LUA_CFLAGS} -isysroot ${CMAKE_OSX_SYSROOT}")
+    set(LUAJIT_CFLAGS "${LUAJIT_CFLAGS} -isysroot ${CMAKE_OSX_SYSROOT}")
+    set(LUAJIT_LDFLAGS "-isysroot ${CMAKE_OSX_SYSROOT}")
   endif ()
 
   if (CMAKE_HOST_APPLE)
     set(MACOSX_TARGET "MACOSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}")
   endif ()
 
+  # We pass C compiler path CC and custom preprocessor definitions via XCFLAGS instead of CFLAGS.
+  # Overriding CFLAGS on the make command line completely discards LuaJIT's own Makefile CFLAGS (such as
+  # -O2 and target-specific compiler optimizations), which breaks the JIT compiler and assembler VM
+  # stack unwinding on platforms like macOS arm64, leading to CPU spin locks or crashes.
   add_custom_target(make_luajit
     COMMAND ${MAKE_COMMAND} libluajit.a ${NINJA_MAKE_JOBS_FLAG}
-      "CFLAGS=${LUA_CFLAGS}" ${MACOSX_TARGET}
+      "CC=${CMAKE_C_COMPILER}" "XCFLAGS=${LUAJIT_CFLAGS}" "LDFLAGS=${LUAJIT_LDFLAGS}" "HOST_LDFLAGS=${LUAJIT_LDFLAGS}" ${MACOSX_TARGET}
     COMMAND ${CMAKE_COMMAND} -E make_directory ${luajit_BINARY_DIR}/include
     COMMAND sh -c "cp ${luajit_SOURCE_DIR}/src/*.h ${luajit_SOURCE_DIR}/src/*.hpp ${luajit_BINARY_DIR}/include/"
     WORKING_DIRECTORY ${luajit_SOURCE_DIR}/src

kvrocks.conf

@@ -492,6 +492,13 @@ txn-context-enabled no
 # Default: no
 lua-strict-key-accessing no
 
+# Max execution time of a Lua script in milliseconds.
+# If the maximum execution time is reached, the server will log it
+# and start rejecting other queries with a BUSY error.
+# Setting it to 0 or a negative value means no timeout.
+# Default: 5000 (5 seconds)
+lua-time-limit 5000
+
 ################################## TLS ###################################
 
 # By default, TLS/SSL is disabled, i.e. `tls-port` is set to 0.

src/commands/cmd_script.cc

@@ -68,7 +68,7 @@ class CommandScript : public Commander {
     // There's a little tricky here since the script command was the write type
     // command but some subcommands like `exists` were readonly, so we want to allow
     // executing on slave here. Maybe we should find other way to do this.
-    if (srv->IsSlave() && subcommand_ != "exists") {
+    if (srv->IsSlave() && subcommand_ != "exists" && subcommand_ != "kill") {
       return {Status::RedisReadOnly, "You can't write against a read only slave"};
     }
 
@@ -84,6 +84,12 @@ class CommandScript : public Commander {
         return s;
       }
       *output = redis::RESP_OK;
+    } else if (args_.size() == 2 && subcommand_ == "kill") {
+      auto s = srv->ScriptKill();
+      if (!s) {
+        return s;
+      }
+      *output = redis::RESP_OK;
     } else if (args_.size() >= 3 && subcommand_ == "exists") {
       *output = redis::MultiLen(args_.size() - 2);
       for (size_t j = 2; j < args_.size(); j++) {

src/config/config.cc

@@ -263,6 +263,7 @@ Config::Config() {
       {"skip-block-cache-deallocation-on-close", false, new YesNoField(&skip_block_cache_deallocation_on_close, false)},
       {"histogram-bucket-boundaries", true, new StringField(&histogram_bucket_boundaries_str_, "")},
       {"lua-strict-key-accessing", false, new YesNoField(&lua_strict_key_accessing, false)},
+      {"lua-time-limit", false, new IntField(&lua_time_limit, 5000, -1, INT_MAX)},
 
       /* rocksdb options */
       {"rocksdb.compression", false,
@@ -834,6 +835,12 @@ void Config::initFieldCallback() {
              }
              return Status::OK();
            }},
+          {"lua-time-limit",
+           [](Server *srv, [[maybe_unused]] const std::string &k, [[maybe_unused]] const std::string &v) -> Status {
+             if (!srv) return Status::OK();
+             srv->ReevaluateScriptTimeout();
+             return Status::OK();
+           }},
       };
   for (const auto &iter : callbacks) {
     auto field_iter = fields_.find(iter.first);

src/config/config.h

@@ -205,6 +205,8 @@ struct Config {
 
   bool lua_strict_key_accessing = false;
 
+  int lua_time_limit = 0;
+
   std::vector<double> histogram_bucket_boundaries;
 
   struct RocksDB {

src/server/redis_connection.cc

@@ -94,15 +94,39 @@ std::string Connection::ToString() {
 }
 
 void Connection::Close() {
-  if (close_cb) close_cb(GetFD());
+  if (is_running_) {
+    EnableFlag(kCloseAsync);
+    if (bev_) {
+      bufferevent_disable(bev_, EV_READ | EV_WRITE);
+      bufferevent_setcb(bev_, nullptr, nullptr, nullptr, nullptr);
+      int fd = bufferevent_getfd(bev_);
+      if (fd != -1) {
+        evutil_closesocket(fd);
+      }
+    }
+    return;
+  }
+
+  if (close_cb) {
+    close_cb(GetFD());
+  }
   owner_->FreeConnection(this);
 }
 
 void Connection::Detach() { owner_->DetachConnection(this); }
 
 void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) {
+  if (is_running_) return;
   is_running_ = true;
-  MakeScopeExit([this] { is_running_ = false; });
+
+  auto exit = MakeScopeExit([this] {
+    is_running_ = false;
+    if (IsFlagEnabled(kCloseAsync)) {
+      Close();
+    } else if (evbuffer_get_length(Input()) > 0) {
+      bufferevent_trigger(bev_, EV_READ, BEV_TRIG_IGNORE_WATERMARKS);
+    }
+  });
 
   SetLastInteraction();
   auto s = req_.Tokenize(Input());
@@ -114,9 +138,6 @@ void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) {
   }
 
   ExecuteCommands(req_.GetCommands());
-  if (IsFlagEnabled(kCloseAsync)) {
-    Close();
-  }
 }
 
 void Connection::OnWrite([[maybe_unused]] bufferevent *bev) {
@@ -433,6 +454,10 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
     to_process_cmds->pop_front();
     if (cmd_tokens.empty()) continue;
 
+    bool is_script_kill = (util::EqualICase(cmd_tokens.front(), "script") && cmd_tokens.size() == 2 &&
+                           util::EqualICase(cmd_tokens[1], "kill"));
+    bool is_shutdown = util::EqualICase(cmd_tokens.front(), "shutdown");
+
     bool is_multi_exec = IsFlagEnabled(Connection::kMultiExec);
     if (IsFlagEnabled(redis::Connection::kCloseAfterReply) && !is_multi_exec) break;
     auto multi_error_exit = MakeScopeExit([&] {
@@ -457,7 +482,6 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
       continue;
     }
     auto current_cmd = std::move(*cmd_s);
-
     const auto &attributes = current_cmd->GetAttributes();
     auto cmd_name = attributes->name;
 
@@ -488,13 +512,23 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
       }
     }
 
+    if (srv_->IsScriptTimedOut()) {
+      if (!is_script_kill && !is_shutdown) {
+        Reply(redis::Error({Status::RedisErrorNoPrefix,
+                            "BUSY Redis is busy running a script. You can only call SCRIPT KILL or SHUTDOWN NOSAVE."}));
+        continue;
+      }
+    }
+
     std::shared_lock<std::shared_mutex> concurrency;  // Allow concurrency
     std::unique_lock<std::shared_mutex> exclusivity;  // Need exclusivity
     // If the command needs to process exclusively, we need to get 'ExclusivityGuard'
     // that can guarantee other threads can't come into critical zone, such as DEBUG,
     // CLUSTER subcommand, CONFIG SET, MULTI, LUA (in the immediate future).
     // Otherwise, we just use 'ConcurrencyGuard' to allow all workers to execute commands at the same time.
-    if (is_multi_exec && !(cmd_flags & kCmdBypassMulti)) {
+    if (is_script_kill || is_shutdown) {
+      // Bypass locks to allow SCRIPT KILL and SHUTDOWN to run even if another thread is stuck in a script
+    } else if (is_multi_exec && !(cmd_flags & kCmdBypassMulti)) {
       // No lock guard, because 'exec' command has acquired 'WorkExclusivityGuard'
     } else if (cmd_flags & kCmdExclusive) {
       exclusivity = srv_->WorkExclusivityGuard();

src/server/server.cc

@@ -101,8 +101,16 @@ Server::Server(engine::Storage *storage, Config *config)
   // init shard pub/sub channels
   pubsub_shard_channels_.resize(config->cluster_enabled ? HASH_SLOTS_SIZE : 1);
 
+  std::vector<int> listen_fds;
   for (int i = 0; i < config->workers; i++) {
-    auto worker = std::make_unique<Worker>(this, config);
+#ifdef __APPLE__
+    auto worker = std::make_unique<Worker>(this, config, listen_fds);
+    if (i == 0) {
+      listen_fds = worker->GetTCPListenFDs();
+    }
+#else
+    auto worker = std::make_unique<Worker>(this, config, std::vector<int>{});
+#endif
     // multiple workers can't listen to the same unix socket, so
     // listen unix socket only from a single worker - the first one
     if (!config->unixsocket.empty() && i == 0) {
@@ -255,6 +263,13 @@ void Server::Stop() {
   if (replication_thread_) replication_thread_->Stop();
   slaveof_mu_.unlock();
 
+  {
+    std::lock_guard<std::mutex> guard(running_scripts_mu_);
+    for (auto *rctx : running_scripts_) {
+      rctx->is_killed = true;
+    }
+  }
+
   for (const auto &worker : worker_threads_) {
     worker->Stop(0 /* immediately terminate  */);
   }
@@ -1879,6 +1894,80 @@ StatusOr<std::unique_ptr<redis::Commander>> Server::LookupAndCreateCommand(const
   return std::move(cmd);
 }
 
+void Server::RegisterRunningScript(lua::ScriptRunCtx *rctx) {
+  std::lock_guard<std::mutex> guard(running_scripts_mu_);
+  running_scripts_.push_back(rctx);
+}
+
+void Server::UnregisterRunningScript(lua::ScriptRunCtx *rctx) {
+  std::lock_guard<std::mutex> guard(running_scripts_mu_);
+  auto it = std::find(running_scripts_.begin(), running_scripts_.end(), rctx);
+  if (it != running_scripts_.end()) {
+    running_scripts_.erase(it);
+  }
+
+  // Re-evaluate if any remaining scripts are timed out, if not, clear the flag
+  bool any_timed_out = false;
+  int limit = config_->lua_time_limit;
+  if (limit > 0 && !running_scripts_.empty()) {
+    uint64_t now_ms = util::GetTimeStampMS();
+    for (const auto *ctx : running_scripts_) {
+      if (now_ms - ctx->start_time_ms >= static_cast<uint64_t>(limit)) {
+        any_timed_out = true;
+        break;
+      }
+    }
+  }
+  is_script_timeout_.store(any_timed_out, std::memory_order_relaxed);
+}
+
+void Server::ReevaluateScriptTimeout() {
+  std::lock_guard<std::mutex> guard(running_scripts_mu_);
+  bool any_timed_out = false;
+  int limit = config_->lua_time_limit;
+  if (limit > 0 && !running_scripts_.empty()) {
+    uint64_t now_ms = util::GetTimeStampMS();
+    for (const auto *ctx : running_scripts_) {
+      if (now_ms - ctx->start_time_ms >= static_cast<uint64_t>(limit)) {
+        any_timed_out = true;
+        break;
+      }
+    }
+  }
+  is_script_timeout_.store(any_timed_out, std::memory_order_relaxed);
+}
+
+bool Server::IsScriptTimedOut() const { return is_script_timeout_.load(std::memory_order_relaxed); }
+
+void Server::SetScriptTimedOut(bool timed_out) { is_script_timeout_.store(timed_out, std::memory_order_relaxed); }
+
+Status Server::ScriptKill() {
+  std::lock_guard<std::mutex> guard(running_scripts_mu_);
+  if (running_scripts_.empty()) {
+    return {Status::NotOK, "NOTBUSY No scripts in execution right now."};
+  }
+
+  bool has_unkillable = false;
+  for (auto *rctx : running_scripts_) {
+    if (rctx->is_write_dirty) {
+      has_unkillable = true;
+    }
+  }
+
+  if (has_unkillable) {
+    return {Status::NotOK,
+            "UNKILLABLE Sorry the script already executed write commands against the dataset. "
+            "You can either wait the script termination or kill the server in a hard way using the SHUTDOWN NOSAVE "
+            "command."};
+  }
+
+  for (auto *rctx : running_scripts_) {
+    rctx->is_killed = true;
+  }
+
+  return Status::OK();
+}
+
 Status Server::ScriptExists(const std::string &sha) const {
   std::string body;
   return ScriptGet(sha, &body);
@@ -2068,8 +2157,14 @@ void Server::AdjustWorkerThreads() {
 }
 
 void Server::increaseWorkerThreads(size_t delta) {
+  std::vector<int> listen_fds;
+#ifdef __APPLE__
+  if (!worker_threads_.empty()) {
+    listen_fds = worker_threads_[0]->GetWorker()->GetTCPListenFDs();
+  }
+#endif
   for (size_t i = 0; i < delta; i++) {
-    auto worker = std::make_unique<Worker>(this, config_);
+    auto worker = std::make_unique<Worker>(this, config_, listen_fds);
     auto worker_thread = std::make_unique<WorkerThread>(std::move(worker));
     worker_thread->Start();
     worker_threads_.emplace_back(std::move(worker_thread));

src/server/server.h

@@ -62,6 +62,10 @@
 
 constexpr const char *REDIS_VERSION = "7.0.0";
 
+namespace lua {
+struct ScriptRunCtx;
+}
+
 struct DBScanInfo {
   // Last scan system clock in seconds
   int64_t last_scan_time_secs = 0;
@@ -315,6 +319,13 @@ class Server {
   void KillClient(int64_t *killed, const std::string &addr, uint64_t id, uint64_t type, bool skipme,
                   redis::Connection *conn);
 
+  void RegisterRunningScript(lua::ScriptRunCtx *rctx);
+  void UnregisterRunningScript(lua::ScriptRunCtx *rctx);
+  void ReevaluateScriptTimeout();
+  bool IsScriptTimedOut() const;
+  void SetScriptTimedOut(bool timed_out);
+  Status ScriptKill();
+
   Status ScriptExists(const std::string &sha) const;
   Status ScriptGet(const std::string &sha, std::string *body) const;
   Status ScriptSet(const std::string &sha, const std::string &body) const;
@@ -478,4 +489,8 @@ class Server {
     uint64_t id;
   };
   std::vector<PausedConnEntry> paused_conns_;
+
+  std::vector<lua::ScriptRunCtx *> running_scripts_;
+  mutable std::mutex running_scripts_mu_;
+  std::atomic<bool> is_script_timeout_{false};
 };