From 31cd4e961db9030243211e27a5339edd038c467c Mon Sep 17 00:00:00 2001 From: xtco3o Date: Thu, 28 May 2026 13:34:37 +0800 Subject: [PATCH 1/3] 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. --- CMakeLists.txt | 7 + cmake/luajit.cmake | 15 ++- kvrocks.conf | 7 + src/commands/cmd_script.cc | 8 +- src/config/config.cc | 1 + src/config/config.h | 2 + src/server/redis_connection.cc | 41 +++++- src/server/server.cc | 81 +++++++++++- src/server/server.h | 14 ++ src/server/worker.cc | 48 ++++++- src/server/worker.h | 5 +- src/storage/scripting.cc | 97 +++++++++++++- src/storage/scripting.h | 9 +- tests/gocase/unit/scripting/scripting_test.go | 121 ++++++++++++++++++ tests/gocase/util/server.go | 4 +- 15 files changed, 442 insertions(+), 18 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index b074b52a819..074b4b6eb02 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/cmake/luajit.cmake b/cmake/luajit.cmake index 83718f1379e..c817f585c13 100644 --- a/cmake/luajit.cmake +++ b/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 diff --git a/kvrocks.conf b/kvrocks.conf index f1541c49676..62e441d46e5 100644 --- a/kvrocks.conf +++ b/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. diff --git a/src/commands/cmd_script.cc b/src/commands/cmd_script.cc index 3f8ee38eb1a..57ed4950ad8 100644 --- a/src/commands/cmd_script.cc +++ b/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++) { diff --git a/src/config/config.cc b/src/config/config.cc index 492eb975bb3..c845edadab1 100644 --- a/src/config/config.cc +++ b/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, diff --git a/src/config/config.h b/src/config/config.h index c17b331e71d..6bf5c7e1c50 100644 --- a/src/config/config.h +++ b/src/config/config.h @@ -205,6 +205,8 @@ struct Config { bool lua_strict_key_accessing = false; + int lua_time_limit = 0; + std::vector histogram_bucket_boundaries; struct RocksDB { diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc index aa301b32ffa..d32d56c34ed 100644 --- a/src/server/redis_connection.cc +++ b/src/server/redis_connection.cc @@ -94,19 +94,35 @@ 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; }); SetLastInteraction(); auto s = req_.Tokenize(Input()); if (!s.IsOK()) { + is_running_ = false; EnableFlag(redis::Connection::kCloseAfterReply); Reply(redis::Error(s)); INFO("[connection] Failed to tokenize the request. Error: {}", s.Msg()); @@ -114,8 +130,12 @@ void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) { } ExecuteCommands(req_.GetCommands()); + is_running_ = false; + if (IsFlagEnabled(kCloseAsync)) { Close(); + } else if (evbuffer_get_length(Input()) > 0) { + bufferevent_trigger(bev_, EV_READ, BEV_TRIG_IGNORE_WATERMARKS); } } @@ -433,6 +453,18 @@ void Connection::ExecuteCommands(std::deque *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 +489,6 @@ void Connection::ExecuteCommands(std::deque *to_process_cmds) { continue; } auto current_cmd = std::move(*cmd_s); - const auto &attributes = current_cmd->GetAttributes(); auto cmd_name = attributes->name; @@ -494,7 +525,9 @@ void Connection::ExecuteCommands(std::deque *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(); diff --git a/src/server/server.cc b/src/server/server.cc index fc763863bf0..69f2ec3649b 100644 --- a/src/server/server.cc +++ b/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 listen_fds; for (int i = 0; i < config->workers; i++) { - auto worker = std::make_unique(this, config); +#ifdef __APPLE__ + auto worker = std::make_unique(this, config, listen_fds); + if (i == 0) { + listen_fds = worker->GetTCPListenFDs(); + } +#else + auto worker = std::make_unique(this, config, std::vector{}); +#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 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,62 @@ StatusOr> Server::LookupAndCreateCommand(const return std::move(cmd); } +void Server::RegisterRunningScript(lua::ScriptRunCtx *rctx) { + std::lock_guard guard(running_scripts_mu_); + running_scripts_.push_back(rctx); +} + +void Server::UnregisterRunningScript(lua::ScriptRunCtx *rctx) { + std::lock_guard 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(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 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 +2139,14 @@ void Server::AdjustWorkerThreads() { } void Server::increaseWorkerThreads(size_t delta) { + std::vector 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(this, config_); + auto worker = std::make_unique(this, config_, listen_fds); auto worker_thread = std::make_unique(std::move(worker)); worker_thread->Start(); worker_threads_.emplace_back(std::move(worker_thread)); diff --git a/src/server/server.h b/src/server/server.h index 63b83e28306..06b61dc7f4c 100644 --- a/src/server/server.h +++ b/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,12 @@ 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; + 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 +488,8 @@ class Server { uint64_t id; }; std::vector paused_conns_; + + std::vector running_scripts_; + mutable std::mutex running_scripts_mu_; + std::atomic is_script_timeout_{false}; }; diff --git a/src/server/worker.cc b/src/server/worker.cc index 45eceb70577..712427c1056 100644 --- a/src/server/worker.cc +++ b/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 &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,36 @@ 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 lock(conns_mu_); + for (const auto &iter : conns_) { + auto *conn = iter.second; + if (conn->IsFlagEnabled(redis::Connection::kCloseAsync)) { + continue; + } + 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> to_be_killed_conns; diff --git a/src/server/worker.h b/src/server/worker.h index 0e3ff7c03ec..abd579ea875 100644 --- a/src/server/worker.h +++ b/src/server/worker.h @@ -44,7 +44,7 @@ class Server; class Worker : EventCallbackBase, EvconnlistenerBase { public: - Worker(Server *srv, Config *config); + Worker(Server *srv, Config *config, const std::vector &tcp_listen_fds = {}); ~Worker(); Worker(const Worker &) = delete; Worker(Worker &&) = delete; @@ -78,8 +78,10 @@ class Worker : EventCallbackBase, EvconnlistenerBase { lua_State *Lua() { return lua_; } void LuaReset(); int64_t GetLuaMemorySize(); + void PollEventLoop(); std::map GetConnections() const { return conns_; } + std::vector GetTCPListenFDs() const { return tcp_listen_fds_; } Server *srv; private: @@ -102,6 +104,7 @@ class Worker : EventCallbackBase, EvconnlistenerBase { struct ev_token_bucket_cfg *rate_limit_group_cfg_ = nullptr; std::atomic lua_; std::atomic is_terminated_ = false; + std::vector tcp_listen_fds_; }; class WorkerThread { diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc index 49727b11238..c1be17df12a 100644 --- a/src/storage/scripting.cc +++ b/src/storage/scripting.cc @@ -31,6 +31,7 @@ #include "commands/commander.h" #include "commands/error_constants.h" +#include "common/logging.h" #include "db_util.h" #include "fmt/format.h" #include "lua.h" @@ -40,9 +41,11 @@ #include "server/redis_connection.h" #include "server/redis_reply.h" #include "server/server.h" +#include "server/worker.h" #include "sha1.h" #include "storage/storage.h" #include "string_util.h" +#include "time_util.h" /* The maximum number of characters needed to represent a long double * as a string (long double has a huge range). @@ -60,6 +63,89 @@ enum { namespace lua { +class ScriptRunCtxGuard { + public: + ScriptRunCtxGuard(Server *srv, lua_State *lua, ScriptRunCtx *rctx) : srv_(srv), lua_(lua), rctx_(rctx) { + rctx_->start_time_ms = util::GetTimeStampMS(); + SaveOnRegistry(lua_, REGISTRY_SCRIPT_RUN_CTX_NAME, rctx_); + srv_->RegisterRunningScript(rctx_); + lua_sethook(lua_, LuaMaskCountHook, LUA_MASKCOUNT, 100000); + } + + ~ScriptRunCtxGuard() { + lua_sethook(lua_, nullptr, 0, 0); + srv_->UnregisterRunningScript(rctx_); + RemoveFromRegistry(lua_, REGISTRY_SCRIPT_RUN_CTX_NAME); + } + + private: + Server *srv_; + lua_State *lua_; + ScriptRunCtx *rctx_; +}; + +static void KillScript(lua_State *lua) { + lua_sethook(lua, LuaMaskCountHook, LUA_MASKLINE, 0); + PushError(lua, "Script killed by user with SCRIPT KILL..."); + RaiseError(lua); +} + +void LuaMaskCountHook(lua_State *lua, [[maybe_unused]] lua_Debug *ar) { + auto *script_run_ctx = GetFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME); + if (script_run_ctx == nullptr) return; + + auto *srv = script_run_ctx->conn->GetServer(); + + if (script_run_ctx->is_killed) { + KillScript(lua); + } + + int limit = srv->GetConfig()->lua_time_limit; + bool is_disconnected = script_run_ctx->conn->IsFlagEnabled(redis::Connection::kCloseAsync); + + int effective_limit = limit; + if (is_disconnected && limit <= 0) { + effective_limit = 60000; // 60 seconds default timeout when disconnected + } + + if (effective_limit > 0) { + uint64_t now_ms = util::GetTimeStampMS(); + if (now_ms - script_run_ctx->start_time_ms >= static_cast(effective_limit)) { + if (limit > 0) { + srv->SetScriptTimedOut(true); + } + + if (!script_run_ctx->slow_logged) { + if (is_disconnected) { + WARN("Slow script detected on disconnected client: still in execution after {} milliseconds. Killing it.", + now_ms - script_run_ctx->start_time_ms); + } else { + WARN( + "Slow script detected: still in execution after {} milliseconds. You can try killing the script using " + "the " + "SCRIPT KILL command.", + now_ms - script_run_ctx->start_time_ms); + } + script_run_ctx->slow_logged = true; + } + + if (is_disconnected) { + if (!script_run_ctx->is_write_dirty) { + KillScript(lua); + } + } else { + // Poll the worker thread's event loop to process SCRIPT KILL or other commands. + auto *worker = script_run_ctx->conn->Owner(); + worker->PollEventLoop(); + + if (script_run_ctx->is_killed) { + KillScript(lua); + } + } + } + } +} + namespace { std::string SanitizeErrorMessage(std::string_view message) { @@ -72,6 +158,8 @@ std::string SanitizeErrorMessage(std::string_view message) { lua_State *CreateState() { lua_State *lua = lua_open(); + // We do not turn off the JIT engine because LUAJIT_ENABLE_CHECKHOOK is defined, + // which allows JIT-compiled code to check hooks safely while preserving performance. LoadLibraries(lua); RemoveUnsupportedFunctions(lua); LoadFuncs(lua); @@ -434,7 +522,7 @@ Status FunctionCall(redis::Connection *conn, engine::Context *ctx, const std::st } lua_pop(lua, 1); - SaveOnRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME, &script_run_ctx); + ScriptRunCtxGuard guard(srv, lua, &script_run_ctx); // save keys on registry the to perform key touching check SaveOnRegistry(lua, REGISTRY_KEYS_NAME, &keys); @@ -451,7 +539,6 @@ Status FunctionCall(redis::Connection *conn, engine::Context *ctx, const std::st } RemoveFromRegistry(lua, REGISTRY_KEYS_NAME); - RemoveFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME); /* Call the Lua garbage collector from time to time to avoid a * full cycle performed by Lua, which adds too latency. @@ -705,7 +792,7 @@ Status EvalGenericCommand(redis::Connection *conn, engine::Context *ctx, const s } lua_pop(lua, 1); - SaveOnRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME, ¤t_script_run_ctx); + ScriptRunCtxGuard guard(srv, lua, ¤t_script_run_ctx); // For the Lua script, should be always run with RESP2 protocol, // unless users explicitly set the protocol version in the script via `redis.setresp`. @@ -737,7 +824,6 @@ Status EvalGenericCommand(redis::Connection *conn, engine::Context *ctx, const s lua_pushnil(lua); lua_setglobal(lua, "ARGV"); RemoveFromRegistry(lua, REGISTRY_KEYS_NAME); - RemoveFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME); /* Call the Lua garbage collector from time to time to avoid a * full cycle performed by Lua, which adds too latency. @@ -888,6 +974,9 @@ int RedisGenericCommand(lua_State *lua, int raise_error) { } std::string output; + if (cmd_flags & redis::kCmdWrite) { + script_run_ctx->is_write_dirty = true; + } s = conn->ExecuteCommand(*script_run_ctx->ctx, cmd_name, args, cmd.get(), &output); if (!s) { PushError(lua, s.Msg().data()); diff --git a/src/storage/scripting.h b/src/storage/scripting.h index bf03f3e104c..07c86efae85 100644 --- a/src/storage/scripting.h +++ b/src/storage/scripting.h @@ -20,6 +20,7 @@ #pragma once +#include #include #include @@ -61,6 +62,7 @@ int RedisErrorReplyCommand(lua_State *lua); int RedisLogCommand(lua_State *lua); int RedisRegisterFunction(lua_State *lua); int RedisSetResp(lua_State *lua); +void LuaMaskCountHook(lua_State *lua, lua_Debug *ar); Status CreateFunction(Server *srv, const std::string &body, std::string *sha, lua_State *lua, bool need_to_store); @@ -148,7 +150,7 @@ using ScriptFlags = uint64_t; /// ScriptRunCtx is used to record context information during the running of Eval scripts and functions. struct ScriptRunCtx { // ScriptFlags - uint64_t flags = 0; + ScriptFlags flags = 0; // current_slot tracks the slot currently accessed by the script // and is used to detect whether there is cross-slot access // between multiple commands in a script or function. @@ -157,6 +159,11 @@ struct ScriptRunCtx { redis::Connection *conn = nullptr; // the storage context engine::Context *ctx = nullptr; + + uint64_t start_time_ms = 0; + std::atomic is_killed{false}; + bool slow_logged = false; + std::atomic is_write_dirty{false}; }; /// SaveOnRegistry saves user-defined data to lua REGISTRY diff --git a/tests/gocase/unit/scripting/scripting_test.go b/tests/gocase/unit/scripting/scripting_test.go index e1318977bf7..e22b46bace0 100644 --- a/tests/gocase/unit/scripting/scripting_test.go +++ b/tests/gocase/unit/scripting/scripting_test.go @@ -24,6 +24,7 @@ import ( _ "embed" "fmt" "math/big" + "strings" "testing" "time" @@ -513,6 +514,126 @@ math.randomseed(ARGV[1]); return tostring(math.random()) r := rdb.Eval(ctx, `redis.setresp(3);`, []string{}) util.ErrorRegexp(t, r.Err(), ".*ERR.*You need set resp3-enabled to yes to enable RESP3.*") }) + + t.Run("SCRIPT KILL and lua-time-limit", func(t *testing.T) { + srv2 := util.StartServer(t, map[string]string{"lua-time-limit": "100"}) + defer srv2.Close() + + rdb2 := srv2.NewClient() + defer func() { require.NoError(t, rdb2.Close()) }() + + require.NoError(t, rdb2.Set(ctx, "k1", "v1", 0).Err()) + + errChan := make(chan error, 1) + go func() { + errChan <- rdb2.Eval(ctx, "while true do end", []string{}).Err() + }() + + time.Sleep(200 * time.Millisecond) + + // Get command might get dispatched to the busy thread and time out, so we retry with a short timeout. + var getErr error + for i := 0; i < 20; i++ { + rdb3 := redis.NewClient(&redis.Options{ + Addr: srv2.HostPort(), + DialTimeout: 500 * time.Millisecond, + ReadTimeout: 500 * time.Millisecond, + WriteTimeout: 500 * time.Millisecond, + }) + getErr = rdb3.Get(ctx, "k1").Err() + _ = rdb3.Close() + if getErr != nil && strings.Contains(getErr.Error(), "BUSY") { + break + } + time.Sleep(100 * time.Millisecond) + } + require.Error(t, getErr) + util.ErrorRegexp(t, getErr, ".*BUSY Redis is busy running a script.*") + + // ScriptKill command might also get dispatched to the busy thread and time out, so we retry it. + var killErr error + for i := 0; i < 20; i++ { + rdb3 := redis.NewClient(&redis.Options{ + Addr: srv2.HostPort(), + DialTimeout: 500 * time.Millisecond, + ReadTimeout: 500 * time.Millisecond, + WriteTimeout: 500 * time.Millisecond, + }) + killErr = rdb3.ScriptKill(ctx).Err() + _ = rdb3.Close() + if killErr == nil { + break + } + time.Sleep(100 * time.Millisecond) + } + require.NoError(t, killErr) + + select { + case err := <-errChan: + util.ErrorRegexp(t, err, ".*Script killed by user with SCRIPT.*") + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for script to be killed") + } + + require.NoError(t, rdb2.Set(ctx, "k1", "v2", 0).Err()) + require.Equal(t, "v2", rdb2.Get(ctx, "k1").Val()) + + go func() { + errChan <- rdb2.Eval(ctx, "redis.call('set', 'k2', 'v2'); while true do end", []string{}).Err() + }() + + time.Sleep(200 * time.Millisecond) + + // ScriptKill command on write script might get dispatched to the busy thread and time out, so we retry it. + for i := 0; i < 20; i++ { + rdb3 := redis.NewClient(&redis.Options{ + Addr: srv2.HostPort(), + DialTimeout: 500 * time.Millisecond, + ReadTimeout: 500 * time.Millisecond, + WriteTimeout: 500 * time.Millisecond, + }) + killErr = rdb3.ScriptKill(ctx).Err() + _ = rdb3.Close() + if killErr != nil && strings.Contains(killErr.Error(), "UNKILLABLE") { + break + } + time.Sleep(100 * time.Millisecond) + } + require.Error(t, killErr) + util.ErrorRegexp(t, killErr, ".*UNKILLABLE Sorry the script already executed write commands.*") + }) + + t.Run("SCRIPT KILL on disconnected client", func(t *testing.T) { + srv2 := util.StartServer(t, map[string]string{"lua-time-limit": "500"}) + defer srv2.Close() + + rdb2 := srv2.NewClient() + require.NoError(t, rdb2.Set(ctx, "k1", "v1", 0).Err()) + + errChan := make(chan error, 1) + go func() { + errChan <- rdb2.Eval(ctx, "while true do end", []string{}).Err() + }() + + // 1. Script is running, elapsed time < 500ms + time.Sleep(100 * time.Millisecond) + rdb3 := srv2.NewClient() + defer func() { require.NoError(t, rdb3.Close()) }() + require.Equal(t, "v1", rdb3.Get(ctx, "k1").Val()) // Should succeed + + // 2. Disconnect Client A (rdb2) + require.NoError(t, rdb2.Close()) + + // 3. Immediately after disconnect, elapsed time is still < 500ms, script should still be running + time.Sleep(100 * time.Millisecond) + require.Equal(t, "v1", rdb3.Get(ctx, "k1").Val()) // Should still succeed + + // 4. Wait for total elapsed time to exceed 500ms, the script should be killed automatically + time.Sleep(500 * time.Millisecond) + + // 5. Verify the script has been killed (we can run Get successfully and server is not BUSY) + require.Equal(t, "v1", rdb3.Get(ctx, "k1").Val()) // Should succeed and not be BUSY + }) } func TestScriptingMasterSlave(t *testing.T) { diff --git a/tests/gocase/util/server.go b/tests/gocase/util/server.go index ea8ee474264..5473f8ccbc3 100644 --- a/tests/gocase/util/server.go +++ b/tests/gocase/util/server.go @@ -28,6 +28,7 @@ import ( "os/exec" "path/filepath" "regexp" + "strings" "sync" "syscall" "testing" @@ -234,7 +235,8 @@ func StartServerWithCLIOptions( dir := *workspace require.NotEmpty(t, dir, "please set the workspace by `-workspace`") - dir, err = os.MkdirTemp(dir, fmt.Sprintf("%s-%d-*", t.Name(), time.Now().UnixMilli())) + name := strings.ReplaceAll(t.Name(), "/", "_") + dir, err = os.MkdirTemp(dir, fmt.Sprintf("%s-%d-*", name, time.Now().UnixMilli())) require.NoError(t, err) configs["dir"] = dir From 9ab96c1c0fb33395a0414b64651d0cfd58a05682 Mon Sep 17 00:00:00 2001 From: xtco3o Date: Thu, 28 May 2026 14:43:42 +0800 Subject: [PATCH 2/3] fix(script): optimize scripting time limit check and event loop polling --- src/server/redis_connection.cc | 17 ++++---- src/server/server.cc | 6 ++- src/storage/scripting.cc | 79 +++++++++++++++++++--------------- src/storage/scripting.h | 1 + 4 files changed, 59 insertions(+), 44 deletions(-) diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc index d32d56c34ed..c41954ce0b2 100644 --- a/src/server/redis_connection.cc +++ b/src/server/redis_connection.cc @@ -119,10 +119,18 @@ void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) { if (is_running_) return; is_running_ = true; + 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()); if (!s.IsOK()) { - is_running_ = false; EnableFlag(redis::Connection::kCloseAfterReply); Reply(redis::Error(s)); INFO("[connection] Failed to tokenize the request. Error: {}", s.Msg()); @@ -130,13 +138,6 @@ void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) { } ExecuteCommands(req_.GetCommands()); - is_running_ = false; - - if (IsFlagEnabled(kCloseAsync)) { - Close(); - } else if (evbuffer_get_length(Input()) > 0) { - bufferevent_trigger(bev_, EV_READ, BEV_TRIG_IGNORE_WATERMARKS); - } } void Connection::OnWrite([[maybe_unused]] bufferevent *bev) { diff --git a/src/server/server.cc b/src/server/server.cc index 69f2ec3649b..22368e6e588 100644 --- a/src/server/server.cc +++ b/src/server/server.cc @@ -1935,8 +1935,6 @@ Status Server::ScriptKill() { for (auto *rctx : running_scripts_) { if (rctx->is_write_dirty) { has_unkillable = true; - } else { - rctx->is_killed = true; } } @@ -1947,6 +1945,10 @@ Status Server::ScriptKill() { "command."}; } + for (auto *rctx : running_scripts_) { + rctx->is_killed = true; + } + return Status::OK(); } diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc index c1be17df12a..bfb16da133b 100644 --- a/src/storage/scripting.cc +++ b/src/storage/scripting.cc @@ -103,44 +103,55 @@ void LuaMaskCountHook(lua_State *lua, [[maybe_unused]] lua_Debug *ar) { int limit = srv->GetConfig()->lua_time_limit; bool is_disconnected = script_run_ctx->conn->IsFlagEnabled(redis::Connection::kCloseAsync); - int effective_limit = limit; - if (is_disconnected && limit <= 0) { - effective_limit = 60000; // 60 seconds default timeout when disconnected + uint64_t now_ms = util::GetTimeStampMS(); + + // If the time limit is reached, we set the script timeout flag and maybe warn + if (limit > 0 && now_ms - script_run_ctx->start_time_ms >= static_cast(limit)) { + srv->SetScriptTimedOut(true); + if (!script_run_ctx->slow_logged) { + WARN( + "Slow script detected: still in execution after {} milliseconds. You can try killing the script using the " + "SCRIPT KILL command.", + now_ms - script_run_ctx->start_time_ms); + script_run_ctx->slow_logged = true; + } } - if (effective_limit > 0) { - uint64_t now_ms = util::GetTimeStampMS(); - if (now_ms - script_run_ctx->start_time_ms >= static_cast(effective_limit)) { - if (limit > 0) { - srv->SetScriptTimedOut(true); - } - - if (!script_run_ctx->slow_logged) { - if (is_disconnected) { - WARN("Slow script detected on disconnected client: still in execution after {} milliseconds. Killing it.", - now_ms - script_run_ctx->start_time_ms); - } else { - WARN( - "Slow script detected: still in execution after {} milliseconds. You can try killing the script using " - "the " - "SCRIPT KILL command.", - now_ms - script_run_ctx->start_time_ms); - } - script_run_ctx->slow_logged = true; - } + // If the client has disconnected, we kill the script (if it hasn't written to the DB). + if (is_disconnected) { + if (!script_run_ctx->slow_logged) { + WARN("Slow script detected on disconnected client: still in execution after {} milliseconds. Killing it.", + now_ms - script_run_ctx->start_time_ms); + script_run_ctx->slow_logged = true; + } + if (!script_run_ctx->is_write_dirty) { + KillScript(lua); + } + } else { + // Determine the polling interval. + // If the limit is enabled, we poll immediately when the limit is exceeded. + // If the limit is disabled, we poll every 5 seconds to check for client disconnection. + uint64_t poll_interval = limit > 0 ? static_cast(limit) : 5000; + + // We only poll if: + // 1. The limit is enabled and the script has exceeded the limit. + // 2. OR the script has been running for at least poll_interval, and we haven't polled in the last poll_interval. + bool should_poll = false; + if (limit > 0 && now_ms - script_run_ctx->start_time_ms >= static_cast(limit)) { + should_poll = (script_run_ctx->last_poll_time_ms == 0 || now_ms - script_run_ctx->last_poll_time_ms >= 100); + } else { + should_poll = + (now_ms - script_run_ctx->start_time_ms >= poll_interval) && + (script_run_ctx->last_poll_time_ms == 0 || now_ms - script_run_ctx->last_poll_time_ms >= poll_interval); + } - if (is_disconnected) { - if (!script_run_ctx->is_write_dirty) { - KillScript(lua); - } - } else { - // Poll the worker thread's event loop to process SCRIPT KILL or other commands. - auto *worker = script_run_ctx->conn->Owner(); - worker->PollEventLoop(); + if (should_poll) { + auto *worker = script_run_ctx->conn->Owner(); + worker->PollEventLoop(); + script_run_ctx->last_poll_time_ms = now_ms; - if (script_run_ctx->is_killed) { - KillScript(lua); - } + if (script_run_ctx->is_killed) { + KillScript(lua); } } } diff --git a/src/storage/scripting.h b/src/storage/scripting.h index 07c86efae85..433829180f7 100644 --- a/src/storage/scripting.h +++ b/src/storage/scripting.h @@ -164,6 +164,7 @@ struct ScriptRunCtx { std::atomic is_killed{false}; bool slow_logged = false; std::atomic is_write_dirty{false}; + uint64_t last_poll_time_ms = 0; }; /// SaveOnRegistry saves user-defined data to lua REGISTRY From 2f1c4750c817e2ea0d5e3ab1e3c9acd364b846bd Mon Sep 17 00:00:00 2001 From: xtco3o Date: Thu, 28 May 2026 16:05:30 +0800 Subject: [PATCH 3/3] fix(script): address review comments and change default time limit --- kvrocks.conf | 4 ++-- src/config/config.cc | 8 +++++++- src/server/redis_connection.cc | 18 +++++++++--------- src/server/server.cc | 16 ++++++++++++++++ src/server/server.h | 1 + src/server/worker.cc | 9 ++++++++- src/storage/scripting.cc | 29 +++++++++++++++++++++++++++++ 7 files changed, 72 insertions(+), 13 deletions(-) diff --git a/kvrocks.conf b/kvrocks.conf index 62e441d46e5..97747e7de61 100644 --- a/kvrocks.conf +++ b/kvrocks.conf @@ -496,8 +496,8 @@ lua-strict-key-accessing no # 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 +# Default: 5000 (5 seconds) +lua-time-limit 5000 ################################## TLS ################################### diff --git a/src/config/config.cc b/src/config/config.cc index c845edadab1..2eaeac30acc 100644 --- a/src/config/config.cc +++ b/src/config/config.cc @@ -263,7 +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)}, + {"lua-time-limit", false, new IntField(&lua_time_limit, 5000, -1, INT_MAX)}, /* rocksdb options */ {"rocksdb.compression", false, @@ -835,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); diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc index c41954ce0b2..ca429dbd211 100644 --- a/src/server/redis_connection.cc +++ b/src/server/redis_connection.cc @@ -454,18 +454,10 @@ void Connection::ExecuteCommands(std::deque *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 && + 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([&] { @@ -520,6 +512,14 @@ void Connection::ExecuteCommands(std::deque *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 concurrency; // Allow concurrency std::unique_lock exclusivity; // Need exclusivity // If the command needs to process exclusively, we need to get 'ExclusivityGuard' diff --git a/src/server/server.cc b/src/server/server.cc index 22368e6e588..f5e4e3307dc 100644 --- a/src/server/server.cc +++ b/src/server/server.cc @@ -1921,6 +1921,22 @@ void Server::UnregisterRunningScript(lua::ScriptRunCtx *rctx) { is_script_timeout_.store(any_timed_out, std::memory_order_relaxed); } +void Server::ReevaluateScriptTimeout() { + std::lock_guard 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(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); } diff --git a/src/server/server.h b/src/server/server.h index 06b61dc7f4c..8888af84071 100644 --- a/src/server/server.h +++ b/src/server/server.h @@ -321,6 +321,7 @@ class Server { void RegisterRunningScript(lua::ScriptRunCtx *rctx); void UnregisterRunningScript(lua::ScriptRunCtx *rctx); + void ReevaluateScriptTimeout(); bool IsScriptTimedOut() const; void SetScriptTimedOut(bool timed_out); Status ScriptKill(); diff --git a/src/server/worker.cc b/src/server/worker.cc index 712427c1056..4c5045ff134 100644 --- a/src/server/worker.cc +++ b/src/server/worker.cc @@ -24,7 +24,9 @@ #include #include +#include #include +#include #include #include @@ -621,7 +623,12 @@ void Worker::PollEventLoop() { } else { auto *output = bufferevent_get_output(bev); if (evbuffer_get_length(output) > 0) { - evbuffer_write(output, conn->GetFD()); + if (evbuffer_write(output, conn->GetFD()) == -1) { + int err = errno; + if (err != EAGAIN && err != EWOULDBLOCK && err != EINTR) { + WARN("[worker] Failed to write to connection output buffer: {}", strerror(err)); + } + } } } } diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc index bfb16da133b..629f12a3a9c 100644 --- a/src/storage/scripting.cc +++ b/src/storage/scripting.cc @@ -23,6 +23,12 @@ #include "scripting.h" +#ifndef _WIN32 +#include +#else +#include +#endif + #include #include @@ -85,6 +91,9 @@ class ScriptRunCtxGuard { }; static void KillScript(lua_State *lua) { + // We set the hook mask to LUA_MASKLINE to ensure that if the script catches the error + // (e.g. using pcall or xpcall) and attempts to continue executing, the hook will be + // triggered again immediately on the very next line to raise the kill error again. lua_sethook(lua, LuaMaskCountHook, LUA_MASKLINE, 0); PushError(lua, "Script killed by user with SCRIPT KILL..."); RaiseError(lua); @@ -103,6 +112,26 @@ void LuaMaskCountHook(lua_State *lua, [[maybe_unused]] lua_Debug *ar) { int limit = srv->GetConfig()->lua_time_limit; bool is_disconnected = script_run_ctx->conn->IsFlagEnabled(redis::Connection::kCloseAsync); + if (!is_disconnected) { + int fd = script_run_ctx->conn->GetFD(); + if (fd >= 0) { + char buf[1]; +#ifdef _WIN32 + int n = recv(fd, buf, 1, MSG_PEEK); + if (n == 0 || (n == -1 && WSAGetLastError() != WSAEWOULDBLOCK && WSAGetLastError() != WSAEINTR)) { + script_run_ctx->conn->EnableFlag(redis::Connection::kCloseAsync); + is_disconnected = true; + } +#else + ssize_t n = recv(fd, buf, 1, MSG_PEEK | MSG_DONTWAIT); + if (n == 0 || (n == -1 && errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR)) { + script_run_ctx->conn->EnableFlag(redis::Connection::kCloseAsync); + is_disconnected = true; + } +#endif + } + } + uint64_t now_ms = util::GetTimeStampMS(); // If the time limit is reached, we set the script timeout flag and maybe warn