diff --git a/be/src/util/brpc_stub_cache.cpp b/be/src/util/brpc_stub_cache.cpp index 031a95c7ceaadd..c3c4d161b825d2 100644 --- a/be/src/util/brpc_stub_cache.cpp +++ b/be/src/util/brpc_stub_cache.cpp @@ -25,6 +25,32 @@ namespace starrocks { +namespace { + +inline int64_t absolute_deadline_us(int64_t ttl_seconds) { + return butil::gettimeofday_us() + ttl_seconds * 1000 * 1000; +} + +} // namespace + +template +void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract) { + std::vector>> tasks; + { + std::lock_guard l(cache->_lock); + cache->_stopping = true; + for (auto& stub : cache->_stub_map) { + tasks.push_back(extract(stub.second)); + } + cache->_stub_map.clear(); + } + if (cache->_pipeline_timer != nullptr) { + for (auto& task : tasks) { + (void)cache->_pipeline_timer->unschedule(task.get()); + } + } +} + BrpcStubCache::BrpcStubCache(ExecEnv* exec_env) : _pipeline_timer(exec_env->pipeline_timer()) { _stub_map.init(239); REGISTER_GAUGE_STARROCKS_METRIC(brpc_endpoint_stub_count, [this]() { @@ -34,20 +60,17 @@ BrpcStubCache::BrpcStubCache(ExecEnv* exec_env) : _pipeline_timer(exec_env->pipe } BrpcStubCache::~BrpcStubCache() { - std::vector> pools_to_cleanup; - { - std::lock_guard l(_lock); - - for (auto& stub : _stub_map) { - pools_to_cleanup.push_back(stub.second); - } - } + wait_clean_tasks_terminate(this, [](const std::shared_ptr& pool) { return pool->_cleanup_task; }); +} - for (auto& pool : pools_to_cleanup) { - (void)_pipeline_timer->unschedule(pool->_cleanup_task.get()); +bool BrpcStubCache::replace_cleanup_task_locked(const butil::EndPoint& endpoint, + std::shared_ptr> task) { + auto pool = _stub_map.seek(endpoint); + if (pool != nullptr) { + (*pool)->_cleanup_task = std::move(task); + return true; } - std::lock_guard l(_lock); - _stub_map.clear(); + return false; } std::shared_ptr BrpcStubCache::get_stub(const butil::EndPoint& endpoint) { @@ -56,17 +79,21 @@ std::shared_ptr BrpcStubCache::get_stub(const auto stub_pool = _stub_map.seek(endpoint); if (stub_pool == nullptr) { auto new_pool = std::make_shared(); - new_pool->_cleanup_task = std::make_shared>(this, endpoint); + new_pool->_cleanup_task = + std::make_shared>(this, endpoint, config::brpc_stub_expire_s); + new_pool->_cleanup_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); _stub_map.insert(endpoint, new_pool); stub_pool = _stub_map.seek(endpoint); - } - if (_pipeline_timer->unschedule((*stub_pool)->_cleanup_task.get()) != TIMER_TASK_RUNNING) { - timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s); + timespec tm = butil::microseconds_to_timespec((*stub_pool)->_cleanup_task->deadline_locked()); auto status = _pipeline_timer->schedule((*stub_pool)->_cleanup_task.get(), tm); if (!status.ok()) { LOG(WARNING) << "Failed to schedule brpc cleanup task: " << endpoint; + _stub_map.erase(endpoint); + return new_pool->get_or_create(endpoint); } + } else { + (*stub_pool)->_cleanup_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); } return (*stub_pool)->get_or_create(endpoint); @@ -96,13 +123,6 @@ std::shared_ptr BrpcStubCache::get_stub(const return get_stub(endpoint); } -void BrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) { - std::lock_guard l(_lock); - - LOG(INFO) << "cleanup brpc stub, endpoint:" << endpoint; - _stub_map.erase(endpoint); -} - BrpcStubCache::StubPool::StubPool() : _idx(-1) { _stubs.reserve(config::brpc_max_connections_per_server); } @@ -143,25 +163,17 @@ HttpBrpcStubCache::~HttpBrpcStubCache() { } void HttpBrpcStubCache::shutdown() { - std::vector>> task_to_cleanup; - - { - std::lock_guard l(_lock); - for (auto& stub : _stub_map) { - task_to_cleanup.push_back(stub.second.second); - } - } + wait_clean_tasks_terminate(this, [](const StubEntry& entry) { return entry.cleanup_task; }); +} - if (_pipeline_timer != nullptr) { - for (auto& task : task_to_cleanup) { - _pipeline_timer->unschedule(task.get()); - } - } - - { - std::lock_guard l(_lock); - _stub_map.clear(); +bool HttpBrpcStubCache::replace_cleanup_task_locked(const butil::EndPoint& endpoint, + std::shared_ptr> task) { + auto entry = _stub_map.seek(endpoint); + if (entry != nullptr) { + entry->cleanup_task = std::move(task); + return true; } + return false; } StatusOr> HttpBrpcStubCache::get_http_stub( @@ -187,33 +199,29 @@ StatusOr> HttpBrpcStubCache::g auto stub_pair_ptr = _stub_map.seek(endpoint); if (stub_pair_ptr == nullptr) { // create - auto new_task = std::make_shared>(this, endpoint); + auto new_task = + std::make_shared>(this, endpoint, config::brpc_stub_expire_s); auto stub = std::make_shared(endpoint, "http"); if (!stub->reset_channel().ok()) { return Status::RuntimeError("init http brpc channel error on " + taddr.hostname + ":" + std::to_string(taddr.port)); } - _stub_map.insert(endpoint, std::make_pair(stub, new_task)); + new_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); + _stub_map.insert(endpoint, StubEntry{stub, new_task}); stub_pair_ptr = _stub_map.seek(endpoint); - } - // schedule clean up task - if (_pipeline_timer->unschedule((*stub_pair_ptr).second.get()) != TIMER_TASK_RUNNING) { - timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s); - auto status = _pipeline_timer->schedule((*stub_pair_ptr).second.get(), tm); + timespec tm = butil::microseconds_to_timespec(stub_pair_ptr->cleanup_task->deadline_locked()); + auto status = _pipeline_timer->schedule(stub_pair_ptr->cleanup_task.get(), tm); if (!status.ok()) { - LOG(WARNING) << "Failed to schedule http brpc cleanup task: " << endpoint; + LOG(WARNING) << "Failed to schedule brpc cleanup task: " << endpoint; + _stub_map.erase(endpoint); + return stub; } + } else { + stub_pair_ptr->cleanup_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); } - return (*stub_pair_ptr).first; -} - -void HttpBrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) { - std::lock_guard l(_lock); - - LOG(INFO) << "cleanup http brpc stub, endpoint:" << endpoint; - _stub_map.erase(endpoint); + return stub_pair_ptr->stub; } #ifndef __APPLE__ @@ -233,25 +241,17 @@ LakeServiceBrpcStubCache::~LakeServiceBrpcStubCache() { } void LakeServiceBrpcStubCache::shutdown() { - std::vector>> task_to_cleanup; - - { - std::lock_guard l(_lock); - for (auto& stub : _stub_map) { - task_to_cleanup.push_back(stub.second.second); - } - } - - if (_pipeline_timer != nullptr) { - for (auto& task : task_to_cleanup) { - _pipeline_timer->unschedule(task.get()); - } - } + wait_clean_tasks_terminate(this, [](const StubEntry& entry) { return entry.cleanup_task; }); +} - { - std::lock_guard l(_lock); - _stub_map.clear(); +bool LakeServiceBrpcStubCache::replace_cleanup_task_locked( + const butil::EndPoint& endpoint, std::shared_ptr> task) { + auto entry = _stub_map.seek(endpoint); + if (entry != nullptr) { + entry->cleanup_task = std::move(task); + return true; } + return false; } DEFINE_FAIL_POINT(get_stub_return_nullptr); @@ -276,32 +276,29 @@ StatusOr> LakeServiceBrp if (stub_pair_ptr == nullptr) { // create auto stub = std::make_shared(endpoint, ""); - auto new_task = std::make_shared>(this, endpoint); + auto new_task = std::make_shared>(this, endpoint, + config::brpc_stub_expire_s); if (!stub->reset_channel().ok()) { return Status::RuntimeError("init lakeService brpc channel error on " + host + ":" + std::to_string(port)); } - _stub_map.insert(endpoint, std::make_pair(stub, new_task)); + new_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); + _stub_map.insert(endpoint, StubEntry{stub, new_task}); stub_pair_ptr = _stub_map.seek(endpoint); - } - // schedule clean up task - if (_pipeline_timer->unschedule((*stub_pair_ptr).second.get()) != TIMER_TASK_RUNNING) { - timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s); - auto status = _pipeline_timer->schedule((*stub_pair_ptr).second.get(), tm); + timespec tm = butil::microseconds_to_timespec(stub_pair_ptr->cleanup_task->deadline_locked()); + auto status = _pipeline_timer->schedule(stub_pair_ptr->cleanup_task.get(), tm); if (!status.ok()) { - LOG(WARNING) << "Failed to schedule lake brpc cleanup task: " << endpoint; + LOG(WARNING) << "Failed to schedule brpc cleanup task: " << endpoint; + _stub_map.erase(endpoint); + return stub; } + } else { + stub_pair_ptr->cleanup_task->renew_deadline_locked(absolute_deadline_us(config::brpc_stub_expire_s)); } - return (*stub_pair_ptr).first; + return stub_pair_ptr->stub; } -void LakeServiceBrpcStubCache::cleanup_expired(const butil::EndPoint& endpoint) { - std::lock_guard l(_lock); - - LOG(INFO) << "cleanup lake service brpc stub, endpoint:" << endpoint; - _stub_map.erase(endpoint); -} #endif -} // namespace starrocks \ No newline at end of file +} // namespace starrocks diff --git a/be/src/util/brpc_stub_cache.h b/be/src/util/brpc_stub_cache.h index 0405f1e244edb9..3cba9129ef08e8 100644 --- a/be/src/util/brpc_stub_cache.h +++ b/be/src/util/brpc_stub_cache.h @@ -59,12 +59,58 @@ class ExecEnv; template class EndpointCleanupTask : public starrocks::pipeline::LightTimerTask { public: - EndpointCleanupTask(StubCacheT* cache, const butil::EndPoint& endpoint) : _cache(cache), _endpoint(endpoint){}; - void Run() override { _cache->cleanup_expired(_endpoint); } + // ttl_seconds is the cache-wide expire window (config::brpc_stub_expire_s). + EndpointCleanupTask(StubCacheT* cache, const butil::EndPoint& endpoint, int64_t ttl_seconds) + : _cache(cache), _endpoint(endpoint), _ttl_seconds(ttl_seconds) {} + // The actual cleanup/renewal decision must run while the cache's _lock is held so + // that _stopping and _deadline are observed atomically with the cache state. + void Run() override { + std::lock_guard l(_cache->_lock); + if (_cache->_stopping) { + return; + } + // Make sure this task is still the authoritative cleanup task for the + // endpoint before rescheduling. If shutdown() cleared the cache or a new + // entry was created for the same endpoint, this task is stale and must + // not schedule anything. + if (!_cache->is_cleanup_task_owner_locked(_endpoint, this)) { + return; + } + int64_t now_us = butil::gettimeofday_us(); + if (now_us >= _deadline) { + LOG(INFO) << "cleanup brpc stub, endpoint:" << _endpoint << ", idle for " + << (now_us - _deadline) / 1000 << " ms past deadline"; + _cache->_stub_map.erase(_endpoint); + return; + } + auto new_task = std::make_shared>(_cache, _endpoint, _ttl_seconds); + new_task->_deadline = _deadline; + if (!_cache->replace_cleanup_task_locked(_endpoint, new_task)) { + return; + } + timespec tm = butil::microseconds_to_timespec(_deadline); + auto status = _cache->_pipeline_timer->schedule(new_task.get(), tm); + if (!status.ok()) { + LOG(WARNING) << "Failed to reschedule brpc cleanup task: " << _endpoint; + // Drop the entry; the next get_*_stub() will recreate it with a fresh task. + _cache->_stub_map.erase(_endpoint); + } + } + + // Reset the absolute deadline (in butil::gettimeofday_us() units) used by the next + // Run() invocation to decide between evict and reschedule. Caller must hold the + // cache lock. + void renew_deadline_locked(int64_t new_deadline) { _deadline = new_deadline; } + int64_t deadline_locked() const { return _deadline; } private: StubCacheT* _cache; butil::EndPoint _endpoint; + // Absolute deadline (in butil::gettimeofday_us() units) used to decide whether a + // firing task should evict the stub or simply reschedule itself. Read/written only + // under the cache's _lock, so it does not need to be atomic. + int64_t _deadline{0}; + int64_t _ttl_seconds{0}; }; class BrpcStubCache { @@ -75,9 +121,22 @@ class BrpcStubCache { std::shared_ptr get_stub(const butil::EndPoint& endpoint); std::shared_ptr get_stub(const TNetworkAddress& taddr); std::shared_ptr get_stub(const std::string& host, int port); - void cleanup_expired(const butil::EndPoint& endpoint); private: + friend class EndpointCleanupTask; + + template + friend void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract); + + bool is_cleanup_task_owner_locked(const butil::EndPoint& endpoint, + const EndpointCleanupTask* task) const { + auto pool = _stub_map.seek(endpoint); + return pool != nullptr && (*pool)->_cleanup_task.get() == task; + } + + bool replace_cleanup_task_locked(const butil::EndPoint& endpoint, + std::shared_ptr> task); + struct StubPool { StubPool(); ~StubPool(); @@ -91,13 +150,13 @@ class BrpcStubCache { SpinLock _lock; butil::FlatMap> _stub_map; pipeline::PipelineTimer* _pipeline_timer; + bool _stopping{false}; }; class HttpBrpcStubCache { public: static HttpBrpcStubCache* getInstance(); StatusOr> get_http_stub(const TNetworkAddress& taddr); - void cleanup_expired(const butil::EndPoint& endpoint); void shutdown(); private: @@ -106,11 +165,29 @@ class HttpBrpcStubCache { HttpBrpcStubCache& operator=(const HttpBrpcStubCache&) = delete; ~HttpBrpcStubCache(); + friend class EndpointCleanupTask; + + template + friend void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract); + + bool is_cleanup_task_owner_locked(const butil::EndPoint& endpoint, + const EndpointCleanupTask* task) const { + auto entry = _stub_map.seek(endpoint); + return entry != nullptr && entry->cleanup_task.get() == task; + } + + bool replace_cleanup_task_locked(const butil::EndPoint& endpoint, + std::shared_ptr> task); + + struct StubEntry { + std::shared_ptr stub; + std::shared_ptr> cleanup_task; + }; + SpinLock _lock; - butil::FlatMap, - std::shared_ptr>>> - _stub_map; + butil::FlatMap _stub_map; pipeline::PipelineTimer* _pipeline_timer; + bool _stopping{false}; }; #ifndef __APPLE__ @@ -118,7 +195,6 @@ class LakeServiceBrpcStubCache { public: static LakeServiceBrpcStubCache* getInstance(); StatusOr> get_stub(const std::string& host, int port); - void cleanup_expired(const butil::EndPoint& endpoint); void shutdown(); private: @@ -127,11 +203,29 @@ class LakeServiceBrpcStubCache { LakeServiceBrpcStubCache& operator=(const LakeServiceBrpcStubCache&) = delete; ~LakeServiceBrpcStubCache(); + friend class EndpointCleanupTask; + + template + friend void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract); + + bool is_cleanup_task_owner_locked(const butil::EndPoint& endpoint, + const EndpointCleanupTask* task) const { + auto entry = _stub_map.seek(endpoint); + return entry != nullptr && entry->cleanup_task.get() == task; + } + + bool replace_cleanup_task_locked(const butil::EndPoint& endpoint, + std::shared_ptr> task); + + struct StubEntry { + std::shared_ptr stub; + std::shared_ptr> cleanup_task; + }; + SpinLock _lock; - butil::FlatMap, - std::shared_ptr>>> - _stub_map; + butil::FlatMap _stub_map; pipeline::PipelineTimer* _pipeline_timer; + bool _stopping{false}; }; #endif diff --git a/be/test/util/brpc_stub_cache_test.cpp b/be/test/util/brpc_stub_cache_test.cpp index bf456c3bbbfb8b..63bacfc7f6728f 100644 --- a/be/test/util/brpc_stub_cache_test.cpp +++ b/be/test/util/brpc_stub_cache_test.cpp @@ -171,4 +171,26 @@ TEST_F(BrpcStubCacheTest, test_http_cleanup) { ASSERT_NE(*stub3, *stub1); } +// Regression test for the lazy-reschedule fix: while an endpoint is being +// accessed within the expire window, the timer fires, sees the stub is still +// active (deadline not yet past), and reschedules instead of evicting, so the +// stub must survive across multiple timer periods. +TEST_F(BrpcStubCacheTest, test_active_access_keeps_stub_alive_across_expire_window) { + config::brpc_stub_expire_s = 2; + BrpcStubCache cache(&_env); + TNetworkAddress address; + address.hostname = "127.0.0.1"; + address.port = 123; + auto stub1 = cache.get_stub(address); + ASSERT_NE(nullptr, stub1); + + // Repeatedly access within the window; each access refreshes the deadline, + // and the single scheduled timer reschedules instead of evicting. + for (int i = 0; i < 3; ++i) { + sleep(1); + auto stub = cache.get_stub(address); + ASSERT_EQ(stub1, stub) << "stub must not be evicted while being accessed"; + } +} + } // namespace starrocks