*: Fix bug that DynamicThreadPool may not destruct tasks in time. (#4210)

dbms/src/Common/DynamicThreadPool.cpp

@@ -77,6 +77,12 @@ void DynamicThreadPool::scheduledToNewDynamicThread(TaskPtr & task)
     t.detach();
 }
 
+void DynamicThreadPool::executeTask(TaskPtr & task)
+{
+    task->execute();
+    task.reset();
+}
+
 void DynamicThreadPool::fixedWork(size_t index)
 {
     Queue * queue = fixed_queues[index].get();
@@ -87,13 +93,13 @@ void DynamicThreadPool::fixedWork(size_t index)
         queue->pop(task);
         if (!task)
             break;
-        task->execute();
+        executeTask(task);
     }
 }
 
 void DynamicThreadPool::dynamicWork(TaskPtr initial_task)
 {
-    initial_task->execute();
+    executeTask(initial_task);
 
     DynamicNode node;
     while (true)
@@ -110,8 +116,7 @@ void DynamicThreadPool::dynamicWork(TaskPtr initial_task)
 
         if (!node.task) // may be timeout or cancelled
             break;
-        node.task->execute();
-        node.task.reset();
+        executeTask(node.task);
     }
     alive_dynamic_threads.fetch_sub(1);
 }

dbms/src/Common/DynamicThreadPool.h

@@ -64,6 +64,8 @@ class DynamicThreadPool
     void fixedWork(size_t index);
     void dynamicWork(TaskPtr initial_task);
 
+    static void executeTask(TaskPtr & task);
+
     const std::chrono::nanoseconds dynamic_auto_shrink_cooldown;
 
     std::vector<std::thread> fixed_threads;

dbms/src/Common/tests/gtest_dynamic_thread_pool.cpp

@@ -3,6 +3,8 @@
 
 namespace DB::tests
 {
+namespace
+{
 class DynamicThreadPoolTest : public ::testing::Test
 {
 };
@@ -158,4 +160,45 @@ try
 }
 CATCH
 
+struct X
+{
+    std::mutex * mu;
+    std::condition_variable * cv;
+    bool * destructed;
+
+    X(std::mutex * mu_, std::condition_variable * cv_, bool * destructed_)
+        : mu(mu_)
+        , cv(cv_)
+        , destructed(destructed_)
+    {}
+
+    ~X()
+    {
+        std::unique_lock lock(*mu);
+        *destructed = true;
+        cv->notify_all();
+    }
+};
+
+TEST_F(DynamicThreadPoolTest, testTaskDestruct)
+try
+{
+    std::mutex mu;
+    std::condition_variable cv;
+    bool destructed = false;
+
+    DynamicThreadPool pool(0, std::chrono::minutes(1));
+    auto tmp = std::make_shared<X>(&mu, &cv, &destructed);
+    pool.schedule(true, [x = tmp] {});
+    tmp.reset();
+
+    {
+        std::unique_lock lock(mu);
+        auto ret = cv.wait_for(lock, std::chrono::seconds(1), [&] { return destructed; });
+        ASSERT_TRUE(ret);
+    }
+}
+CATCH
+
+} // namespace
 } // namespace DB::tests