feat: add thread init and destroy callbacks to TC_ThreadPool

Author: ruanshudong

unit-test/util/test_tc_thread_pool.cpp

@@ -55,4 +55,118 @@ TEST_F(UtilThreadPollTest, testDecontructor)
 	}
 
 	tpool.stop();
+}
+
+TEST_F(UtilThreadPollTest, testThreadCallbacks)
+{
+	std::atomic<int> initCount(0);
+	std::atomic<int> destroyCount(0);
+	
+	TC_ThreadPool pool;
+	pool.init(4);
+	
+	pool.setThreadInitCallback([&initCount](){
+		initCount++;
+	});
+	
+	pool.setThreadDestroyCallback([&destroyCount](){
+		destroyCount++;
+	});
+	
+	pool.start();
+	
+	// 执行一些任务
+	for (int i = 0; i < 10; ++i)
+	{
+		pool.exec([](){ 
+			std::this_thread::sleep_for(std::chrono::milliseconds(10));
+		});
+	}
+	
+	pool.waitForAllDone();
+	pool.stop();
+	
+	// 验证回调被调用了正确的次数
+	EXPECT_EQ(initCount.load(), 4);
+	EXPECT_EQ(destroyCount.load(), 4);
+}
+
+TEST_F(UtilThreadPollTest, testThreadLocalStorage)
+{
+	std::mutex mtx;
+	std::set<int> values;
+	
+	TC_ThreadPool pool;
+	pool.init(4);
+	
+	pool.setThreadInitCallback([](){
+		static std::atomic<int> counter(1000);
+		thread_local int g_value = 0;
+		g_value = counter.fetch_add(1);
+	});
+	
+	pool.start();
+	
+	for (int i = 0; i < 100; ++i)
+	{
+		pool.exec([&mtx, &values](){
+			static thread_local int g_value = 0;
+			std::lock_guard<std::mutex> lock(mtx);
+			if (g_value != 0) {
+				values.insert(g_value);
+			}
+		});
+	}
+	
+	pool.waitForAllDone();
+	pool.stop();
+}
+
+TEST_F(UtilThreadPollTest, testCallbackException)
+{
+	std::atomic<int> taskCount(0);
+	
+	TC_ThreadPool pool;
+	pool.init(2);
+	
+	// 初始化回调抛出异常
+	pool.setThreadInitCallback([](){
+		throw std::runtime_error("Init error");
+	});
+	
+	pool.start();
+	
+	// 任务应该仍然能够执行
+	for (int i = 0; i < 10; ++i)
+	{
+		pool.exec([&taskCount](){ 
+			taskCount++;
+		});
+	}
+	
+	pool.waitForAllDone();
+	pool.stop();
+	
+	// 验证任务都被执行了
+	EXPECT_EQ(taskCount.load(), 10);
+}
+
+TEST_F(UtilThreadPollTest, testSetCallbackAfterStart)
+{
+	TC_ThreadPool pool;
+	pool.init(2);
+	pool.start();
+	
+	// 应该抛出异常
+	EXPECT_THROW(
+		pool.setThreadInitCallback([](){}),
+		TC_ThreadPool_Exception
+	);
+	
+	EXPECT_THROW(
+		pool.setThreadDestroyCallback([](){}),
+		TC_ThreadPool_Exception
+	);
+	
+	pool.stop();
 }

util/include/util/tc_thread_pool.h

@@ -23,24 +23,31 @@ namespace tars
  * 使用说明:
  * TC_ThreadPool tpool;
  * tpool.init(5);   //初始化线程池线程数
- * //启动线程有两种方式
- * //第一种, 直接启动
+ * 
+ * //设置线程生命周期回调（可选）
+ * tpool.setThreadInitCallback([](){
+ *     //线程初始化代码，每个工作线程启动时执行一次
+ *     //可以初始化 thread_local 变量
+ * });
+ * tpool.setThreadDestroyCallback([](){
+ *     //线程清理代码，每个工作线程退出前执行一次
+ *     //可以清理 thread_local 资源
+ * });
+ * 
+ * //启动线程
  * tpool.start();
- * //第二种, 启动时指定初始化函数, 比如定义函数
- * void testFunction(int i)
- * {
- *     cout << i << endl;
- * }
- * tpool.start(testFunction, 5);    //start的第一函数是std::bind返回的函数(std::function), 后面跟参数
+ * 
  * //将任务丢到线程池中
  * tpool.exec(testFunction, 10);    //参数和start相同
+ * 
  * //等待线程池结束, 有两种方式:
  * //第一种等待线程池中无任务
  * tpool.waitForAllDone(1000);      //参数<0时, 表示无限等待(注意有人调用stop也会推出)
  * //第二种等待外部有人调用线程池的stop函数
  * tpool.waitForStop(1000);
  * //此时: 外部需要结束线程池是调用
  * tpool.stop();
+ * 
  * 注意:
  * TC_ThreadPool::exec执行任务返回的是个future, 因此可以通过future异步获取结果, 比如:
  * int testInt(int i)
@@ -91,6 +98,18 @@ class UTIL_DLL_API TC_ThreadPool
     };
     typedef shared_ptr<TaskFunc> TaskFuncPtr;
 public:
+    /**
+     * @brief 线程初始化回调函数类型
+     * 在工作线程启动后、执行任务前调用（每个线程仅调用一次）
+     */
+    typedef std::function<void()> ThreadInitCallback;
+    
+    /**
+     * @brief 线程销毁回调函数类型
+     * 在工作线程退出前调用（每个线程仅调用一次）
+     */
+    typedef std::function<void()> ThreadDestroyCallback;
+
     /**
     * @brief 构造函数
     *
@@ -142,6 +161,30 @@ class UTIL_DLL_API TC_ThreadPool
     */
     void start();
 
+    /**
+     * @brief 设置线程初始化回调函数
+     * 
+     * @param callback 初始化回调函数，在每个工作线程启动时调用
+     * 
+     * 注意：
+     * - 必须在 start() 之前调用
+     * - 回调函数在工作线程中执行，可以安全地访问 TLS 变量
+     * - 如果回调函数抛出异常，异常会被捕获并忽略，线程继续运行
+     */
+    void setThreadInitCallback(const ThreadInitCallback& callback);
+    
+    /**
+     * @brief 设置线程销毁回调函数
+     * 
+     * @param callback 销毁回调函数，在每个工作线程退出前调用
+     * 
+     * 注意：
+     * - 必须在 start() 之前调用
+     * - 回调函数在工作线程中执行，可以安全地访问 TLS 变量
+     * - 如果回调函数抛出异常，异常会被捕获并忽略
+     */
+    void setThreadDestroyCallback(const ThreadDestroyCallback& callback);
+
     /**
     * @brief 用线程池启用任务(F是function, Args是参数)
     *
@@ -230,6 +273,10 @@ class UTIL_DLL_API TC_ThreadPool
     bool                      _bTerminate;
 
     std::atomic<int>          _atomic{ 0 };
+
+    ThreadInitCallback        _threadInitCallback;
+    
+    ThreadDestroyCallback     _threadDestroyCallback;
 };
 
 
@@ -260,6 +307,24 @@ class UTIL_DLL_API TC_ThreadPoolHash
         */
     void start();
 
+    /**
+     * @brief 设置所有线程池的初始化回调
+     * 
+     * @param callback 初始化回调函数
+     * 
+     * 注意：必须在 start() 之前调用
+     */
+    void setThreadInitCallback(const TC_ThreadPool::ThreadInitCallback& callback);
+    
+    /**
+     * @brief 设置所有线程池的销毁回调
+     * 
+     * @param callback 销毁回调函数
+     * 
+     * 注意：必须在 start() 之前调用
+     */
+    void setThreadDestroyCallback(const TC_ThreadPool::ThreadDestroyCallback& callback);
+
     /**
     * @brief 用线程池启用任务(F是function, Args是参数)
     *
@@ -302,7 +367,9 @@ class UTIL_DLL_API TC_ThreadPoolHash
 protected:
 private:
     vector<TC_ThreadPool*> _pools;
-
+    
+    TC_ThreadPool::ThreadInitCallback    _threadInitCallback;
+    TC_ThreadPool::ThreadDestroyCallback _threadDestroyCallback;
 };
 
 }

util/src/tc_thread_pool.cpp

@@ -90,6 +90,32 @@ void TC_ThreadPool::start()
     }
 }
 
+void TC_ThreadPool::setThreadInitCallback(const ThreadInitCallback& callback)
+{
+    std::unique_lock<std::mutex> lock(_mutex);
+    
+    if (!_threads.empty())
+    {
+        throw TC_ThreadPool_Exception(
+            "[TC_ThreadPool::setThreadInitCallback] cannot set callback after thread pool started!");
+    }
+    
+    _threadInitCallback = callback;
+}
+
+void TC_ThreadPool::setThreadDestroyCallback(const ThreadDestroyCallback& callback)
+{
+    std::unique_lock<std::mutex> lock(_mutex);
+    
+    if (!_threads.empty())
+    {
+        throw TC_ThreadPool_Exception(
+            "[TC_ThreadPool::setThreadDestroyCallback] cannot set callback after thread pool started!");
+    }
+    
+    _threadDestroyCallback = callback;
+}
+
 bool TC_ThreadPool::get(TaskFuncPtr& task)
 {
     std::unique_lock<std::mutex> lock(_mutex);
@@ -116,6 +142,19 @@ bool TC_ThreadPool::get(TaskFuncPtr& task)
 
 void TC_ThreadPool::run()
 {
+    // 线程初始化回调
+    if (_threadInitCallback)
+    {
+        try
+        {
+            _threadInitCallback();
+        }
+        catch (...)
+        {
+            // 捕获异常，防止线程启动失败
+        }
+    }
+
     //调用处理部分
     while (!isTerminate())
     {
@@ -149,6 +188,19 @@ void TC_ThreadPool::run()
             }
         }
     }
+
+    // 线程退出回调
+    if (_threadDestroyCallback)
+    {
+        try
+        {
+            _threadDestroyCallback();
+        }
+        catch (...)
+        {
+            // 捕获异常，确保线程正常退出
+        }
+    }
 }
 
 bool TC_ThreadPool::waitForAllDone(int millsecond)
@@ -187,10 +239,39 @@ void TC_ThreadPoolHash::init(size_t num)
     {
         TC_ThreadPool* p = new TC_ThreadPool();
         p->init(1);
+        
+        // 应用回调
+        if (_threadInitCallback)
+        {
+            p->setThreadInitCallback(_threadInitCallback);
+        }
+        if (_threadDestroyCallback)
+        {
+            p->setThreadDestroyCallback(_threadDestroyCallback);
+        }
+        
         _pools.push_back(p);
     }
 }
 
+void TC_ThreadPoolHash::setThreadInitCallback(const TC_ThreadPool::ThreadInitCallback& callback)
+{
+    _threadInitCallback = callback;
+    for (auto pool : _pools)
+    {
+        pool->setThreadInitCallback(callback);
+    }
+}
+
+void TC_ThreadPoolHash::setThreadDestroyCallback(const TC_ThreadPool::ThreadDestroyCallback& callback)
+{
+    _threadDestroyCallback = callback;
+    for (auto pool : _pools)
+    {
+        pool->setThreadDestroyCallback(callback);
+    }
+}
+
 TC_ThreadPool* TC_ThreadPoolHash::getThread(size_t index)
 {
     if (_pools.empty() || (index + 1) > _pools.size())