This is an automated email from the ASF dual-hosted git repository.
git-hulk 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 4d7c1a6c7 feat(server): support the client-output-buffer-limit
configuration (#3560)
4d7c1a6c7 is described below
commit 4d7c1a6c75b1eda28b976c8276b1237824c26e90
Author: hulk <[email protected]>
AuthorDate: Sun Jul 19 09:44:40 2026 +0800
feat(server): support the client-output-buffer-limit configuration (#3560)
Add the client-output-buffer-limit configuration to disconnect clients
that are not reading data from the server fast enough, following the
same semantics as Redis: a client is disconnected immediately once the
hard limit is reached, or when it stays over the soft limit for more
than the configured seconds continuously.
Unlike Redis, all the classes are configured in a single line, and only
the specified classes are changed. The limits are disabled by default
for all classes of clients, e.g. to protect against slow Pub/Sub
subscribers with the same limits as Redis:
```
client-output-buffer-limit pubsub 32m 8m 60
```
The limit takes effect on the normal and pubsub classes since the
replication stream is written to the socket directly instead of going
through the connection output buffer. The slave class is accepted for
compatibility but takes no effect: slow replicas are still handled by
max-replication-lag and replication-send-timeout-ms.
The check runs every time data is appended to the connection output
buffer. A client over the limit is closed asynchronously by manually
triggering the deferred write callback, since the write callback of a
stuck client may never fire on its own; the callback then runs in the
owner worker's event loop where freeing the connection is safe.
Also expose the client_output_buffer_limit_disconnections counter in
INFO stats, and classify clients subscribed to only shard channels as
pubsub clients so that they are covered by the pubsub class as well.
Part of #2284
Assistant-By Claude Fable 5
---
kvrocks.conf | 33 +++
src/commands/blocking_commander.h | 13 +
src/commands/cmd_stream.cc | 24 ++
src/config/config.cc | 69 ++++++
src/config/config.h | 24 ++
src/server/redis_connection.cc | 76 +++++-
src/server/redis_connection.h | 17 +-
src/server/server.cc | 2 +
src/server/worker.cc | 25 +-
src/stats/stats.h | 4 +
.../unit/limits/client_output_buffer_limit_test.go | 261 +++++++++++++++++++++
11 files changed, 537 insertions(+), 11 deletions(-)
diff --git a/kvrocks.conf b/kvrocks.conf
index b2ab11fbd..d13a13cd1 100644
--- a/kvrocks.conf
+++ b/kvrocks.conf
@@ -116,6 +116,39 @@ persist-cluster-nodes-enabled yes
#
maxclients 10000
+# The client output buffer limits can be used to force disconnection of clients
+# that are not reading data from the server fast enough for some reason (a
+# common reason is that a Pub/Sub client can't consume messages as fast as the
+# publisher can produce them).
+#
+# The limit can be set differently for the three different classes of clients:
+#
+# normal -> normal clients including MONITOR clients
+# slave -> slave clients (accepted for compatibility with Redis but takes no
+# effect since the replication stream doesn't go through the
+# connection output buffer; use max-replication-lag and
+# replication-send-timeout-ms to handle slow replicas instead)
+# pubsub -> clients subscribed to at least one channel, pattern or shard
channel
+#
+# The format of every class is the following:
+#
+# <class> <hard limit> <soft limit> <soft seconds>
+#
+# A client is immediately disconnected once the hard limit is reached, or if
+# the soft limit is reached and remains reached for the specified number of
+# seconds (continuously). Setting a limit to 0 disables it, and the soft limit
+# should be less than the hard limit when both are set.
+#
+# Unlike Redis, all the classes should be configured in a single line, and only
+# the specified classes are changed. The hard/soft limits accept the units
+# k/m/g (power of 1024).
+#
+# The limits are disabled by default for all classes of clients. For example,
+# to protect against slow Pub/Sub subscribers with the same limits as Redis:
+#
+# client-output-buffer-limit pubsub 32m 8m 60
+client-output-buffer-limit normal 0 0 0 slave 0 0 0 pubsub 0 0 0
+
# Require clients to issue AUTH <PASSWORD> before processing any other
# commands. This might be useful in environments in which you do not trust
# others with access to the host running kvrocks.
diff --git a/src/commands/blocking_commander.h
b/src/commands/blocking_commander.h
index 619c5dbf1..fc668554e 100644
--- a/src/commands/blocking_commander.h
+++ b/src/commands/blocking_commander.h
@@ -69,6 +69,19 @@ class BlockingCommander : public Commander,
}
void OnWrite(bufferevent *bev) {
+ // The connection might be scheduled to close while it's blocked, e.g. it
+ // exceeded the output buffer limit or was killed by the CLIENT KILL
command.
+ // The manually triggered write callback lands here instead of
+ // Connection::OnWrite since the blocking command replaced the bufferevent
+ // callbacks, so the close must be handled here as well, the same way as
+ // the EOF handling in OnEvent.
+ if (conn_->IsFlagEnabled(Connection::kCloseAsync) ||
conn_->IsFlagEnabled(Connection::kCloseAfterReply)) {
+ if (timer_) timer_.reset();
+ UnblockKeys();
+ conn_->Close();
+ return;
+ }
+
bool done{false};
{
// The blocking command should not be executed when the server is in
exclusive state,
diff --git a/src/commands/cmd_stream.cc b/src/commands/cmd_stream.cc
index 73a5279c8..f2facb982 100644
--- a/src/commands/cmd_stream.cc
+++ b/src/commands/cmd_stream.cc
@@ -1340,6 +1340,18 @@ class CommandXRead : public Commander,
}
void OnWrite(bufferevent *bev) {
+ // The connection might be scheduled to close while it's blocked, e.g. it
+ // exceeded the output buffer limit or was killed by the CLIENT KILL
command.
+ // The manually triggered write callback lands here instead of
+ // Connection::OnWrite since the blocking command replaced the bufferevent
+ // callbacks, so the close must be handled here as well.
+ if (conn_->IsFlagEnabled(Connection::kCloseAsync) ||
conn_->IsFlagEnabled(Connection::kCloseAfterReply)) {
+ if (timer_) timer_.reset();
+ unblockAll();
+ conn_->Close();
+ return;
+ }
+
if (timer_ != nullptr) {
timer_.reset();
}
@@ -1643,6 +1655,18 @@ class CommandXReadGroup : public Commander,
}
void OnWrite(bufferevent *bev) {
+ // The connection might be scheduled to close while it's blocked, e.g. it
+ // exceeded the output buffer limit or was killed by the CLIENT KILL
command.
+ // The manually triggered write callback lands here instead of
+ // Connection::OnWrite since the blocking command replaced the bufferevent
+ // callbacks, so the close must be handled here as well.
+ if (conn_->IsFlagEnabled(Connection::kCloseAsync) ||
conn_->IsFlagEnabled(Connection::kCloseAfterReply)) {
+ if (timer_) timer_.reset();
+ unblockAll();
+ conn_->Close();
+ return;
+ }
+
if (timer_ != nullptr) {
timer_.reset();
}
diff --git a/src/config/config.cc b/src/config/config.cc
index 3932574cd..0de6a2c6f 100644
--- a/src/config/config.cc
+++ b/src/config/config.cc
@@ -180,6 +180,8 @@ Config::Config() {
{"timeout", false, new IntField(&timeout, 0, 0, INT_MAX)},
{"tcp-backlog", true, new IntField(&backlog, 511, 0, INT_MAX)},
{"maxclients", false, new IntField(&maxclients, 10240, 0, INT_MAX)},
+ {"client-output-buffer-limit", false,
+ new StringField(&client_output_buffer_limit_str_, "normal 0 0 0 slave 0
0 0 pubsub 0 0 0")},
{"max-backup-to-keep", false, new IntField(&max_backup_to_keep, 1, 0,
1)},
{"max-backup-keep-hours", false, new IntField(&max_backup_keep_hours, 0,
0, INT_MAX)},
{"master-use-repl-port", false, new YesNoField(&master_use_repl_port,
false)},
@@ -444,6 +446,55 @@ void Config::initFieldValidator() {
// The callback function would be invoked after the field was set,
// it may change related fields or re-format the field. for example,
// when the 'dir' was set, the db-dir or backup-dir should be reset as well.
+// Parses a client-output-buffer-limit spec like "normal 0 0 0 pubsub 32m 8m
60"
+// and applies it to the limits of the specified client kinds. The full spec is
+// parsed before applying anything, so a malformed quadruple cannot leave the
+// limits partially updated.
+Status Config::parseClientOutputBufferLimits(const std::string &v) {
+ std::vector<std::string> args = util::Split(v, " \t");
+ if (args.empty() || args.size() % 4 != 0) {
+ return {Status::NotOK, "should be in the format of <class> <hard limit>
<soft limit> <soft seconds> ..."};
+ }
+
+ struct ParsedLimit {
+ ClientKind kind;
+ uint64_t hard_limit_bytes, soft_limit_bytes;
+ int64_t soft_limit_seconds;
+ };
+ std::vector<ParsedLimit> parsed;
+ for (size_t i = 0; i < args.size(); i += 4) {
+ ClientKind kind = ClientKind::kNormal;
+ if (util::EqualICase(args[i], "normal")) {
+ kind = ClientKind::kNormal;
+ } else if (util::EqualICase(args[i], "slave") || util::EqualICase(args[i],
"replica")) {
+ kind = ClientKind::kSlave;
+ } else if (util::EqualICase(args[i], "pubsub")) {
+ kind = ClientKind::kPubsub;
+ } else {
+ return {Status::NotOK, fmt::format("unknown client kind '{}'", args[i])};
+ }
+ auto hard = GET_OR_RET(ParseSizeAndUnit(args[i + 1]).Prefixed("invalid
hard limit"));
+ auto soft = GET_OR_RET(ParseSizeAndUnit(args[i + 2]).Prefixed("invalid
soft limit"));
+ auto secs = GET_OR_RET(
+ ParseInt<int64_t>(args[i + 3], NumericRange<int64_t>{0, INT64_MAX},
10).Prefixed("invalid soft seconds"));
+ if (hard > INT64_MAX || soft > INT64_MAX) {
+ return {Status::NotOK, fmt::format("the hard limit and soft limit should
be no greater than {}", INT64_MAX)};
+ }
+ if (hard != 0 && soft >= hard) {
+ return {Status::NotOK, "the soft limit should be less than the hard
limit"};
+ }
+ parsed.push_back({kind, hard, soft, secs});
+ }
+
+ for (const auto &p : parsed) {
+ auto &limit = GetClientOutputBufferLimit(p.kind);
+ limit.hard_limit_bytes.store(p.hard_limit_bytes,
std::memory_order_relaxed);
+ limit.soft_limit_bytes.store(p.soft_limit_bytes,
std::memory_order_relaxed);
+ limit.soft_limit_seconds.store(p.soft_limit_seconds,
std::memory_order_relaxed);
+ }
+ return Status::OK();
+}
+
void Config::initFieldCallback() {
auto set_db_option_cb = [](Server *srv, const std::string &k, const
std::string &v) -> Status {
if (!srv) return Status::OK(); // srv is nullptr when load config from
file
@@ -551,6 +602,24 @@ void Config::initFieldCallback() {
srv->AdjustOpenFilesLimit();
return Status::OK();
}},
+ {"client-output-buffer-limit",
+ [this]([[maybe_unused]] Server *srv, [[maybe_unused]] const
std::string &k, const std::string &v) -> Status {
+ if (auto s = parseClientOutputBufferLimits(v); !s.IsOK()) return
s;
+
+ // Canonicalize the stored string so that CONFIG GET/REWRITE
always
+ // report all classes, even if only a subset was specified.
+ constexpr const char *client_kinds[] = {"normal", "slave",
"pubsub"};
+ std::string canonical;
+ for (size_t i = 0; i < static_cast<size_t>(ClientKind::kCount);
i++) {
+ const auto &limit = client_output_buffer_limits[i];
+ canonical += fmt::format("{}{} {} {} {}", i == 0 ? "" : " ",
client_kinds[i],
+
limit.hard_limit_bytes.load(std::memory_order_relaxed),
+
limit.soft_limit_bytes.load(std::memory_order_relaxed),
+
limit.soft_limit_seconds.load(std::memory_order_relaxed));
+ }
+ client_output_buffer_limit_str_ = std::move(canonical);
+ return Status::OK();
+ }},
{"slaveof", replicaof_cb},
{"replicaof", replicaof_cb},
{"profiling-sample-commands",
diff --git a/src/config/config.h b/src/config/config.h
index 62c13c491..a6dd19f21 100644
--- a/src/config/config.h
+++ b/src/config/config.h
@@ -23,6 +23,7 @@
#include <rocksdb/options.h>
#include <sys/resource.h>
+#include <atomic>
#include <map>
#include <memory>
#include <set>
@@ -45,6 +46,22 @@ constexpr const uint32_t PORT_LIMIT = 65535;
enum SupervisedMode { kSupervisedNone = 0, kSupervisedAutoDetect,
kSupervisedSystemd, kSupervisedUpStart };
+enum class ClientKind {
+ kNormal = 0,
+ kSlave = 1,
+ kPubsub = 2,
+ kCount = 3,
+};
+
+// Limits are read from reply paths that may run outside the worker threads
+// (e.g. WAIT wakeups from the feed-replica thread), so keep the fields atomic
+// instead of guarding them with the command exclusivity guard.
+struct ClientOutputBufferLimit {
+ std::atomic<uint64_t> hard_limit_bytes = 0;
+ std::atomic<uint64_t> soft_limit_bytes = 0;
+ std::atomic<int64_t> soft_limit_seconds = 0;
+};
+
constexpr const char *TLS_AUTH_CLIENTS_NO = "no";
constexpr const char *TLS_AUTH_CLIENTS_OPTIONAL = "optional";
@@ -109,6 +126,11 @@ struct Config {
spdlog::level::level_enum log_level = spdlog::level::info;
int backlog = 511;
int maxclients = 10000;
+ ClientOutputBufferLimit
client_output_buffer_limits[static_cast<size_t>(ClientKind::kCount)];
+
+ ClientOutputBufferLimit &GetClientOutputBufferLimit(ClientKind kind) {
+ return client_output_buffer_limits[static_cast<size_t>(kind)];
+ }
int max_backup_to_keep = 1;
int max_backup_keep_hours = 24;
int slowlog_log_slower_than = 100000;
@@ -292,6 +314,7 @@ struct Config {
std::string compaction_checker_range_str_;
std::string compaction_checker_cron_str_;
std::string profiling_sample_commands_str_;
+ std::string client_output_buffer_limit_str_;
std::map<std::string, std::unique_ptr<ConfigField>> fields_;
std::vector<std::string> rename_command_;
std::string histogram_bucket_boundaries_str_;
@@ -299,6 +322,7 @@ struct Config {
void initFieldValidator();
void initFieldCallback();
+ Status parseClientOutputBufferLimits(const std::string &v);
Status parseConfigFromPair(const std::pair<std::string, std::string> &input,
int line_number);
Status parseConfigFromString(const std::string &input, int line_number);
bool checkFieldValueIsDefault(const std::string &key, const std::string
&value) const;
diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc
index aa301b32f..65061c017 100644
--- a/src/server/redis_connection.cc
+++ b/src/server/redis_connection.cc
@@ -93,11 +93,60 @@ std::string Connection::ToString() {
evbuffer_get_length(Input()), evbuffer_get_length(Output()), last_cmd_,
set_info_.lib_name, set_info_.lib_ver);
}
-void Connection::Close() {
+void Connection::Close(bool is_async) {
+ if (is_async) {
+ // Only the first caller should schedule the close since concurrent reply
+ // paths (e.g. publishers on other workers) may race here.
+ if (flags_.fetch_or(kCloseAsync) & kCloseAsync) return;
+
+ // The write callback of a stuck client may never be invoked since its
output
+ // buffer cannot drain, so trigger the callback manually instead of waiting
+ // for it. Ignoring watermarks is required because the write callback is
only
+ // triggered when the output buffer size is not larger than the low
watermark.
+ // The callback is deferred(BEV_OPT_DEFER_CALLBACKS) and runs in the owner
+ // worker's event loop, where it's safe to free the connection.
+ bufferevent_trigger(bev_, EV_WRITE, BEV_TRIG_IGNORE_WATERMARKS |
BEV_TRIG_DEFER_CALLBACKS);
+ return;
+ }
+
if (close_cb) close_cb(GetFD());
owner_->FreeConnection(this);
}
+bool Connection::IsExceedOutputBufferLimit() {
+ // Connections that are already scheduled to close don't need to be checked
+ // again. The replication stream is written to the socket directly instead of
+ // going through the connection output buffer, so the slave kind is not
+ // applicable here: slow replicas are handled by max-replication-lag and
+ // replication-send-timeout-ms.
+ if (IsFlagEnabled(kCloseAsync) || IsFlagEnabled(kCloseAfterReply) ||
IsFlagEnabled(kSlave)) return false;
+
+ auto kind = GetClientType() == kTypePubsub ? ClientKind::kPubsub :
ClientKind::kNormal;
+ const auto &limit = srv_->GetConfig()->GetClientOutputBufferLimit(kind);
+ uint64_t hard_limit_bytes =
limit.hard_limit_bytes.load(std::memory_order_relaxed);
+ uint64_t soft_limit_bytes =
limit.soft_limit_bytes.load(std::memory_order_relaxed);
+ if (hard_limit_bytes == 0 && soft_limit_bytes == 0) return false;
+
+ uint64_t used_bytes = evbuffer_get_length(Output());
+ if (hard_limit_bytes != 0 && used_bytes >= hard_limit_bytes) return true;
+
+ if (soft_limit_bytes != 0) {
+ if (used_bytes >= soft_limit_bytes) {
+ int64_t soft_limit_seconds =
limit.soft_limit_seconds.load(std::memory_order_relaxed);
+ int64_t now = util::GetTimeStamp();
+ int64_t reached_time =
obuf_soft_limit_reached_time_.load(std::memory_order_relaxed);
+ if (reached_time == 0) {
+ obuf_soft_limit_reached_time_.store(now, std::memory_order_relaxed);
+ } else if (now - reached_time > soft_limit_seconds) {
+ return true;
+ }
+ } else {
+ obuf_soft_limit_reached_time_.store(0, std::memory_order_relaxed);
+ }
+ }
+ return false;
+}
+
void Connection::Detach() { owner_->DetachConnection(this); }
void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) {
@@ -152,6 +201,11 @@ void Connection::OnEvent(bufferevent *bev, int16_t events)
{
}
void Connection::Reply(const std::string &msg) {
+ // Connections scheduled to be closed asynchronously don't need any more
+ // replies, the pending output buffer is dropped when the connection is
freed.
+ if (IsFlagEnabled(kCloseAsync)) {
+ return;
+ }
if (reply_mode_ == ReplyMode::SKIP) {
reply_mode_ = ReplyMode::ON;
return;
@@ -165,6 +219,13 @@ void Connection::Reply(const std::string &msg) {
queued_replies_.push_back(msg);
} else {
redis::Reply(bufferevent_get_output(bev_), msg);
+ if (IsExceedOutputBufferLimit()) {
+ WARN(
+ "[connection] Client {} (id={}) scheduled to be closed ASAP for
overcoming of output buffer limits, obuf: {}",
+ addr_, id_, evbuffer_get_length(Output()));
+ srv_->stats.IncrClientOutputBufferLimitDisconnections();
+ Close(true /* is_async */);
+ }
}
}
@@ -198,7 +259,8 @@ uint64_t Connection::GetIdleTime() const { return
static_cast<uint64_t>(util::Ge
uint64_t Connection::GetClientType() const {
if (IsFlagEnabled(kSlave)) return kTypeSlave;
- if (!subscribe_channels_.empty() || !subscribe_patterns_.empty()) return
kTypePubsub;
+ if (!subscribe_channels_.empty() || !subscribe_patterns_.empty() ||
!subscribe_shard_channels_.empty())
+ return kTypePubsub;
return kTypeNormal;
}
@@ -222,9 +284,10 @@ void Connection::DisableFlag(Flag flag) { flags_ &=
(~flag); }
bool Connection::IsFlagEnabled(Flag flag) const { return (flags_ & flag) > 0; }
bool Connection::CanMigrate() const {
- return !is_running_ //
reading or writing
- && !IsFlagEnabled(redis::Connection::kCloseAfterReply) //
close after reply
- && saved_current_command_ == nullptr //
not executing blocking command like BLPOP
+ return !is_running_ // reading or
writing
+ && !IsFlagEnabled(redis::Connection::kCloseAfterReply) // close
after reply
+ && !IsFlagEnabled(redis::Connection::kCloseAsync) // async close
might be pending on the current event base
+ && saved_current_command_ == nullptr // not executing
blocking command like BLPOP
&& subscribe_channels_.empty() && subscribe_patterns_.empty(); //
not subscribing any channel
}
@@ -434,7 +497,8 @@ void Connection::ExecuteCommands(std::deque<CommandTokens>
*to_process_cmds) {
if (cmd_tokens.empty()) continue;
bool is_multi_exec = IsFlagEnabled(Connection::kMultiExec);
- if (IsFlagEnabled(redis::Connection::kCloseAfterReply) && !is_multi_exec)
break;
+ if ((IsFlagEnabled(Connection::kCloseAfterReply) ||
IsFlagEnabled(Connection::kCloseAsync)) && !is_multi_exec)
+ break;
auto multi_error_exit = MakeScopeExit([&] {
if (is_multi_exec) multi_error_ = true;
});
diff --git a/src/server/redis_connection.h b/src/server/redis_connection.h
index 909ce8759..c4bbd630c 100644
--- a/src/server/redis_connection.h
+++ b/src/server/redis_connection.h
@@ -62,7 +62,11 @@ class Connection : public EvbufCallbackBase<Connection> {
Connection(const Connection &) = delete;
Connection &operator=(const Connection &) = delete;
- void Close();
+ // Closes the connection immediately by default. Pass is_async=true to
+ // schedule the close on the owner worker's event loop instead, which is
+ // required when the connection is being closed from another thread or
+ // while it's still executing commands.
+ void Close(bool is_async = false);
void Detach();
void OnRead(bufferevent *bev);
void OnWrite(bufferevent *bev);
@@ -70,6 +74,13 @@ class Connection : public EvbufCallbackBase<Connection> {
void SendFile(int fd);
std::string ToString();
+ // Returns true if the connection output buffer size exceeds the configured
+ // client-output-buffer-limit of its client kind, following the same
hard/soft
+ // limit semantics as Redis. It should be checked every time data is appended
+ // to the connection output buffer, and the caller is responsible for closing
+ // the connection when it returns true.
+ bool IsExceedOutputBufferLimit();
+
void Reply(const std::string &msg);
const std::vector<std::string> &GetQueuedReplies() const;
void ClearQueuedReplies() { queued_replies_.clear(); }
@@ -246,6 +257,10 @@ class Connection : public EvbufCallbackBase<Connection> {
std::vector<std::string> queued_replies_;
bool is_paused_ = false;
+
+ // The first time the output buffer size was found to exceed the soft limit
+ // of client-output-buffer-limit, or 0 if it's currently below the limit.
+ std::atomic<int64_t> obuf_soft_limit_reached_time_ = 0;
};
} // namespace redis
diff --git a/src/server/server.cc b/src/server/server.cc
index 6ebc2c3c4..78c578abd 100644
--- a/src/server/server.cc
+++ b/src/server/server.cc
@@ -1406,6 +1406,8 @@ Server::InfoEntries Server::GetStatsInfo() {
entries.emplace_back("sync_full", stats.fullsync_count.load());
entries.emplace_back("sync_partial_ok", stats.psync_ok_count.load());
entries.emplace_back("sync_partial_err", stats.psync_err_count.load());
+ entries.emplace_back("client_output_buffer_limit_disconnections",
+ stats.client_output_buffer_limit_disconnections.load());
auto db_stats = storage->GetDBStats();
entries.emplace_back("keyspace_hits", db_stats->keyspace_hits.load());
diff --git a/src/server/worker.cc b/src/server/worker.cc
index 150d27aab..fcd9343a4 100644
--- a/src/server/worker.cc
+++ b/src/server/worker.cc
@@ -495,13 +495,30 @@ void Worker::UnpauseConnection(int fd, uint64_t id) {
Status Worker::Reply(int fd, const std::string &reply) {
std::unique_lock<std::mutex> lock(conns_mu_);
auto iter = conns_.find(fd);
- if (iter != conns_.end()) {
- iter->second->SetLastInteraction();
- redis::Reply(iter->second->Output(), reply);
+ if (iter == conns_.end()) {
+ return {Status::NotOK, "connection doesn't exist"};
+ }
+
+ // A connection that is scheduled to close doesn't need any more replies
+ // since its pending output will be dropped when it's freed, but it's still
+ // counted as a receiver like Redis does, as it remains subscribed until
+ // the connection is freed.
+ if (iter->second->IsFlagEnabled(redis::Connection::kCloseAsync)) {
return Status::OK();
}
- return {Status::NotOK, "connection doesn't exist"};
+ iter->second->SetLastInteraction();
+ redis::Reply(iter->second->Output(), reply);
+ if (iter->second->IsExceedOutputBufferLimit()) {
+ WARN("[worker] Client {} (id={}) scheduled to be closed ASAP for
overcoming of output buffer limits, obuf: {}",
+ iter->second->GetAddr(), iter->second->GetID(),
evbuffer_get_length(iter->second->Output()));
+ srv->stats.IncrClientOutputBufferLimitDisconnections();
+ // The message was already appended to the output buffer before the
+ // connection was scheduled to close, so the subscriber is still counted
+ // as a receiver of the message, the same as Redis.
+ iter->second->Close(true /* is_async */);
+ }
+ return Status::OK();
}
void Worker::BecomeMonitorConn(redis::Connection *conn) {
diff --git a/src/stats/stats.h b/src/stats/stats.h
index 0bae042d6..961298831 100644
--- a/src/stats/stats.h
+++ b/src/stats/stats.h
@@ -76,6 +76,7 @@ class Stats {
std::atomic<uint64_t> fullsync_count = {0};
std::atomic<uint64_t> psync_err_count = {0};
std::atomic<uint64_t> psync_ok_count = {0};
+ std::atomic<uint64_t> client_output_buffer_limit_disconnections = {0};
std::map<std::string, CommandStat> commands_stats;
using BucketBoundaries = std::vector<double>;
@@ -91,6 +92,9 @@ class Stats {
void IncrFullSyncCount() { fullsync_count.fetch_add(1,
std::memory_order_relaxed); }
void IncrPSyncErrCount() { psync_err_count.fetch_add(1,
std::memory_order_relaxed); }
void IncrPSyncOKCount() { psync_ok_count.fetch_add(1,
std::memory_order_relaxed); }
+ void IncrClientOutputBufferLimitDisconnections() {
+ client_output_buffer_limit_disconnections.fetch_add(1,
std::memory_order_relaxed);
+ }
static int64_t GetMemoryRSS();
void TrackInstantaneousMetric(int metric, uint64_t current_reading);
uint64_t GetInstantaneousMetric(int metric) const;
diff --git a/tests/gocase/unit/limits/client_output_buffer_limit_test.go
b/tests/gocase/unit/limits/client_output_buffer_limit_test.go
new file mode 100644
index 000000000..0b2e37cf2
--- /dev/null
+++ b/tests/gocase/unit/limits/client_output_buffer_limit_test.go
@@ -0,0 +1,261 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package limits
+
+import (
+ "context"
+ "strconv"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/apache/kvrocks/tests/gocase/util"
+ "github.com/stretchr/testify/require"
+)
+
+func getClientOutputBufferLimitDisconnections(t *testing.T, srv
*util.KvrocksServer) int {
+ t.Helper()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+ count, err := strconv.Atoi(util.FindInfoEntry(rdb,
"client_output_buffer_limit_disconnections"))
+ require.NoError(t, err)
+ return count
+}
+
+func TestClientOutputBufferLimitConfig(t *testing.T) {
+ srv := util.StartServer(t, map[string]string{})
+ defer srv.Close()
+
+ ctx := context.Background()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+
+ t.Run("limits are disabled by default", func(t *testing.T) {
+ v, err := rdb.ConfigGet(ctx,
"client-output-buffer-limit").Result()
+ require.NoError(t, err)
+ require.Equal(t, "normal 0 0 0 slave 0 0 0 pubsub 0 0 0",
v["client-output-buffer-limit"])
+ })
+
+ t.Run("partial spec only changes the specified classes", func(t
*testing.T) {
+ require.NoError(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", "pubsub 256k 128k 30").Err())
+ v, err := rdb.ConfigGet(ctx,
"client-output-buffer-limit").Result()
+ require.NoError(t, err)
+ require.Equal(t, "normal 0 0 0 slave 0 0 0 pubsub 262144 131072
30", v["client-output-buffer-limit"])
+
+ // 'replica' is an alias of the slave class
+ require.NoError(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", "replica 1m 512k 10").Err())
+ v, err = rdb.ConfigGet(ctx,
"client-output-buffer-limit").Result()
+ require.NoError(t, err)
+ require.Equal(t, "normal 0 0 0 slave 1048576 524288 10 pubsub
262144 131072 30", v["client-output-buffer-limit"])
+ })
+
+ t.Run("invalid spec is rejected and keeps the old value", func(t
*testing.T) {
+ require.NoError(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", "normal 100m 0 0").Err())
+ for _, spec := range []string{
+ "",
+ "normal 0 0",
+ "unknown 0 0 0",
+ "normal x 0 0",
+ "normal 0 x 0",
+ "normal 0 0 x",
+ "normal 0 0 0 pubsub 1m 512k",
+ // hard/soft limits and soft seconds are limited to
INT64_MAX
+ "normal 9223372036854775808 0 0",
+ "normal 0 9223372036854775808 0",
+ "normal 0 0 9223372036854775808",
+ "normal 0 0 -1",
+ // the soft limit should be less than the hard limit
+ "normal 1m 2m 0",
+ "normal 1m 1m 0",
+ } {
+ require.Error(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", spec).Err())
+ }
+ v, err := rdb.ConfigGet(ctx,
"client-output-buffer-limit").Result()
+ require.NoError(t, err)
+ require.Equal(t, "normal 104857600 0 0 slave 1048576 524288 10
pubsub 262144 131072 30",
+ v["client-output-buffer-limit"])
+ require.NoError(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", "normal 0 0 0").Err())
+ })
+
+ t.Run("config rewrite and restart keep the value", func(t *testing.T) {
+ require.NoError(t, rdb.ConfigSet(ctx,
"client-output-buffer-limit", "pubsub 4m 1m 20").Err())
+ require.NoError(t, rdb.ConfigRewrite(ctx).Err())
+ srv.Restart()
+
+ c := srv.NewClient()
+ defer func() { require.NoError(t, c.Close()) }()
+ v, err := c.ConfigGet(ctx,
"client-output-buffer-limit").Result()
+ require.NoError(t, err)
+ require.Equal(t, "normal 0 0 0 slave 1048576 524288 10 pubsub
4194304 1048576 20",
+ v["client-output-buffer-limit"])
+ })
+}
+
+func TestClientOutputBufferLimitPubsubHardLimit(t *testing.T) {
+ srv := util.StartServer(t, map[string]string{
+ "client-output-buffer-limit": "normal 0 0 0 slave 0 0 0 pubsub
1m 0 0",
+ })
+ defer srv.Close()
+
+ ctx := context.Background()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+
+ t.Run("slow subscriber is disconnected once the hard limit is reached",
func(t *testing.T) {
+ sub := srv.NewTCPClient()
+ defer func() { require.NoError(t, sub.Close()) }()
+ require.NoError(t, sub.WriteArgs("SUBSCRIBE", "ch"))
+ sub.MustRead(t, "*3")
+ sub.MustRead(t, "$9")
+ sub.MustRead(t, "subscribe")
+ sub.MustRead(t, "$2")
+ sub.MustRead(t, "ch")
+ sub.MustRead(t, ":1")
+
+ // A single message bigger than the hard limit must kill the
subscriber
+ // right when it's appended to the output buffer. The message
is made
+ // much bigger than the limit so that the kernel socket buffers
cannot
+ // concurrently drain the output buffer below it.
+ payload := strings.Repeat("x", 8*1024*1024)
+ res := rdb.Publish(ctx, "ch", payload)
+ require.NoError(t, res.Err())
+ // The message was already appended to the output buffer of the
+ // subscriber before it was scheduled to close, so the
subscriber is
+ // still counted as a receiver of the message, the same as
Redis.
+ require.EqualValues(t, 1, res.Val())
+
+ require.Eventually(t, func() bool {
+ return getClientOutputBufferLimitDisconnections(t, srv)
== 1
+ }, 5*time.Second, 100*time.Millisecond)
+ })
+}
+
+func TestClientOutputBufferLimitPubsubSoftLimit(t *testing.T) {
+ srv := util.StartServer(t, map[string]string{
+ "client-output-buffer-limit": "normal 0 0 0 slave 0 0 0 pubsub
0 512k 1",
+ })
+ defer srv.Close()
+
+ ctx := context.Background()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+
+ t.Run("slow subscriber is disconnected after staying over the soft
limit", func(t *testing.T) {
+ sub := srv.NewTCPClient()
+ defer func() { require.NoError(t, sub.Close()) }()
+ require.NoError(t, sub.WriteArgs("SUBSCRIBE", "ch"))
+ sub.MustRead(t, "*3")
+ sub.MustRead(t, "$9")
+ sub.MustRead(t, "subscribe")
+ sub.MustRead(t, "$2")
+ sub.MustRead(t, "ch")
+ sub.MustRead(t, ":1")
+
+ // The subscriber stops reading: the first big message crosses
the soft
+ // limit and starts the clock, but must not kill the connection
yet.
+ // The message is made much bigger than the limit so that the
kernel
+ // socket buffers cannot drain the output buffer below it in
between.
+ payload := strings.Repeat("x", 8*1024*1024)
+ require.NoError(t, rdb.Publish(ctx, "ch", payload).Err())
+ require.Equal(t, 0, getClientOutputBufferLimitDisconnections(t,
srv))
+
+ // Still over the soft limit after more than soft-seconds: the
next
+ // append must disconnect the subscriber.
+ time.Sleep(2 * time.Second)
+ require.NoError(t, rdb.Publish(ctx, "ch", "ping").Err())
+
+ require.Eventually(t, func() bool {
+ return getClientOutputBufferLimitDisconnections(t, srv)
== 1
+ }, 5*time.Second, 100*time.Millisecond)
+ })
+}
+
+func TestBlockedClientOutputBufferLimit(t *testing.T) {
+ srv := util.StartServer(t, map[string]string{
+ "client-output-buffer-limit": "normal 0 0 0 slave 0 0 0 pubsub
1m 0 0",
+ })
+ defer srv.Close()
+
+ ctx := context.Background()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+
+ t.Run("subscriber blocked on BLPOP is disconnected once the hard limit
is reached", func(t *testing.T) {
+ sub := srv.NewTCPClient()
+ defer func() { require.NoError(t, sub.Close()) }()
+ require.NoError(t, sub.WriteArgs("SUBSCRIBE", "ch"))
+ sub.MustRead(t, "*3")
+ sub.MustRead(t, "$9")
+ sub.MustRead(t, "subscribe")
+ sub.MustRead(t, "$2")
+ sub.MustRead(t, "ch")
+ sub.MustRead(t, ":1")
+
+ // Block the subscriber on a key that will never be written, so
that the
+ // blocking command replaces the bufferevent callbacks of the
connection.
+ require.NoError(t, sub.WriteArgs("BLPOP", "blocked-key", "0"))
+ require.Eventually(t, func() bool {
+ return strings.Contains(rdb.ClientList(ctx).Val(),
"cmd=blpop")
+ }, 5*time.Second, 100*time.Millisecond)
+ require.Equal(t, "1", util.FindInfoEntry(rdb,
"blocked_clients"))
+
+ // The blocked client must be closed instead of only being
counted as
+ // disconnected while lingering with a full output buffer.
+ payload := strings.Repeat("x", 8*1024*1024)
+ require.NoError(t, rdb.Publish(ctx, "ch", payload).Err())
+
+ require.Eventually(t, func() bool {
+ return !strings.Contains(rdb.ClientList(ctx).Val(),
"cmd=blpop")
+ }, 5*time.Second, 100*time.Millisecond)
+ require.Equal(t, 1, getClientOutputBufferLimitDisconnections(t,
srv))
+
+ // The disconnected client must be unblocked as well, so it
should no
+ // longer be counted as a blocked client.
+ require.Equal(t, "0", util.FindInfoEntry(rdb,
"blocked_clients"))
+
+ // The stale registration of the disconnected client should not
consume
+ // the pushed element: it must still be delivered to a new
consumer.
+ require.NoError(t, rdb.LPush(ctx, "blocked-key", "value").Err())
+ require.Equal(t, []string{"blocked-key", "value"},
rdb.BLPop(ctx, time.Second, "blocked-key").Val())
+ })
+}
+
+func TestNormalClientOutputBufferLimit(t *testing.T) {
+ srv := util.StartServer(t, map[string]string{
+ "client-output-buffer-limit": "normal 1m 0 0 slave 0 0 0 pubsub
0 0 0",
+ })
+ defer srv.Close()
+
+ ctx := context.Background()
+ rdb := srv.NewClient()
+ defer func() { require.NoError(t, rdb.Close()) }()
+
+ t.Run("normal client exceeding the hard limit is disconnected", func(t
*testing.T) {
+ require.NoError(t, rdb.Set(ctx, "big", strings.Repeat("x",
2*1024*1024), 0).Err())
+
+ c := srv.NewTCPClient()
+ defer func() { require.NoError(t, c.Close()) }()
+ require.NoError(t, c.WriteArgs("GET", "big"))
+
+ require.Eventually(t, func() bool {
+ return getClientOutputBufferLimitDisconnections(t, srv)
== 1
+ }, 5*time.Second, 100*time.Millisecond)
+ })
+}
