From fc6d9b99c36229d180a5ce0acfd53c4495a62a2a Mon Sep 17 00:00:00 2001 From: Winchell Date: Fri, 28 Aug 2026 11:29:47 +0800 Subject: [PATCH] Fix SelectiveChannel response copy for non-reflectable messages --- src/brpc/nonreflectable_message.h | 40 ++++++++++++++++++++++++++--- src/brpc/selective_channel.cpp | 11 ++++++-- test/brpc_channel_unittest.cpp | 42 +++++++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 5 deletions(-) diff --git a/src/brpc/nonreflectable_message.h b/src/brpc/nonreflectable_message.h index fab6301249..e7ded50e8f 100644 --- a/src/brpc/nonreflectable_message.h +++ b/src/brpc/nonreflectable_message.h @@ -25,6 +25,11 @@ namespace brpc { +class NonreflectableMessageBase : public ::google::protobuf::Message { +public: + virtual bool CopyFromSameType(const ::google::protobuf::Message& other) = 0; +}; + // // In bRPC, some non-Protobuf based protocol messages are also designed to implement // Protobuf Message interfaces, to provide a unified protocol message. @@ -38,7 +43,7 @@ namespace brpc { // and use only #if version_check #endif, to make maintenance easier. // template -class NonreflectableMessage : public ::google::protobuf::Message { +class NonreflectableMessage : public NonreflectableMessageBase { public: inline NonreflectableMessage() = default; inline NonreflectableMessage(const NonreflectableMessage&) : NonreflectableMessage() {} @@ -97,6 +102,18 @@ class NonreflectableMessage : public ::google::protobuf::Message { MergeFrom(other); } + bool CopyFromSameType(const ::google::protobuf::Message& other) override { + if (&other == this) { + return true; + } + if (other.GetDescriptor() != descriptor()) { + return false; + } + Clear(); + MergeFrom(static_cast(other)); + return true; + } + void MergeFrom(const ::google::protobuf::Message& other) PB_526_OVERRIDE { if (&other == this) { return; @@ -229,6 +246,17 @@ class NonreflectableMessage : public ::google::protobuf::Message { private: static T _instance; +#if GOOGLE_PROTOBUF_VERSION >= 5029000 + static void* PlacementNew_(const void*, void* mem, + ::google::protobuf::Arena* arena) { + T* message = ::new (mem) T(); + if (arena != nullptr) { + arena->OwnDestructor(message); + } + return message; + } +#endif + #if GOOGLE_PROTOBUF_VERSION >= 5027000 struct NonreflectableMessageClassData : ClassDataFull { constexpr NonreflectableMessageClassData() @@ -239,7 +267,10 @@ class NonreflectableMessage : public ::google::protobuf::Message { nullptr, // tc_table nullptr, // is_initialized nullptr, // merge_to_from - ::google::protobuf::internal::MessageCreator(), // message_creator + ::google::protobuf::internal::MessageCreator( + &NonreflectableMessage::PlacementNew_, + sizeof(T), + static_cast(alignof(T))), // message_creator 0, // cached_size_offset false, // is_lite }, @@ -253,7 +284,10 @@ class NonreflectableMessage : public ::google::protobuf::Message { nullptr, // on_demand_register_arena_dtor nullptr, // is_initialized nullptr, // merge_to_from - ::google::protobuf::internal::MessageCreator(), // message_creator + ::google::protobuf::internal::MessageCreator( + &NonreflectableMessage::PlacementNew_, + sizeof(T), + static_cast(alignof(T))), // message_creator 0, // cached_size_offset false, // is_lite }, diff --git a/src/brpc/selective_channel.cpp b/src/brpc/selective_channel.cpp index c53ac13b33..cedbc1d729 100644 --- a/src/brpc/selective_channel.cpp +++ b/src/brpc/selective_channel.cpp @@ -22,6 +22,7 @@ #include "brpc/socket.h" // SocketUser #include "brpc/load_balancer.h" // LoadBalancer #include "brpc/details/controller_private_accessor.h" // RPCSender +#include "brpc/nonreflectable_message.h" #include "brpc/selective_channel.h" #include "brpc/global.h" @@ -394,8 +395,14 @@ void SubDone::Run() { main_cntl->_error_code = _cntl._error_code; } else { if (_cntl._response != main_cntl->_response) { - main_cntl->_response->GetReflection()->Swap( - main_cntl->_response, _cntl._response); + NonreflectableMessageBase* nr_msg = + dynamic_cast(main_cntl->_response); + if (nr_msg != nullptr) { + CHECK(nr_msg->CopyFromSameType(*_cntl._response)); + } else { + main_cntl->_response->GetReflection()->Swap( + main_cntl->_response, _cntl._response); + } } } const Controller::CompletionInfo info = { _cid, true }; diff --git a/test/brpc_channel_unittest.cpp b/test/brpc_channel_unittest.cpp index b6125e26d8..27f711d483 100644 --- a/test/brpc_channel_unittest.cpp +++ b/test/brpc_channel_unittest.cpp @@ -37,6 +37,7 @@ #include "brpc/channel.h" #include "brpc/details/load_balancer_with_naming.h" #include "brpc/parallel_channel.h" +#include "brpc/redis.h" #include "brpc/selective_channel.h" #include "brpc/socket_map.h" #include "brpc/controller.h" @@ -232,6 +233,16 @@ class DelayedCloseEchoService : public ::test::EchoService { "Close connection after delay"); } }; +class PingCommandHandler : public brpc::RedisCommandHandler { +public: + brpc::RedisCommandHandlerResult Run( + const std::vector&, + brpc::RedisReply* output, + bool) override { + output->SetStatus("PONG"); + return brpc::REDIS_CMD_HANDLED; + } +}; pthread_once_t register_mock_protocol = PTHREAD_ONCE_INIT; @@ -2645,6 +2656,37 @@ TEST_F(ChannelTest, empty_selective_channel) { CallMethod(&channel, &cntl, &req, &res, false); EXPECT_EQ(ENODATA, cntl.ErrorCode()) << cntl.ErrorText(); } +TEST_F(ChannelTest, selective_channel_supports_nonreflectable_response) { + PingCommandHandler ping_handler; + std::unique_ptr redis_service(new brpc::RedisService); + ASSERT_TRUE(redis_service->AddCommandHandler("ping", &ping_handler)); + + brpc::Server server; + brpc::ServerOptions server_options; + server_options.redis_service = redis_service.release(); + ASSERT_EQ(0, server.Start("127.0.0.1:0", &server_options)); + + brpc::ChannelOptions channel_options; + channel_options.protocol = brpc::PROTOCOL_REDIS; + std::unique_ptr sub_channel(new brpc::Channel); + ASSERT_EQ(0, sub_channel->Init(server.listen_address(), &channel_options)); + + brpc::Controller cntl; + brpc::RedisRequest request; + brpc::RedisResponse response; + ASSERT_TRUE(request.AddCommand("ping")); + + brpc::SelectiveChannel channel; + ASSERT_EQ(0, channel.Init("rr", nullptr)); + ASSERT_EQ(0, channel.AddChannel(sub_channel.release(), nullptr)); + + cntl.set_timeout_ms(1000); + channel.CallMethod(nullptr, &cntl, &request, &response, nullptr); + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + ASSERT_EQ(1, response.reply_size()); + ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type()); + ASSERT_EQ("PONG", response.reply(0).data()); +} class BadCall : public brpc::CallMapper { brpc::SubCall Map(int,