From d5001f898ce3f73c1c46412f4d78c84655939c19 Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Fri, 28 Aug 2026 04:11:33 +0000 Subject: [PATCH 1/3] fix: include cstdint in thread-local test --- test/baidu_thread_local_unittest.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/test/baidu_thread_local_unittest.cpp b/test/baidu_thread_local_unittest.cpp index 01209b6b..4cc629f3 100644 --- a/test/baidu_thread_local_unittest.cpp +++ b/test/baidu_thread_local_unittest.cpp @@ -17,6 +17,7 @@ #include #include +#include #include "butil/thread_local.h" namespace { From c8de3fbca0b89951367acaa22a919ad86c0f7d28 Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Fri, 28 Aug 2026 04:35:03 +0000 Subject: [PATCH 2/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 | 39 +++++++++- src/brpc/server.h | 9 +++ test/brpc_server_unittest.cpp | 137 ++++++++++++++++++++++++++++++++++ 7 files changed, 280 insertions(+), 6 deletions(-) diff --git a/docs/cn/server.md b/docs/cn/server.md index 81469b9f..a46f4405 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 52ed0d88..39260008 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 38f8b6dc..1b007dec 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 @@ -39,6 +41,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(NULL) , _use_rdma(false) { @@ -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; @@ -85,6 +98,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. @@ -198,9 +212,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, @@ -277,7 +327,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; @@ -309,6 +364,9 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { } #endif 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; } @@ -370,6 +428,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 c82cdcc1..51430c2a 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" @@ -57,6 +58,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. @@ -71,6 +76,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); @@ -92,6 +101,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; @@ -107,6 +119,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 07bb0b27..40ddc223 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -132,6 +132,7 @@ ServerOptions::ServerOptions() , server_owns_interceptor(false) , num_threads(8) , max_concurrency(0) + , redis_max_connections(0) , session_local_data_factory(NULL) , reserved_session_local_data(0) , thread_local_data_factory(NULL) @@ -594,6 +595,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) && !whitelist.erase(protocols[i].name)) { @@ -772,6 +780,28 @@ int Server::StartInternal(const butil::EndPoint& endpoint, } return -1; } + + 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 == NULL || + real_opt.enabled_protocols != "redis" || + real_opt.has_builtin_services || + service_count() != 0 || + real_opt.nshead_service != NULL || + real_opt.thrift_service != NULL || + real_opt.mongo_service_adaptor != NULL || + real_opt.http_master_service != NULL || + real_opt.rtmp_service != NULL)) { + 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 (opt) { _options = *opt; } else { @@ -1053,7 +1083,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; } @@ -1094,7 +1125,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; } @@ -1665,8 +1697,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 844d554e..6c41811e 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -128,6 +128,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; @@ -272,6 +280,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 c22b6b53..14d9c7d4 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -49,6 +49,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 "echo.pb.h" @@ -151,6 +152,23 @@ class EchoServiceImpl : public test::EchoService { butil::atomic count; }; +class ConnectionLimitRedisService : public brpc::RedisService { +public: + std::unique_ptr NewConnectionContext( + brpc::Socket*) const override { + return std::unique_ptr( + new brpc::ConnectionContext); + } + + brpc::RedisCommandHandlerResult DispatchCommand( + brpc::ConnectionContext*, + const std::vector&, + brpc::RedisReply*, + bool) const override { + return brpc::REDIS_CMD_HANDLED; + } +}; + // An evil service that fakes its `ServiceDescriptor' class EvilService : public test::EchoService { public: @@ -1247,6 +1265,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 ConnectionLimitRedisService; + 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 = NULL; + } + EXPECT_EQ(-1, rc); +} + +TEST_F(ServerTest, reject_redis_connections_over_limit) { + brpc::Server server; + brpc::ServerOptions opt; + opt.redis_service = new ConnectionLimitRedisService; + 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, NULL)); + 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", NULL)); + SendSleepRPC(rpc_server.listen_address(), 0, true); + + butil::fd_guard rejected_client(tcp_connect(ep, NULL)); + 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 ConnectionLimitRedisService; + 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, NULL)); + 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, NULL)); + 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 6d56abecbb9ce6c6b862a1bb6389f5bd920b34aa Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Fri, 28 Aug 2026 04:47:32 +0000 Subject: [PATCH 3/3] feat: update Redis connection limit at runtime --- docs/cn/server.md | 2 ++ docs/en/server.md | 2 ++ src/brpc/acceptor.cpp | 12 ++++++--- src/brpc/acceptor.h | 3 ++- src/brpc/server.cpp | 49 +++++++++++++++++++++++++---------- src/brpc/server.h | 8 ++++++ test/brpc_server_unittest.cpp | 37 +++++++++++++++++++++++++- 7 files changed, 94 insertions(+), 19 deletions(-) diff --git a/docs/cn/server.md b/docs/cn/server.md index a46f4405..f197ef6f 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 39260008..2b4be43f 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. Lowering the limit does not close existing connections; new connections are rejected until the count falls below the new value. Raising it takes effect on subsequent admission checks, and setting it to 0 disables the limit. A Redis-only Server started with an unlimited value can enable the limit later through this method. + ## 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 1b007dec..a049b9bb 100644 --- a/src/brpc/acceptor.cpp +++ b/src/brpc/acceptor.cpp @@ -98,7 +98,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. @@ -222,8 +222,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( @@ -231,6 +232,11 @@ bool Acceptor::TryAcquireRedisConnectionSlot() { return true; } +void Acceptor::SetRedisMaxConnections(size_t max_connections) { + _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 51430c2a..5be62b68 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -103,6 +103,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; @@ -125,7 +126,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 40ddc223..c8899aa0 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -574,6 +574,19 @@ bool is_http_protocol(const char* name) { return strcmp(name, "http") == 0 || strcmp(name, "h2") == 0; } +bool is_redis_only_public_listener(const ServerOptions& opt, + size_t user_service_count) { + return opt.redis_service != NULL && + opt.enabled_protocols == "redis" && + !opt.has_builtin_services && + user_service_count == 0 && + opt.nshead_service == NULL && + opt.thrift_service == NULL && + opt.mongo_service_adaptor == NULL && + opt.http_master_service == NULL && + opt.rtmp_service == NULL; +} + Acceptor* Server::BuildAcceptor() { std::set whitelist; for (butil::StringSplitter sp(_options.enabled_protocols.c_str(), ' '); @@ -590,16 +603,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 == NULL) { // 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 && @@ -788,15 +802,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 == NULL || - real_opt.enabled_protocols != "redis" || - real_opt.has_builtin_services || - service_count() != 0 || - real_opt.nshead_service != NULL || - real_opt.thrift_service != NULL || - real_opt.mongo_service_adaptor != NULL || - real_opt.http_master_service != NULL || - real_opt.rtmp_service != NULL)) { + !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)"; @@ -1710,6 +1716,21 @@ void Server::GetStat(ServerStatistics* stat) const { stat->builtin_service_count = builtin_service_count(); } +int Server::SetRedisMaxConnections(size_t max_connections) { + if (!IsRunning() || _am == NULL) { + LOG(WARNING) << "SetRedisMaxConnections is only allowed for 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 6c41811e..1f023b7e 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -133,6 +133,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 it while running. // Default: 0 (unlimited) size_t redis_max_connections; @@ -511,6 +512,13 @@ 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 lowering the limit; + // new connections are rejected until the count falls below it. A value of + // 0 disables the limit. Returns 0 on success, -1 if the 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 14d9c7d4..cc98ab8d 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -1288,10 +1288,11 @@ TEST_F(ServerTest, reject_redis_connections_over_limit) { brpc::Server server; brpc::ServerOptions opt; opt.redis_service = new ConnectionLimitRedisService; - 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, NULL)); @@ -1314,6 +1315,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", NULL)); + EXPECT_EQ(-1, rpc_server.SetRedisMaxConnections(1)); SendSleepRPC(rpc_server.listen_address(), 0, true); butil::fd_guard rejected_client(tcp_connect(ep, NULL)); @@ -1331,7 +1333,40 @@ 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, NULL)); + 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 does not close either existing connection, but it + // immediately prevents another connection from being admitted. + ASSERT_EQ(0, server.SetRedisMaxConnections(1)); + server.GetStat(&stat); + EXPECT_EQ(2ul, stat.connection_count); + butil::fd_guard lowered_rejected_client(tcp_connect(ep, NULL)); + ASSERT_GT(lowered_rejected_client, 0); + ASSERT_EQ(0, setsockopt(lowered_rejected_client, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + char lowered_response[64]; + const ssize_t lowered_nr = recv(lowered_rejected_client, + lowered_response, + sizeof(lowered_response), + 0); + ASSERT_EQ(expected.size(), (size_t)lowered_nr); + EXPECT_EQ(expected, std::string(lowered_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) {