diff --git a/src/brpc/selective_channel.cpp b/src/brpc/selective_channel.cpp index 5d2fbe898e..c6ab236f70 100644 --- a/src/brpc/selective_channel.cpp +++ b/src/brpc/selective_channel.cpp @@ -17,6 +17,7 @@ #include +#include #include #include "bthread/bthread.h" // bthread_id_xx #include "brpc/socket.h" // SocketUser @@ -286,11 +287,30 @@ inline int ChannelBalancer::SelectChannel(const LoadBalancer::SelectIn& in, } int ChannelBalancer::CheckHealth() { - BAIDU_SCOPED_LOCK(_mutex); - for (ChannelToIdMap::const_iterator it = _chan_map.begin(); - it != _chan_map.end(); ++it) { - if (!it->second->Failed() && - it->first->CheckHealth() == 0) { + std::vector ids; + { + BAIDU_SCOPED_LOCK(_mutex); + ids.reserve(_chan_map.size()); + for (ChannelToIdMap::const_iterator it = _chan_map.begin(); + it != _chan_map.end(); ++it) { + SubChannel* sub = static_cast(it->second->user()); + if (sub->ownership == DOESNT_OWN_CHANNEL) { + // The caller may delete a non-owning channel as soon as + // RemoveAndDestroyChannel() returns. Keep the mutex while + // using it, so removal cannot return during this call. + if (!it->second->Failed() && sub->chan->CheckHealth() == 0) { + return 0; + } + } else { + ids.push_back(it->second->id()); + } + } + } + for (size_t i = 0; i < ids.size(); ++i) { + SocketUniquePtr ptr; + if (Socket::AddressFailedAsWell(ids[i], &ptr) >= 0 && + !ptr->Failed() && + static_cast(ptr->user())->chan->CheckHealth() == 0) { return 0; } } @@ -299,17 +319,41 @@ int ChannelBalancer::CheckHealth() { void ChannelBalancer::Describe(std::ostream& os, const DescribeOptions& options) { - BAIDU_SCOPED_LOCK(_mutex); if (!options.verbose) { + BAIDU_SCOPED_LOCK(_mutex); os << _chan_map.size(); return; } - for (ChannelToIdMap::const_iterator it = _chan_map.begin(); - it != _chan_map.end(); ++it) { - if (it != _chan_map.begin()) { - os << ' '; + std::vector ids; + bool first = true; + { + BAIDU_SCOPED_LOCK(_mutex); + ids.reserve(_chan_map.size()); + for (ChannelToIdMap::const_iterator it = _chan_map.begin(); + it != _chan_map.end(); ++it) { + SubChannel* sub = static_cast(it->second->user()); + if (sub->ownership == DOESNT_OWN_CHANNEL) { + // Synchronization with removal protects the externally owned + // channel while Describe() is using it. + if (!first) { + os << ' '; + } + sub->chan->Describe(os, options); + first = false; + } else { + ids.push_back(it->second->id()); + } + } + } + for (size_t i = 0; i < ids.size(); ++i) { + SocketUniquePtr ptr; + if (Socket::AddressFailedAsWell(ids[i], &ptr) >= 0) { + if (!first) { + os << ' '; + } + static_cast(ptr->user())->chan->Describe(os, options); + first = false; } - it->first->Describe(os, options); } } diff --git a/test/brpc_channel_unittest.cpp b/test/brpc_channel_unittest.cpp index f6155cbdb5..22a637c374 100644 --- a/test/brpc_channel_unittest.cpp +++ b/test/brpc_channel_unittest.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -122,6 +123,230 @@ class DeferredSelectiveChannel : public brpc::ChannelBase { std::vector* _calls; }; +class BlockingDescribeChannel : public brpc::ChannelBase { +public: + BlockingDescribeChannel() : _started(false), _release(false) {} + + void CallMethod(const google::protobuf::MethodDescriptor*, + google::protobuf::RpcController*, + const google::protobuf::Message*, + google::protobuf::Message*, + google::protobuf::Closure*) override {} + int CheckHealth() override { return 0; } + int Weight() override { return 1; } + void Describe(std::ostream&, const brpc::DescribeOptions&) const override { + _started.store(true, butil::memory_order_release); + while (!_release.load(butil::memory_order_acquire)) { + bthread_usleep(1000); + } + } + bool started() const { + return _started.load(butil::memory_order_acquire); + } + void Release() { _release.store(true, butil::memory_order_release); } + +private: + mutable butil::atomic _started; + mutable butil::atomic _release; +}; + +class BlockingHealthChannel : public brpc::ChannelBase { +public: + explicit BlockingHealthChannel(butil::atomic* destroyed) + : _destroyed(destroyed), _started(false), _release(false) {} + + ~BlockingHealthChannel() override { + _destroyed->fetch_add(1, butil::memory_order_relaxed); + } + + void CallMethod(const google::protobuf::MethodDescriptor*, + google::protobuf::RpcController*, + const google::protobuf::Message*, + google::protobuf::Message*, + google::protobuf::Closure*) override {} + int CheckHealth() override { + _started.store(true, butil::memory_order_release); + while (!_release.load(butil::memory_order_acquire)) { + bthread_usleep(1000); + } + return 0; + } + int Weight() override { return 1; } + bool started() const { + return _started.load(butil::memory_order_acquire); + } + void Release() { _release.store(true, butil::memory_order_release); } + +private: + butil::atomic* _destroyed; + butil::atomic _started; + butil::atomic _release; +}; + +struct SelectiveDescribeArgs { + brpc::SelectiveChannel* channel; + butil::atomic done; + SelectiveDescribeArgs() : channel(nullptr), done(false) {} +}; + +void* RunSelectiveDescribe(void* arg) { + SelectiveDescribeArgs* args = static_cast(arg); + std::ostringstream os; + brpc::DescribeOptions options; + options.verbose = true; + args->channel->Describe(os, options); + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +struct SelectiveHealthArgs { + brpc::ChannelBase* channel; + butil::atomic done; + int rc; + + SelectiveHealthArgs() : channel(nullptr), done(false), rc(-1) {} +}; + +void* RunSelectiveHealth(void* arg) { + SelectiveHealthArgs* args = static_cast(arg); + args->rc = args->channel->CheckHealth(); + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +struct SelectiveAddArgs { + brpc::SelectiveChannel* channel; + brpc::SelectiveChannel::ChannelHandle handle; + butil::atomic done; + int rc; + SelectiveAddArgs() : channel(nullptr), done(false), rc(-1) {} +}; + +struct SelectiveRemoveArgs { + brpc::SelectiveChannel* channel; + brpc::SelectiveChannel::ChannelHandle handle; + BlockingDescribeChannel* sub_channel; + butil::atomic started; + butil::atomic done; + + SelectiveRemoveArgs() + : channel(nullptr), sub_channel(nullptr), started(false), done(false) {} +}; + +class ScopedSelectiveChannelThreads { +public: + explicit ScopedSelectiveChannelThreads(BlockingDescribeChannel* channel) + : _describe_started(false), _add_started(false), _remove_started(false), + _channel(channel) {} + + ~ScopedSelectiveChannelThreads() { + _channel->Release(); + if (_describe_started) { + pthread_join(_describe_tid, nullptr); + } + if (_add_started) { + pthread_join(_add_tid, nullptr); + } + if (_remove_started) { + pthread_join(_remove_tid, nullptr); + } + } + + pthread_t _describe_tid; + pthread_t _add_tid; + pthread_t _remove_tid; + bool _describe_started; + bool _add_started; + bool _remove_started; + +private: + BlockingDescribeChannel* _channel; +}; + +void* RunSelectiveAdd(void* arg) { + SelectiveAddArgs* args = static_cast(arg); + args->rc = args->channel->AddChannel(new DeferredSelectiveChannel(nullptr), + &args->handle); + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +void* RunSelectiveRemove(void* arg) { + SelectiveRemoveArgs* args = static_cast(arg); + args->started.store(true, butil::memory_order_release); + args->channel->RemoveAndDestroyChannel(args->handle); + delete args->sub_channel; + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +struct SelectiveHealthRemoveArgs { + brpc::SelectiveChannel* channel; + brpc::SelectiveChannel::ChannelHandle handle; + BlockingHealthChannel* sub_channel; + butil::atomic started; + butil::atomic done; + + SelectiveHealthRemoveArgs() + : channel(nullptr), sub_channel(nullptr), started(false), done(false) {} +}; + +void* RunSelectiveHealthRemove(void* arg) { + SelectiveHealthRemoveArgs* args = + static_cast(arg); + args->started.store(true, butil::memory_order_release); + args->channel->RemoveAndDestroyChannel(args->handle); + delete args->sub_channel; + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +struct SelectiveOwningRemoveArgs { + brpc::SelectiveChannel* channel; + brpc::SelectiveChannel::ChannelHandle handle; + butil::atomic done; + + SelectiveOwningRemoveArgs() : channel(nullptr), done(false) {} +}; + +void* RunSelectiveOwningRemove(void* arg) { + SelectiveOwningRemoveArgs* args = + static_cast(arg); + args->channel->RemoveAndDestroyChannel(args->handle); + args->done.store(true, butil::memory_order_release); + return nullptr; +} + +class ScopedSelectiveHealthThreads { +public: + explicit ScopedSelectiveHealthThreads(BlockingHealthChannel* channel) + : _health_started(false), _add_started(false), _remove_started(false), + _channel(channel) {} + + ~ScopedSelectiveHealthThreads() { + _channel->Release(); + if (_health_started) { + pthread_join(_health_tid, nullptr); + } + if (_add_started) { + pthread_join(_add_tid, nullptr); + } + if (_remove_started) { + pthread_join(_remove_tid, nullptr); + } + } + + pthread_t _health_tid; + pthread_t _add_tid; + pthread_t _remove_tid; + bool _health_started; + bool _add_started; + bool _remove_started; + +private: + BlockingHealthChannel* _channel; +}; + class SelectiveRetryPolicy : public brpc::RetryPolicy { public: bool DoRetry(const brpc::Controller* cntl) const override { @@ -4086,4 +4311,232 @@ TEST_F(ChannelTest, selective_channel_add_failure_ownership) { ASSERT_EQ(0, CountingSubChannel::s_alive); } +TEST(SelectiveChannelTest, describe_does_not_block_add_channel) { + brpc::SelectiveChannel channel; + ASSERT_EQ(0, channel.Init("rr", nullptr)); + + BlockingDescribeChannel* blocking = new BlockingDescribeChannel; + brpc::SelectiveChannel::ChannelHandle first_handle; + ASSERT_EQ(0, channel.AddChannel(blocking, &first_handle)); + SelectiveDescribeArgs describe_args; + describe_args.channel = &channel; + SelectiveAddArgs add_args; + + { + ScopedSelectiveChannelThreads threads(blocking); + ASSERT_EQ(0, pthread_create(&threads._describe_tid, nullptr, + RunSelectiveDescribe, &describe_args)); + threads._describe_started = true; + for (int i = 0; i < 1000 && !blocking->started(); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(blocking->started()); + + if (blocking->started()) { + add_args.channel = &channel; + ASSERT_EQ(0, pthread_create(&threads._add_tid, nullptr, + RunSelectiveAdd, &add_args)); + threads._add_started = true; + for (int i = 0; i < 100 && !add_args.done.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + const bool add_finished = add_args.done.load( + butil::memory_order_acquire); + EXPECT_TRUE(add_finished); + if (add_finished) { + EXPECT_EQ(0, add_args.rc); + } + } + } + + if (add_args.rc == 0) { + channel.RemoveAndDestroyChannel(add_args.handle); + } + channel.RemoveAndDestroyChannel(first_handle); +} + +TEST(SelectiveChannelTest, describe_protects_non_owning_channel_removal) { + brpc::SelectiveChannel channel; + ASSERT_EQ(0, channel.Init("rr", nullptr)); + + BlockingDescribeChannel* blocking = new BlockingDescribeChannel; + brpc::SelectiveChannel::SubChannelOptions options; + options.ownership = brpc::DOESNT_OWN_CHANNEL; + brpc::SelectiveChannel::ChannelHandle handle; + ASSERT_EQ(0, channel.AddChannel(blocking, options, &handle)); + + SelectiveDescribeArgs describe_args; + describe_args.channel = &channel; + SelectiveRemoveArgs remove_args; + remove_args.channel = &channel; + remove_args.handle = handle; + remove_args.sub_channel = blocking; + + { + ScopedSelectiveChannelThreads threads(blocking); + const int describe_rc = pthread_create(&threads._describe_tid, nullptr, + RunSelectiveDescribe, &describe_args); + EXPECT_EQ(0, describe_rc); + if (describe_rc == 0) { + threads._describe_started = true; + for (int i = 0; i < 1000 && !blocking->started(); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(blocking->started()); + + if (blocking->started()) { + const int remove_rc = pthread_create(&threads._remove_tid, nullptr, + RunSelectiveRemove, &remove_args); + EXPECT_EQ(0, remove_rc); + if (remove_rc == 0) { + threads._remove_started = true; + for (int i = 0; i < 1000 && !remove_args.started.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(remove_args.started.load( + butil::memory_order_acquire)); + for (int i = 0; i < 100 && !remove_args.done.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + EXPECT_FALSE(remove_args.done.load( + butil::memory_order_acquire)); + } + } + } + } + + if (!remove_args.done.load(butil::memory_order_acquire)) { + channel.RemoveAndDestroyChannel(handle); + delete blocking; + } +} + +TEST(SelectiveChannelTest, health_check_does_not_block_owning_channel_changes) { + butil::atomic destroyed(0); + brpc::SelectiveChannel channel; + ASSERT_EQ(0, channel.Init("rr", nullptr)); + BlockingHealthChannel* blocking = new BlockingHealthChannel(&destroyed); + brpc::SelectiveChannel::ChannelHandle blocking_handle; + ASSERT_EQ(0, channel.AddChannel(blocking, &blocking_handle)); + SelectiveHealthArgs health_args; + health_args.channel = static_cast(&channel); + SelectiveAddArgs add_args; + add_args.channel = &channel; + SelectiveOwningRemoveArgs remove_args; + remove_args.channel = &channel; + remove_args.handle = blocking_handle; + + { + ScopedSelectiveHealthThreads threads(blocking); + ASSERT_EQ(0, pthread_create(&threads._health_tid, nullptr, + RunSelectiveHealth, &health_args)); + threads._health_started = true; + for (int i = 0; i < 1000 && !blocking->started(); ++i) { + bthread_usleep(1000); + } + ASSERT_TRUE(blocking->started()); + + ASSERT_EQ(0, pthread_create(&threads._add_tid, nullptr, + RunSelectiveAdd, &add_args)); + threads._add_started = true; + for (int i = 0; i < 100 && !add_args.done.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + const bool add_finished = add_args.done.load( + butil::memory_order_acquire); + EXPECT_TRUE(add_finished); + if (add_finished) { + EXPECT_EQ(0, add_args.rc); + } + + ASSERT_EQ(0, pthread_create(&threads._remove_tid, nullptr, + RunSelectiveOwningRemove, &remove_args)); + threads._remove_started = true; + for (int i = 0; i < 100 && !remove_args.done.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(remove_args.done.load(butil::memory_order_acquire)); + EXPECT_EQ(0, destroyed.load(butil::memory_order_acquire)); + } + + if (add_args.rc == 0) { + channel.RemoveAndDestroyChannel(add_args.handle); + } + if (!remove_args.done.load(butil::memory_order_acquire)) { + channel.RemoveAndDestroyChannel(blocking_handle); + } + EXPECT_TRUE(health_args.done.load(butil::memory_order_acquire)); + EXPECT_EQ(0, health_args.rc); + EXPECT_EQ(1, destroyed.load(butil::memory_order_acquire)); +} + +TEST(SelectiveChannelTest, health_check_protects_non_owning_channel_removal) { + butil::atomic destroyed(0); + brpc::SelectiveChannel channel; + ASSERT_EQ(0, channel.Init("rr", nullptr)); + BlockingHealthChannel* blocking = new BlockingHealthChannel(&destroyed); + brpc::SelectiveChannel::SubChannelOptions options; + options.ownership = brpc::DOESNT_OWN_CHANNEL; + brpc::SelectiveChannel::ChannelHandle handle; + ASSERT_EQ(0, channel.AddChannel(blocking, options, &handle)); + + SelectiveHealthArgs health_args; + health_args.channel = static_cast(&channel); + SelectiveHealthRemoveArgs remove_args; + remove_args.channel = &channel; + remove_args.handle = handle; + remove_args.sub_channel = blocking; + + { + ScopedSelectiveHealthThreads threads(blocking); + const int health_rc = pthread_create(&threads._health_tid, nullptr, + RunSelectiveHealth, &health_args); + EXPECT_EQ(0, health_rc); + if (health_rc == 0) { + threads._health_started = true; + for (int i = 0; i < 1000 && !blocking->started(); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(blocking->started()); + + if (blocking->started()) { + const int remove_rc = pthread_create( + &threads._remove_tid, nullptr, + RunSelectiveHealthRemove, &remove_args); + EXPECT_EQ(0, remove_rc); + if (remove_rc == 0) { + threads._remove_started = true; + for (int i = 0; i < 1000 && !remove_args.started.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + EXPECT_TRUE(remove_args.started.load( + butil::memory_order_acquire)); + for (int i = 0; i < 100 && !remove_args.done.load( + butil::memory_order_acquire); ++i) { + bthread_usleep(1000); + } + EXPECT_FALSE(remove_args.done.load( + butil::memory_order_acquire)); + EXPECT_EQ(0, destroyed.load(butil::memory_order_acquire)); + } + } + } + } + + if (!remove_args.done.load(butil::memory_order_acquire)) { + channel.RemoveAndDestroyChannel(handle); + delete blocking; + } + EXPECT_TRUE(health_args.done.load(butil::memory_order_acquire)); + EXPECT_EQ(0, health_args.rc); + EXPECT_TRUE(remove_args.done.load(butil::memory_order_acquire)); + EXPECT_EQ(1, destroyed.load(butil::memory_order_acquire)); +} + } //namespace