From 9c5358cc6246c958a7ae78dfe25afac6d7edaa79 Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Fri, 28 Aug 2026 04:23:18 +0000 Subject: [PATCH 1/3] feat: limit Redis server connections --- docs/cn/server.md | 6 ++ docs/en/server.md | 6 ++ src/brpc/acceptor.cpp | 69 +++++++++++++++++-- src/brpc/acceptor.h | 20 ++++++ src/brpc/server.cpp | 38 ++++++++++- src/brpc/server.h | 9 +++ test/brpc_server_unittest.cpp | 120 ++++++++++++++++++++++++++++++++++ 7 files changed, 262 insertions(+), 6 deletions(-) diff --git a/docs/cn/server.md b/docs/cn/server.md index 7a9c47a4ef..5ddfbdd568 100644 --- a/docs/cn/server.md +++ b/docs/cn/server.md @@ -378,6 +378,12 @@ Server.set_version(...)可以为server设置一个名称+版本，可通过/vers | ------------------------- | ----- | ---------------------------------------- | ------------------- | | log_idle_connection_close | false | Print log when an idle connection is closed | src/brpc/socket.cpp | +## 限制Redis连接数 + +设置`ServerOptions.redis_max_connections`可以限制Redis专用公网监听端口上的并发连接数。默认值为0，表示不限制。非零值要求设置`redis_service`，将`enabled_protocols`严格设置为`"redis"`，关闭内置服务，并且不能在同一个Server上注册RPC或其他协议服务。 + +Acceptor会在创建brpc Socket前预留连接名额，因此空闲连接也计入上限，并发accept不会突破限制。超过限制的明文连接会收到`-ERR max number of clients reached`；启用SSL的监听端口会在TLS握手前直接关闭连接。内部监听端口和其他Server实例不受影响。`ServerStatistics.rejected_redis_connection_count`记录累计拒绝的连接数。 + ## pid_file 如果设置了此字段，Server启动时会创建一个同名文件，内容为进程号。默认为空。 diff --git a/docs/en/server.md b/docs/en/server.md index 80806db583..7af6109fe6 100644 --- a/docs/en/server.md +++ b/docs/en/server.md @@ -375,6 +375,12 @@ If [-log_idle_connection_close](http://brpc.baidu.com:8765/flags/log_idle_connec | ------------------------- | ----- | ---------------------------------------- | ------------------- | | log_idle_connection_close | false | Print log when an idle connection is closed | src/brpc/socket.cpp | +## Limit Redis connections + +Set `ServerOptions.redis_max_connections` to limit simultaneous connections on a Redis-only public listener. The default value is 0, which disables the limit. A non-zero value requires `redis_service` to be set, `enabled_protocols` to be exactly `"redis"`, builtin services to be disabled, and no RPC or other protocol services to share the Server. + +The acceptor reserves a slot before creating a brpc Socket, so idle connections count toward the limit and concurrent accepts cannot exceed it. An over-limit plaintext connection receives `-ERR max number of clients reached`; an SSL-enabled listener closes it before starting a TLS handshake. Internal listeners and other Server instances are unaffected. `ServerStatistics.rejected_redis_connection_count` reports the cumulative number of rejected connections. + ## pid_file If this field is non-empty, Server creates a file named so at start-up, with pid as the content. Empty by default. diff --git a/src/brpc/acceptor.cpp b/src/brpc/acceptor.cpp index 8333a17127..f5d23c802b 100644 --- a/src/brpc/acceptor.cpp +++ b/src/brpc/acceptor.cpp @@ -16,7 +16,9 @@ // under the License. +#include #include +#include #include #include "butil/fd_guard.h" // fd_guard #include "butil/fd_utility.h" // make_close_on_exec @@ -38,6 +40,9 @@ Acceptor::Acceptor(bthread_keytable_pool_t* pool) , _listened_fd(-1) , _acception_id(0) , _empty_cond(&_map_mutex) + , _connection_count(0) + , _rejected_redis_connection_count(0) + , _redis_max_connections(0) , _force_ssl(false) , _ssl_ctx(nullptr) , _socket_mode(SOCKET_MODE_TCP) @@ -52,6 +57,14 @@ Acceptor::~Acceptor() { int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, const std::shared_ptr& ssl_ctx, bool force_ssl) { + return StartAccept( + listened_fd, idle_timeout_sec, ssl_ctx, force_ssl, 0); +} + +int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, + const std::shared_ptr& ssl_ctx, + bool force_ssl, + size_t redis_max_connections) { if (listened_fd < 0) { LOG(FATAL) << "Invalid listened_fd=" << listened_fd; return -1; @@ -87,6 +100,7 @@ int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, _idle_timeout_sec = idle_timeout_sec; _force_ssl = force_ssl; _ssl_ctx = ssl_ctx; + _redis_max_connections = redis_max_connections; // Creation of _acception_id is inside lock so that OnNewConnections // (which may run immediately) should see sane fields set below. @@ -200,9 +214,45 @@ void Acceptor::Join() { } size_t Acceptor::ConnectionCount() const { - // Notice that _socket_map may be modified concurrently. This actually - // assumes that size() is safe to call concurrently. - return _socket_map.size(); + return _connection_count.load(butil::memory_order_relaxed); +} + +size_t Acceptor::RejectedRedisConnectionCount() const { + return _rejected_redis_connection_count.load(butil::memory_order_relaxed); +} + +bool Acceptor::TryAcquireRedisConnectionSlot() { + size_t count = _connection_count.load(butil::memory_order_relaxed); + do { + if (_redis_max_connections != 0 && + count >= _redis_max_connections) { + return false; + } + } while (!_connection_count.compare_exchange_weak( + count, count + 1, butil::memory_order_relaxed)); + return true; +} + +void Acceptor::RejectRedisConnection(int fd) { + _rejected_redis_connection_count.fetch_add( + 1, butil::memory_order_relaxed); + + // Reject SSL-capable listeners before doing any TLS work. Plaintext here + // would violate the TLS record protocol and could trigger an expensive + // handshake in a higher layer. + if (_ssl_ctx) { + return; + } + + static const char response[] = + "-ERR max number of clients reached\r\n"; + ssize_t nwritten; + do { + nwritten = send(fd, + response, + sizeof(response) - 1, + MSG_DONTWAIT | MSG_NOSIGNAL); + } while (nwritten < 0 && errno == EINTR); } void Acceptor::ListConnections(std::vector* conn_list, @@ -275,7 +325,12 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { acception->SetFailed(EINVAL, "Impossible! acception->user() MUST be Acceptor"); return; } - + + if (!am->TryAcquireRedisConnectionSlot()) { + am->RejectRedisConnection(in_fd); + continue; + } + SocketId socket_id; SocketOptions options; options.keytable_pool = am->_keytable_pool; @@ -288,6 +343,9 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { options.socket_mode = am->_socket_mode; options.bthread_tag = am->_bthread_tag; if (Socket::Create(options, &socket_id) != 0) { + const size_t previous = am->_connection_count.fetch_sub( + 1, butil::memory_order_relaxed); + CHECK_GT(previous, 0u); LOG(ERROR) << "Fail to create Socket"; continue; } @@ -349,6 +407,9 @@ void Acceptor::BeforeRecycle(Socket* sock) { // If a Socket could not be addressed shortly after its creation, it // was not added into `_socket_map'. _socket_map.erase(sock->id()); + const size_t previous = + _connection_count.fetch_sub(1, butil::memory_order_relaxed); + CHECK_GT(previous, 0u); if (_socket_map.empty()) { _empty_cond.Broadcast(); } diff --git a/src/brpc/acceptor.h b/src/brpc/acceptor.h index 2f138b30ea..e365e2a944 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -19,6 +19,7 @@ #define BRPC_ACCEPTOR_H #include "bthread/bthread.h" // bthread_t +#include "butil/atomicops.h" // butil::atomic #include "butil/synchronization/condition_variable.h" #include "butil/containers/flat_map.h" #include "brpc/input_messenger.h" @@ -58,6 +59,10 @@ friend class Server; int StartAccept(int listened_fd, int idle_timeout_sec, const std::shared_ptr& ssl_ctx, bool force_ssl); + int StartAccept(int listened_fd, int idle_timeout_sec, + const std::shared_ptr& ssl_ctx, + bool force_ssl, + size_t redis_max_connections); // [thread-safe] Stop accepting connections. // `closewait_ms' is not used anymore. @@ -72,6 +77,10 @@ friend class Server; // Get number of existing connections. size_t ConnectionCount() const; + // Get the cumulative number of connections rejected by the Redis-only + // listener's connection limit. + size_t RejectedRedisConnectionCount() const; + // Clear `conn_list' and append all connections into it. void ListConnections(std::vector* conn_list); @@ -93,6 +102,9 @@ friend class Server; // Remove the accepted socket `sock' from inside void BeforeRecycle(Socket* sock) override; + bool TryAcquireRedisConnectionSlot(); + void RejectRedisConnection(int fd); + bthread_keytable_pool_t* _keytable_pool; // owned by Server Status _status; int _idle_timeout_sec; @@ -108,6 +120,14 @@ friend class Server; // The map containing all the accepted sockets SocketMap _socket_map; + // A slot is reserved before Socket::Create(), closing the race where a + // socket starts processing before it is inserted into _socket_map. These + // atomics protect only numeric admission and publish no socket state, so + // relaxed memory ordering is sufficient. + butil::atomic _connection_count; + butil::atomic _rejected_redis_connection_count; + size_t _redis_max_connections; + bool _force_ssl; std::shared_ptr _ssl_ctx; diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp index 0852cfc105..31c8ea0ad8 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -137,6 +137,7 @@ ServerOptions::ServerOptions() , server_owns_interceptor(false) , num_threads(8) , max_concurrency(0) + , redis_max_connections(0) , session_local_data_factory(nullptr) , reserved_session_local_data(0) , thread_local_data_factory(nullptr) @@ -635,6 +636,13 @@ Acceptor* Server::BuildAcceptor() { // The protocol does not support server-side. continue; } + if (_options.redis_max_connections != 0 && + strcmp(protocols[i].name, "redis") != 0) { + // A pre-TLS connection limit cannot inspect the protocol. The + // validated Redis-only listener must therefore install no RPC or + // HTTP parser that could make this port a shared interface. + continue; + } if (has_whitelist && !is_http_protocol(protocols[i].name) && !is_rdma_handshake_protocol(protocols[i].name) && @@ -867,6 +875,27 @@ int Server::StartInternal(const butil::EndPoint& endpoint, const ServerOptions default_opt; const ServerOptions& real_opt = opt ? *opt : default_opt; + // Admission happens before protocol parsing (and, importantly, before a + // TLS handshake), so it is only safe on a listener dedicated to Redis. + // Reject ambiguous configurations instead of accidentally limiting RPCs + // sharing the public port. + if (real_opt.redis_max_connections != 0 && + (real_opt.redis_service == nullptr || + real_opt.enabled_protocols != "redis" || + real_opt.has_builtin_services || + service_count() != 0 || + real_opt.nshead_service != nullptr || + real_opt.thrift_service != nullptr || + real_opt.mongo_service_adaptor != nullptr || + real_opt.baidu_master_service != nullptr || + real_opt.http_master_service != nullptr || + real_opt.rtmp_service != nullptr)) { + LOG(ERROR) << "redis_max_connections requires a Redis-only public " + "listener (redis_service set, enabled_protocols=redis, " + "no RPC or builtin services)"; + return -1; + } + if (!real_opt.h2_settings.IsValid(true/*log_error*/)) { LOG(ERROR) << "Invalid h2_settings"; return -1; @@ -1164,7 +1193,8 @@ int Server::StartInternal(const butil::EndPoint& endpoint, // Pass ownership of `sockfd' to `_am' if (_am->StartAccept(sockfd, _options.idle_timeout_sec, _default_ssl_ctx, - _options.force_ssl) != 0) { + _options.force_ssl, + _options.redis_max_connections) != 0) { LOG(ERROR) << "Fail to start acceptor"; return -1; } @@ -1206,7 +1236,8 @@ int Server::StartInternal(const butil::EndPoint& endpoint, // Pass ownership of `sockfd' to `_internal_am' if (_internal_am->StartAccept(sockfd, _options.idle_timeout_sec, _default_ssl_ctx, - false) != 0) { + false, + 0) != 0) { LOG(ERROR) << "Fail to start internal_acceptor"; return -1; } @@ -1788,8 +1819,11 @@ google::protobuf::Service* Server::FindServiceByName( void Server::GetStat(ServerStatistics* stat) const { stat->connection_count = 0; + stat->rejected_redis_connection_count = 0; if (_am) { stat->connection_count += _am->ConnectionCount(); + stat->rejected_redis_connection_count += + _am->RejectedRedisConnectionCount(); } if (_internal_am) { stat->connection_count += _internal_am->ConnectionCount(); diff --git a/src/brpc/server.h b/src/brpc/server.h index 30846f750d..c530fec5e4 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -132,6 +132,14 @@ struct ServerOptions { // Default: 0 (unlimited) int max_concurrency; + // Maximum number of connections accepted by a Redis-only public + // listener. This option is rejected unless redis_service is configured, + // enabled_protocols is exactly "redis", no protobuf/RPC services are + // registered, and builtin services are disabled. The internal listener + // and other Server instances are never subject to this limit. + // Default: 0 (unlimited) + size_t redis_max_connections; + // Default value of method-level max concurrencies, // Overridable by Server.MaxConcurrencyOf(). AdaptiveMaxConcurrency method_max_concurrency; @@ -303,6 +311,7 @@ struct ServerOptions { // server. But bvar contains more stats and is more convenient. struct ServerStatistics { size_t connection_count; + size_t rejected_redis_connection_count; int user_service_count; int builtin_service_count; }; diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index ed0268e853..64990a51a6 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -50,6 +50,7 @@ #include "brpc/server.h" #include "brpc/restful.h" #include "brpc/channel.h" +#include "brpc/redis.h" #include "brpc/socket_map.h" #include "brpc/controller.h" #include "brpc/compress.h" @@ -1290,6 +1291,125 @@ TEST_F(ServerTest, close_idle_connections) { ASSERT_EQ(0ul, stat.connection_count); } +TEST_F(ServerTest, redis_connection_limit_requires_dedicated_listener) { + brpc::Server server; + EchoServiceImpl echo_service; + ASSERT_EQ(0, server.AddService( + &echo_service, brpc::SERVER_DOESNT_OWN_SERVICE)); + + brpc::ServerOptions opt; + opt.redis_service = new brpc::RedisService; + opt.redis_max_connections = 1; + opt.enabled_protocols = "redis"; + opt.has_builtin_services = false; + const int rc = server.Start("127.0.0.1:0", &opt); + if (rc != 0) { + delete opt.redis_service; + opt.redis_service = nullptr; + } + EXPECT_EQ(-1, rc); +} + +TEST_F(ServerTest, reject_redis_connections_over_limit) { + brpc::Server server; + brpc::ServerOptions opt; + opt.redis_service = new brpc::RedisService; + opt.redis_max_connections = 1; + opt.enabled_protocols = "redis"; + opt.has_builtin_services = false; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + + const butil::EndPoint ep = server.listen_address(); + butil::fd_guard first_client(tcp_connect(ep, nullptr)); + ASSERT_GT(first_client, 0); + + brpc::ServerStatistics stat; + for (int retry = 0; retry < 100; ++retry) { + server.GetStat(&stat); + if (stat.connection_count == 1) { + break; + } + usleep(1000); + } + ASSERT_EQ(1ul, stat.connection_count); + + // The Redis limit belongs to this acceptor, not to the process. A separate + // RPC Server must remain reachable while the Redis listener is full. + EchoServiceImpl rpc_service; + brpc::Server rpc_server; + ASSERT_EQ(0, rpc_server.AddService( + &rpc_service, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, rpc_server.Start("127.0.0.1:0", nullptr)); + SendSleepRPC(rpc_server.listen_address(), 0, true); + + butil::fd_guard rejected_client(tcp_connect(ep, nullptr)); + ASSERT_GT(rejected_client, 0); + struct timeval timeout = {1, 0}; + ASSERT_EQ(0, setsockopt(rejected_client, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + char response[64]; + const ssize_t nr = recv(rejected_client, response, sizeof(response), 0); + const std::string expected = "-ERR max number of clients reached\r\n"; + ASSERT_EQ(expected.size(), (size_t)nr); + EXPECT_EQ(expected, std::string(response, (size_t)nr)); + + server.GetStat(&stat); + EXPECT_EQ(1ul, stat.connection_count); + EXPECT_EQ(1ul, stat.rejected_redis_connection_count); + + first_client.reset(-1); + for (int retry = 0; retry < 100; ++retry) { + server.GetStat(&stat); + if (stat.connection_count == 0) { + break; + } + usleep(1000); + } + EXPECT_EQ(0ul, stat.connection_count); +} + +TEST_F(ServerTest, reject_redis_connection_before_tls_handshake) { + brpc::Server server; + brpc::ServerOptions opt; + opt.redis_service = new brpc::RedisService; + opt.redis_max_connections = 1; + opt.enabled_protocols = "redis"; + opt.has_builtin_services = false; + opt.force_ssl = true; + brpc::CertInfo& cert = opt.mutable_ssl_options()->default_cert; + cert.certificate = "cert1.crt"; + cert.private_key = "cert1.key"; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + + const butil::EndPoint ep = server.listen_address(); + // Holding an idle TCP socket consumes the only slot without initiating a + // TLS handshake. + butil::fd_guard first_client(tcp_connect(ep, nullptr)); + ASSERT_GT(first_client, 0); + brpc::ServerStatistics stat; + for (int retry = 0; retry < 100; ++retry) { + server.GetStat(&stat); + if (stat.connection_count == 1) { + break; + } + usleep(1000); + } + ASSERT_EQ(1ul, stat.connection_count); + + butil::fd_guard rejected_client(tcp_connect(ep, nullptr)); + ASSERT_GT(rejected_client, 0); + struct timeval timeout = {1, 0}; + ASSERT_EQ(0, setsockopt(rejected_client, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + char response; + // EOF without sending a ClientHello proves that rejection happened in + // accept, before brpc created a Socket or entered TLS authentication. + EXPECT_EQ(0, recv(rejected_client, &response, sizeof(response), 0)); + + server.GetStat(&stat); + EXPECT_EQ(1ul, stat.rejected_redis_connection_count); +} + TEST_F(ServerTest, logoff_and_multiple_start) { butil::Timer timer; butil::EndPoint ep; From bca102309f7e26421c1f9dcf77f42412f120ecda Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Fri, 28 Aug 2026 04:51:40 +0000 Subject: [PATCH 2/3] feat: update Redis connection limit at runtime --- docs/cn/server.md | 2 ++ docs/en/server.md | 2 ++ src/brpc/acceptor.cpp | 14 +++++++--- src/brpc/acceptor.h | 3 ++- src/brpc/server.cpp | 50 ++++++++++++++++++++++++----------- src/brpc/server.h | 7 +++++ test/brpc_server_unittest.cpp | 36 ++++++++++++++++++++++++- 7 files changed, 94 insertions(+), 20 deletions(-) diff --git a/docs/cn/server.md b/docs/cn/server.md index 5ddfbdd568..7d51875212 100644 --- a/docs/cn/server.md +++ b/docs/cn/server.md @@ -384,6 +384,8 @@ Server.set_version(...)可以为server设置一个名称+版本，可通过/vers Acceptor会在创建brpc Socket前预留连接名额，因此空闲连接也计入上限，并发accept不会突破限制。超过限制的明文连接会收到`-ERR max number of clients reached`；启用SSL的监听端口会在TLS握手前直接关闭连接。内部监听端口和其他Server实例不受影响。`ServerStatistics.rejected_redis_connection_count`记录累计拒绝的连接数。 +运行中的Redis专用Server可以调用`Server::SetRedisMaxConnections()`原子更新上限。调高上限会影响后续连接准入；调低上限不会断开已有连接，活跃连接数降到新上限以下后才会重新接受新连接。设置为0会关闭限制。以无限制值启动的Redis专用Server也可以稍后动态开启限制。 + ## pid_file 如果设置了此字段，Server启动时会创建一个同名文件，内容为进程号。默认为空。 diff --git a/docs/en/server.md b/docs/en/server.md index 7af6109fe6..1b665d1b93 100644 --- a/docs/en/server.md +++ b/docs/en/server.md @@ -381,6 +381,8 @@ Set `ServerOptions.redis_max_connections` to limit simultaneous connections on a The acceptor reserves a slot before creating a brpc Socket, so idle connections count toward the limit and concurrent accepts cannot exceed it. An over-limit plaintext connection receives `-ERR max number of clients reached`; an SSL-enabled listener closes it before starting a TLS handshake. Internal listeners and other Server instances are unaffected. `ServerStatistics.rejected_redis_connection_count` reports the cumulative number of rejected connections. +Call `Server::SetRedisMaxConnections()` to atomically update the limit on a running Redis-only Server. Raising the limit affects subsequent admission checks. Lowering it does not close existing connections; new connections are accepted again after the active count falls below the limit. Set the limit to 0 to disable it. A Redis-only Server started with an unlimited value can enable the limit later. + ## pid_file If this field is non-empty, Server creates a file named so at start-up, with pid as the content. Empty by default. diff --git a/src/brpc/acceptor.cpp b/src/brpc/acceptor.cpp index f5d23c802b..ed41aa2c4d 100644 --- a/src/brpc/acceptor.cpp +++ b/src/brpc/acceptor.cpp @@ -100,7 +100,7 @@ int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, _idle_timeout_sec = idle_timeout_sec; _force_ssl = force_ssl; _ssl_ctx = ssl_ctx; - _redis_max_connections = redis_max_connections; + SetRedisMaxConnections(redis_max_connections); // Creation of _acception_id is inside lock so that OnNewConnections // (which may run immediately) should see sane fields set below. @@ -224,8 +224,9 @@ size_t Acceptor::RejectedRedisConnectionCount() const { bool Acceptor::TryAcquireRedisConnectionSlot() { size_t count = _connection_count.load(butil::memory_order_relaxed); do { - if (_redis_max_connections != 0 && - count >= _redis_max_connections) { + const size_t max_connections = + _redis_max_connections.load(butil::memory_order_relaxed); + if (max_connections != 0 && count >= max_connections) { return false; } } while (!_connection_count.compare_exchange_weak( @@ -233,6 +234,13 @@ bool Acceptor::TryAcquireRedisConnectionSlot() { return true; } +void Acceptor::SetRedisMaxConnections(size_t max_connections) { + // The limit controls only future numeric admission decisions and does not + // publish socket state, so a relaxed store is sufficient. + _redis_max_connections.store( + max_connections, butil::memory_order_relaxed); +} + void Acceptor::RejectRedisConnection(int fd) { _rejected_redis_connection_count.fetch_add( 1, butil::memory_order_relaxed); diff --git a/src/brpc/acceptor.h b/src/brpc/acceptor.h index e365e2a944..d1e15a96c2 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -104,6 +104,7 @@ friend class Server; bool TryAcquireRedisConnectionSlot(); void RejectRedisConnection(int fd); + void SetRedisMaxConnections(size_t max_connections); bthread_keytable_pool_t* _keytable_pool; // owned by Server Status _status; @@ -126,7 +127,7 @@ friend class Server; // relaxed memory ordering is sufficient. butil::atomic _connection_count; butil::atomic _rejected_redis_connection_count; - size_t _redis_max_connections; + butil::atomic _redis_max_connections; bool _force_ssl; std::shared_ptr _ssl_ctx; diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp index 31c8ea0ad8..d69b0cb78b 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -619,6 +619,20 @@ BUTIL_FORCE_INLINE bool is_rdma_handshake_protocol(const char* name) { return strcmp(name, "rdma_handshake") == 0; } +bool is_redis_only_public_listener(const ServerOptions& opt, + size_t user_service_count) { + return opt.redis_service != nullptr && + opt.enabled_protocols == "redis" && + !opt.has_builtin_services && + user_service_count == 0 && + opt.nshead_service == nullptr && + opt.thrift_service == nullptr && + opt.mongo_service_adaptor == nullptr && + opt.baidu_master_service == nullptr && + opt.http_master_service == nullptr && + opt.rtmp_service == nullptr; +} + Acceptor* Server::BuildAcceptor() { std::set whitelist; for (butil::StringSplitter sp(_options.enabled_protocols.c_str(), ' '); @@ -631,16 +645,17 @@ Acceptor* Server::BuildAcceptor() { InputMessageHandler handler; std::vector protocols; ListProtocols(&protocols); + const bool redis_only = + is_redis_only_public_listener(_options, service_count()); for (size_t i = 0; i < protocols.size(); ++i) { if (protocols[i].process_request == nullptr) { // The protocol does not support server-side. continue; } - if (_options.redis_max_connections != 0 && - strcmp(protocols[i].name, "redis") != 0) { - // A pre-TLS connection limit cannot inspect the protocol. The - // validated Redis-only listener must therefore install no RPC or - // HTTP parser that could make this port a shared interface. + if (redis_only && strcmp(protocols[i].name, "redis") != 0) { + // This dedicated listener may enable its connection limit at + // runtime. Install no RPC or HTTP parser that could make pre-TLS + // admission affect a shared interface. continue; } if (has_whitelist && @@ -880,16 +895,7 @@ int Server::StartInternal(const butil::EndPoint& endpoint, // Reject ambiguous configurations instead of accidentally limiting RPCs // sharing the public port. if (real_opt.redis_max_connections != 0 && - (real_opt.redis_service == nullptr || - real_opt.enabled_protocols != "redis" || - real_opt.has_builtin_services || - service_count() != 0 || - real_opt.nshead_service != nullptr || - real_opt.thrift_service != nullptr || - real_opt.mongo_service_adaptor != nullptr || - real_opt.baidu_master_service != nullptr || - real_opt.http_master_service != nullptr || - real_opt.rtmp_service != nullptr)) { + !is_redis_only_public_listener(real_opt, service_count())) { LOG(ERROR) << "redis_max_connections requires a Redis-only public " "listener (redis_service set, enabled_protocols=redis, " "no RPC or builtin services)"; @@ -1832,6 +1838,20 @@ void Server::GetStat(ServerStatistics* stat) const { stat->builtin_service_count = builtin_service_count(); } +int Server::SetRedisMaxConnections(size_t max_connections) { + if (!IsRunning() || _am == nullptr) { + LOG(WARNING) << "SetRedisMaxConnections requires a running Server"; + return -1; + } + if (!is_redis_only_public_listener(_options, service_count())) { + LOG(WARNING) << "SetRedisMaxConnections requires a Redis-only public " + "listener"; + return -1; + } + _am->SetRedisMaxConnections(max_connections); + return 0; +} + void Server::ListServices(std::vector *services) { if (!services) { return; diff --git a/src/brpc/server.h b/src/brpc/server.h index c530fec5e4..6357aabac8 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -137,6 +137,7 @@ struct ServerOptions { // enabled_protocols is exactly "redis", no protobuf/RPC services are // registered, and builtin services are disabled. The internal listener // and other Server instances are never subject to this limit. + // Use Server::SetRedisMaxConnections() to update the limit at runtime. // Default: 0 (unlimited) size_t redis_max_connections; @@ -548,6 +549,12 @@ class Server { // Get statistics of this server void GetStat(ServerStatistics* stat) const; + // Atomically update the connection limit of a running Redis-only public + // listener. Existing connections are not closed when the limit is lowered. + // Set to 0 to disable the limit. Returns 0 on success, -1 if this Server is + // not running or its public listener is not dedicated to Redis. + int SetRedisMaxConnections(size_t max_connections); + // Get the options passed to Start(). const ServerOptions& options() const { return _options; } diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index 64990a51a6..846662569e 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -1314,10 +1314,11 @@ TEST_F(ServerTest, reject_redis_connections_over_limit) { brpc::Server server; brpc::ServerOptions opt; opt.redis_service = new brpc::RedisService; - opt.redis_max_connections = 1; + opt.redis_max_connections = 0; opt.enabled_protocols = "redis"; opt.has_builtin_services = false; ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + ASSERT_EQ(0, server.SetRedisMaxConnections(1)); const butil::EndPoint ep = server.listen_address(); butil::fd_guard first_client(tcp_connect(ep, nullptr)); @@ -1340,6 +1341,7 @@ TEST_F(ServerTest, reject_redis_connections_over_limit) { ASSERT_EQ(0, rpc_server.AddService( &rpc_service, brpc::SERVER_DOESNT_OWN_SERVICE)); ASSERT_EQ(0, rpc_server.Start("127.0.0.1:0", nullptr)); + EXPECT_EQ(-1, rpc_server.SetRedisMaxConnections(1)); SendSleepRPC(rpc_server.listen_address(), 0, true); butil::fd_guard rejected_client(tcp_connect(ep, nullptr)); @@ -1357,7 +1359,39 @@ TEST_F(ServerTest, reject_redis_connections_over_limit) { EXPECT_EQ(1ul, stat.connection_count); EXPECT_EQ(1ul, stat.rejected_redis_connection_count); + ASSERT_EQ(0, server.SetRedisMaxConnections(2)); + butil::fd_guard second_client(tcp_connect(ep, nullptr)); + ASSERT_GT(second_client, 0); + for (int retry = 0; retry < 100; ++retry) { + server.GetStat(&stat); + if (stat.connection_count == 2) { + break; + } + usleep(1000); + } + ASSERT_EQ(2ul, stat.connection_count); + + // Lowering the limit only gates future accepts; it does not disconnect + // the two clients that are already established. + ASSERT_EQ(0, server.SetRedisMaxConnections(1)); + server.GetStat(&stat); + EXPECT_EQ(2ul, stat.connection_count); + + butil::fd_guard lowered_limit_client(tcp_connect(ep, nullptr)); + ASSERT_GT(lowered_limit_client, 0); + ASSERT_EQ(0, setsockopt(lowered_limit_client, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + const ssize_t lowered_nr = + recv(lowered_limit_client, response, sizeof(response), 0); + ASSERT_EQ(expected.size(), (size_t)lowered_nr); + EXPECT_EQ(expected, std::string(response, (size_t)lowered_nr)); + + server.GetStat(&stat); + EXPECT_EQ(2ul, stat.connection_count); + EXPECT_EQ(2ul, stat.rejected_redis_connection_count); + first_client.reset(-1); + second_client.reset(-1); for (int retry = 0; retry < 100; ++retry) { server.GetStat(&stat); if (stat.connection_count == 0) { From 6a6f39859d002e48917d5a81251e96e92bc7309b Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Sat, 29 Aug 2026 07:22:45 +0000 Subject: [PATCH 3/3] fix: handle partial Redis rejection writes --- src/brpc/acceptor.cpp | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/brpc/acceptor.cpp b/src/brpc/acceptor.cpp index ed41aa2c4d..982723f40f 100644 --- a/src/brpc/acceptor.cpp +++ b/src/brpc/acceptor.cpp @@ -254,13 +254,21 @@ void Acceptor::RejectRedisConnection(int fd) { static const char response[] = "-ERR max number of clients reached\r\n"; - ssize_t nwritten; - do { - nwritten = send(fd, - response, - sizeof(response) - 1, - MSG_DONTWAIT | MSG_NOSIGNAL); - } while (nwritten < 0 && errno == EINTR); + const size_t response_size = sizeof(response) - 1; + size_t offset = 0; + while (offset < response_size) { + const ssize_t nwritten = send(fd, + response + offset, + response_size - offset, + MSG_DONTWAIT | MSG_NOSIGNAL); + if (nwritten > 0) { + offset += nwritten; + } else if (nwritten < 0 && errno == EINTR) { + continue; + } else { + break; + } + } } void Acceptor::ListConnections(std::vector* conn_list,