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 7dc92edb7 Fix socket reference leak during removal and revival (#3691)
7dc92edb7 is described below
commit 7dc92edb7b2cbfdbd6548faacbdf761fa6dbe6df
Author: Bright Chen <[email protected]>
AuthorDate: Sat Oct 31 19:01:42 2026 +0800
Fix socket reference leak during removal and revival (#3691)
---
src/brpc/load_balancer.cpp | 6 +--
src/brpc/selective_channel.cpp | 21 ++++----
src/brpc/versioned_ref_with_id.h | 57 +++++++++++++++------
test/brpc_channel_unittest.cpp | 44 ++++++++++++++++
test/brpc_socket_unittest.cpp | 106 +++++++++++++++++++++++++++++++++++++++
5 files changed, 205 insertions(+), 29 deletions(-)
diff --git a/src/brpc/load_balancer.cpp b/src/brpc/load_balancer.cpp
index 0544a4836..725e1a62c 100644
--- a/src/brpc/load_balancer.cpp
+++ b/src/brpc/load_balancer.cpp
@@ -83,17 +83,17 @@ int SharedLoadBalancer::Init(const char* lb_protocol) {
std::string lb_name;
butil::StringPiece lb_params;
if (!ParseParameters(lb_protocol, &lb_name, &lb_params)) {
- LOG(FATAL) << "Fail to parse this load balancer protocol '" <<
lb_protocol << '\'';
+ LOG(ERROR) << "Fail to parse this load balancer protocol '" <<
lb_protocol << '\'';
return -1;
}
const LoadBalancer* lb = LoadBalancerExtension()->Find(lb_name.c_str());
if (lb == nullptr) {
- LOG(FATAL) << "Fail to find LoadBalancer by `" << lb_name << "'";
+ LOG(ERROR) << "Fail to find LoadBalancer by `" << lb_name << "'";
return -1;
}
_lb = lb->New(lb_params);
if (_lb == nullptr) {
- LOG(FATAL) << "Fail to new LoadBalancer";
+ LOG(ERROR) << "Fail to new LoadBalancer";
return -1;
}
if (FLAGS_show_lb_in_vars && !_exposed) {
diff --git a/src/brpc/selective_channel.cpp b/src/brpc/selective_channel.cpp
index 5d2fbe898..5143fb0eb 100644
--- a/src/brpc/selective_channel.cpp
+++ b/src/brpc/selective_channel.cpp
@@ -228,21 +228,19 @@ int ChannelBalancer::AddChannel(ChannelBase* sub_channel,
// OWNS_CHANNEL).
ptr->set_user(sub_chan);
cleanup_subchan.dismiss();
- if (rc > 0 && !ptr->HCEnabled()) {
- LOG(ERROR) << "Health check of SocketId="
- << sock_id << " is disabled";
+ auto cleanup_socket = butil::MakeScopeGuard([&]() noexcept {
// sub_chan will be deleted when the socket is recycled.
ptr->SetFailed();
- // Cancel health checking.
+ // Permanently prevent revival before releasing the HC reference.
+ ptr->ReleaseAdditionalReference();
ptr->ReleaseHCRelatedReference();
+ });
+ if (rc > 0 && !ptr->HCEnabled()) {
+ LOG(ERROR) << "Health check of SocketId=" << sock_id << " is disabled";
return -1;
}
if (!AddServer(ServerId(sock_id, subopt.tag))) {
LOG(ERROR) << "Duplicated sub_channel=" << sub_channel;
- // sub_chan will be deleted when the socket is recycled.
- ptr->SetFailed();
- // Cancel health checking.
- ptr->ReleaseHCRelatedReference();
return -1;
}
// The health-check-related reference has been held on created.
@@ -251,6 +249,7 @@ int ChannelBalancer::AddChannel(ChannelBase* sub_channel,
handle->id = sock_id;
handle->tag = subopt.tag;
}
+ cleanup_socket.dismiss();
return 0;
}
@@ -266,9 +265,9 @@ void ChannelBalancer::RemoveAndDestroyChannel(const
SelectiveChannel::ChannelHan
BAIDU_SCOPED_LOCK(_mutex);
CHECK_EQ(1UL, _chan_map.erase(sub->chan));
}
- if (rc == 0) {
- ptr->ReleaseAdditionalReference();
- }
+ // Even a failed socket must be marked as permanently removed before
+ // releasing the HC reference, otherwise a concurrent HC can revive it.
+ ptr->ReleaseAdditionalReference();
// Cancel health checking.
ptr->ReleaseHCRelatedReference();
}
diff --git a/src/brpc/versioned_ref_with_id.h b/src/brpc/versioned_ref_with_id.h
index b7213c349..9fc5f5266 100644
--- a/src/brpc/versioned_ref_with_id.h
+++ b/src/brpc/versioned_ref_with_id.h
@@ -301,7 +301,8 @@ public:
// Release the additional reference which added inside `Create'
// before so that `VersionedRefWithId' will be recycled automatically
- // once on one is addressing it.
+ // once no one is addressing it. Unlike SetFailed(), this permanently
+ // prevents revival, even if SetFailed() already released the reference.
int ReleaseAdditionalReference();
VRefId id() const { return _this_id; }
@@ -312,14 +313,16 @@ public:
protected:
friend void DereferenceVersionedRefWithId<>(T* r);
- // Status flag used to mark that
+ // State of the additional reference.
enum AdditionalRefStatus {
// 1. Additional reference has been increased;
ADDITIONAL_REF_USING,
// 2. Additional reference is increasing;
ADDITIONAL_REF_REVIVING,
- // 3. Additional reference has been decreased.
- ADDITIONAL_REF_RECYCLED
+ // 3. SetFailed() released the reference; Revive() may add it back.
+ ADDITIONAL_REF_RELEASED_ON_FAILURE,
+ // 4. The owner has released the reference and forbidden revival.
+ ADDITIONAL_REF_RELEASED
};
AdditionalRefStatus additional_ref_status() const {
@@ -352,6 +355,8 @@ friend void DereferenceVersionedRefWithId<>(T* r);
private:
typedef butil::ResourceId<T> resource_id_t;
+ int ReleaseAdditionalReferenceImpl(bool permanent);
+
// 1. When `failed_as_well=true', returns 0 on success,
// 1 on failed object, -1 on recycled.
// 2. When `failed_as_well=true', returns 0 on success,
@@ -384,8 +389,10 @@ private:
// decreased, or is increasing.
// additional ref status:
// constructor / `Create': REF_USING
- // `SetFailed': REF_USING -> REF_RECYCLED
- // `Revive' REF_RECYCLED -> REF_REVIVING -> REF_USING
+ // `SetFailed`: REF_USING -> REF_RELEASED_ON_FAILURE
+ // `Revive`: REF_RELEASED_ON_FAILURE -> REF_REVIVING -> REF_USING
+ // `ReleaseAdditionalReference`: REF_USING/REF_RELEASED_ON_FAILURE ->
REF_RELEASED
+ // REF_RELEASED is terminal until the slot is reused by Create().
butil::atomic<AdditionalRefStatus> _additional_ref_status;
};
@@ -538,7 +545,7 @@ int VersionedRefWithId<T>::SetFailedImpl(Args&&... args) {
CallOnFailed(static_cast<T*>(this), std::forward<Args>(args)...);
// Deref additionally which is added at creation so that this
// queue's reference will hit 0(recycle) when no one addresses it.
- ReleaseAdditionalReference();
+ ReleaseAdditionalReferenceImpl(false);
// NOTE: This object may be recycled at this point, don't
// touch anything.
return 0;
@@ -548,12 +555,17 @@ int VersionedRefWithId<T>::SetFailedImpl(Args&&... args) {
template<typename T>
int VersionedRefWithId<T>::ReleaseAdditionalReference() {
+ return ReleaseAdditionalReferenceImpl(true);
+}
+
+template<typename T>
+int VersionedRefWithId<T>::ReleaseAdditionalReferenceImpl(bool permanent) {
do {
AdditionalRefStatus expect = ADDITIONAL_REF_USING;
if (_additional_ref_status.compare_exchange_strong(
- expect, ADDITIONAL_REF_RECYCLED,
- butil::memory_order_relaxed,
- butil::memory_order_relaxed)) {
+ expect, permanent ? ADDITIONAL_REF_RELEASED
+ : ADDITIONAL_REF_RELEASED_ON_FAILURE,
+ butil::memory_order_relaxed, butil::memory_order_relaxed)) {
CallBeforeAdditionalRefReleased(static_cast<T*>(this));
return Dereference();
}
@@ -561,8 +573,17 @@ int VersionedRefWithId<T>::ReleaseAdditionalReference() {
if (expect == ADDITIONAL_REF_REVIVING) {
// sched_yield to wait until status is not REF_REVIVING.
sched_yield();
+ } else if (permanent && expect == ADDITIONAL_REF_RELEASED_ON_FAILURE) {
+ // SetFailed() already dropped the reference, but the owner's
+ // removal must still prevent Revive() from adding it back.
+ if (_additional_ref_status.compare_exchange_strong(
+ expect, ADDITIONAL_REF_RELEASED,
+ butil::memory_order_relaxed,
+ butil::memory_order_relaxed)) {
+ return -1;
+ }
} else {
- return -1; // REF_RECYCLED
+ return -1; // REF_RELEASED_ON_FAILURE or REF_RELEASED
}
} while (true);
}
@@ -628,18 +649,24 @@ int VersionedRefWithId<T>::Dereference() {
template<typename T>
void VersionedRefWithId<T>::Revive(int32_t at_least_nref) {
+ AdditionalRefStatus expect = ADDITIONAL_REF_RELEASED_ON_FAILURE;
+ if (!_additional_ref_status.compare_exchange_strong(
+ expect, ADDITIONAL_REF_REVIVING,
+ butil::memory_order_relaxed,
+ butil::memory_order_relaxed)) {
+ CHECK_EQ(ADDITIONAL_REF_RELEASED, expect);
+ return;
+ }
const uint32_t id_ver = VersionOfVRefId(_this_id);
uint64_t vref = _versioned_ref.load(butil::memory_order_relaxed);
- _additional_ref_status.store(
- ADDITIONAL_REF_REVIVING, butil::memory_order_relaxed);
while (true) {
CHECK_EQ(id_ver + 1, VersionOfVRef(vref)) << "id=" << id();
int32_t nref = NRefOfVRef(vref);
if (nref < at_least_nref) {
- // Set the status to REF_RECYCLED since no one uses this object
+ // Revival was abandoned without reacquiring the reference.
_additional_ref_status.store(
- ADDITIONAL_REF_RECYCLED, butil::memory_order_relaxed);
+ ADDITIONAL_REF_RELEASED_ON_FAILURE,
butil::memory_order_relaxed);
CHECK_EQ(1, nref);
LOG(WARNING) << description() << " was abandoned during revival";
return;
diff --git a/test/brpc_channel_unittest.cpp b/test/brpc_channel_unittest.cpp
index f6155cbdb..3366aa1f3 100644
--- a/test/brpc_channel_unittest.cpp
+++ b/test/brpc_channel_unittest.cpp
@@ -2576,6 +2576,25 @@ TEST_F(ChannelTest, init_using_unknown_naming_service) {
ASSERT_EQ(-1, channel.Init("unknown://unknown", "unknown", nullptr));
}
+TEST_F(ChannelTest, selective_channel_invalid_load_balancer) {
+ GFLAGS_NAMESPACE::FlagSaver flag_saver;
+#if !BRPC_WITH_GLOG
+ // Invalid configuration must return an error even when fatal logs abort.
+ ASSERT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption(
+ "crash_on_fatal_log", "true").empty());
+#endif
+ const char* invalid_options[] = {
+ "", "unknown", "rr:invalid=1", "random:invalid=1"
+ };
+ for (const char* option : invalid_options) {
+ SCOPED_TRACE(option);
+ brpc::SelectiveChannel channel;
+ EXPECT_EQ(-1, channel.Init(option, nullptr));
+ // A failed initialization must leave the channel usable for retry.
+ EXPECT_EQ(0, channel.Init("rr", nullptr));
+ }
+}
+
TEST_F(ChannelTest, init_using_unexist_fns) {
brpc::Channel channel;
ASSERT_EQ(-1, channel.Init("fiLe://no_such_file", "rr", nullptr));
@@ -3977,6 +3996,31 @@ public:
};
int CountingSubChannel::s_alive = 0;
+TEST_F(ChannelTest, selective_channel_removal_prevents_revival) {
+ brpc::SelectiveChannel channel;
+ ASSERT_EQ(0, channel.Init("rr", nullptr));
+ brpc::SelectiveChannel::ChannelHandle handle;
+ ASSERT_EQ(0, channel.AddChannel(new CountingSubChannel, &handle));
+ ASSERT_EQ(1, CountingSubChannel::s_alive);
+ brpc::SocketUniquePtr ptr;
+ ASSERT_EQ(0, brpc::Socket::Address(handle.id, &ptr));
+
+ // Keep the HC reference but disable automatic scheduling: the retained
+ // pointer below simulates the health check task resuming after removal.
+ ptr->_is_hc_related_ref_held = false;
+ EXPECT_EQ(0, ptr->SetFailed());
+ ptr->_is_hc_related_ref_held = true;
+ channel.RemoveAndDestroyChannel(handle);
+ EXPECT_FALSE(ptr->HCEnabled());
+ ptr->Revive(1);
+ EXPECT_TRUE(ptr->Failed());
+ EXPECT_EQ(1, ptr->nref());
+ // Also clean up the leaked additional reference if this test regresses.
+ ptr->SetFailed();
+ ptr.reset();
+ EXPECT_EQ(0, CountingSubChannel::s_alive);
+}
+
TEST_F(ChannelTest, selective_channel_add_failure_ownership) {
// nullptr sub channel is rejected without touching ownership.
{
diff --git a/test/brpc_socket_unittest.cpp b/test/brpc_socket_unittest.cpp
index ee9cb31e6..48e49a854 100644
--- a/test/brpc_socket_unittest.cpp
+++ b/test/brpc_socket_unittest.cpp
@@ -1087,6 +1087,112 @@ TEST_F(SocketTest,
derived_socket_keeps_parent_keepalive_options) {
}
}
+class QuietRevival : public brpc::SocketUser {
+ void AfterRevived(brpc::Socket*) override {}
+};
+
+TEST_F(SocketTest, removal_racing_with_revival) {
+ QuietRevival user;
+ // Exercise both orderings without depending on the health check timer.
+ for (int revive_first = 0; revive_first < 2; ++revive_first) {
+ SCOPED_TRACE(revive_first);
+ brpc::SocketOptions options;
+ options.user = &user;
+ brpc::SocketId id;
+ ASSERT_EQ(0, brpc::Socket::Create(options, &id));
+ BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); };
+ // These two references stand in for the owner's HC reference and
+ // the health check task's reference. The timer itself is disabled.
+ brpc::SocketUniquePtr owner;
+ brpc::SocketUniquePtr health_check;
+ ASSERT_EQ(0, brpc::Socket::Address(id, &owner));
+ ASSERT_EQ(0, brpc::Socket::Address(id, &health_check));
+ ASSERT_EQ(0, owner->SetFailed());
+ ASSERT_EQ(2, owner->nref());
+
+ if (revive_first) {
+ health_check->Revive(2);
+ EXPECT_FALSE(owner->Failed());
+ EXPECT_EQ(3, owner->nref());
+ EXPECT_EQ(0, owner->ReleaseAdditionalReference());
+ } else {
+ EXPECT_EQ(-1, owner->ReleaseAdditionalReference());
+ // The owner's HC reference has not been released yet: Revive(2)
+ // must still reject this socket after permanent removal.
+ health_check->Revive(2);
+ EXPECT_TRUE(owner->Failed());
+ }
+ EXPECT_EQ(2, owner->nref());
+ EXPECT_EQ(-1, owner->ReleaseAdditionalReference());
+ EXPECT_EQ(2, owner->nref());
+ owner.reset();
+ health_check.reset();
+ EXPECT_EQ(-1, brpc::Socket::Status(id));
+ }
+}
+
+TEST_F(SocketTest, removal_before_failure_prevents_revival) {
+ QuietRevival user;
+ brpc::SocketOptions options;
+ options.user = &user;
+ brpc::SocketId id;
+ ASSERT_EQ(0, brpc::Socket::Create(options, &id));
+ BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); };
+ brpc::SocketUniquePtr ptr;
+ ASSERT_EQ(0, brpc::Socket::Address(id, &ptr));
+ ASSERT_EQ(0, ptr->ReleaseAdditionalReference());
+ ASSERT_EQ(0, ptr->SetFailed());
+ ptr->Revive(1);
+ EXPECT_TRUE(ptr->Failed());
+ EXPECT_EQ(1, ptr->nref());
+ ptr.reset();
+ EXPECT_EQ(-1, brpc::Socket::Status(id));
+}
+
+struct RemovalRace {
+ explicit RemovalRace(brpc::Socket* s) : socket(s), ready(2) {}
+ brpc::Socket* socket;
+ bthread::CountdownEvent ready;
+};
+
+static void* ReleaseAdditionalReferenceInThread(void* arg) {
+ RemovalRace* race = static_cast<RemovalRace*>(arg);
+ race->ready.signal();
+ race->ready.wait();
+ race->socket->ReleaseAdditionalReference();
+ return nullptr;
+}
+
+TEST_F(SocketTest, concurrent_removal_and_revival) {
+ QuietRevival user;
+ for (int i = 0; i < 200; ++i) {
+ SCOPED_TRACE(i);
+ brpc::SocketOptions options;
+ options.user = &user;
+ brpc::SocketId id;
+ ASSERT_EQ(0, brpc::Socket::Create(options, &id));
+ BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); };
+ brpc::SocketUniquePtr owner;
+ brpc::SocketUniquePtr health_check;
+ ASSERT_EQ(0, brpc::Socket::Address(id, &owner));
+ ASSERT_EQ(0, brpc::Socket::Address(id, &health_check));
+ ASSERT_EQ(0, owner->SetFailed());
+ RemovalRace race(owner.get());
+ pthread_t tid;
+ ASSERT_EQ(0, pthread_create(&tid, nullptr,
+ ReleaseAdditionalReferenceInThread,
+ &race));
+ race.ready.signal();
+ race.ready.wait();
+ health_check->Revive(2);
+ ASSERT_EQ(0, pthread_join(tid, nullptr));
+ EXPECT_EQ(2, owner->nref());
+ owner.reset();
+ health_check.reset();
+ EXPECT_EQ(-1, brpc::Socket::Status(id));
+ }
+}
+
TEST_F(SocketTest, client_binding_ipv6_connect) {
butil::EndPoint point;
ASSERT_EQ(0, butil::str2endpoint("[::1]:0", &point));
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
