Fixed undefined behaviour in ThreadPool

dbms/src/Common/ThreadPool.h

@@ -10,6 +10,7 @@
 #include <optional>
 #include <ext/singleton.h>
 
+#include <Poco/Event.h>
 #include <Common/ThreadStatus.h>
 
 
@@ -133,18 +134,19 @@ class ThreadFromGlobalPool
 
     template <typename Function, typename... Args>
     explicit ThreadFromGlobalPool(Function && func, Args &&... args)
+        : state(std::make_shared<Poco::Event>())
     {
-        mutex = std::make_shared<std::mutex>();
-
-        /// The function object must be copyable, so we wrap lock_guard in shared_ptr.
+        /// NOTE: If this will throw an exception, the descructor won't be called.
         GlobalThreadPool::instance().scheduleOrThrow([
-            mutex = mutex,
-            lock = std::make_shared<std::lock_guard<std::mutex>>(*mutex),
+            state = state,
             func = std::forward<Function>(func),
             args = std::make_tuple(std::forward<Args>(args)...)]
         {
-            DB::ThreadStatus thread_status;
-            std::apply(func, args);
+            {
+                DB::ThreadStatus thread_status;
+                std::apply(func, args);
+            }
+            state->set();
         });
     }
 
@@ -157,7 +159,7 @@ class ThreadFromGlobalPool
     {
         if (joinable())
             std::terminate();
-        mutex = std::move(rhs.mutex);
+        state = std::move(rhs.state);
         return *this;
     }
 
@@ -171,26 +173,26 @@ class ThreadFromGlobalPool
     {
         if (!joinable())
             std::terminate();
-        {
-            std::lock_guard lock(*mutex);
-        }
-        mutex.reset();
+
+        state->wait();
+        state.reset();
     }
 
     void detach()
     {
         if (!joinable())
             std::terminate();
-        mutex.reset();
+        state.reset();
     }
 
     bool joinable() const
     {
-        return static_cast<bool>(mutex);
+        return state != nullptr;
     }
 
 private:
-    std::shared_ptr<std::mutex> mutex;  /// Object must be moveable.
+    /// The state used in this object and inside the thread job.
+    std::shared_ptr<Poco::Event> state;
 };
 
 

dbms/src/Common/tests/thread_pool_3.cpp

@@ -1,4 +1,4 @@
-#include <atomic>
+#include <mutex>
 #include <iostream>
 #include <Common/ThreadPool.h>
 
@@ -10,8 +10,9 @@ void test()
 {
     Pool pool(10, 2, 10);
 
+    std::mutex mutex;
     for (size_t i = 0; i < 10; ++i)
-        pool.schedule([]{ std::cerr << '.'; });
+        pool.schedule([&]{ std::lock_guard lock(mutex); std::cerr << '.'; });
     pool.wait();
 }