diff --git a/docs/cn/server.md b/docs/cn/server.md index 7a9c47a4ef..7d51875212 100644 --- a/docs/cn/server.md +++ b/docs/cn/server.md @@ -378,6 +378,14 @@ 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`记录累计拒绝的连接数。 + +运行中的Redis专用Server可以调用`Server::SetRedisMaxConnections()`原子更新上限。调高上限会影响后续连接准入；调低上限不会断开已有连接，活跃连接数降到新上限以下后才会重新接受新连接。设置为0会关闭限制。以无限制值启动的Redis专用Server也可以稍后动态开启限制。 + ## pid_file 如果设置了此字段，Server启动时会创建一个同名文件，内容为进程号。默认为空。 diff --git a/docs/en/server.md b/docs/en/server.md index 80806db583..1b665d1b93 100644 --- a/docs/en/server.md +++ b/docs/en/server.md @@ -375,6 +375,14 @@ 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. + +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 8333a17127..982723f40f 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; + SetRedisMaxConnections(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,61 @@ 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 { + 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( + count, count + 1, butil::memory_order_relaxed)); + 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); + + // 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"; + 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, @@ -275,7 +341,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 +359,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 +423,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..d1e15a96c2 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,10 @@ friend class Server; // Remove the accepted socket `sock' from inside void BeforeRecycle(Socket* sock) override; + bool TryAcquireRedisConnectionSlot(); + void RejectRedisConnection(int fd); + void SetRedisMaxConnections(size_t max_connections); + bthread_keytable_pool_t* _keytable_pool; // owned by Server Status _status; int _idle_timeout_sec; @@ -108,6 +121,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; + 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 0852cfc105..d69b0cb78b 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) @@ -618,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(), ' '); @@ -630,11 +645,19 @@ 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 (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 && !is_http_protocol(protocols[i].name) && !is_rdma_handshake_protocol(protocols[i].name) && @@ -867,6 +890,18 @@ 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 && + !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)"; + return -1; + } + if (!real_opt.h2_settings.IsValid(true/*log_error*/)) { LOG(ERROR) << "Invalid h2_settings"; return -1; @@ -1164,7 +1199,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 +1242,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 +1825,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(); @@ -1798,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 30846f750d..6357aabac8 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -132,6 +132,15 @@ 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. + // Use Server::SetRedisMaxConnections() to update the limit at runtime. + // 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 +312,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; }; @@ -539,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 ed0268e853..846662569e 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,159 @@ 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 = 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)); + 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)); + EXPECT_EQ(-1, rpc_server.SetRedisMaxConnections(1)); + 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); + + 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) { + 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;