From b8307adf8d39693f5342f53f4af4380abc8adde3 Mon Sep 17 00:00:00 2001 From: hui lai Date: Mon, 21 Sep 2026 16:27:33 +0800 Subject: [PATCH] [fix](workload group) Drain adaptive flush callbacks before pool teardown (#68264) ### What problem does this PR solve? Adaptive flush callbacks can outlive the ThreadPool objects they borrow. Workload group ID changes make teardown cancel a different registration name, and timer cancellation can miss a rearmed timer or free its state before an already-started callback acquires its mutex. These races can crash the backend during workload group changes or shutdown. - Save a stable, per-pool registration key and cancel it even when adaptive adjustment has subsequently been disabled. - Read the final timer ID under the callback mutex, then wait for brpc to finish any running callback before releasing its state. - Serialize registration, cancellation and stop; drain duplicate registrations and reject registration after stop. - Stop adaptive callbacks before backend pool teardown, and cancel registrations on direct WG scheduler destruction. ### Release note Fix backend crashes caused by adaptive flush callbacks accessing released thread pools during workload group changes or backend shutdown. --- be/src/runtime/exec_env_init.cpp | 7 + .../runtime/workload_group/workload_group.cpp | 51 +++--- .../runtime/workload_group/workload_group.h | 4 + .../workload_group/workload_group_manager.cpp | 9 +- .../adaptive_thread_pool_controller.cpp | 57 ++++-- .../storage/adaptive_thread_pool_controller.h | 23 +-- .../workload_group_manager_test.cpp | 127 ++++++++++++++ .../adaptive_thread_pool_controller_test.cpp | 162 ++++++++++++++++++ 8 files changed, 393 insertions(+), 47 deletions(-) diff --git a/be/src/runtime/exec_env_init.cpp b/be/src/runtime/exec_env_init.cpp index a3cf2d5db6fc66..c282382e863c03 100644 --- a/be/src/runtime/exec_env_init.cpp +++ b/be/src/runtime/exec_env_init.cpp @@ -101,6 +101,7 @@ #include "service/backend_options.h" #include "service/backend_service.h" #include "service/point_query_executor.h" +#include "storage/adaptive_thread_pool_controller.h" #include "storage/cache/ann_index_ivf_list_cache.h" #include "storage/cache/page_cache.h" #include "storage/id_manager.h" @@ -850,6 +851,12 @@ void ExecEnv::destroy() { // _routine_load_task_executor should be stopped before _new_load_stream_mgr. SAFE_STOP(_routine_load_task_executor); SAFE_STOP(_stream_load_recorder_manager); + // Adaptive callbacks borrow WG/global flush pools and the S3 upload pool. + // Drain them before any of these dependencies can be destroyed. + if (_storage_engine) { + _storage_engine->adaptive_thread_controller()->stop(); + } + // stop workload scheduler SAFE_STOP(_workload_sched_mgr); // Stop workload group execution threads before FragmentMgr. Running pipeline tasks can still diff --git a/be/src/runtime/workload_group/workload_group.cpp b/be/src/runtime/workload_group/workload_group.cpp index e051eda337ed84..28c522bfd701fa 100644 --- a/be/src/runtime/workload_group/workload_group.cpp +++ b/be/src/runtime/workload_group/workload_group.cpp @@ -29,6 +29,7 @@ #include "cloud/config.h" #include "common/config.h" #include "common/logging.h" +#include "cpp/sync_point.h" #include "exec/pipeline/task_queue.h" #include "exec/pipeline/task_scheduler.h" #include "exec/scan/scanner_scheduler.h" @@ -570,10 +571,15 @@ Status WorkloadGroup::upsert_thread_pool_no_lock(WorkloadGroupInfo* wg_info, std::make_unique(pipeline_exec_thread_num, blocking_exec_thread_num, "p_" + wg_name, cg_cpu_ctl_ptr); - Status ret = pipeline_task_scheduler->start(); + Status ret = SYNC_POINT_HOOK_RETURN_VALUE( + pipeline_task_scheduler->start(), + "WorkloadGroup::upsert_thread_pool_no_lock::task_scheduler_start"); if (ret.ok()) { _task_sched = std::move(pipeline_task_scheduler); } else { + // A failed start may leave only some schedulers running. Stop all of + // them before destruction, which requires both schedulers to be shut down. + pipeline_task_scheduler->stop(); upsert_ret = ret; LOG(INFO) << "[upsert wg thread pool] task scheduler start failed, gid= " << wg_id; } @@ -636,17 +642,7 @@ Status WorkloadGroup::upsert_thread_pool_no_lock(WorkloadGroupInfo* wg_info, LOG(INFO) << "[upsert wg thread pool] create " + pool_name + " succ, gid=" << wg_id << ", max thread num=" << max_flush_thread_num << ", min thread num=" << min_flush_thread_num; - // Register the new pool with adaptive thread controller - if (config::enable_adaptive_flush_threads) { - auto* controller = - ExecEnv::GetInstance()->storage_engine().adaptive_thread_controller(); - auto* flush_pool = _memtable_flush_pool.get(); - controller->add("flush_wg_" + std::to_string(_id), {flush_pool}, - AdaptiveThreadPoolController::make_flush_adjust_func(controller, - flush_pool), - config::max_flush_thread_num_per_cpu, - config::min_flush_thread_num_per_cpu); - } + register_adaptive_flush_no_lock(); } else { upsert_ret = ret; LOG(INFO) << "[upsert wg thread pool] create " + pool_name + " failed, gid=" << wg_id; @@ -780,21 +776,36 @@ void WorkloadGroup::stop_schedulers_no_lock() { _remote_scan_task_sched->stop(); } if (_memtable_flush_pool) { - // Unregister from adaptive controller before destroying the pool to avoid UAF: - // the adjustment loop holds raw ThreadPool* pointers and must not access them - // after the pool is gone. - if (config::enable_adaptive_flush_threads) { - auto* controller = - ExecEnv::GetInstance()->storage_engine().adaptive_thread_controller(); - controller->cancel("flush_wg_" + std::to_string(_id)); - } + cancel_adaptive_flush_no_lock(); _memtable_flush_pool->shutdown(); _memtable_flush_pool->wait(); } } +void WorkloadGroup::register_adaptive_flush_no_lock() { + if (config::enable_adaptive_flush_threads) { + auto* controller = ExecEnv::GetInstance()->storage_engine().adaptive_thread_controller(); + auto* flush_pool = _memtable_flush_pool.get(); + _adaptive_flush_key = fmt::format("flush_wg_{}_{}", _id, fmt::ptr(flush_pool)); + controller->add( + _adaptive_flush_key, {flush_pool}, + AdaptiveThreadPoolController::make_flush_adjust_func(controller, flush_pool), + config::max_flush_thread_num_per_cpu, config::min_flush_thread_num_per_cpu); + } +} + +void WorkloadGroup::cancel_adaptive_flush_no_lock() { + if (!_adaptive_flush_key.empty()) { + auto* controller = ExecEnv::GetInstance()->storage_engine().adaptive_thread_controller(); + // A runtime config change must not skip cancellation of an existing registration. + controller->cancel(_adaptive_flush_key); + _adaptive_flush_key.clear(); + } +} + void WorkloadGroup::destroy_schedulers() { std::lock_guard wlock(_task_sched_lock); + cancel_adaptive_flush_no_lock(); _task_sched.reset(); _scan_task_sched.reset(); _remote_scan_task_sched.reset(); diff --git a/be/src/runtime/workload_group/workload_group.h b/be/src/runtime/workload_group/workload_group.h index 2a310bf66a0f32..2263af386cc73c 100644 --- a/be/src/runtime/workload_group/workload_group.h +++ b/be/src/runtime/workload_group/workload_group.h @@ -219,6 +219,8 @@ class WorkloadGroup : public std::enable_shared_from_this { void upsert_cgroup_cpu_ctl_no_lock(WorkloadGroupInfo* wg_info); Status upsert_thread_pool_no_lock(WorkloadGroupInfo* wg_info, std::shared_ptr cg_cpu_ctl_ptr); + void register_adaptive_flush_no_lock(); + void cancel_adaptive_flush_no_lock(); void stop_schedulers_no_lock(); void destroy_schedulers(); @@ -262,6 +264,8 @@ class WorkloadGroup : public std::enable_shared_from_this { std::unique_ptr _scan_task_sched {nullptr}; std::unique_ptr _remote_scan_task_sched {nullptr}; std::unique_ptr _memtable_flush_pool {nullptr}; + // Registration identity must survive normal WG ID changes and ID reuse. + std::string _adaptive_flush_key; std::map> _scan_io_throttle_map; std::shared_ptr _remote_scan_io_throttle {nullptr}; diff --git a/be/src/runtime/workload_group/workload_group_manager.cpp b/be/src/runtime/workload_group/workload_group_manager.cpp index 019620a586d06a..9c50c3c5fb6538 100644 --- a/be/src/runtime/workload_group/workload_group_manager.cpp +++ b/be/src/runtime/workload_group/workload_group_manager.cpp @@ -964,7 +964,14 @@ Status WorkloadGroupMgr::create_internal_wg() { WorkloadGroupInfo wg_info = WorkloadGroupInfo::parse_topic_info(twg_info); auto normal_wg = std::make_shared(wg_info); - RETURN_IF_ERROR(normal_wg->upsert_task_scheduler(&wg_info)); + auto status = normal_wg->upsert_task_scheduler(&wg_info); + if (!status.ok()) { + // A later pool may have started and registered an adaptive callback even + // when an earlier scheduler failed. This WG is not owned by the manager + // yet, so drain its callbacks before the local shared_ptr releases it. + normal_wg->try_stop_schedulers(); + return status; + } { std::lock_guard w_lock(_group_mutex); diff --git a/be/src/storage/adaptive_thread_pool_controller.cpp b/be/src/storage/adaptive_thread_pool_controller.cpp index bcb7b089389e1b..405d8fcf0cc24c 100644 --- a/be/src/storage/adaptive_thread_pool_controller.cpp +++ b/be/src/storage/adaptive_thread_pool_controller.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include "cloud/config.h" @@ -27,6 +28,7 @@ #include "common/logging.h" #include "common/metrics/system_metrics.h" #include "common/status.h" +#include "cpp/sync_point.h" #include "util/threadpool.h" #include "util/time.h" @@ -49,14 +51,13 @@ int AdaptiveThreadPoolController::PoolGroup::get_min_threads() const { void AdaptiveThreadPoolController::_on_timer(void* raw) { auto* arg = static_cast(raw); - // Hold mu for the entire callback (fire + re-registration). - // cancel() acquires mu after bthread_timer_del, so this provides - // cancel-with-wait semantics without a dedicated thread. + TEST_SYNC_POINT("AdaptiveThreadPoolController::callback_entered"); + + // Keep registration and adjustment serialized with cancellation. std::lock_guard lk(arg->mu); if (arg->stopped.load(std::memory_order_acquire)) { - // cancel() set stopped before we took the lock. - // cancel() owns arg and will delete it after taking mu. + // cancel() joins this timer before deleting arg. return; } @@ -66,6 +67,8 @@ void AdaptiveThreadPoolController::_on_timer(void* raw) { return; // cancel() will clean up } + TEST_SYNC_POINT("AdaptiveThreadPoolController::before_rearm"); + // Re-register the next one-shot timer. bthread_timer_t tid; if (bthread_timer_add(&tid, butil::milliseconds_from_now(arg->interval_ms), _on_timer, arg) == @@ -83,6 +86,8 @@ void AdaptiveThreadPoolController::init(SystemMetrics* system_metrics, } void AdaptiveThreadPoolController::stop() { + std::lock_guard lifecycle_lock(_lifecycle_mutex); + _stopped = true; std::vector names; { std::lock_guard lk(_mutex); @@ -91,13 +96,19 @@ void AdaptiveThreadPoolController::stop() { } } for (const auto& name : names) { - cancel(name); + _cancel(name); } } void AdaptiveThreadPoolController::add(std::string name, std::vector pools, AdjustFunc adjust_func, double max_threads_per_cpu, double min_threads_per_cpu, int64_t interval_ms) { + std::lock_guard lifecycle_lock(_lifecycle_mutex); + if (_stopped) { + return; + } + _cancel(name); + PoolGroup group; group.name = name; group.pools = std::move(pools); @@ -114,6 +125,8 @@ void AdaptiveThreadPoolController::add(std::string name, std::vectorname = name; arg->interval_ms = interval_ms; + // Even an immediately due timer must not run before its ID and group are published. + std::lock_guard timer_lock(arg->mu); bthread_timer_t tid; if (bthread_timer_add(&tid, butil::milliseconds_from_now(interval_ms), _on_timer, arg) == 0) { arg->timer_id.store(tid, std::memory_order_release); @@ -133,6 +146,11 @@ void AdaptiveThreadPoolController::add(std::string name, std::vector lifecycle_lock(_lifecycle_mutex); + _cancel(name); +} + +void AdaptiveThreadPoolController::_cancel(const std::string& name) { TimerArg* arg = nullptr; { std::lock_guard lk(_mutex); @@ -149,18 +167,24 @@ void AdaptiveThreadPoolController::cancel(const std::string& name) { // Signal the callback to stop re-registering. arg->stopped.store(true, std::memory_order_release); + // Once per removed registration, after stopped is visible. Cancelling a + // missing registration must not reach this synchronization point. + TEST_SYNC_POINT("AdaptiveThreadPoolController::cancel_stopped"); - // Try to cancel a pending (not yet fired) timer. Read timer_id after - // setting stopped so any re-registration in a concurrent callback has - // already stored the latest id by now (it holds mu, which we haven't - // taken yet). - bthread_timer_t tid = arg->timer_id.load(std::memory_order_acquire); - bthread_timer_del(tid); // returns non-zero if already fired; that's fine + // A callback may have passed its stopped check and still be re-registering. + // Take mu before reading the final ID, rather than cancelling a stale ID. + bthread_timer_t tid; + { + std::lock_guard lk(arg->mu); + tid = arg->timer_id.load(std::memory_order_acquire); + } - // Wait for any in-flight callback to finish. The callback holds mu while - // running _fire_group and re-registering, so acquiring mu here ensures - // we don't free arg while the callback is still executing. - { std::lock_guard lk(arg->mu); } + // The timer can already be running without having acquired mu. Joining via + // brpc's running state covers that window too. Do not hold mu while waiting: + // such a callback must acquire it, observe stopped and return. + while (tid != 0 && bthread_timer_del(tid) == 1) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } delete arg; LOG(INFO) << "Adaptive: cancelled pool group '" << name << "'"; @@ -198,6 +222,7 @@ void AdaptiveThreadPoolController::_fire_group(const std::string& name) { // Fire all groups once regardless of schedule. For testing. void AdaptiveThreadPoolController::adjust_once() { + std::lock_guard lifecycle_lock(_lifecycle_mutex); std::vector names; { std::lock_guard lk(_mutex); diff --git a/be/src/storage/adaptive_thread_pool_controller.h b/be/src/storage/adaptive_thread_pool_controller.h index ae78b7883dd92f..792259fa98b01f 100644 --- a/be/src/storage/adaptive_thread_pool_controller.h +++ b/be/src/storage/adaptive_thread_pool_controller.h @@ -40,18 +40,14 @@ struct TimerArg { std::string name; int64_t interval_ms; - // Set by cancel() before calling bthread_timer_del. The callback checks - // this flag after acquiring `mu` and skips re-registration when true. + // Set before cancel() acquires mu, preventing further adjustment/re-registration. std::atomic stopped {false}; - // Tracks the most recently registered timer id. Updated under `mu` by the - // callback after each re-registration; read by cancel() to call - // bthread_timer_del on the latest pending timer. + // Updated and read under mu, including the initial registration in add(). std::atomic timer_id {0}; - // Held for the entire duration of the callback (fire + re-registration). - // cancel() acquires it after bthread_timer_del to wait for any in-flight - // invocation to complete before freeing `this`. + // Serializes initial registration, adjustment, re-registration and cancellation. + // Taking this lock alone does not join a callback that has not acquired it yet. std::mutex mu; }; @@ -88,10 +84,11 @@ class AdaptiveThreadPoolController { // Initialize with system-level dependencies. void init(SystemMetrics* system_metrics, ThreadPool* s3_file_upload_pool); - // Cancel all registered pool groups. Must be called before the pools are destroyed. + // Permanently stop registration and cancel all groups before pools are destroyed. void stop(); - // Register a pool group and start a recurring bthread_timer_add chain. + // Register a timer chain, draining an existing registration with the same name. + // Lifecycle methods must not be called from an AdjustFunc. void add(std::string name, std::vector pools, AdjustFunc adjust_func, double max_threads_per_cpu, double min_threads_per_cpu, int64_t interval_ms = kDefaultIntervalMs); @@ -137,10 +134,16 @@ class AdaptiveThreadPoolController { void _apply_thread_count(PoolGroup& group, int target_threads, const std::string& reason); + // Requires _lifecycle_mutex. + void _cancel(const std::string& name); + private: SystemMetrics* _system_metrics = nullptr; ThreadPool* _s3_file_upload_pool = nullptr; + // Serializes add/cancel/stop so concurrent teardown also waits for cancellation. + std::mutex _lifecycle_mutex; + bool _stopped = false; mutable std::mutex _mutex; mutable std::mutex _metrics_state_mutex; std::map _pool_groups; diff --git a/be/test/runtime/workload_group/workload_group_manager_test.cpp b/be/test/runtime/workload_group/workload_group_manager_test.cpp index 97cbaaf01f8c5b..8a10c619bfd249 100644 --- a/be/test/runtime/workload_group/workload_group_manager_test.cpp +++ b/be/test/runtime/workload_group/workload_group_manager_test.cpp @@ -29,9 +29,11 @@ #include #include #include +#include #include "common/config.h" #include "common/status.h" +#include "cpp/sync_point.h" #include "exec/pipeline/dependency.h" #include "exec/pipeline/pipeline_tracing.h" #include "exec/spill/spill_file_manager.h" @@ -40,10 +42,13 @@ #include "runtime/runtime_query_statistics_mgr.h" #include "runtime/thread_context.h" #include "runtime/workload_group/workload_group.h" +#include "storage/adaptive_thread_pool_controller.h" #include "storage/olap_define.h" +#include "storage/storage_engine.h" #include "testutil/mock/mock_query_task_controller.h" #include "util/defer_op.h" #include "util/mem_info.h" +#include "util/threadpool.h" namespace doris { @@ -1233,4 +1238,126 @@ TEST_F(WorkloadGroupManagerTest, phase4_skips_cancelled_query_memory_exceeded) { live_query->query_mem_tracker()->consume(-1024 * 4); } +TEST_F(WorkloadGroupManagerTest, FailedInternalGroupSetupCancelsAdaptiveFlush) { + auto* env = ExecEnv::GetInstance(); + auto saved_engine = std::move(env->_storage_engine); + const auto saved_config = std::make_tuple( + config::enable_adaptive_flush_threads, config::enable_task_executor_in_internal_table, + config::enable_task_executor_in_external_table, config::pipeline_executor_size, + config::blocking_pipeline_executor_size, config::doris_scanner_thread_pool_thread_num, + config::doris_max_remote_scanner_thread_pool_thread_num, + config::doris_scanner_min_thread_pool_thread_num, config::min_active_scan_threads, + config::min_active_file_scan_threads, config::flush_thread_num_per_store); + Defer restore {[&] { + env->set_storage_engine(std::move(saved_engine)); + std::tie(config::enable_adaptive_flush_threads, + config::enable_task_executor_in_internal_table, + config::enable_task_executor_in_external_table, config::pipeline_executor_size, + config::blocking_pipeline_executor_size, + config::doris_scanner_thread_pool_thread_num, + config::doris_max_remote_scanner_thread_pool_thread_num, + config::doris_scanner_min_thread_pool_thread_num, config::min_active_scan_threads, + config::min_active_file_scan_threads, config::flush_thread_num_per_store) = + saved_config; + }}; + env->set_storage_engine(std::make_unique(EngineOptions {})); + auto* controller = env->storage_engine().adaptive_thread_controller(); + config::enable_adaptive_flush_threads = true; + config::enable_task_executor_in_internal_table = false; + config::enable_task_executor_in_external_table = false; + config::pipeline_executor_size = 1; + config::blocking_pipeline_executor_size = 1; + config::doris_scanner_thread_pool_thread_num = 1; + config::doris_max_remote_scanner_thread_pool_thread_num = 1; + config::doris_scanner_min_thread_pool_thread_num = 1; + config::min_active_scan_threads = 1; + config::min_active_file_scan_threads = 1; + config::flush_thread_num_per_store = 1; + + auto* sp = SyncPoint::get_instance(); + sp->enable_processing(); + Defer disable_sync_points {[&] { sp->disable_processing(); }}; + int failed_starts = 0; + int cancelled_registrations = 0; + SyncPoint::CallbackGuard start_guard; + SyncPoint::CallbackGuard cancel_guard; + sp->set_call_back( + "WorkloadGroup::upsert_thread_pool_no_lock::task_scheduler_start", + [&](auto&& args) { + auto* result = try_any_cast_ret(args); + result->first = Status::InternalError("injected pipeline scheduler failure"); + result->second = true; + ++failed_starts; + }, + &start_guard); + sp->set_call_back( + "AdaptiveThreadPoolController::cancel_stopped", + [&](auto&&) { ++cancelled_registrations; }, &cancel_guard); + + const auto status = _wg_manager->create_internal_wg(); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("injected pipeline scheduler failure"), std::string::npos); + EXPECT_EQ(failed_starts, 1); + EXPECT_TRUE(_wg_manager->_workload_groups.empty()); + // The flush pool was registered despite the earlier scheduler failure, and + // must be cancelled even though the WG never entered the manager's map. + EXPECT_EQ(cancelled_registrations, 1); + { + std::lock_guard lock(controller->_mutex); + EXPECT_TRUE(controller->_pool_groups.empty()); + } + // Also clean up if an assertion above detects a missing cancellation. + controller->stop(); +} + +// Exercise the actual registration/cancellation paths without starting query schedulers. +TEST_F(WorkloadGroupManagerTest, AdaptiveFlushRegistrationSurvivesIdChangeAndReuse) { + auto* env = ExecEnv::GetInstance(); + auto saved_engine = std::move(env->_storage_engine); + const bool saved_adaptive = config::enable_adaptive_flush_threads; + Defer restore {[&] { + env->set_storage_engine(std::move(saved_engine)); + config::enable_adaptive_flush_threads = saved_adaptive; + }}; + env->set_storage_engine(std::make_unique(EngineOptions {})); + auto* controller = env->storage_engine().adaptive_thread_controller(); + config::enable_adaptive_flush_threads = true; + + auto wg = _wg_manager->get_or_create_workload_group({.id = 1, .name = "normal"}); + ASSERT_TRUE(ThreadPoolBuilder("wg_flush_test") + .set_min_threads(1) + .set_max_threads(2) + .build(&wg->_memtable_flush_pool) + .ok()); + wg->register_adaptive_flush_no_lock(); + const auto key = wg->_adaptive_flush_key; + ASSERT_GT(controller->get_current_threads(key), 0); + _wg_manager->reset_workload_group_id("normal", 100); + EXPECT_EQ(wg->id(), 100); + EXPECT_EQ(wg->_adaptive_flush_key, key); + + // A second WG using the original ID must not replace the first registration. + auto reused = _wg_manager->get_or_create_workload_group({.id = 1, .name = "reused"}); + ASSERT_TRUE(ThreadPoolBuilder("wg_flush_reused") + .set_min_threads(1) + .set_max_threads(2) + .build(&reused->_memtable_flush_pool) + .ok()); + reused->register_adaptive_flush_no_lock(); + const auto reused_key = reused->_adaptive_flush_key; + EXPECT_NE(key, reused_key); + + // Disabling adjustment must not disable cleanup of already registered pools. + config::enable_adaptive_flush_threads = false; + wg->try_stop_schedulers(); + wg->destroy_schedulers(); + EXPECT_EQ(controller->get_current_threads(key), 0); + EXPECT_GT(controller->get_current_threads(reused_key), 0); + // Direct scheduler destruction must also drain the registration. + reused->destroy_schedulers(); + EXPECT_EQ(controller->get_current_threads(reused_key), 0); + config::enable_adaptive_flush_threads = true; + controller->adjust_once(); +} + } // namespace doris diff --git a/be/test/storage/adaptive_thread_pool_controller_test.cpp b/be/test/storage/adaptive_thread_pool_controller_test.cpp index 8e3b34841f8e34..1590ff5411d2ba 100644 --- a/be/test/storage/adaptive_thread_pool_controller_test.cpp +++ b/be/test/storage/adaptive_thread_pool_controller_test.cpp @@ -20,12 +20,16 @@ #include #include +#include +#include #include #include "common/config.h" #include "common/metrics/metrics.h" #include "common/metrics/system_metrics.h" +#include "cpp/sync_point.h" #include "testutil/test_util.h" +#include "util/defer_op.h" #include "util/threadpool.h" namespace doris { @@ -68,6 +72,71 @@ class AdaptiveThreadPoolControllerTest : public testing::Test { if (_pool2) _pool2->shutdown(); } + void check_cancel_race(const std::string& point) { + config::enable_adaptive_flush_threads = true; + auto* sp = SyncPoint::get_instance(); + sp->enable_processing(); + Defer disable_sync_points {[&] { sp->disable_processing(); }}; + std::promise entered; + std::promise release; + std::promise cancelling; + auto entered_future = entered.get_future(); + auto release_future = release.get_future().share(); + auto cancelling_future = cancelling.get_future(); + std::atomic entered_calls {0}; + std::atomic cancellations {0}; + // Remove callbacks before destroying the state they capture. + SyncPoint::CallbackGuard entered_guard; + SyncPoint::CallbackGuard cancelling_guard; + sp->set_call_back( + point, + [&](auto&&) { + // A timeout can release the callback before cancellation starts. + // Report repeated entries through the count instead of throwing. + if (entered_calls.fetch_add(1) == 0) { + entered.set_value(); + } + release_future.wait(); + }, + &entered_guard); + sp->set_call_back( + "AdaptiveThreadPoolController::cancel_stopped", + [&](auto&&) { + if (cancellations.fetch_add(1) == 0) { + cancelling.set_value(); + } + }, + &cancelling_guard); + + AdaptiveThreadPoolController controller; + controller.add( + "race", {_pool.get()}, + AdaptiveThreadPoolController::make_flush_adjust_func(&controller, _pool.get()), 4, + 0.5, 1); + // Always release the callback before joining/stopping, including on test failure. + if (entered_future.wait_for(std::chrono::seconds(5)) != std::future_status::ready) { + release.set_value(); + controller.stop(); + FAIL() << "Timer did not reach " << point; + } + auto cancelled = std::async(std::launch::async, [&] { controller.cancel("race"); }); + auto cancelling_status = cancelling_future.wait_for(std::chrono::seconds(5)); + EXPECT_EQ(cancelling_status, std::future_status::ready); + EXPECT_EQ(cancelled.wait_for(std::chrono::milliseconds(20)), std::future_status::timeout); + auto second_cancel = std::async(std::launch::async, [&] { controller.cancel("race"); }); + EXPECT_EQ(second_cancel.wait_for(std::chrono::milliseconds(20)), + std::future_status::timeout); + release.set_value(); + cancelled.get(); + second_cancel.get(); + EXPECT_EQ(entered_calls.load(), 1); + EXPECT_EQ(cancellations.load(), 1); + EXPECT_EQ(controller.get_current_threads("race"), 0); + _pool.reset(); + // A timer rearmed after the final stopped check must have been cancelled too. + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + bool _original_enable_adaptive; std::unique_ptr _pool; std::unique_ptr _pool2; @@ -354,4 +423,97 @@ TEST_F(AdaptiveThreadPoolControllerTest, TestCancel) { EXPECT_EQ(controller.get_current_threads("test"), 0); } +TEST_F(AdaptiveThreadPoolControllerTest, CancelJoinsCallbackBeforeLock) { + check_cancel_race("AdaptiveThreadPoolController::callback_entered"); +} + +TEST_F(AdaptiveThreadPoolControllerTest, CancelJoinsRearmingCallback) { + check_cancel_race("AdaptiveThreadPoolController::before_rearm"); +} + +TEST_F(AdaptiveThreadPoolControllerTest, StopRejectsNewRegistrations) { + AdaptiveThreadPoolController controller; + controller.stop(); + controller.add("late", {_pool.get()}, + AdaptiveThreadPoolController::make_flush_adjust_func(&controller, _pool.get()), + 4, 0.5, 1); + EXPECT_EQ(controller.get_current_threads("late"), 0); +} + +TEST_F(AdaptiveThreadPoolControllerTest, ReplacingRegistrationDrainsOldTimer) { + config::enable_adaptive_flush_threads = true; + auto* sp = SyncPoint::get_instance(); + sp->enable_processing(); + Defer disable_sync_points {[&] { sp->disable_processing(); }}; + std::promise entered; + std::promise release; + std::promise cancelling; + auto entered_future = entered.get_future(); + auto release_future = release.get_future().share(); + auto cancelling_future = cancelling.get_future(); + std::atomic old_calls {0}; + std::atomic new_calls {0}; + std::atomic old_callback_finished {false}; + std::atomic cancellations {0}; + SyncPoint::CallbackGuard cancelling_guard; + sp->set_call_back( + "AdaptiveThreadPoolController::cancel_stopped", + [&](auto&&) { + if (cancellations.fetch_add(1) == 0) { + cancelling.set_value(); + } + }, + &cancelling_guard); + + AdaptiveThreadPoolController controller; + controller.add( + "same", {_pool.get()}, + [&, pool = _pool.get()](int current, int, int, std::string&) { + if (old_calls.fetch_add(1) == 0) { + entered.set_value(); + } + release_future.wait(); + EXPECT_EQ(pool->get_queue_size(), 0); + old_callback_finished.store(true); + return current; + }, + 4, 0.5, 1); + if (entered_future.wait_for(std::chrono::seconds(5)) != std::future_status::ready) { + release.set_value(); + controller.stop(); + FAIL() << "Old timer did not enter its adjustment callback"; + } + + auto replaced = std::async(std::launch::async, [&] { + controller.add( + "same", {_pool2.get()}, + [&](int, int min_t, int, std::string&) { + new_calls.fetch_add(1); + return min_t; + }, + 4, 0.5, 60000); + EXPECT_TRUE(old_callback_finished.load()) + << "Replacement returned before the old callback finished"; + }); + // Wait for replacement to actually start cancellation, not merely for its + // worker to be scheduled. The old AdjustFunc stays blocked until released. + EXPECT_EQ(cancelling_future.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_EQ(replaced.wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); + EXPECT_FALSE(old_callback_finished.load()); + release.set_value(); + replaced.get(); + + EXPECT_TRUE(old_callback_finished.load()); + EXPECT_EQ(old_calls.load(), 1); + EXPECT_EQ(new_calls.load(), 0); + _pool.reset(); + controller.adjust_once(); + EXPECT_EQ(new_calls.load(), 1); + EXPECT_EQ(old_calls.load(), 1); + controller.cancel("same"); + // Check registration removal; the barriers above verify callback completion. + EXPECT_EQ(controller.get_current_threads("same"), 0); + _pool2.reset(); +} + } // namespace doris