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

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

Author: duanyyyyyyy

be/src/util/brpc_stub_cache.cpp

@@ -65,7 +65,7 @@ std::shared_ptr<PInternalService_RecoverableStub> BrpcStubCache::get_stub(const
     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;
+        LOG(WARNING) << "Failed to schedule brpc cleanup task: " << endpoint;
     }
 
     return (*stub_pool)->get_or_create(endpoint);
@@ -99,10 +99,7 @@ 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);
-    }
+    _stub_map.erase(endpoint);
 }
 
 BrpcStubCache::StubPool::StubPool() : _idx(-1) {
@@ -137,6 +134,26 @@ HttpBrpcStubCache* HttpBrpcStubCache::getInstance() {
 
 HttpBrpcStubCache::HttpBrpcStubCache() {
     _stub_map.init(500);
+    _task_map.init(500);
+    _pipeline_timer = ExecEnv::GetInstance()->pipeline_timer();
+}
+
+HttpBrpcStubCache::~HttpBrpcStubCache() {
+    std::vector<std::shared_ptr<HttpEndpointCleanupTask>> task_to_cleanup;
+
+    {
+        std::lock_guard<SpinLock> l(_lock);
+        for (auto& stub : _task_map) {
+            task_to_cleanup.push_back(stub.second);
+        }
+    }
+
+    for (auto& task : task_to_cleanup) {
+        task->unschedule(_pipeline_timer);
+    }
+
+    _stub_map.clear();
+    _task_map.clear();
 }
 
 StatusOr<std::shared_ptr<PInternalService_RecoverableStub>> HttpBrpcStubCache::get_http_stub(
@@ -158,6 +175,21 @@ StatusOr<std::shared_ptr<PInternalService_RecoverableStub>> HttpBrpcStubCache::g
     }
     // get is exist
     std::lock_guard<SpinLock> l(_lock);
+
+    // schedule clean up task
+    auto task = _task_map.seek(endpoint);
+    if (task == nullptr) {
+        auto new_task = std::make_shared<HttpEndpointCleanupTask>(this, endpoint);
+        _task_map.insert(endpoint, new_task);
+        task = _task_map.seek(endpoint);
+    }
+    _pipeline_timer->unschedule((*task).get());
+    timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s);
+    auto status = _pipeline_timer->schedule((*task).get(), tm);
+    if (!status.ok()) {
+        LOG(WARNING) << "Failed to schedule http brpc cleanup task: " << endpoint;
+    }
+
     auto stub_ptr = _stub_map.seek(endpoint);
     if (stub_ptr != nullptr) {
         return *stub_ptr;
@@ -172,13 +204,41 @@ StatusOr<std::shared_ptr<PInternalService_RecoverableStub>> HttpBrpcStubCache::g
     return stub;
 }
 
+void HttpBrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) {
+    std::lock_guard<SpinLock> l(_lock);
+
+    LOG(INFO) << "cleanup http stubs from endpoint:" << endpoint;
+    _stub_map.erase(endpoint);
+    _task_map.erase(endpoint);
+}
+
 LakeServiceBrpcStubCache* LakeServiceBrpcStubCache::getInstance() {
     static LakeServiceBrpcStubCache cache;
     return &cache;
 }
 
 LakeServiceBrpcStubCache::LakeServiceBrpcStubCache() {
     _stub_map.init(500);
+    _task_map.init(500);
+    _pipeline_timer = ExecEnv::GetInstance()->pipeline_timer();
+}
+
+LakeServiceBrpcStubCache::~LakeServiceBrpcStubCache() {
+    std::vector<std::shared_ptr<LakeEndpointCleanupTask>> task_to_cleanup;
+
+    {
+        std::lock_guard<SpinLock> l(_lock);
+        for (auto& stub : _task_map) {
+            task_to_cleanup.push_back(stub.second);
+        }
+    }
+
+    for (auto& task : task_to_cleanup) {
+        task->unschedule(_pipeline_timer);
+    }
+
+    _stub_map.clear();
+    _task_map.clear();
 }
 
 DEFINE_FAIL_POINT(get_stub_return_nullptr);
@@ -197,22 +257,54 @@ StatusOr<std::shared_ptr<starrocks::LakeService_RecoverableStub>> LakeServiceBrp
     }
     // get if exist
     std::lock_guard<SpinLock> l(_lock);
+
+    // schedule clean up task
+    auto task = _task_map.seek(endpoint);
+    if (task == nullptr) {
+        auto new_task = std::make_shared<LakeEndpointCleanupTask>(this, endpoint);
+        _task_map.insert(endpoint, new_task);
+        task = _task_map.seek(endpoint);
+    }
+    _pipeline_timer->unschedule((*task).get());
+    timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s);
+    auto status = _pipeline_timer->schedule((*task).get(), tm);
+    if (!status.ok()) {
+        LOG(WARNING) << "Failed to schedule lake brpc cleanup task: " << endpoint;
+    }
+
     auto stub_ptr = _stub_map.seek(endpoint);
     FAIL_POINT_TRIGGER_EXECUTE(get_stub_return_nullptr, { stub_ptr = nullptr; });
     if (stub_ptr != nullptr) {
         return *stub_ptr;
     }
     // create
     auto stub = std::make_shared<starrocks::LakeService_RecoverableStub>(endpoint, "");
+
     if (!stub->reset_channel().ok()) {
         return Status::RuntimeError("init brpc http channel error on " + host + ":" + std::to_string(port));
     }
     _stub_map.insert(endpoint, stub);
     return stub;
 }
 
+void LakeServiceBrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) {
+    std::lock_guard<SpinLock> l(_lock);
+
+    LOG(INFO) << "cleanup lake service stubs from endpoint:" << endpoint;
+    _stub_map.erase(endpoint);
+    _task_map.erase(endpoint);
+}
+
 void EndpointCleanupTask::Run() {
     _cache->cleanup_expired(_endpoint);
 }
 
+void HttpEndpointCleanupTask::Run() {
+    _cache->cleanup_expired(_endpoint);
+}
+
+void LakeEndpointCleanupTask::Run() {
+    _cache->cleanup_expired(_endpoint);
+}
+
 } // namespace starrocks

be/src/util/brpc_stub_cache.h

@@ -52,6 +52,8 @@ namespace starrocks {
 
 class ExecEnv;
 class EndpointCleanupTask;
+class HttpEndpointCleanupTask;
+class LakeEndpointCleanupTask;
 
 class BrpcStubCache {
 public:
@@ -83,28 +85,36 @@ class HttpBrpcStubCache {
 public:
     static HttpBrpcStubCache* getInstance();
     StatusOr<std::shared_ptr<PInternalService_RecoverableStub>> get_http_stub(const TNetworkAddress& taddr);
+    void cleanup_expired(const butil::EndPoint& endpoint);
 
 private:
     HttpBrpcStubCache();
     HttpBrpcStubCache(const HttpBrpcStubCache&) = delete;
     HttpBrpcStubCache& operator=(const HttpBrpcStubCache&) = delete;
+    ~HttpBrpcStubCache();
 
     SpinLock _lock;
     butil::FlatMap<butil::EndPoint, std::shared_ptr<PInternalService_RecoverableStub>> _stub_map;
+    butil::FlatMap<butil::EndPoint, std::shared_ptr<HttpEndpointCleanupTask>> _task_map;
+    pipeline::PipelineTimer* _pipeline_timer;
 };
 
 class LakeServiceBrpcStubCache {
 public:
     static LakeServiceBrpcStubCache* getInstance();
     StatusOr<std::shared_ptr<starrocks::LakeService_RecoverableStub>> get_stub(const std::string& host, int port);
+    void cleanup_expired(const butil::EndPoint& endpoint);
 
 private:
     LakeServiceBrpcStubCache();
     LakeServiceBrpcStubCache(const LakeServiceBrpcStubCache&) = delete;
     LakeServiceBrpcStubCache& operator=(const LakeServiceBrpcStubCache&) = delete;
+    ~LakeServiceBrpcStubCache();
 
     SpinLock _lock;
     butil::FlatMap<butil::EndPoint, std::shared_ptr<LakeService_RecoverableStub>> _stub_map;
+    butil::FlatMap<butil::EndPoint, std::shared_ptr<LakeEndpointCleanupTask>> _task_map;
+    pipeline::PipelineTimer* _pipeline_timer;
 };
 
 class EndpointCleanupTask : public starrocks::pipeline::PipelineTimerTask {
@@ -117,4 +127,26 @@ class EndpointCleanupTask : public starrocks::pipeline::PipelineTimerTask {
     butil::EndPoint _endpoint;
 };
 
+class HttpEndpointCleanupTask : public starrocks::pipeline::PipelineTimerTask {
+public:
+    HttpEndpointCleanupTask(HttpBrpcStubCache* cache, const butil::EndPoint& endpoint)
+            : _cache(cache), _endpoint(endpoint){};
+    void Run() override;
+
+private:
+    HttpBrpcStubCache* _cache;
+    butil::EndPoint _endpoint;
+};
+
+class LakeEndpointCleanupTask : public starrocks::pipeline::PipelineTimerTask {
+public:
+    LakeEndpointCleanupTask(LakeServiceBrpcStubCache* cache, const butil::EndPoint& endpoint)
+            : _cache(cache), _endpoint(endpoint){};
+    void Run() override;
+
+private:
+    LakeServiceBrpcStubCache* _cache;
+    butil::EndPoint _endpoint;
+};
+
 } // namespace starrocks

be/test/util/brpc_stub_cache_test.cpp

@@ -138,4 +138,37 @@ TEST_F(BrpcStubCacheTest, test_cleanup) {
     ASSERT_NE(stub3, stub1);
 }
 
+TEST_F(BrpcStubCacheTest, test_lake_cleanup) {
+    config::brpc_stub_expire_s = 1;
+    LakeServiceBrpcStubCache cache;
+    std::string hostname = "127.0.0.1";
+    int32_t port = 123;
+    auto stub1 = cache.get_stub(hostname, port);
+    ASSERT_TRUE(stub1.ok());
+    ASSERT_NE(nullptr, *stub1);
+    auto stub2 = cache.get_stub(hostname, port);
+    ASSERT_TRUE(stub1.ok());
+    ASSERT_EQ(*stub2, *stub1);
+
+    sleep(2);
+    auto stub3 = cache.get_stub(hostname, port);
+    ASSERT_NE(*stub3, *stub1);
+}
+
+TEST_F(BrpcStubCacheTest, test_http_cleanup) {
+    config::brpc_stub_expire_s = 1;
+    HttpBrpcStubCache cache;
+    TNetworkAddress address;
+    address.hostname = "127.0.0.1";
+    address.port = 123;
+    auto stub1 = cache.get_http_stub(address);
+    ASSERT_NE(nullptr, *stub1);
+    auto stub2 = cache.get_http_stub(address);
+    ASSERT_EQ(*stub2, *stub1);
+
+    sleep(2);
+    auto stub3 = cache.get_http_stub(address);
+    ASSERT_NE(*stub3, *stub1);
+}
+
 } // namespace starrocks

docs/en/administration/management/BE_configuration.md

@@ -256,6 +256,15 @@ curl http://<BE_IP>:<BE_HTTP_PORT>/varz
 - Description: The maximum body size of a bRPC.
 - Introduced in: -
 
+##### brpc_stub_expire_s
+
+- Default: 3600
+- Type: Int
+- Unit: Seconds
+- Is mutable: Yes
+- Description: The expire time of BRPC stub cache, default 60 minutes.
+- Introduced in: -
+
 <!--
 ##### brpc_socket_max_unwritten_bytes