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

Signed-off-by: duanyyyyyyy <yan.duan9759@gmail.com>

Author: duanyyyyyyy

be/src/common/config.h

@@ -58,6 +58,9 @@ CONF_String(ssl_private_key_path, "");
 // The max number of single connections maintained by the brpc client and each server.
 // 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 expire configurations
+// The expire time of BRPC stub cache, default 60 minutes.
+CONF_mInt32(brpc_stub_expire_s, "3600"); // 60 minutes
 
 // Declare a selection strategy for those servers have many ips.
 // Note that there should at most one ip match this list.

be/src/runtime/exec_env.cpp

@@ -497,7 +497,7 @@ Status ExecEnv::init(const std::vector<StorePath>& store_paths, bool as_cn) {
 #endif
     _load_channel_mgr = new LoadChannelMgr();
     _load_stream_mgr = new LoadStreamMgr();
-    _brpc_stub_cache = new BrpcStubCache();
+    _brpc_stub_cache = new BrpcStubCache(this);
     _stream_load_executor = new StreamLoadExecutor(this);
     _stream_context_mgr = new StreamContextMgr();
     _transaction_mgr = new TransactionMgr(this);

be/src/util/brpc_stub_cache.cpp

@@ -14,6 +14,8 @@
 
 #include "util/brpc_stub_cache.h"
 
+#include <runtime/exec_env.h>
+
 #include "common/config.h"
 #include "gen_cpp/internal_service.pb.h"
 #include "gen_cpp/lake_service.pb.h"
@@ -22,28 +24,50 @@
 
 namespace starrocks {
 
-BrpcStubCache::BrpcStubCache() {
+BrpcStubCache::BrpcStubCache(ExecEnv* exec_env) {
     _stub_map.init(239);
+    _pipeline_timer = exec_env->pipeline_timer();
     REGISTER_GAUGE_STARROCKS_METRIC(brpc_endpoint_stub_count, [this]() {
         std::lock_guard<SpinLock> l(_lock);
         return _stub_map.size();
     });
 }
 
 BrpcStubCache::~BrpcStubCache() {
-    for (auto& stub : _stub_map) {
-        delete stub.second;
+    std::vector<std::shared_ptr<StubPool>> pools_to_cleanup;
+    {
+        std::lock_guard<SpinLock> l(_lock);
+
+        for (auto& stub : _stub_map) {
+            pools_to_cleanup.push_back(stub.second);
+        }
+    }
+
+    for (auto& pool : pools_to_cleanup) {
+        pool->_cleanup_task->unschedule(_pipeline_timer);
     }
+
+    _stub_map.clear();
 }
 
 std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::get_stub(const butil::EndPoint& endpoint) {
     std::lock_guard<SpinLock> l(_lock);
+
     auto stub_pool = _stub_map.seek(endpoint);
     if (stub_pool == nullptr) {
-        StubPool* pool = new StubPool();
-        _stub_map.insert(endpoint, pool);
-        return pool->get_or_create(endpoint);
+        auto new_pool = std::make_shared<StubPool>();
+        new_pool->_cleanup_task = new EndpointCleanupTask(this, endpoint);
+        _stub_map.insert(endpoint, new_pool);
+        stub_pool = _stub_map.seek(endpoint);
+    }
+
+    _pipeline_timer->unschedule((*stub_pool)->_cleanup_task);
+    timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s);
+    auto status = _pipeline_timer->schedule((*stub_pool)->_cleanup_task, tm);
+    if (!status.ok()) {
+        LOG(WARNING) << "Failed to schedule endpoint: " << endpoint;
     }
+
     return (*stub_pool)->get_or_create(endpoint);
 }
 
@@ -71,10 +95,25 @@ std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::get_stub(const
     return get_stub(endpoint);
 }
 
+void BrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) {
+    std::lock_guard<SpinLock> l(_lock);
+
+    LOG(INFO) << "cleanup stubs from endpoint:" << endpoint;
+    auto pool = _stub_map.seek(endpoint);
+    if (pool != nullptr) {
+        _stub_map.erase(endpoint);
+    }
+}
+
 BrpcStubCache::StubPool::StubPool() : _idx(-1) {
     _stubs.reserve(config::brpc_max_connections_per_server);
 }
 
+BrpcStubCache::StubPool::~StubPool() {
+    _stubs.clear();
+    SAFE_DELETE(_cleanup_task);
+}
+
 std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::StubPool::get_or_create(
         const butil::EndPoint& endpoint) {
     if (UNLIKELY(_stubs.size() < config::brpc_max_connections_per_server)) {
@@ -172,4 +211,8 @@ StatusOr<std::shared_ptr<starrocks::LakeService_RecoverableStub>> LakeServiceBrp
     return stub;
 }
 
+void EndpointCleanupTask::Run() {
+    _cache->cleanup_expired(_endpoint);
+}
+
 } // namespace starrocks

be/src/util/brpc_stub_cache.h

@@ -34,6 +34,8 @@
 
 #pragma once
 
+#include <exec/pipeline/schedule/pipeline_timer.h>
+
 #include <memory>
 #include <mutex>
 #include <vector>
@@ -48,26 +50,33 @@
 
 namespace starrocks {
 
+class ExecEnv;
+class EndpointCleanupTask;
+
 class BrpcStubCache {
 public:
-    BrpcStubCache();
+    BrpcStubCache(ExecEnv* exec_env);
     ~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 cleanup_expired(const butil::EndPoint& endpoint);
 
 private:
     struct StubPool {
         StubPool();
+        ~StubPool();
         std::shared_ptr<PInternalService_RecoverableStub> get_or_create(const butil::EndPoint& endpoint);
 
         std::vector<std::shared_ptr<PInternalService_RecoverableStub>> _stubs;
         int64_t _idx;
+        EndpointCleanupTask* _cleanup_task = nullptr;
     };
 
     SpinLock _lock;
-    butil::FlatMap<butil::EndPoint, StubPool*> _stub_map;
+    butil::FlatMap<butil::EndPoint, std::shared_ptr<StubPool>> _stub_map;
+    pipeline::PipelineTimer* _pipeline_timer;
 };
 
 class HttpBrpcStubCache {
@@ -98,4 +107,14 @@ class LakeServiceBrpcStubCache {
     butil::FlatMap<butil::EndPoint, std::shared_ptr<LakeService_RecoverableStub>> _stub_map;
 };
 
+class EndpointCleanupTask : public starrocks::pipeline::PipelineTimerTask {
+public:
+    EndpointCleanupTask(BrpcStubCache* cache, const butil::EndPoint& endpoint) : _cache(cache), _endpoint(endpoint){};
+    void Run() override;
+
+private:
+    BrpcStubCache* _cache;
+    butil::EndPoint _endpoint;
+};
+
 } // namespace starrocks

be/test/http/stream_load_test.cpp

@@ -91,7 +91,7 @@ class StreamLoadActionTest : public testing::Test {
         config::streaming_load_max_mb = 1;
 
         _env._load_stream_mgr = new LoadStreamMgr();
-        _env._brpc_stub_cache = new BrpcStubCache();
+        _env._brpc_stub_cache = new BrpcStubCache(&_env);
         _env._stream_load_executor = new StreamLoadExecutor(&_env);
 
         _evhttp_req = evhttp_request_new(nullptr, nullptr);

be/test/http/transaction_stream_load_test.cpp

@@ -68,7 +68,7 @@ class TransactionStreamLoadActionTest : public testing::Test {
         config::streaming_load_max_mb = 1;
 
         _env._load_stream_mgr = new LoadStreamMgr();
-        _env._brpc_stub_cache = new BrpcStubCache();
+        _env._brpc_stub_cache = new BrpcStubCache(&_env);
         _env._stream_load_executor = new StreamLoadExecutor(&_env);
         _env._stream_context_mgr = new StreamContextMgr();
         _env._transaction_mgr = new TransactionMgr(&_env);

be/test/util/brpc_stub_cache_test.cpp

@@ -18,7 +18,9 @@
 #include "util/brpc_stub_cache.h"
 
 #include <gtest/gtest.h>
+#include <testutil/assert.h>
 
+#include "runtime/exec_env.h"
 #include "util/failpoint/fail_point.h"
 
 namespace starrocks {
@@ -27,10 +29,22 @@ class BrpcStubCacheTest : public testing::Test {
 public:
     BrpcStubCacheTest() = default;
     ~BrpcStubCacheTest() override = default;
+    void SetUp() override {
+        _env._pipeline_timer = new pipeline::PipelineTimer();
+        ASSERT_OK(_env._pipeline_timer->start());
+    }
+    void TearDown() override {
+        delete _env._pipeline_timer;
+        _env._pipeline_timer = nullptr;
+        config::brpc_stub_expire_s = 3600;
+    }
+
+private:
+    ExecEnv _env;
 };
 
 TEST_F(BrpcStubCacheTest, normal) {
-    BrpcStubCache cache;
+    BrpcStubCache cache(&_env);
     TNetworkAddress address;
     address.hostname = "127.0.0.1";
     address.port = 123;
@@ -46,7 +60,7 @@ TEST_F(BrpcStubCacheTest, normal) {
 }
 
 TEST_F(BrpcStubCacheTest, invalid) {
-    BrpcStubCache cache;
+    BrpcStubCache cache(&_env);
     TNetworkAddress address;
     address.hostname = "invalid.cm.invalid";
     address.port = 123;
@@ -55,7 +69,7 @@ TEST_F(BrpcStubCacheTest, invalid) {
 }
 
 TEST_F(BrpcStubCacheTest, reset) {
-    BrpcStubCache cache;
+    BrpcStubCache cache(&_env);
     TNetworkAddress address;
     address.hostname = "127.0.0.1";
     address.port = 123;
@@ -108,4 +122,20 @@ TEST_F(BrpcStubCacheTest, test_http_stub) {
     ASSERT_EQ(nullptr, *stub4);
 }
 
+TEST_F(BrpcStubCacheTest, test_cleanup) {
+    config::brpc_stub_expire_s = 1;
+    BrpcStubCache cache(&_env);
+    TNetworkAddress address;
+    address.hostname = "127.0.0.1";
+    address.port = 123;
+    auto stub1 = cache.get_stub(address);
+    ASSERT_NE(nullptr, stub1);
+    auto stub2 = cache.get_stub(address);
+    ASSERT_EQ(stub2, stub1);
+
+    sleep(2);
+    auto stub3 = cache.get_stub(address);
+    ASSERT_NE(stub3, stub1);
+}
+
 } // namespace starrocks