diff --git a/src/brpc/channel.cpp b/src/brpc/channel.cpp index 83fc37b077..3c07036b31 100644 --- a/src/brpc/channel.cpp +++ b/src/brpc/channel.cpp @@ -43,6 +43,20 @@ namespace brpc { +struct Channel::SingleServerState { + explicit SingleServerState(const SocketMapKey& key_in) : key(key_in) {} + + SocketMapKey key; +}; + +struct ChannelInitOptions { + ChannelOptions options; + Protocol::SerializeRequest serialize_request; + Protocol::PackRequest pack_request; + Protocol::GetMethodName get_method_name; + int preferred_index; +}; + DECLARE_bool(enable_rpcz); DECLARE_bool(usercode_in_pthread); DEFINE_string(health_check_path, "", "Http path of health check call." @@ -174,77 +188,82 @@ Channel::Channel(ProfilerLinker) } Channel::~Channel() { - if (_server_id != INVALID_SOCKET_ID) { - const ChannelSignature sig = ComputeChannelSignature(_options); - SocketMapRemove(SocketMapKey(_server_address, sig)); - } + ResetSingleServer(); } - -int Channel::InitChannelOptions(const ChannelOptions* options) { - if (options) { // Override default options if user provided one. - _options = *options; +void Channel::ResetSingleServer() { + if (_single_server_state) { + SocketMapRemove(_single_server_state->key); + _single_server_state.reset(); } - const Protocol* protocol = FindProtocol(_options.protocol); + _server_id = INVALID_SOCKET_ID; +} + +static int InitChannelOptions(const ChannelOptions& source, + ChannelInitOptions* output) { + output->options = source; + const Protocol* protocol = FindProtocol(output->options.protocol); if (NULL == protocol || !protocol->support_client()) { LOG(ERROR) << "Channel does not support the protocol"; return -1; } - if (_options.hc_option.health_check_path.empty()) { - _options.hc_option.health_check_path = FLAGS_health_check_path; - _options.hc_option.health_check_timeout_ms = FLAGS_health_check_timeout_ms; + if (output->options.hc_option.health_check_path.empty()) { + output->options.hc_option.health_check_path = FLAGS_health_check_path; + output->options.hc_option.health_check_timeout_ms = FLAGS_health_check_timeout_ms; } - auto ret = TransportFactory::ContextInitOrDie(_options.socket_mode, false, &_options); + auto ret = TransportFactory::ContextInitOrDie( + output->options.socket_mode, false, &output->options); if (ret != 0) { LOG(ERROR) << "Fail to initialize transport context for channel, ret=" << ret; return -1; } - _serialize_request = protocol->serialize_request; - _pack_request = protocol->pack_request; - _get_method_name = protocol->get_method_name; + output->serialize_request = protocol->serialize_request; + output->pack_request = protocol->pack_request; + output->get_method_name = protocol->get_method_name; // Check connection_type - if (_options.connection_type == CONNECTION_TYPE_UNKNOWN) { + if (output->options.connection_type == CONNECTION_TYPE_UNKNOWN) { // Save has_error which will be overriden in later assignments to // connection_type. - const bool has_error = _options.connection_type.has_error(); + const bool has_error = output->options.connection_type.has_error(); if (protocol->supported_connection_type & CONNECTION_TYPE_SINGLE) { - _options.connection_type = CONNECTION_TYPE_SINGLE; + output->options.connection_type = CONNECTION_TYPE_SINGLE; } else if (protocol->supported_connection_type & CONNECTION_TYPE_POOLED) { - _options.connection_type = CONNECTION_TYPE_POOLED; + output->options.connection_type = CONNECTION_TYPE_POOLED; } else { - _options.connection_type = CONNECTION_TYPE_SHORT; + output->options.connection_type = CONNECTION_TYPE_SHORT; } if (has_error) { - LOG(ERROR) << "Channel=" << this << " chose connection_type=" - << _options.connection_type.name() << " for protocol=" - << _options.protocol.name(); + LOG(ERROR) << "Channel chose connection_type=" + << output->options.connection_type.name() << " for protocol=" + << output->options.protocol.name(); } } else { - if (!(_options.connection_type & protocol->supported_connection_type)) { + if (!(output->options.connection_type & protocol->supported_connection_type)) { LOG(ERROR) << protocol->name << " does not support connection_type=" - << ConnectionTypeToString(_options.connection_type); + << ConnectionTypeToString(output->options.connection_type); return -1; } } - _preferred_index = get_client_side_messenger()->FindProtocolIndex(_options.protocol); - if (_preferred_index < 0) { + output->preferred_index = + get_client_side_messenger()->FindProtocolIndex(output->options.protocol); + if (output->preferred_index < 0) { LOG(ERROR) << "Fail to get index for protocol=" - << _options.protocol.name(); + << output->options.protocol.name(); return -1; } - if (_options.protocol == PROTOCOL_ESP) { - if (_options.auth == NULL) { - _options.auth = policy::global_esp_authenticator(); + if (output->options.protocol == PROTOCOL_ESP) { + if (output->options.auth == NULL) { + output->options.auth = policy::global_esp_authenticator(); } } // Normalize connection_group - std::string& cg = _options.connection_group; + std::string& cg = output->options.connection_group; if (!cg.empty() && (::isspace(cg.front()) || ::isspace(cg.back()))) { butil::TrimWhitespace(cg, butil::TRIM_ALL, &cg); } @@ -252,6 +271,20 @@ int Channel::InitChannelOptions(const ChannelOptions* options) { return 0; } +int Channel::InitChannelOptions(const ChannelOptions* options) { + ChannelInitOptions initialized_options; + if (::brpc::InitChannelOptions(options ? *options : _options, + &initialized_options) != 0) { + return -1; + } + _options = initialized_options.options; + _serialize_request = initialized_options.serialize_request; + _pack_request = initialized_options.pack_request; + _get_method_name = initialized_options.get_method_name; + _preferred_index = initialized_options.preferred_index; + return 0; +} + int Channel::Init(const char* server_addr_and_port, const ChannelOptions* options) { GlobalInitializeOrDie(); @@ -359,47 +392,76 @@ int Channel::InitSingle(const butil::EndPoint& server_addr_and_port, const ChannelOptions* options, int raw_port) { GlobalInitializeOrDie(); - if (InitChannelOptions(options) != 0) { + ChannelInitOptions initialized_options; + if (::brpc::InitChannelOptions(options ? *options : _options, + &initialized_options) != 0) { return -1; } int* port_out = raw_port == -1 ? &raw_port: NULL; - ParseURL(raw_server_address, &_scheme, &_service_name, port_out); - const std::string host = _service_name; + std::string scheme; + std::string service_name; + ParseURL(raw_server_address, &scheme, &service_name, port_out); + const std::string host = service_name; if (raw_port != -1) { - _service_name.append(":").append(std::to_string(raw_port)); + service_name.append(":").append(std::to_string(raw_port)); } - if (_options.protocol == brpc::PROTOCOL_HTTP && _scheme == "https") { - SetHttpsPeerName(host, &_options); + if (initialized_options.options.protocol == brpc::PROTOCOL_HTTP && + scheme == "https") { + SetHttpsPeerName(host, &initialized_options.options); } + _options = initialized_options.options; + _serialize_request = initialized_options.serialize_request; + _pack_request = initialized_options.pack_request; + _get_method_name = initialized_options.get_method_name; + _preferred_index = initialized_options.preferred_index; const int port = server_addr_and_port.port; if (port < 0) { LOG(ERROR) << "Invalid port=" << port; return -1; } butil::EndPoint client_endpoint; - if (!_options.client_host.empty() && - butil::str2ip(_options.client_host.c_str(), &client_endpoint.ip) != 0 && - butil::hostname2ip(_options.client_host.c_str(), &client_endpoint.ip) != 0) { - LOG(ERROR) << "Invalid client host=`" << _options.client_host << '\''; + if (!initialized_options.options.client_host.empty() && + butil::str2ip(initialized_options.options.client_host.c_str(), + &client_endpoint.ip) != 0 && + butil::hostname2ip(initialized_options.options.client_host.c_str(), + &client_endpoint.ip) != 0) { + LOG(ERROR) << "Invalid client host=`" + << initialized_options.options.client_host << '\''; return -1; } - _server_address = server_addr_and_port; - const ChannelSignature sig = ComputeChannelSignature(_options); + const ChannelSignature sig = + ComputeChannelSignature(initialized_options.options); std::shared_ptr ssl_ctx; - if (CreateSocketSSLContext(_options, &ssl_ctx) != 0) { + if (CreateSocketSSLContext(initialized_options.options, &ssl_ctx) != 0) { return -1; } SocketOptions opt; opt.local_side = client_endpoint; opt.initial_ssl_ctx = ssl_ctx; - opt.socket_mode = _options.socket_mode; - opt.hc_option = _options.hc_option; - opt.device_name = _options.device_name; - if (SocketMapInsert(SocketMapKey(server_addr_and_port, sig), - &_server_id, opt) != 0) { + opt.socket_mode = initialized_options.options.socket_mode; + opt.hc_option = initialized_options.options.hc_option; + opt.device_name = initialized_options.options.device_name; + const SocketMapKey socket_map_key(server_addr_and_port, sig); + SocketId server_id; + if (SocketMapInsert(socket_map_key, &server_id, opt) != 0) { LOG(ERROR) << "Fail to insert into SocketMap"; return -1; } + std::unique_ptr single_server_state( + new (std::nothrow) SingleServerState(socket_map_key)); + if (!single_server_state) { + SocketMapRemove(socket_map_key); + LOG(ERROR) << "Fail to allocate single-server state"; + return -1; + } + + ResetSingleServer(); + _scheme.swap(scheme); + _service_name.swap(service_name); + _server_address = server_addr_and_port; + _server_id = server_id; + _single_server_state.swap(single_server_state); + _lb.reset(); return 0; } @@ -411,23 +473,36 @@ int Channel::Init(const char* ns_url, return Init(ns_url, options); } GlobalInitializeOrDie(); - if (InitChannelOptions(options) != 0) { + ChannelInitOptions initialized_options; + if (::brpc::InitChannelOptions(options ? *options : _options, + &initialized_options) != 0) { return -1; } int raw_port = -1; - ParseURL(ns_url, &_scheme, &_service_name, &raw_port); - const std::string host = _service_name; + std::string scheme; + std::string service_name; + ParseURL(ns_url, &scheme, &service_name, &raw_port); + const std::string host = service_name; if (raw_port != -1) { - _service_name.append(":").append(std::to_string(raw_port)); + service_name.append(":").append(std::to_string(raw_port)); } - if (_options.protocol == brpc::PROTOCOL_HTTP && _scheme == "https") { - SetHttpsPeerName(host, &_options); + if (initialized_options.options.protocol == brpc::PROTOCOL_HTTP && + scheme == "https") { + SetHttpsPeerName(host, &initialized_options.options); } + _options = initialized_options.options; + _serialize_request = initialized_options.serialize_request; + _pack_request = initialized_options.pack_request; + _get_method_name = initialized_options.get_method_name; + _preferred_index = initialized_options.preferred_index; butil::EndPoint client_endpoint; - if (!_options.client_host.empty() && - butil::str2ip(_options.client_host.c_str(), &client_endpoint.ip) != 0 && - butil::hostname2ip(_options.client_host.c_str(), &client_endpoint.ip) != 0) { - LOG(ERROR) << "Invalid client host=`" << _options.client_host << '\''; + if (!initialized_options.options.client_host.empty() && + butil::str2ip(initialized_options.options.client_host.c_str(), + &client_endpoint.ip) != 0 && + butil::hostname2ip(initialized_options.options.client_host.c_str(), + &client_endpoint.ip) != 0) { + LOG(ERROR) << "Invalid client host=`" + << initialized_options.options.client_host << '\''; return -1; } std::unique_ptr lb(new (std::nothrow) @@ -437,21 +512,26 @@ int Channel::Init(const char* ns_url, return -1; } GetNamingServiceThreadOptions ns_opt; - ns_opt.succeed_without_server = _options.succeed_without_server; - ns_opt.log_succeed_without_server = _options.log_succeed_without_server; - ns_opt.socket_option.socket_mode = _options.socket_mode; - ns_opt.channel_signature = ComputeChannelSignature(_options); - ns_opt.socket_option.hc_option = _options.hc_option; + ns_opt.succeed_without_server = initialized_options.options.succeed_without_server; + ns_opt.log_succeed_without_server = + initialized_options.options.log_succeed_without_server; + ns_opt.socket_option.socket_mode = initialized_options.options.socket_mode; + ns_opt.channel_signature = ComputeChannelSignature(initialized_options.options); + ns_opt.socket_option.hc_option = initialized_options.options.hc_option; ns_opt.socket_option.local_side = client_endpoint; - ns_opt.socket_option.device_name = _options.device_name; - if (CreateSocketSSLContext(_options, + ns_opt.socket_option.device_name = initialized_options.options.device_name; + if (CreateSocketSSLContext(initialized_options.options, &ns_opt.socket_option.initial_ssl_ctx) != 0) { return -1; } - if (lb->Init(ns_url, lb_name, _options.ns_filter, &ns_opt) != 0) { + if (lb->Init(ns_url, lb_name, initialized_options.options.ns_filter, &ns_opt) != 0) { LOG(ERROR) << "Fail to initialize LoadBalancerWithNaming"; return -1; } + + ResetSingleServer(); + _scheme.swap(scheme); + _service_name.swap(service_name); _lb.reset(lb.release()); return 0; } diff --git a/src/brpc/channel.h b/src/brpc/channel.h index 28a17ac8ea..994a4cb3c0 100644 --- a/src/brpc/channel.h +++ b/src/brpc/channel.h @@ -22,6 +22,7 @@ // To brpc developers: This is a header included by user, don't depend // on internal structures, use opaque pointers instead. +#include // std::unique_ptr #include // std::ostream #include "bthread/errno.h" // Redefine errno #include "butil/intrusive_ptr.hpp" // butil::intrusive_ptr @@ -245,7 +246,10 @@ friend class SelectiveChannel; // therefore destroy the `controller' inside `done' static void CallMethodImpl(Controller* controller, SharedLoadBalancer* lb); + struct SingleServerState; + int InitChannelOptions(const ChannelOptions* options); + void ResetSingleServer(); int InitSingle(const butil::EndPoint& server_addr_and_port, const char* raw_server_address, const ChannelOptions* options, @@ -265,6 +269,7 @@ friend class SelectiveChannel; butil::intrusive_ptr _lb; ChannelOptions _options; int _preferred_index; + std::unique_ptr _single_server_state; }; enum ChannelOwnership { diff --git a/test/brpc_channel_unittest.cpp b/test/brpc_channel_unittest.cpp index 6f4540d6a2..9edc8ef601 100644 --- a/test/brpc_channel_unittest.cpp +++ b/test/brpc_channel_unittest.cpp @@ -2316,6 +2316,47 @@ TEST_F(ChannelTest, init_as_single_server) { } } +TEST_F(ChannelTest, reinit_single_server_releases_socket_map_references) { + butil::EndPoint first_endpoint; + butil::EndPoint second_endpoint; + ASSERT_EQ(0, str2endpoint("127.0.0.1:59347", &first_endpoint)); + ASSERT_EQ(0, str2endpoint("127.0.0.1:59348", &second_endpoint)); + + { + brpc::Channel channel; + ASSERT_EQ(0, channel.Init(first_endpoint, NULL)); + ASSERT_EQ(0, channel.Init(first_endpoint, NULL)); + ASSERT_EQ(0, channel.Init(second_endpoint, NULL)); + } + + brpc::SocketId id; + EXPECT_NE(0, brpc::SocketMapFind(brpc::SocketMapKey(first_endpoint), &id)); + EXPECT_NE(0, brpc::SocketMapFind(brpc::SocketMapKey(second_endpoint), &id)); +} + +TEST_F(ChannelTest, failed_reinit_releases_existing_socket_map_reference) { + butil::EndPoint endpoint; + ASSERT_EQ(0, str2endpoint("127.0.0.1:59349", &endpoint)); + + { + brpc::Channel channel; + ASSERT_EQ(0, channel.Init(endpoint, NULL)); + const brpc::SocketId original_id = channel._server_id; + + brpc::ChannelOptions invalid_options; + invalid_options.client_host = "not a valid client host"; + ASSERT_EQ(-1, channel.Init(endpoint, &invalid_options)); + + EXPECT_EQ(endpoint, channel._server_address); + EXPECT_EQ(original_id, channel._server_id); + EXPECT_EQ(invalid_options.client_host, channel.options().client_host); + EXPECT_TRUE(channel.SingleServer()); + } + + brpc::SocketId id; + EXPECT_NE(0, brpc::SocketMapFind(brpc::SocketMapKey(endpoint), &id)); +} + TEST_F(ChannelTest, init_using_unknown_naming_service) { brpc::Channel channel; ASSERT_EQ(-1, channel.Init("unknown://unknown", "unknown", NULL));