diff --git a/be/src/common/brpc/brpc_stub_cache.cpp b/be/src/common/brpc/brpc_stub_cache.cpp index db146ef3823bec..344c3488ee4082 100644 --- a/be/src/common/brpc/brpc_stub_cache.cpp +++ b/be/src/common/brpc/brpc_stub_cache.cpp @@ -16,6 +16,7 @@ #include "base/failpoint/fail_point.h" #include "base/metrics.h" +#include "base/time/time.h" #include "common/config_network_fwd.h" #include "gen_cpp/internal_service.pb.h" #ifndef __APPLE__ @@ -40,8 +41,41 @@ std::mutex& singleton_cache_mutex() { return mutex; } +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; + BthreadTimer* timer = nullptr; + { + std::lock_guard l(cache->_lock); + cache->_stopping = true; + timer = cache->_timer; + cache->_timer = nullptr; + for (auto& stub : cache->_stub_map) { + tasks.push_back(extract(stub.second)); + } + cache->_stub_map.clear(); + } + if (timer != nullptr) { + for (auto& task : tasks) { + task->unschedule_and_join(timer); + } + } +} + +template +void reset_state_for_rebind(CacheT* cache, BthreadTimer* timer) { + DCHECK(timer != nullptr); + std::lock_guard l(cache->_lock); + DCHECK(cache->_stub_map.empty() || cache->_timer == timer); + cache->_stopping = false; + cache->_timer = timer; +} + struct BrpcStubCache::Metrics { Metrics(MetricRegistry* metric_registry, BrpcStubCache* cache) : registry(metric_registry), cache(cache) { DCHECK(registry != nullptr); @@ -71,21 +105,17 @@ BrpcStubCache::BrpcStubCache(BthreadTimer* timer, MetricRegistry* metric_registr BrpcStubCache::~BrpcStubCache() { _metrics.reset(); + wait_clean_tasks_terminate(this, [](const std::shared_ptr& pool) { return pool->_cleanup_task; }); +} - std::vector> pools_to_cleanup; - { - std::lock_guard l(_lock); - - for (auto& stub : _stub_map) { - pools_to_cleanup.push_back(stub.second); - } - } - - for (auto& pool : pools_to_cleanup) { - (void)_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) { @@ -94,17 +124,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 (_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 = _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); @@ -134,13 +168,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() { _stubs.reserve(config::brpc_max_connections_per_server); } @@ -190,31 +217,21 @@ HttpBrpcStubCache::~HttpBrpcStubCache() { } void HttpBrpcStubCache::bind_timer(BthreadTimer* timer) { - DCHECK(timer != nullptr); - std::lock_guard l(_lock); - DCHECK(_stub_map.empty() || _timer == timer); - _timer = timer; + reset_state_for_rebind(this, timer); } void HttpBrpcStubCache::shutdown() { - std::vector>> task_to_cleanup; - BthreadTimer* timer = nullptr; - - { - std::lock_guard l(_lock); - timer = _timer; - _timer = nullptr; - for (auto& stub : _stub_map) { - task_to_cleanup.push_back(stub.second.second); - } - _stub_map.clear(); - } + wait_clean_tasks_terminate(this, [](const StubEntry& entry) { return entry.cleanup_task; }); +} - if (timer != nullptr) { - for (auto& task : task_to_cleanup) { - timer->unschedule(task.get()); - } +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( @@ -243,33 +260,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 (_timer->unschedule((*stub_pair_ptr).second.get()) != TIMER_TASK_RUNNING) { - timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s); - auto status = _timer->schedule((*stub_pair_ptr).second.get(), tm); + timespec tm = butil::microseconds_to_timespec(stub_pair_ptr->cleanup_task->deadline_locked()); + auto status = _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__ @@ -298,31 +311,21 @@ LakeServiceBrpcStubCache::~LakeServiceBrpcStubCache() { } void LakeServiceBrpcStubCache::bind_timer(BthreadTimer* timer) { - DCHECK(timer != nullptr); - std::lock_guard l(_lock); - DCHECK(_stub_map.empty() || _timer == timer); - _timer = timer; + reset_state_for_rebind(this, timer); } void LakeServiceBrpcStubCache::shutdown() { - std::vector>> task_to_cleanup; - BthreadTimer* timer = nullptr; - - { - std::lock_guard l(_lock); - timer = _timer; - _timer = nullptr; - for (auto& stub : _stub_map) { - task_to_cleanup.push_back(stub.second.second); - } - _stub_map.clear(); - } + wait_clean_tasks_terminate(this, [](const StubEntry& entry) { return entry.cleanup_task; }); +} - if (timer != nullptr) { - for (auto& task : task_to_cleanup) { - timer->unschedule(task.get()); - } +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); @@ -350,32 +353,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 (_timer->unschedule((*stub_pair_ptr).second.get()) != TIMER_TASK_RUNNING) { - timespec tm = butil::seconds_from_now(config::brpc_stub_expire_s); - auto status = _timer->schedule((*stub_pair_ptr).second.get(), tm); + timespec tm = butil::microseconds_to_timespec(stub_pair_ptr->cleanup_task->deadline_locked()); + auto status = _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 diff --git a/be/src/common/brpc/brpc_stub_cache.h b/be/src/common/brpc/brpc_stub_cache.h index 2ba92f57289e5f..da3833fb6471d8 100644 --- a/be/src/common/brpc/brpc_stub_cache.h +++ b/be/src/common/brpc/brpc_stub_cache.h @@ -41,8 +41,10 @@ #include "base/brpc/brpc.h" #include "base/concurrency/spinlock.h" #include "base/network/network_util.h" +#include "base/time/time.h" #include "common/brpc/internal_service_recoverable_stub.h" #include "common/bthread_timer.h" +#include "common/logging.h" #include "common/statusor.h" #include "gen_cpp/Types_types.h" // TNetworkAddress @@ -57,14 +59,60 @@ class MetricRegistry; constexpr int TIMER_TASK_RUNNING = 1; template -class EndpointCleanupTask : public BthreadLightTimerTask { +class EndpointCleanupTask : public BthreadTimerTask { 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->_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 +123,25 @@ 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); + + template + friend void reset_state_for_rebind(CacheT* cache, BthreadTimer* timer); + + 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 Metrics; struct StubPool { StubPool(); @@ -93,6 +157,7 @@ class BrpcStubCache { butil::FlatMap> _stub_map; BthreadTimer* _timer; std::unique_ptr _metrics; + bool _stopping{false}; }; class HttpBrpcStubCache { @@ -103,19 +168,38 @@ class HttpBrpcStubCache { static void initialize(BthreadTimer* timer); static HttpBrpcStubCache* getInstance(); StatusOr> get_http_stub(const TNetworkAddress& taddr); - void cleanup_expired(const butil::EndPoint& endpoint); void shutdown(); private: explicit HttpBrpcStubCache(BthreadTimer* timer); ~HttpBrpcStubCache(); void bind_timer(BthreadTimer* timer); + friend class EndpointCleanupTask; + + template + friend void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract); + + template + friend void reset_state_for_rebind(CacheT* cache, BthreadTimer* timer); + + 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; BthreadTimer* _timer; + bool _stopping{false}; }; #ifndef __APPLE__ @@ -127,19 +211,38 @@ class LakeServiceBrpcStubCache { static void initialize(BthreadTimer* timer); static LakeServiceBrpcStubCache* getInstance(); StatusOr> get_stub(const std::string& host, int port); - void cleanup_expired(const butil::EndPoint& endpoint); void shutdown(); private: explicit LakeServiceBrpcStubCache(BthreadTimer* timer); ~LakeServiceBrpcStubCache(); void bind_timer(BthreadTimer* timer); + friend class EndpointCleanupTask; + + template + friend void wait_clean_tasks_terminate(CacheT* cache, ExtractFn extract); + + template + friend void reset_state_for_rebind(CacheT* cache, BthreadTimer* timer); + + 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; BthreadTimer* _timer; + bool _stopping{false}; }; #endif diff --git a/be/test/common/brpc/brpc_stub_cache_test.cpp b/be/test/common/brpc/brpc_stub_cache_test.cpp index 4ab56fe524dd65..a7bd080a3ff2e3 100644 --- a/be/test/common/brpc/brpc_stub_cache_test.cpp +++ b/be/test/common/brpc/brpc_stub_cache_test.cpp @@ -181,6 +181,88 @@ TEST_F(BrpcStubCacheTest, test_http_cleanup) { ASSERT_NE(*stub3, *stub1); } +// Regression test: destroying BrpcStubCache while a cleanup task is scheduled +// (and possibly firing) must join the task before the cache state is torn down. +TEST_F(BrpcStubCacheTest, test_destructor_joins_inflight_cleanup_tasks) { + config::brpc_stub_expire_s = 1; + auto cache = std::make_unique(_timer.get()); + TNetworkAddress address; + address.hostname = "127.0.0.1"; + address.port = 123; + auto stub = cache->get_stub(address); + ASSERT_NE(nullptr, stub); + + // Trigger ~BrpcStubCache() while the cleanup task is (or will be) in flight. + // Drain the cache and assert the unique_ptr release returns cleanly. + cache.reset(); + + // Reacquire the endpoint through a fresh cache; the slot must have been + // cleanly torn down without leaking the previous task. + auto cache2 = std::make_unique(_timer.get()); + auto fresh_stub = cache2->get_stub(address); + ASSERT_NE(nullptr, fresh_stub); + cache2.reset(); +} + +// 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 (idle < brpc_stub_expire_s), 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(_timer.get()); + 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 last-access + // time, 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"; + } +} + +TEST_F(BrpcStubCacheTest, test_http_active_access_keeps_stub_alive_across_expire_window) { + config::brpc_stub_expire_s = 2; + HttpBrpcStubCache cache(_timer.get()); + TNetworkAddress address; + address.hostname = "127.0.0.1"; + address.port = 123; + auto stub1 = cache.get_http_stub(address); + ASSERT_NE(nullptr, *stub1); + + for (int i = 0; i < 3; ++i) { + sleep(1); + auto stub = cache.get_http_stub(address); + ASSERT_NE(nullptr, *stub); + ASSERT_EQ(*stub1, *stub) << "http stub must not be evicted while being accessed"; + } +} + +#ifndef __APPLE__ +TEST_F(BrpcStubCacheTest, test_lake_active_access_keeps_stub_alive_across_expire_window) { + config::brpc_stub_expire_s = 2; + LakeServiceBrpcStubCache cache(_timer.get()); + 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); + + for (int i = 0; i < 3; ++i) { + sleep(1); + auto stub = cache.get_stub(hostname, port); + ASSERT_TRUE(stub.ok()); + ASSERT_NE(nullptr, *stub); + ASSERT_EQ(*stub1, *stub) << "lake stub must not be evicted while being accessed"; + } +} +#endif + TEST_F(BrpcStubCacheTest, http_singleton_reinitialize_rebinds_pipeline_timer) { auto timer2 = std::make_unique(); ASSERT_OK(timer2->start()); @@ -188,7 +270,6 @@ TEST_F(BrpcStubCacheTest, http_singleton_reinitialize_rebinds_pipeline_timer) { HttpBrpcStubCache::initialize(_timer.get()); auto* cache = HttpBrpcStubCache::getInstance(); ASSERT_NE(nullptr, cache); - ASSERT_EQ(_timer.get(), cache->_timer); TNetworkAddress address; address.hostname = "127.0.0.1"; @@ -198,10 +279,9 @@ TEST_F(BrpcStubCacheTest, http_singleton_reinitialize_rebinds_pipeline_timer) { ASSERT_NE(nullptr, *stub); cache->shutdown(); - ASSERT_EQ(nullptr, cache->_timer); + ASSERT_FALSE(cache->get_http_stub(address).ok()); HttpBrpcStubCache::initialize(timer2.get()); - ASSERT_EQ(timer2.get(), cache->_timer); auto rebound_stub = cache->get_http_stub(address); ASSERT_TRUE(rebound_stub.ok()); @@ -218,17 +298,15 @@ TEST_F(BrpcStubCacheTest, lake_singleton_reinitialize_rebinds_pipeline_timer) { LakeServiceBrpcStubCache::initialize(_timer.get()); auto* cache = LakeServiceBrpcStubCache::getInstance(); ASSERT_NE(nullptr, cache); - ASSERT_EQ(_timer.get(), cache->_timer); auto stub = cache->get_stub("127.0.0.1", 123); ASSERT_TRUE(stub.ok()); ASSERT_NE(nullptr, *stub); cache->shutdown(); - ASSERT_EQ(nullptr, cache->_timer); + ASSERT_FALSE(cache->get_stub("127.0.0.1", 123).ok()); LakeServiceBrpcStubCache::initialize(timer2.get()); - ASSERT_EQ(timer2.get(), cache->_timer); auto rebound_stub = cache->get_stub("127.0.0.1", 123); ASSERT_TRUE(rebound_stub.ok());