Copilot commented on code in PR #3498:
URL: https://github.com/apache/brpc/pull/3498#discussion_r3975231898
##########
src/brpc/acceptor.cpp:
##########
@@ -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;
+ }
Review Comment:
RejectRedisConnection() does not close/shutdown `fd` (and in the SSL case
returns immediately). If `fd` lifetime is intentionally managed by an outer
`fd_guard` in the accept loop, consider documenting that expectation here
explicitly; otherwise, strongly consider closing the socket inside
RejectRedisConnection() to make the pre-TLS rejection invariant self-contained
and harder to regress.
##########
src/brpc/acceptor.cpp:
##########
@@ -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;
+ }
+ }
Review Comment:
The send() loop breaks on EAGAIN/EWOULDBLOCK (and any non-EINTR error),
which can result in clients not receiving the documented plaintext error
response under load/backpressure. If best-effort delivery is acceptable, the
docs/tests should tolerate partial/no response; otherwise, consider a bounded
retry strategy (e.g., poll/select with a short deadline) or using an existing
utility that handles partial writes robustly without blocking indefinitely.
##########
test/brpc_server_unittest.cpp:
##########
@@ -1424,6 +1425,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_GE(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_GE(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));
Review Comment:
This test assumes a single recv() returns the full error line, but TCP does
not preserve message boundaries; recv() may return a partial payload even when
the server writes it in one call. To avoid flaky failures, read in a loop until
either the expected bytes are collected, EOF is hit, or the timeout expires
(and then compare the accumulated buffer to `expected`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
