branch-4.0: [fix](be) Fix time-sharing executor queued task count

be/src/vec/exec/executor/time_sharing/time_sharing_task_executor.cpp

@@ -305,7 +305,7 @@ TimeSharingTaskExecutor::~TimeSharingTaskExecutor() {
         }
         {
             std::unique_lock<std::mutex> l(_lock);
-            _tokenless->_entries->remove_all(splits_to_destroy);
+            _remove_queued_splits_unlocked(splits_to_destroy);
         }
     }
 
@@ -450,7 +450,7 @@ Status TimeSharingTaskExecutor::_do_submit(std::shared_ptr<PrioritizedSplitRunne
     DCHECK(state == SplitThreadPoolToken::State::IDLE ||
            state == SplitThreadPoolToken::State::RUNNING);
     split->submit_time_watch().start();
-    _tokenless->_entries->offer(std::move(split));
+    _offer_split_unlocked(std::move(split));
     if (state == SplitThreadPoolToken::State::IDLE) {
         _tokenless->transition(SplitThreadPoolToken::State::RUNNING);
     }
@@ -462,8 +462,6 @@ Status TimeSharingTaskExecutor::_do_submit(std::shared_ptr<PrioritizedSplitRunne
     //    1. If it is a SERIAL token, and there are unsubmitted tasks, submit them to the queue.
     //    2. If it is a CONCURRENT token, and there are still unsubmitted tasks, and the upper limit of concurrency is not reached,
     //       then submitted to the queue.
-    _total_queued_tasks++;
-
     // Wake up an idle thread for this task. Choosing the thread at the front of
     // the list ensures LIFO semantics as idling threads are also added to the front.
     //
@@ -599,7 +597,8 @@ void TimeSharingTaskExecutor::_dispatch_thread() {
                         lock.unlock();
                         l.lock();
                         if (_tokenless->state() == SplitThreadPoolToken::State::RUNNING) {
-                            _tokenless->_entries->offer(split);
+                            split->submit_time_watch().reset();
+                            _offer_split_unlocked(split);
                         }
                         l.unlock();
                     } else {
@@ -615,7 +614,8 @@ void TimeSharingTaskExecutor::_dispatch_thread() {
                                 split->reset_level_priority();
                                 std::unique_lock<std::mutex> l(_lock);
                                 if (_tokenless->state() == SplitThreadPoolToken::State::RUNNING) {
-                                    _tokenless->_entries->offer(split);
+                                    split->submit_time_watch().reset();
+                                    _offer_split_unlocked(split);
                                 }
                             } else {
                                 LOG(WARNING) << "blocked split is failed, split_id: "
@@ -800,7 +800,7 @@ Status TimeSharingTaskExecutor::remove_task(std::shared_ptr<TaskHandle> task_han
         }
         {
             std::unique_lock<std::mutex> l(_lock);
-            _tokenless->_entries->remove_all(splits_to_destroy);
+            _remove_queued_splits_unlocked(splits_to_destroy);
         }
     }
 
@@ -880,6 +880,36 @@ Status TimeSharingTaskExecutor::re_enqueue_split(std::shared_ptr<TaskHandle> tas
     return _do_submit(prioritized_split);
 }
 
+void TimeSharingTaskExecutor::_offer_split_unlocked(std::shared_ptr<PrioritizedSplitRunner> split) {
+    _tokenless->_entries->offer(std::move(split));
+    ++_total_queued_tasks;
+}
+
+void TimeSharingTaskExecutor::_remove_queued_splits_unlocked(
+        const std::vector<std::shared_ptr<PrioritizedSplitRunner>>& splits) {
+    if (splits.empty()) {
+        return;
+    }
+
+    const size_t queue_size_before = _tokenless->_entries->size();
+    _tokenless->_entries->remove_all(splits);
+    const size_t queue_size_after = _tokenless->_entries->size();
+    DCHECK_GE(queue_size_before, queue_size_after);
+
+    const auto removed = static_cast<int>(queue_size_before - queue_size_after);
+    DCHECK_GE(_total_queued_tasks, removed);
+    _total_queued_tasks -= removed;
+
+    if (_tokenless->state() == SplitThreadPoolToken::State::RUNNING &&
+        _tokenless->_active_threads == 0 && _tokenless->_entries->size() == 0) {
+        _tokenless->transition(SplitThreadPoolToken::State::IDLE);
+    }
+
+    if (_total_queued_tasks == 0 && _active_threads == 0) {
+        _idle_cond.notify_all();
+    }
+}
+
 void TimeSharingTaskExecutor::_split_finished(std::shared_ptr<PrioritizedSplitRunner> split,
                                               const Status& status) {
     _completed_splits_per_level[split->priority().level()]++;

be/src/vec/exec/executor/time_sharing/time_sharing_task_executor.h

@@ -295,6 +295,15 @@ class TimeSharingTaskExecutor : public TaskExecutor {
     // // Submits a task to be run via token.
     Status _do_submit(std::shared_ptr<PrioritizedSplitRunner> split);
 
+    // Offer a split to the executor queue and keep _total_queued_tasks consistent.
+    // REQUIRES: _lock is held.
+    void _offer_split_unlocked(std::shared_ptr<PrioritizedSplitRunner> split);
+
+    // Remove queued splits and keep _total_queued_tasks/token state consistent.
+    // REQUIRES: _lock is held.
+    void _remove_queued_splits_unlocked(
+            const std::vector<std::shared_ptr<PrioritizedSplitRunner>>& splits);
+
     //NOTE: not thread safe, caller should keep it thread-safe by using lock
     Status _try_create_thread(int thread_num, std::lock_guard<std::mutex>&);
 

be/test/vec/exec/executor/time_sharing/time_sharing_task_executor_test.cpp

@@ -29,6 +29,7 @@
 #include <string>
 #include <thread>
 
+#include "common/exception.h"
 #include "vec/exec/executor/ticker.h"
 #include "vec/exec/executor/time_sharing/time_sharing_task_handle.h"
 
@@ -291,6 +292,28 @@ class TestingSplitRunner : public SplitRunner {
     ListenableFuture<Void> _completion_future {};
 };
 
+class ThrowingSplitRunner : public SplitRunner {
+public:
+    explicit ThrowingSplitRunner(Status status) : _status(std::move(status)) {}
+
+    Status init() override { return Status::OK(); }
+
+    Result<SharedListenableFuture<Void>> process_for(std::chrono::nanoseconds) override {
+        throw Exception(_status);
+    }
+
+    void close(const Status& status) override {}
+
+    bool is_finished() override { return false; }
+
+    Status finished_status() override { return _status; }
+
+    std::string get_info() const override { return ""; }
+
+private:
+    Status _status;
+};
+
 class QueueOnlySplitRunner : public SplitRunner {
 public:
     Status init() override { return Status::OK(); }
@@ -363,6 +386,58 @@ class TimeSharingTaskExecutorTest : public testing::Test {
     }
 };
 
+TEST_F(TimeSharingTaskExecutorTest, test_remove_task_clears_queued_task_count) {
+    auto ticker = std::make_shared<TestingTicker>();
+
+    TimeSharingTaskExecutor::ThreadConfig thread_config;
+    thread_config.thread_name = "leak_repro";
+    thread_config.workload_group = "normal";
+    thread_config.max_thread_num = 0;
+    thread_config.min_thread_num = 0;
+    thread_config.max_queue_size = 2;
+    TimeSharingTaskExecutor executor(thread_config, 0, 1, 1, ticker);
+    ASSERT_TRUE(executor.init().ok());
+    ASSERT_TRUE(executor.start().ok());
+
+    try {
+        for (int i = 0; i < thread_config.max_queue_size; ++i) {
+            auto task_handle = TEST_TRY(executor.create_task(
+                    TaskId("removed_task_" + std::to_string(i)), []() { return 0.0; }, 1,
+                    std::chrono::milliseconds(1), std::optional<int>(1)));
+            auto split = std::make_shared<QueueOnlySplitRunner>();
+
+            auto enqueue_result = executor.enqueue_splits(task_handle, false, {split});
+            ASSERT_TRUE(enqueue_result.has_value()) << enqueue_result.error();
+            EXPECT_EQ(executor.waiting_splits_size(), 1);
+
+            ASSERT_TRUE(executor.remove_task(task_handle).ok());
+            EXPECT_FALSE(split->is_started());
+            EXPECT_EQ(executor.waiting_splits_size(), 0);
+            EXPECT_EQ(executor.get_queue_size(), 0);
+        }
+
+        EXPECT_EQ(executor.num_active_threads(), 0);
+        EXPECT_EQ(executor.waiting_splits_size(), 0);
+        EXPECT_EQ(executor.get_queue_size(), 0);
+
+        auto task_handle = TEST_TRY(executor.create_task(
+                TaskId("next_task"), []() { return 0.0; }, 1, std::chrono::milliseconds(1),
+                std::optional<int>(1)));
+        auto split = std::make_shared<QueueOnlySplitRunner>();
+
+        auto enqueue_result = executor.enqueue_splits(task_handle, false, {split});
+        ASSERT_TRUE(enqueue_result.has_value()) << enqueue_result.error();
+        EXPECT_EQ(executor.waiting_splits_size(), 1);
+        EXPECT_EQ(executor.get_queue_size(), 1);
+
+        static_cast<void>(executor.remove_task(task_handle));
+    } catch (...) {
+        executor.stop();
+        throw;
+    }
+    executor.stop();
+}
+
 TEST_F(TimeSharingTaskExecutorTest, test_invalid_task_handle_returns_error) {
     auto ticker = std::make_shared<TestingTicker>();