feat: limit Redis server connections

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`，随后关闭连接。背压或套接字错误可能导致响应部分或全部未送达；accept循环不会等待客户端连接变为可写。启用SSL的监听端口会在TLS握手前直接关闭连接。内部监听端口和其他Server实例不受影响。`ServerStatistics.rejected_redis_connection_count`记录累计拒绝的连接数。
+
+运行中的Redis专用Server可以调用`Server::SetRedisMaxConnections()`原子更新上限。调高上限会影响后续连接准入；调低上限不会断开已有连接，活跃连接数降到新上限以下后才会重新接受新连接。设置为0会关闭限制。以无限制值启动的Redis专用Server也可以稍后动态开启限制。
+
 ## pid_file
 
 如果设置了此字段，Server启动时会创建一个同名文件，内容为进程号。默认为空。

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. For an over-limit plaintext connection, the server makes a nonblocking, best-effort attempt to send `-ERR max number of clients reached` and then closes the connection. Backpressure or a socket error may prevent delivery of some or all of the response; the accept loop does not wait for the client to become writable. An SSL-enabled listener closes the connection 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.

src/brpc/acceptor.cpp

@@ -16,7 +16,9 @@
 // under the License.
 
 
+#include <errno.h>
 #include <inttypes.h>
+#include <sys/socket.h>
 #include <gflags/gflags.h>
 #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<SocketSSLContext>& 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<SocketSSLContext>& 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,79 @@ 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) {
+    // Borrowed fd: OnNewConnectionsUntilEAGAIN() retains ownership in its
+    // fd_guard. After this returns, the caller's continue destroys the guard
+    // and closes fd on both the SSL and plaintext paths. Do not close it here.
+    _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";
+    // Delivery is best-effort: handle short writes, but never wait for a slow
+    // peer to become writable and stall admission for other connections.
+    int send_flags = MSG_DONTWAIT;
+#if defined(MSG_NOSIGNAL)
+    send_flags |= MSG_NOSIGNAL;
+#elif defined(SO_NOSIGPIPE)
+    const int enabled = 1;
+    if (setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE,
+                   &enabled, sizeof(enabled)) != 0) {
+        return;
+    }
+#else
+    // Omit the optional error if SIGPIPE cannot be suppressed for this send.
+    return;
+#endif
+    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,
+                                      send_flags);
+        if (nwritten > 0) {
+            offset += nwritten;
+        } else if (nwritten < 0 && errno == EINTR) {
+            continue;
+        } else {
+            break;
+        }
+    }
 }
 
 void Acceptor::ListConnections(std::vector<SocketId>* conn_list,
@@ -275,7 +359,13 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) {
             acception->SetFailed(EINVAL, "Impossible! acception->user() MUST be Acceptor");
             return;
         }
-
+
+        if (!am->TryAcquireRedisConnectionSlot()) {
+            am->RejectRedisConnection(in_fd);
+            // in_fd still owns the fd; leaving this iteration closes it.
+            continue;
+        }
+
         SocketId socket_id;
         SocketOptions options;
         options.keytable_pool = am->_keytable_pool;
@@ -288,6 +378,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 +442,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();
     }

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<SocketSSLContext>& ssl_ctx,
                     bool force_ssl);
+    int StartAccept(int listened_fd, int idle_timeout_sec,
+                    const std::shared_ptr<SocketSSLContext>& 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<SocketId>* 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<size_t> _connection_count;
+    butil::atomic<size_t> _rejected_redis_connection_count;
+    butil::atomic<size_t> _redis_max_connections;
+
     bool _force_ssl;
     std::shared_ptr<SocketSSLContext> _ssl_ctx;
 

src/brpc/server.cpp

@@ -138,6 +138,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)
@@ -619,6 +620,25 @@ BUTIL_FORCE_INLINE bool is_rdma_handshake_protocol(const char* name) {
     return strcmp(name, "rdma_handshake") == 0;
 }
 
+static const char kRedisOnlyPublicListenerRequirements[] =
+    "a Redis-only public listener (redis_service set, enabled_protocols exactly "
+    "\"redis\", has_builtin_services=false, no registered RPC services, "
+    "and all other protocol service pointers null)";
+
+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::unordered_set<std::string> whitelist;
     for (butil::StringSplitter sp(_options.enabled_protocols.c_str(), ' ');
@@ -631,6 +651,8 @@ Acceptor* Server::BuildAcceptor() {
     InputMessageHandler handler;
     std::vector<Protocol> 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.
@@ -640,6 +662,12 @@ Acceptor* Server::BuildAcceptor() {
         // rdma_handshake are always served, but they are still
         // valid names for the whitelist.
         bool in_whitelist = (whitelist.erase(protocols[i].name) != 0);
+        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 && !in_whitelist &&
             !is_http_protocol(protocols[i].name) &&
             !is_rdma_handshake_protocol(protocols[i].name)) {
@@ -870,6 +898,17 @@ 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 "
+                   << kRedisOnlyPublicListenerRequirements;
+        return -1;
+    }
+
     if (!real_opt.h2_settings.IsValid(true/*log_error*/)) {
         LOG(ERROR) << "Invalid h2_settings";
         return -1;
@@ -1167,7 +1206,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;
         }
@@ -1209,7 +1249,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;
         }
@@ -1791,8 +1832,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();
@@ -1801,6 +1845,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 "
+                     << kRedisOnlyPublicListenerRequirements;
+        return -1;
+    }
+    _am->SetRedisMaxConnections(max_connections);
+    return 0;
+}
+
 void Server::ListServices(std::vector<google::protobuf::Service*> *services) {
     if (!services) {
         return;