xtco3o commented on code in PR #3505:
URL: https://github.com/apache/kvrocks/pull/3505#discussion_r3316250522
##########
src/server/redis_connection.cc:
##########
@@ -433,6 +454,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;
+ }
+ }
Review Comment:
Fixed. I've tightened the matching condition to check for exactly two tokens
(`cmd_tokens.size() == 2`) for `SCRIPT KILL` commands. Furthermore, I moved the
`IsScriptTimedOut()` check to occur after the authentication and namespace
checks, but still before acquiring any concurrency/exclusivity guards,
preventing unauthenticated clients from bypassing the BUSY check.
##########
src/server/worker.cc:
##########
@@ -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<std::mutex> 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());
+ }
Review Comment:
Fixed. I've updated the logic to check the return value of `evbuffer_write`
and log non-recoverable errors (i.e. those other than `EAGAIN`, `EWOULDBLOCK`,
and `EINTR`) to prevent silent failures. Bypassing write watermarks, callbacks,
and rate limiting is acceptable here because client bufferevents in Kvrocks do
not utilize these features.
##########
src/storage/scripting.cc:
##########
@@ -60,6 +63,100 @@ 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);
Review Comment:
Clarified with comments. We keep the `LUA_MASKLINE` mask active on the hook
so that if a script attempts to catch the raised error using `pcall` or
`xpcall` to continue running, the hook will immediately trigger again on the
next line and re-raise the error. This ensures the script is terminated. Added
inline documentation to clarify this intent.
##########
src/server/server.cc:
##########
@@ -1879,6 +1894,64 @@ 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);
+}
+
+bool Server::IsScriptTimedOut() const { return
is_script_timeout_.load(std::memory_order_relaxed); }
+
Review Comment:
Fixed. Added `ReevaluateScriptTimeout()` method to `Server` and registered a
configuration callback for `lua-time-limit` in `config.cc`. When
`lua-time-limit` is changed via `CONFIG SET`, `ReevaluateScriptTimeout()` is
called to dynamically update `is_script_timeout_`. Also adjusted the default
value of `lua-time-limit` to 5000 in both `config.cc` and `kvrocks.conf`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
