thweetkomputer commented on code in PR #3498:
URL: https://github.com/apache/brpc/pull/3498#discussion_r3975316236
##########
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:
Agreed that the ownership contract should be explicit here. Added comments
in f05f8f61 at both `RejectRedisConnection()` and its call site: the function
borrows the fd; the caller retains it in `butil::fd_guard`, and `continue`
destroys the guard and closes the connection on both SSL and plaintext paths.
No additional close is needed inside the helper. The plaintext test now also
checks EOF, complementing the existing pre-TLS EOF test. Rebuilt
`brpc_server_unittest`; all four focused tests passed.
##########
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:
Best-effort delivery is intentional: waiting for writability in the accept
loop would stall admission for other clients. In f05f8f61, both server guides
now explicitly describe the nonblocking attempt, subsequent close, and the
possibility of a partial or missing response under backpressure or socket
errors. Added the same rationale beside the send loop. The loopback test still
requires the complete error because these are fresh connections with no induced
send backpressure; weakening that assertion to accept an empty response would
miss a broken normal send path. This test assumption is now documented. No
blocking retry was introduced. All four focused tests passed.
##########
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:
Fixed in f05f8f61: both plaintext rejection checks now accumulate reads
until the expected response has been collected, EOF occurs, or a read
fails/times out. EINTR is retried and the existing receive timeout remains in
effect. The test deliberately uses a seven-byte buffer so the full response
must span multiple recv calls, then verifies EOF after the error. Rebuilt
`brpc_server_unittest`; all four focused tests passed.
--
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]
