[Enhancement] Add a cleaner for BrpcStubCache to cleanup unused connections

Author: duanyyyyyyy

be/src/common/config.h

@@ -59,6 +59,10 @@ CONF_String(ssl_private_key_path, "");
 // These connections are created during the first few access and will be used thereafter
 CONF_Int32(brpc_max_connections_per_server, "1");
 
+// BRPC stub cache cleanup configurations
+// Interval in seconds to run cleanup for expired BRPC stubs
+CONF_mInt32(brpc_stub_cleanup_interval_s, "600"); // 10 minutes
+
 // Declare a selection strategy for those servers have many ips.
 // Note that there should at most one ip match this list.
 // this is a list in semicolon-delimited format, in CIDR notation, e.g. 10.10.10.0/24

be/src/util/brpc_stub_cache.cpp

@@ -17,6 +17,7 @@
 #include "common/config.h"
 #include "gen_cpp/internal_service.pb.h"
 #include "gen_cpp/lake_service.pb.h"
+#include "util/misc.h"
 #include "util/failpoint/fail_point.h"
 #include "util/starrocks_metrics.h"
 
@@ -28,6 +29,17 @@ BrpcStubCache::BrpcStubCache() {
         std::lock_guard<SpinLock> l(_lock);
         return _stub_map.size();
     });
+    _cleanup_thread = std::thread([this] {
+#ifdef GOOGLE_PROFILER
+    ProfilerRegisterThread();
+#endif
+        while (!_stopped.load()) {
+            LOG(INFO) << "Start to clean up expired brpc stubs";
+            check_and_cleanup_unhealthy_stubs();
+            nap_sleep(config::brpc_stub_cleanup_interval_s, [this] { return _stopped.load(); });
+        }
+    });
+    Thread::set_thread_name(_cleanup_thread, "brpc_cleanup_thread");
 }
 
 BrpcStubCache::~BrpcStubCache() {
@@ -71,6 +83,30 @@ std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::get_stub(const
     return get_stub(endpoint);
 }
 
+void BrpcStubCache::check_and_cleanup_unhealthy_stubs() {
+    std::vector<butil::EndPoint> snapshot;
+    {
+        std::lock_guard<SpinLock> l(_lock);
+
+        for (auto& entry : _stub_map) {
+            snapshot.push_back(entry.first);
+        }
+    }
+
+    for (auto endpoint : snapshot) {
+        auto stub_pool = get_stub_pool_readonly(endpoint);
+        if (stub_pool && stub_pool->is_unhealthy()) {
+            {
+                std::lock_guard<SpinLock> l(_lock);
+
+                LOG(INFO) << "cleanup stubs from endpoint:" << endpoint;
+                _stub_map.erase(endpoint);
+                delete stub_pool;
+            }
+        }
+    }
+}
+
 BrpcStubCache::StubPool::StubPool() : _idx(-1) {
     _stubs.reserve(config::brpc_max_connections_per_server);
 }
@@ -91,6 +127,28 @@ std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::StubPool::get_o
     return _stubs[_idx];
 }
 
+bool BrpcStubCache::StubPool::is_unhealthy() {
+    if (_stubs.size() == 0) {
+        // When the vector is empty do not need to clean
+         return false;
+    }
+
+    auto& first_stub = _stubs[0];
+    if (!first_stub) {
+        return true;
+    }
+
+    auto status = first_stub->check_health();
+    return !status.ok();
+}
+
+BrpcStubCache::StubPool* BrpcStubCache::get_stub_pool_readonly(
+        const butil::EndPoint& endpoint) {
+    std::lock_guard<SpinLock> l(_lock);
+    auto stub_ptr = _stub_map.seek(endpoint);
+    return stub_ptr ? *stub_ptr : nullptr;
+}
+
 HttpBrpcStubCache* HttpBrpcStubCache::getInstance() {
     static HttpBrpcStubCache cache;
     return &cache;

be/src/util/brpc_stub_cache.h

@@ -37,6 +37,7 @@
 #include <memory>
 #include <mutex>
 #include <vector>
+#include <atomic>
 
 #include "common/statusor.h"
 #include "gen_cpp/Types_types.h" // TNetworkAddress
@@ -45,6 +46,7 @@
 #include "util/lake_service_recoverable_stub.h"
 #include "util/network_util.h"
 #include "util/spinlock.h"
+#include "util/thread.h"
 
 namespace starrocks {
 
@@ -56,18 +58,26 @@ class BrpcStubCache {
     std::shared_ptr<PInternalService_RecoverableStub> get_stub(const butil::EndPoint& endpoint);
     std::shared_ptr<PInternalService_RecoverableStub> get_stub(const TNetworkAddress& taddr);
     std::shared_ptr<PInternalService_RecoverableStub> get_stub(const std::string& host, int port);
+    void check_and_cleanup_unhealthy_stubs();
 
 private:
     struct StubPool {
         StubPool();
+        ~StubPool();
+
         std::shared_ptr<PInternalService_RecoverableStub> get_or_create(const butil::EndPoint& endpoint);
+        bool is_unhealthy();
 
         std::vector<std::shared_ptr<PInternalService_RecoverableStub>> _stubs;
         int64_t _idx;
     };
 
+    StubPool* get_stub_pool_readonly(const butil::EndPoint& endpoint);
+
     SpinLock _lock;
     butil::FlatMap<butil::EndPoint, StubPool*> _stub_map;
+    std::atomic<bool> _stopped{false};
+    std::thread _cleanup_thread;
 };
 
 class HttpBrpcStubCache {

be/src/util/internal_service_recoverable_stub.cpp

@@ -60,6 +60,70 @@ Status PInternalService_RecoverableStub::reset_channel(int64_t next_connection_g
     return Status::OK();
 }
 
+Status PInternalService_RecoverableStub::check_health() {
+
+#ifdef BE_TEST
+    if (_endpoint.port == 125) {
+        LOG(INFO) << "BE_TEST mode: Port 125 simulated as unhealthy for " << _endpoint;
+        return Status::InternalError("Simulated failure for port 111");
+    }
+    LOG(INFO) << "BE_TEST mode: Endpoint " << _endpoint << " simulated as healthy";
+    return Status::OK();
+#endif
+
+    // extract host
+    std::string host = butil::endpoint2str(_endpoint).c_str();
+    size_t colon_pos = host.find(':');
+    if (colon_pos != std::string::npos) {
+        host = host.substr(0, colon_pos);
+    }
+
+    // generate endpoint
+    butil::EndPoint http_endpoint;
+    if (butil::str2endpoint((host + ":" + std::to_string(config::be_http_port)).c_str(), &http_endpoint) != 0) {
+        LOG(WARNING) << "Invalid HTTP endpoint: " << host << ":" << config::be_http_port;
+        return Status::InternalError("Invalid HTTP endpoint");
+    }
+
+    // generate channel options
+    brpc::ChannelOptions options;
+    options.connect_timeout_ms = config::rpc_connect_timeout_ms;
+    options.timeout_ms = config::rpc_connect_timeout_ms;
+    options.max_retry = 1;
+    options.protocol = "http";
+
+    // init channel
+    brpc::Channel channel;
+    if (channel.Init(http_endpoint, &options) != 0) {
+        LOG(WARNING) << "Failed to init HTTP channel for " << http_endpoint;
+        return Status::InternalError("Failed to init HTTP channel");
+    }
+
+    // generate controller
+    brpc::Controller cntl;
+    cntl.http_request().set_method(brpc::HTTP_METHOD_GET);
+    cntl.http_request().uri() = "/api/health";
+
+    // trigger health API
+    channel.CallMethod(nullptr, &cntl, nullptr, nullptr, nullptr);
+
+    // get and check response
+    if (cntl.Failed()) {
+        LOG(WARNING) << "Health check failed for endpoint " << _endpoint 
+                     << ", error: " << cntl.ErrorText();
+        return Status::InternalError("Health check request failed");
+    }
+    int status_code = cntl.http_response().status_code();
+    if (status_code == 200) {
+        LOG(INFO) << "Health check passed for endpoint " << _endpoint << " (HTTP " << status_code << ")";
+        return Status::OK();
+    } else {
+        LOG(WARNING) << "Health check failed for endpoint " << _endpoint 
+                     << ", HTTP status: " << status_code;
+        return Status::InternalError("Health check failed");
+    }
+}
+
 void PInternalService_RecoverableStub::tablet_writer_open(::google::protobuf::RpcController* controller,
                                                           const ::starrocks::PTabletWriterOpenRequest* request,
                                                           ::starrocks::PTabletWriterOpenResult* response,

be/src/util/internal_service_recoverable_stub.h

@@ -28,6 +28,7 @@ class PInternalService_RecoverableStub : public PInternalService,
     ~PInternalService_RecoverableStub();
 
     Status reset_channel(int64_t next_connection_group = 0);
+    Status check_health();
 
     std::shared_ptr<starrocks::PInternalService_Stub> stub() const {
         std::shared_lock l(_mutex);

be/test/util/brpc_stub_cache_test.cpp

@@ -108,4 +108,25 @@ TEST_F(BrpcStubCacheTest, test_http_stub) {
     ASSERT_EQ(nullptr, *stub4);
 }
 
+TEST_F(BrpcStubCacheTest, cleanup) {
+    BrpcStubCache cache;
+    TNetworkAddress address;
+    address.hostname = "127.0.0.1";
+    // healthy port
+    address.port = 123;
+    auto stub1 = cache.get_stub(address);
+    ASSERT_NE(nullptr, stub1);
+    cache.check_and_cleanup_unhealthy_stubs();
+    auto stub2 = cache.get_stub(address);
+    ASSERT_EQ(stub2, stub1);
+
+    // unhealthy port
+    address.port = 125;
+    auto stub3 = cache.get_stub(address);
+    ASSERT_NE(nullptr, stub1);
+    cache.check_and_cleanup_unhealthy_stubs();
+    auto stub4 = cache.get_stub(address);
+    ASSERT_NE(stub4, stub3);
+}
+
 } // namespace starrocks