Skip to content

Fix SelectiveChannel response copy for non-reflectable messages #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cw20050111-prog wants to merge 1 commit into
base: SelectiveChannel
Choose a base branch
from
+88 −5
Open

Fix SelectiveChannel response copy for non-reflectable messages #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 37 additions & 3 deletions src/brpc/nonreflectable_message.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -38,7 +43,7 @@ namespace brpc {
// and use only #if version_check #endif, to make maintenance easier.
//
template <typename T>
class NonreflectableMessage : public ::google::protobuf::Message {
class NonreflectableMessage : public NonreflectableMessageBase {
public:
inline NonreflectableMessage() = default;
inline NonreflectableMessage(const NonreflectableMessage&) : NonreflectableMessage() {}
Expand Down Expand Up @@ -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<const T&>(other));
return true;
}

void MergeFrom(const ::google::protobuf::Message& other) PB_526_OVERRIDE {
if (&other == this) {
return;
Expand Down Expand Up @@ -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()
Expand All @@ -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<uint8_t>(alignof(T))), // message_creator
0, // cached_size_offset
false, // is_lite
},
Expand All @@ -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<uint8_t>(alignof(T))), // message_creator
0, // cached_size_offset
false, // is_lite
},
Expand Down
11 changes: 9 additions & 2 deletions src/brpc/selective_channel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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<NonreflectableMessageBase*>(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 };
Expand Down
42 changes: 42 additions & 0 deletions test/brpc_channel_unittest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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<butil::StringPiece>&,
brpc::RedisReply* output,
bool) override {
output->SetStatus("PONG");
return brpc::REDIS_CMD_HANDLED;
}
};

pthread_once_t register_mock_protocol = PTHREAD_ONCE_INIT;

Expand Down Expand Up @@ -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<brpc::RedisService> 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<brpc::Channel> 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,
Expand Down