Bugfix: fixed the issue where the timer memory managed by PeripheryTaskScheduler was not released correctly when the program exited

Author: chhy2009

trpc/future/future_test.cc

@@ -641,15 +641,13 @@ TEST(Future, ThenCopyableExecutor) {
 // Test mltiple threads executor.
 TEST(Future, TestExecuteOkInDifferentThread) {
   static int exec_count = 0;
-  int loop_times = 50000;
+  int loop_times = 10000;
   for (int i = 0; i < loop_times; i++) {
     Promise<int> pr;
     auto fut = pr.GetFuture();
     std::thread t([pr = std::move(pr)]() mutable {
-      std::this_thread::sleep_for(std::chrono::nanoseconds(1));
       pr.SetValue(1);
     });
-    std::this_thread::sleep_for(std::chrono::nanoseconds(10000));
     fut.Then([](int&& val) {
       exec_count++;
       return MakeReadyFuture<>();

trpc/metrics/prometheus/prometheus_metrics.cc

@@ -88,7 +88,7 @@ void PrometheusMetrics::Start() noexcept {
 
 void PrometheusMetrics::Stop() noexcept {
   if (push_gateway_task_id_ != 0) {
-    PeripheryTaskScheduler::GetInstance()->StopInnerTask(push_gateway_task_id_);
+    PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(push_gateway_task_id_);
     PeripheryTaskScheduler::GetInstance()->JoinInnerTask(push_gateway_task_id_);
     push_gateway_task_id_ = 0;
   }

trpc/naming/domain/selector_domain.cc

@@ -306,7 +306,7 @@ void SelectorDomain::Start() noexcept {
 
 void SelectorDomain::Stop() noexcept {
   if (task_id_) {
-    PeripheryTaskScheduler::GetInstance()->StopInnerTask(task_id_);
+    PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(task_id_);
     task_id_ = 0;
   }
 }

trpc/rpcz/collector.cc

@@ -84,7 +84,7 @@ void RpczCollectorTask::Destroy() {
 
 void RpczCollectorTask::Stop() {
   if (task_id_) {
-    trpc::PeripheryTaskScheduler::GetInstance()->StopInnerTask(task_id_);
+    trpc::PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(task_id_);
     task_id_ = 0;
   }
 }

trpc/runtime/common/heartbeat/heartbeat_report.cc

@@ -117,7 +117,7 @@ void HeartBeatReport::Stop() {
 
   enable_ = false;
 
-  PeripheryTaskScheduler::GetInstance()->StopInnerTask(task_id_);
+  PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(task_id_);
   PeripheryTaskScheduler::GetInstance()->JoinInnerTask(task_id_);
 
   task_id_ = 0;

trpc/runtime/common/periphery_task_scheduler.cc

@@ -47,9 +47,9 @@ void PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::Stop() {
 }
 
 void PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::Join() {
-  for (unsigned i = 0; i < thread_num_; ++i) {
-    if (workers_[i].joinable()) {
-      workers_[i].join();
+  for (auto & worker : workers_) {
+    if (worker.joinable()) {
+      worker.join();
     }
   }
   workers_.clear();
@@ -61,6 +61,11 @@ void PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::Join() {
   }
 
   while (!tasks_.empty()) {
+    // Compatibility handling for scenarios where the user does not call DetachTask/RemoveTask: 
+    // actively call Deref when program exit.
+    if (tasks_.top()->UnsafeRefCount() > 1) {
+      tasks_.top()->Deref();
+    }
     tasks_.pop();
   }
 }
@@ -135,6 +140,19 @@ bool PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::StopTaskImpl(uint64_t t
   return StopAndDestroyTask(task_id, false);
 }
 
+bool PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::DetachTaskImpl(uint64_t task_id) {
+  if (TRPC_UNLIKELY(exited_.load(std::memory_order_relaxed))) {
+    return false;
+  }
+
+  TaskPtr task_ptr = GetTaskPtr(task_id);
+  if (task_ptr.Get() == nullptr) {
+    return false;
+  }
+
+  return true;
+}
+
 bool PeripheryTaskScheduler::PeripheryTaskSchedulerImpl::JoinTaskImpl(uint64_t task_id) {
   if (TRPC_UNLIKELY(exited_.load(std::memory_order_relaxed))) {
     TRPC_FMT_ERROR("PeripheryTaskScheduler is already exited.");
@@ -323,6 +341,13 @@ bool PeripheryTaskScheduler::RemoveTask(uint64_t task_id) {
   return scheduler_->RemoveTaskImpl(task_id);
 }
 
+bool PeripheryTaskScheduler::DetachTask(std::uint64_t task_id) {
+  if (!scheduler_) {
+    return false;
+  }
+  return scheduler_->DetachTaskImpl(task_id);
+}
+
 bool PeripheryTaskScheduler::StopTask(uint64_t task_id) {
   if (!scheduler_) {
     return false;
@@ -368,6 +393,13 @@ bool PeripheryTaskScheduler::RemoveInnerTask(uint64_t task_id) {
   return inner_scheduler_->RemoveTaskImpl(task_id);
 }
 
+bool PeripheryTaskScheduler::DetachInnerTask(std::uint64_t task_id) {
+  if (!inner_scheduler_) {
+    return false;
+  }
+  return inner_scheduler_->DetachTaskImpl(task_id);
+}
+
 bool PeripheryTaskScheduler::StopInnerTask(uint64_t task_id) {
   if (!inner_scheduler_) {
     return false;

trpc/runtime/common/periphery_task_scheduler.h

@@ -102,9 +102,16 @@ class PeripheryTaskScheduler {
   /// @brief Remove task, used in scenarios where it is not necessary to wait for tasks to complete before exiting.
   /// @param task_id task id
   /// @return on success, return true. on error, return false
-  /// @note This interface can only be called once with the same ID.
+  /// @note This interface can only be called once with the same ID. After call it, the task_id can't be used again
   bool RemoveTask(std::uint64_t task_id);
 
+  /// @brief Detach task, after calling this interface, the lifecycle of this task will be managed by the scheduler, 
+  //         users no longer need to concern themselves with the release of the task.
+  /// @param task_id task id
+  /// @return on success return true, otherwise return false
+  /// @note the task_id can't be used again after calling this interface
+  bool DetachTask(std::uint64_t task_id);
+
   /// @brief Same as 'SubmitTask', but is used only internally by the framework.
   std::uint64_t SubmitInnerTask(Function<void()>&& task, const std::string& name = "");
 
@@ -119,10 +126,13 @@ class PeripheryTaskScheduler {
   /// @brief Same as 'RemoveTask', but is used only internally by the framework.
   bool RemoveInnerTask(std::uint64_t task_id);
 
-  /// @brief Same as 'Stoptask', but is used only internally by the framework.
+  /// @brief Same as 'DetachTask', but is used only internally by the framework.
+  bool DetachInnerTask(std::uint64_t task_id);
+
+  /// @brief Same as 'StopTask', but is used only internally by the framework.
   bool StopInnerTask(std::uint64_t task_id);
 
-  /// @brief Same as 'Jointask', but is used only internally by the framework.
+  /// @brief Same as 'JoinTask', but is used only internally by the framework.
   bool JoinInnerTask(std::uint64_t task_id);
 
   /// @brief Used to destroy resources accessed by scheduled tasks after all scheduled task execution threads have
@@ -158,6 +168,7 @@ class PeripheryTaskScheduler {
     bool StopTaskImpl(std::uint64_t task_id);
     bool JoinTaskImpl(std::uint64_t task_id);
     bool RemoveTaskImpl(std::uint64_t task_id);
+    bool DetachTaskImpl(uint64_t task_id);
 
     void Init(size_t thread_num);
     void Start();

trpc/runtime/common/periphery_task_scheduler_test.cc

@@ -74,6 +74,13 @@ void TestRemoveTask() {
   ASSERT_TRUE(PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(task_id));
 }
 
+void TestDetachTask() {
+  std::uint64_t task_id = PeripheryTaskScheduler::GetInstance()->SubmitInnerTask([]() {});
+  ASSERT_TRUE(task_id > 0);
+  ASSERT_FALSE(PeripheryTaskScheduler::GetInstance()->DetachInnerTask(task_id + 1));
+  ASSERT_TRUE(PeripheryTaskScheduler::GetInstance()->DetachInnerTask(task_id));
+}
+
 void TestSubmitPeriodicTask() {
   int count = 0;
   Latch latch(1);
@@ -268,6 +275,8 @@ TEST_F(PeripheryTaskSchedulerTest, SubmitTaskTest) { TestSubmitTask(); }
 
 TEST_F(PeripheryTaskSchedulerTest, RemoveTaskTest) { TestRemoveTask(); }
 
+TEST_F(PeripheryTaskSchedulerTest, DetachTaskTest) { TestDetachTask(); }
+
 TEST_F(PeripheryTaskSchedulerTest, SubmitPeriodicTaskTest) { TestSubmitPeriodicTask(); }
 
 TEST_F(PeripheryTaskSchedulerTest, RemoveTaskAdvanceTest) { TestRemoveTaskAdvance(); }

trpc/runtime/common/runtime_info_report/runtime_info_reporter.cc

@@ -173,7 +173,7 @@ void StopReportRuntimeInfo() {
     return;
   }
 
-  PeripheryTaskScheduler::GetInstance()->StopInnerTask(report_task_id);
+  PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(report_task_id);
   PeripheryTaskScheduler::GetInstance()->JoinInnerTask(report_task_id);
 
   report_task_id = 0;

trpc/runtime/common/stats/frame_stats.cc

@@ -37,7 +37,7 @@ void FrameStats::Start() {
 
 void FrameStats::Stop() {
   if (task_id_) {
-    PeripheryTaskScheduler::GetInstance()->StopInnerTask(task_id_);
+    PeripheryTaskScheduler::GetInstance()->RemoveInnerTask(task_id_);
     PeripheryTaskScheduler::GetInstance()->JoinInnerTask(task_id_);
     task_id_ = 0;
   }