branch-4.1: [fix](workload group) Drain adaptive flush callbacks before pool teardown #68264

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

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<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;
         }
@@ -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<std::shared_mutex> wlock(_task_sched_lock);
+    cancel_adaptive_flush_no_lock();
     _task_sched.reset();
     _scan_task_sched.reset();
     _remote_scan_task_sched.reset();

be/src/runtime/workload_group/workload_group.h

@@ -219,6 +219,8 @@ class WorkloadGroup : public std::enable_shared_from_this<WorkloadGroup> {
     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();
 
@@ -262,6 +264,8 @@ class WorkloadGroup : public std::enable_shared_from_this<WorkloadGroup> {
     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};

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<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);

be/src/storage/adaptive_thread_pool_controller.cpp

@@ -20,13 +20,15 @@
 #include <butil/time.h>
 
 #include <algorithm>
+#include <chrono>
 #include <thread>
 
 #include "cloud/config.h"
 #include "common/config.h"
 #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);

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 @@ 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<ThreadPool*> 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<std::string, PoolGroup> _pool_groups;