[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

@@ -134,7 +134,6 @@ HttpBrpcStubCache* HttpBrpcStubCache::getInstance() {
 
 HttpBrpcStubCache::HttpBrpcStubCache() {
     _stub_map.init(500);
-    _task_map.init(500);
     _pipeline_timer = ExecEnv::GetInstance()->pipeline_timer();
 }
 
@@ -143,8 +142,8 @@ HttpBrpcStubCache::~HttpBrpcStubCache() {
 
     {
         std::lock_guard<SpinLock> l(_lock);
-        for (auto& stub : _task_map) {
-            task_to_cleanup.push_back(stub.second);
+        for (auto& stub : _stub_map) {
+            task_to_cleanup.push_back(stub.second.second);
         }
     }
 
@@ -153,7 +152,6 @@ HttpBrpcStubCache::~HttpBrpcStubCache() {
     }
 
     _stub_map.clear();
-    _task_map.clear();
 }
 
 StatusOr<std::shared_ptr<PInternalService_RecoverableStub>> HttpBrpcStubCache::get_http_stub(
@@ -176,40 +174,34 @@ 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 stub_pair_ptr = _stub_map.seek(endpoint);
+    if (stub_pair_ptr == nullptr) {
+        // create
         auto new_task = std::make_shared<HttpEndpointCleanupTask>(this, endpoint);
-        _task_map.insert(endpoint, new_task);
-        task = _task_map.seek(endpoint);
+        auto stub = std::make_shared<PInternalService_RecoverableStub>(endpoint, "http");
+        if (!stub->reset_channel().ok()) {
+            return Status::RuntimeError("init brpc http channel error on " + taddr.hostname + ":" +
+                                        std::to_string(taddr.port));
+        }
+        _stub_map.insert(endpoint, std::make_pair(stub, new_task));
+        stub_pair_ptr = _stub_map.seek(endpoint);
     }
-    _pipeline_timer->unschedule((*task).get());
+
+    // schedule clean up task
+    _pipeline_timer->unschedule((*stub_pair_ptr).second.get());
     timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s);
-    auto status = _pipeline_timer->schedule((*task).get(), tm);
+    auto status = _pipeline_timer->schedule((*stub_pair_ptr).second.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;
-    }
-    // create
-    auto stub = std::make_shared<PInternalService_RecoverableStub>(endpoint, "http");
-    if (!stub->reset_channel().ok()) {
-        return Status::RuntimeError("init brpc http channel error on " + taddr.hostname + ":" +
-                                    std::to_string(taddr.port));
-    }
-    _stub_map.insert(endpoint, stub);
-    return stub;
+    return (*stub_pair_ptr).first;
 }
 
 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() {
@@ -219,7 +211,6 @@ LakeServiceBrpcStubCache* LakeServiceBrpcStubCache::getInstance() {
 
 LakeServiceBrpcStubCache::LakeServiceBrpcStubCache() {
     _stub_map.init(500);
-    _task_map.init(500);
     _pipeline_timer = ExecEnv::GetInstance()->pipeline_timer();
 }
 
@@ -228,8 +219,8 @@ LakeServiceBrpcStubCache::~LakeServiceBrpcStubCache() {
 
     {
         std::lock_guard<SpinLock> l(_lock);
-        for (auto& stub : _task_map) {
-            task_to_cleanup.push_back(stub.second);
+        for (auto& stub : _stub_map) {
+            task_to_cleanup.push_back(stub.second.second);
         }
     }
 
@@ -238,7 +229,6 @@ LakeServiceBrpcStubCache::~LakeServiceBrpcStubCache() {
     }
 
     _stub_map.clear();
-    _task_map.clear();
 }
 
 DEFINE_FAIL_POINT(get_stub_return_nullptr);
@@ -258,41 +248,34 @@ 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 stub_pair_ptr = _stub_map.seek(endpoint);
+    FAIL_POINT_TRIGGER_EXECUTE(get_stub_return_nullptr, { stub_pair_ptr = nullptr; });
+    if (stub_pair_ptr == nullptr) {
+        // create
+        auto stub = std::make_shared<starrocks::LakeService_RecoverableStub>(endpoint, "");
         auto new_task = std::make_shared<LakeEndpointCleanupTask>(this, endpoint);
-        _task_map.insert(endpoint, new_task);
-        task = _task_map.seek(endpoint);
+        if (!stub->reset_channel().ok()) {
+            return Status::RuntimeError("init brpc lake channel error on " + host + ":" + std::to_string(port));
+        }
+        _stub_map.insert(endpoint, std::make_pair(stub, new_task));
+        stub_pair_ptr = _stub_map.seek(endpoint);
     }
-    _pipeline_timer->unschedule((*task).get());
+
+    // schedule clean up task
+    _pipeline_timer->unschedule((*stub_pair_ptr).second.get());
     timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s);
-    auto status = _pipeline_timer->schedule((*task).get(), tm);
+    auto status = _pipeline_timer->schedule((*stub_pair_ptr).second.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;
+    return (*stub_pair_ptr).first;
 }
 
 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() {

be/src/util/brpc_stub_cache.h

@@ -94,8 +94,9 @@ class HttpBrpcStubCache {
     ~HttpBrpcStubCache();
 
     SpinLock _lock;
-    butil::FlatMap<butil::EndPoint, std::shared_ptr<PInternalService_RecoverableStub>> _stub_map;
-    butil::FlatMap<butil::EndPoint, std::shared_ptr<HttpEndpointCleanupTask>> _task_map;
+    butil::FlatMap<butil::EndPoint, std::pair<std::shared_ptr<PInternalService_RecoverableStub>,
+                                              std::shared_ptr<HttpEndpointCleanupTask>>>
+            _stub_map;
     pipeline::PipelineTimer* _pipeline_timer;
 };
 
@@ -112,8 +113,9 @@ class LakeServiceBrpcStubCache {
     ~LakeServiceBrpcStubCache();
 
     SpinLock _lock;
-    butil::FlatMap<butil::EndPoint, std::shared_ptr<LakeService_RecoverableStub>> _stub_map;
-    butil::FlatMap<butil::EndPoint, std::shared_ptr<LakeEndpointCleanupTask>> _task_map;
+    butil::FlatMap<butil::EndPoint,
+                   std::pair<std::shared_ptr<LakeService_RecoverableStub>, std::shared_ptr<LakeEndpointCleanupTask>>>
+            _stub_map;
     pipeline::PipelineTimer* _pipeline_timer;
 };