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

This commit merges all changes of the current dev branch compared to unstable into a single commit. It includes implementation of the lua-time-limit configuration option, SCRIPT KILL command to safely interrupt long-running scripts, connection guard optimizations, worker event loop recursive event handling, and macOS TCP listener socket sharing.

Author: x-at-01

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,26 @@ 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")
   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}")
   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}" ${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: 0 (disabled)
+lua-time-limit 0
+
 ################################## 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, 0, -1, INT_MAX)},
 
       /* rocksdb options */
       {"rocksdb.compression", false,

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

@@ -433,6 +433,18 @@ 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");
+
+    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;
+      }
+    }
+
     bool is_multi_exec = IsFlagEnabled(Connection::kMultiExec);
     if (IsFlagEnabled(redis::Connection::kCloseAfterReply) && !is_multi_exec) break;
     auto multi_error_exit = MakeScopeExit([&] {
@@ -457,7 +469,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;
 
@@ -494,7 +505,9 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
     // 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,61 @@ 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);
+  running_script_count_.fetch_add(1, std::memory_order_relaxed);
+}
+
+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);
+    running_script_count_.fetch_sub(1, std::memory_order_relaxed);
+  }
+}
+
+bool Server::IsScriptTimedOut() const {
+  int limit = config_->lua_time_limit;
+  if (limit <= 0) return false;
+  if (running_script_count_.load(std::memory_order_relaxed) == 0) return false;
+
+  uint64_t now_ms = util::GetTimeStampMS();
+  std::lock_guard<std::mutex> guard(running_scripts_mu_);
+  for (const auto *rctx : running_scripts_) {
+    if (now_ms - rctx->start_time_ms >= static_cast<uint64_t>(limit)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+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;
+    } else {
+      rctx->is_killed = 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."};
+  }
+
+  return Status::OK();
+}
+
 Status Server::ScriptExists(const std::string &sha) const {
   std::string body;
   return ScriptGet(sha, &body);
@@ -2068,8 +2138,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,11 @@ 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);
+  bool IsScriptTimedOut() const;
+  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 +487,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<size_t> running_script_count_{0};
 };

src/server/worker.cc

@@ -54,7 +54,8 @@
 #include "server.h"
 #include "storage/scripting.h"
 
-Worker::Worker(Server *srv, Config *config) : srv(srv), base_(event_base_new()) {
+Worker::Worker(Server *srv, Config *config, const std::vector<int> &tcp_listen_fds)
+    : srv(srv), base_(event_base_new()) {
   if (!base_) throw std::runtime_error{"event base failed to be created"};
 
   timer_.reset(NewEvent(base_, -1, EV_PERSIST));
@@ -66,6 +67,19 @@ Worker::Worker(Server *srv, Config *config) : srv(srv), base_(event_base_new())
       ERROR("[worker] Failed to listen to socket with fd: {}, Error: {}", config->socket_fd, s.Msg());
       exit(1);
     }
+  } else if (!tcp_listen_fds.empty()) {
+    for (int fd : tcp_listen_fds) {
+      int dup_fd = dup(fd);
+      if (dup_fd == -1) {
+        ERROR("[worker] Failed to dup fd: {}, Error: {}", fd, strerror(errno));
+        exit(1);
+      }
+      evconnlistener *lev = NewEvconnlistener<&Worker::newTCPConnection>(
+          base_, LEV_OPT_THREADSAFE | LEV_OPT_CLOSE_ON_FREE, config->backlog, dup_fd);
+      listen_events_.emplace_back(lev);
+      tcp_listen_fds_.push_back(dup_fd);
+      INFO("[worker] Listening on shared TCP fd: {} (original: {})", dup_fd, fd);
+    }
   } else {
     const uint32_t ports[3] = {config->port, config->tls_port, 0};
 
@@ -241,6 +255,7 @@ Status Worker::listenFD(int fd, uint32_t expected_port, int backlog) {
   evconnlistener *lev =
       NewEvconnlistener<&Worker::newTCPConnection>(base_, LEV_OPT_THREADSAFE | LEV_OPT_CLOSE_ON_FREE, backlog, dup_fd);
   listen_events_.emplace_back(lev);
+  tcp_listen_fds_.push_back(dup_fd);
   INFO("[worker] Listening on dup'ed fd: {}", dup_fd);
   return Status::OK();
 }
@@ -285,6 +300,7 @@ Status Worker::listenTCP(const std::string &host, uint32_t port, int backlog) {
     auto lev =
         NewEvconnlistener<&Worker::newTCPConnection>(base_, LEV_OPT_THREADSAFE | LEV_OPT_CLOSE_ON_FREE, backlog, fd);
     listen_events_.emplace_back(lev);
+    tcp_listen_fds_.push_back(fd);
   }
 
   return Status::OK();
@@ -582,6 +598,33 @@ void Worker::LuaReset() {
 
 int64_t Worker::GetLuaMemorySize() { return (int64_t)lua_gc(lua_, LUA_GCCOUNT, 0) * 1024; }
 
+void Worker::PollEventLoop() {
+  event_base_loop(base_, EVLOOP_ONCE | EVLOOP_NONBLOCK);
+
+  // Flush all active connections' output buffers to their sockets
+  std::lock_guard<std::mutex> lock(conns_mu_);
+  for (const auto &iter : conns_) {
+    auto *conn = iter.second;
+    auto *bev = conn->GetBufferEvent();
+    if (bev) {
+      bool is_tls = false;
+#ifdef ENABLE_OPENSSL
+      if (bufferevent_openssl_get_ssl(bev) != nullptr) {
+        is_tls = true;
+      }
+#endif
+      if (is_tls) {
+        bufferevent_flush(bev, EV_WRITE, BEV_FLUSH);
+      } else {
+        auto *output = bufferevent_get_output(bev);
+        if (evbuffer_get_length(output) > 0) {
+          evbuffer_write(output, conn->GetFD());
+        }
+      }
+    }
+  }
+}
+
 void Worker::KickoutIdleClients(int timeout) {
   std::vector<std::pair<int, uint64_t>> to_be_killed_conns;