From 61ee4ebeef0d920512c6fff10642d54143e6403b Mon Sep 17 00:00:00 2001
From: tqcq <99722391+tqcq@users.noreply.github.com>
Date: Wed, 1 May 2024 14:14:57 +0800
Subject: [PATCH] feat support any safety task

---
 .../task_queue/pending_task_safety_flag.h     | 28 ++++++++++++++++---
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/sled/task_queue/pending_task_safety_flag.h b/src/sled/task_queue/pending_task_safety_flag.h
index 8c3b6e4..00b326b 100644
--- a/src/sled/task_queue/pending_task_safety_flag.h
+++ b/src/sled/task_queue/pending_task_safety_flag.h
@@ -8,6 +8,7 @@
 #define SLED_TASK_QUEUE_PENDING_TASK_SAFETY_FLAG_H
 
 #pragma once
+#include "sled/meta/type_traits.h"
 #include "sled/ref_counted_base.h"
 #include "sled/scoped_refptr.h"
 #include "sled/synchronization/sequence_checker.h"
@@ -53,11 +54,30 @@ private:
     scoped_refptr<PendingTaskSafetyFlag> flag_;
 };
 
-inline std::function<void()>
-SafeTask(scoped_refptr<PendingTaskSafetyFlag> flag, std::function<void()> task)
+// inline std::function<void()>
+// SafeTask(scoped_refptr<PendingTaskSafetyFlag> flag, std::function<void()> task)
+// {
+//     return [flag, task]() mutable {
+//         if (flag->alive()) { std::move(task)(); }
+//     };
+// }
+
+template<typename F, typename... Args, typename = EnableIfT<!std::is_void<InvokeResultT<F, Args...>>::value>>
+std::function<InvokeResultT<F, Args...>(Args &&...)>
+SafeTask(scoped_refptr<PendingTaskSafetyFlag> flag, F &&f)
 {
-    return [flag, task]() mutable {
-        if (flag->alive()) { std::move(task)(); }
+    return [flag, f](Args &&...args) mutable -> InvokeResultT<F, Args...> {
+        if (flag->alive()) { return f(std::forward<Args>(args)...); }
+        return {};
+    };
+}
+
+template<typename F, typename... Args, typename = EnableIfT<std::is_void<InvokeResultT<F, Args...>>::value>>
+std::function<void(Args &&...)>
+SafeTask(scoped_refptr<PendingTaskSafetyFlag> flag, F &&f)
+{
+    return [flag, f](Args &&...args) mutable -> void {
+        if (flag->alive()) { f(std::forward<Args>(args)...); }
     };
 }