This is an automated email from the ASF dual-hosted git repository.
sollhui pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new bf127c2d9ca [fix](workload group) Drain adaptive flush callbacks
before pool teardown (#68264)
bf127c2d9ca is described below
commit bf127c2d9ca1cdbb70c604320120fea1f1e33002
Author: hui lai <[email protected]>
AuthorDate: Mon Sep 21 16:27:33 2026 +0800
[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 +
be/src/runtime/workload_group/workload_group.cpp | 51 ++++---
be/src/runtime/workload_group/workload_group.h | 4 +
.../workload_group/workload_group_manager.cpp | 9 +-
be/src/storage/adaptive_thread_pool_controller.cpp | 57 ++++++--
be/src/storage/adaptive_thread_pool_controller.h | 23 +--
.../workload_group/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 9ff7cc5f4e6..4833cbbf5eb 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"
@@ -870,6 +871,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 76ca0fbeaaf..ecaa59972e6 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"
@@ -568,10 +569,15 @@ Status
WorkloadGroup::upsert_thread_pool_no_lock(WorkloadGroupInfo* wg_info,
std::make_unique<HybridTaskScheduler>(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;
}
@@ -634,17 +640,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;
@@ -778,21 +774,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<std::shared_mutex> 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 bff74bdeef5..a82097c14c8 100644
--- a/be/src/runtime/workload_group/workload_group.h
+++ b/be/src/runtime/workload_group/workload_group.h
@@ -220,6 +220,8 @@ private:
void upsert_cgroup_cpu_ctl_no_lock(WorkloadGroupInfo* wg_info);
Status upsert_thread_pool_no_lock(WorkloadGroupInfo* wg_info,
std::shared_ptr<CgroupCpuCtl>
cg_cpu_ctl_ptr);
+ void register_adaptive_flush_no_lock();
+ void cancel_adaptive_flush_no_lock();
void stop_schedulers_no_lock();
void destroy_schedulers();
@@ -263,6 +265,8 @@ private:
std::unique_ptr<ScannerScheduler> _scan_task_sched {nullptr};
std::unique_ptr<ScannerScheduler> _remote_scan_task_sched {nullptr};
std::unique_ptr<ThreadPool> _memtable_flush_pool {nullptr};
+ // Registration identity must survive normal WG ID changes and ID reuse.
+ std::string _adaptive_flush_key;
std::map<std::string, std::shared_ptr<IOThrottle>> _scan_io_throttle_map;
std::shared_ptr<IOThrottle> _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 2a283f29972..ed4ce5bd0d0 100644
--- a/be/src/runtime/workload_group/workload_group_manager.cpp
+++ b/be/src/runtime/workload_group/workload_group_manager.cpp
@@ -992,7 +992,14 @@ Status WorkloadGroupMgr::create_internal_wg() {
WorkloadGroupInfo wg_info = WorkloadGroupInfo::parse_topic_info(twg_info);
auto normal_wg = std::make_shared<WorkloadGroup>(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<std::shared_mutex> 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 bcb7b089389..405d8fcf0cc 100644
--- a/be/src/storage/adaptive_thread_pool_controller.cpp
+++ b/be/src/storage/adaptive_thread_pool_controller.cpp
@@ -20,6 +20,7 @@
#include <butil/time.h>
#include <algorithm>
+#include <chrono>
#include <thread>
#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<TimerArg*>(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<std::mutex> 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<std::mutex> lifecycle_lock(_lifecycle_mutex);
+ _stopped = true;
std::vector<std::string> names;
{
std::lock_guard<std::mutex> 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<ThreadPool*> pools,
AdjustFunc adjust_func, double
max_threads_per_cpu,
double min_threads_per_cpu, int64_t
interval_ms) {
+ std::lock_guard<std::mutex> 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::vector<ThreadPool*
arg->name = name;
arg->interval_ms = interval_ms;
+ // Even an immediately due timer must not run before its ID and group are
published.
+ std::lock_guard<std::mutex> 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<ThreadPool*
}
void AdaptiveThreadPoolController::cancel(const std::string& name) {
+ std::lock_guard<std::mutex> lifecycle_lock(_lifecycle_mutex);
+ _cancel(name);
+}
+
+void AdaptiveThreadPoolController::_cancel(const std::string& name) {
TimerArg* arg = nullptr;
{
std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> 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<std::mutex> lifecycle_lock(_lifecycle_mutex);
std::vector<std::string> names;
{
std::lock_guard<std::mutex> lk(_mutex);
diff --git a/be/src/storage/adaptive_thread_pool_controller.h
b/be/src/storage/adaptive_thread_pool_controller.h
index ae78b7883dd..792259fa98b 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<bool> 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<bthread_timer_t> 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 @@ public:
// 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<ThreadPool*> pools, AdjustFunc
adjust_func,
double max_threads_per_cpu, double min_threads_per_cpu,
int64_t interval_ms = kDefaultIntervalMs);
@@ -137,10 +134,16 @@ private:
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<std::string, PoolGroup> _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 bf7bc9e57d1..03a9cedae3d 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 <filesystem>
#include <memory>
#include <sstream>
+#include <tuple>
#include "common/config.h"
#include "common/status.h"
+#include "cpp/sync_point.h"
#include "exec/pipeline/dependency.h"
#include "exec/spill/spill_file_manager.h"
#include "load/memtable/memtable_memory_limiter.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 {
@@ -1231,4 +1236,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<StorageEngine>(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<Status>(args);
+ result->first = Status::InternalError<false>("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<std::mutex> 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<StorageEngine>(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 8e3b34841f8..1590ff5411d 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 <gtest/gtest.h>
#include <algorithm>
+#include <atomic>
+#include <future>
#include <thread>
#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 @@ protected:
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<void> entered;
+ std::promise<void> release;
+ std::promise<void> cancelling;
+ auto entered_future = entered.get_future();
+ auto release_future = release.get_future().share();
+ auto cancelling_future = cancelling.get_future();
+ std::atomic<int> entered_calls {0};
+ std::atomic<int> 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<ThreadPool> _pool;
std::unique_ptr<ThreadPool> _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<void> entered;
+ std::promise<void> release;
+ std::promise<void> cancelling;
+ auto entered_future = entered.get_future();
+ auto release_future = release.get_future().share();
+ auto cancelling_future = cancelling.get_future();
+ std::atomic<int> old_calls {0};
+ std::atomic<int> new_calls {0};
+ std::atomic<bool> old_callback_finished {false};
+ std::atomic<int> 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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
