This is an automated email from the ASF dual-hosted git repository.
wwbmmm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/brpc.git
The following commit(s) were added to refs/heads/master by this push:
new e4b2287d Limit streams accepted per request (#3493)
e4b2287d is described below
commit e4b2287db066e6eb0b80224d4eef53d90e9a93c3
Author: Xiaofeng Wang <[email protected]>
AuthorDate: Thu Aug 27 16:18:41 2026 +0800
Limit streams accepted per request (#3493)
* Limit streams accepted per request
Reject a streaming RPC request when its declared stream count exceeds a
configurable limit. Check the count before creating any response streams.
Cover the limit boundary, rejection, and a subsequent valid request through
one client/server connection.
* Propagate stream limit errors
Return the configured limit to the RPC caller.
Synchronize test state across server and client threads.
* Match stream limit format types
Cast the configured limit to the type required by the format string.
---
src/brpc/stream.cpp | 20 ++++++++++
test/brpc_streaming_rpc_unittest.cpp | 77 ++++++++++++++++++++++++++++++++++++
2 files changed, 97 insertions(+)
diff --git a/src/brpc/stream.cpp b/src/brpc/stream.cpp
index 7e032b69..42cb0eab 100644
--- a/src/brpc/stream.cpp
+++ b/src/brpc/stream.cpp
@@ -40,6 +40,11 @@ DECLARE_int64(socket_max_streams_unconsumed_bytes);
DEFINE_uint64(stream_write_max_segment_size, 512 * 1024 * 1024,
"Stream message exceeding this size will be automatically split
into smaller segments");
BRPC_VALIDATE_GFLAG(stream_write_max_segment_size, PositiveInteger);
+DEFINE_int32(stream_max_streams_per_request, 64,
+ "Maximum number of streams that StreamAccept creates for one "
+ "request. Raise this flag if an application needs to accept "
+ "more streams per request");
+BRPC_VALIDATE_GFLAG(stream_max_streams_per_request, PositiveInteger);
const static butil::IOBuf *TIMEOUT_TASK = (butil::IOBuf*)-1L;
@@ -1034,6 +1039,21 @@ int StreamAccept(StreamIds& response_streams,
Controller& cntl,
LOG(ERROR) << "No stream along with this request";
return -1;
}
+ const int64_t stream_count = static_cast<int64_t>(
+ cntl._remote_stream_settings->extra_stream_ids_size()) + 1;
+ if (stream_count > FLAGS_stream_max_streams_per_request) {
+ cntl.SetFailed(EREQUEST,
+ "Reject %" PRId64 " streams in one request, exceeding "
+ "-stream_max_streams_per_request=%" PRId64,
+ stream_count,
+ static_cast<int64_t>(
+ FLAGS_stream_max_streams_per_request));
+ LOG(ERROR) << "Reject " << stream_count
+ << " streams in one request, exceeding "
+ "-stream_max_streams_per_request="
+ << FLAGS_stream_max_streams_per_request;
+ return -1;
+ }
StreamOptions opt;
if (options != nullptr) {
opt = *options;
diff --git a/test/brpc_streaming_rpc_unittest.cpp
b/test/brpc_streaming_rpc_unittest.cpp
index 12c2ff3c..7d1dc366 100644
--- a/test/brpc_streaming_rpc_unittest.cpp
+++ b/test/brpc_streaming_rpc_unittest.cpp
@@ -1249,6 +1249,83 @@ private:
int _adjustment;
};
+class MyServiceWithStreamCountLimit : public test::EchoService {
+public:
+ void Echo(::google::protobuf::RpcController* controller,
+ const ::test::EchoRequest* request,
+ ::test::EchoResponse* response,
+ ::google::protobuf::Closure* done) override {
+ brpc::ClosureGuard done_guard(done);
+ brpc::Controller* cntl = static_cast<brpc::Controller*>(controller);
+ response->set_message(request->message());
+
+ brpc::StreamIds response_streams;
+ accept_result.store(
+ brpc::StreamAccept(response_streams, *cntl, nullptr),
+ std::memory_order_release);
+ accepted_streams.store(
+ response_streams.size(), std::memory_order_release);
+ }
+
+ std::atomic<int> accept_result{0};
+ std::atomic<size_t> accepted_streams{0};
+};
+
+TEST_F(StreamingRpcTest, limit_streams_accepted_per_request) {
+ std::string old_stream_limit;
+ ASSERT_TRUE(GFLAGS_NAMESPACE::GetCommandLineOption(
+ "stream_max_streams_per_request", &old_stream_limit));
+ BRPC_SCOPE_EXIT {
+ GFLAGS_NAMESPACE::SetCommandLineOption(
+ "stream_max_streams_per_request", old_stream_limit.c_str());
+ };
+ ASSERT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption(
+ "stream_max_streams_per_request", "2").empty());
+
+ brpc::Server server;
+ MyServiceWithStreamCountLimit service;
+ ASSERT_EQ(0, server.AddService(
+ &service, brpc::SERVER_DOESNT_OWN_SERVICE));
+ ASSERT_EQ(0, server.Start(0, nullptr));
+
+ brpc::Channel channel;
+ ASSERT_EQ(0, channel.Init(server.listen_address(), nullptr));
+
+ for (size_t stream_count : {2u, 3u, 2u}) {
+ brpc::Controller cntl;
+ brpc::StreamIds request_streams;
+ ASSERT_EQ(0, brpc::StreamCreate(
+ request_streams, stream_count, cntl, nullptr));
+ ASSERT_EQ(stream_count, request_streams.size());
+
+ test::EchoService_Stub stub(&channel);
+ stub.Echo(&cntl, &request, &response, nullptr);
+ if (stream_count == 2) {
+ ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
+ ASSERT_EQ(0, service.accept_result.load(
+ std::memory_order_acquire));
+ ASSERT_EQ(stream_count, service.accepted_streams.load(
+ std::memory_order_acquire));
+ } else {
+ ASSERT_TRUE(cntl.Failed());
+ ASSERT_EQ(brpc::EREQUEST, cntl.ErrorCode());
+ ASSERT_NE(std::string::npos, cntl.ErrorText().find(
+
"stream_max_streams_per_request"));
+ ASSERT_EQ(-1, service.accept_result.load(
+ std::memory_order_acquire));
+ ASSERT_EQ(0u, service.accepted_streams.load(
+ std::memory_order_acquire));
+ }
+
+ for (brpc::StreamId stream_id : request_streams) {
+ brpc::StreamClose(stream_id);
+ }
+ }
+
+ server.Stop(0);
+ server.Join();
+}
+
TEST_F(StreamingRpcTest, reject_mismatched_returned_stream_identifiers) {
const size_t STREAM_COUNT = 3;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
