Add Shutdown method to TaskRunner.

PiperOrigin-RevId: 628448253

internal/platform/task_runner.h

@@ -29,13 +29,19 @@ class TaskRunner {
   // Posts a task to task runner. The task runs immediately or not depends on
   // the implementation of class. If the implementation supports multiple
   // threads, posted tasks could run concurrently.
+  // Returns false if TashRunner has been shutdown.
   virtual bool PostTask(absl::AnyInvocable<void()> task) = 0;
 
   // Posts a task to run with delay. Multiple tasks can be scheduled. Tasks will
   // execute in the order of their delay expiring, not in the order they were
   // posted.
+  // Returns false if TashRunner has been shutdown.
   virtual bool PostDelayedTask(absl::Duration delay,
                                absl::AnyInvocable<void()> task) = 0;
+
+  // Shutdown this TaskRunner so that scheduled tasks are no longer executed.
+  // New tasks posted after Shutdown all will be ignored.
+  virtual void Shutdown() = 0;
 };
 
 }  // namespace nearby

internal/platform/task_runner_impl.cc

@@ -20,6 +20,7 @@
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
 #include "internal/platform/implementation/crypto.h"
 #include "internal/platform/multi_thread_executor.h"
@@ -38,6 +39,16 @@ TaskRunnerImpl::TaskRunnerImpl(uint32_t runner_count) {
 }
 
 TaskRunnerImpl::~TaskRunnerImpl() {
+  {
+    absl::MutexLock lock(&mutex_);
+    if (closed_) {
+      return;
+    }
+  }
+  Shutdown();
+}
+
+void TaskRunnerImpl::Shutdown() {
   absl::flat_hash_map<uint64_t, std::unique_ptr<Timer>> timers;
   {
     absl::MutexLock lock(&mutex_);
@@ -53,38 +64,48 @@ TaskRunnerImpl::~TaskRunnerImpl() {
 }
 
 bool TaskRunnerImpl::PostTask(absl::AnyInvocable<void()> task) {
+  {
+    absl::MutexLock lock(&mutex_);
+    if (closed_) {
+      return false;
+    }
+  }
+
   if (task) {
     // Because of cannot get the executor status from platform API, just returns
     // true after calling the Execute method.
     executor_->Execute(std::move(task));
   }
-
   return true;
 }
 
 bool TaskRunnerImpl::PostDelayedTask(absl::Duration delay,
                                      absl::AnyInvocable<void()> task) {
-  if (!task) {
-    return true;
-  }
-
   absl::MutexLock lock(&mutex_);
   if (closed_) {
     return false;
   }
+  if (!task) {
+    return true;
+  }
   uint64_t id = GenerateId();
   std::unique_ptr<Timer> timer = std::make_unique<TimerImpl>();
   if (timer->Start(absl::ToInt64Milliseconds(delay), 0,
                    [this, id, task = std::move(task)]() mutable {
-                     absl::MutexLock lock(&mutex_);
-                     if (closed_) {
-                       return;
+                     std::unique_ptr<Timer> timer;
+                     {
+                      absl::MutexLock lock(&mutex_);
+                      if (closed_) {
+                        return;
+                      }
+                      timer = std::move(timers_map_.extract(id).mapped());
                      }
                      PostTask(std::move(task));
                      // We can't destroy the timer directly from the timer
                      // callback.
-                     auto timer = timers_map_.extract(id);
-                     PostTask([timer = std::move(timer)]() {});
+                     if (timer) {
+                      PostTask([timer = std::move(timer)]() {});
+                     }
                    })) {
     timers_map_.emplace(id, std::move(timer));
     return true;

internal/platform/task_runner_impl.h

@@ -23,13 +23,14 @@
 #undef UNICODE
 #endif
 
+#include <cstdint>
 #include <memory>
 
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/synchronization/mutex.h"
-#include "internal/platform/atomic_boolean.h"
+#include "absl/time/time.h"
 #include "internal/platform/submittable_executor.h"
 #include "internal/platform/task_runner.h"
 #include "internal/platform/timer.h"
@@ -41,11 +42,14 @@ class TaskRunnerImpl : public TaskRunner {
   explicit TaskRunnerImpl(uint32_t runner_count);
   ~TaskRunnerImpl() override;
 
-  bool PostTask(absl::AnyInvocable<void()> task) override;
+  bool PostTask(absl::AnyInvocable<void()> task) override
+      ABSL_LOCKS_EXCLUDED(mutex_);
   bool PostDelayedTask(absl::Duration delay,
                        absl::AnyInvocable<void()> task) override
       ABSL_LOCKS_EXCLUDED(mutex_);
 
+  void Shutdown() override ABSL_LOCKS_EXCLUDED(mutex_);
+
  private:
   uint64_t GenerateId();
 

internal/platform/task_runner_impl_test.cc

@@ -21,6 +21,7 @@
 #include "gtest/gtest.h"
 #include "absl/synchronization/notification.h"
 #include "absl/time/clock.h"
+#include "absl/time/time.h"
 #include "internal/platform/count_down_latch.h"
 
 namespace nearby {
@@ -132,6 +133,20 @@ TEST_P(TaskRunnerImplTest, PostEmptyTask) {
   EXPECT_TRUE(task_runner.PostDelayedTask(absl::Milliseconds(100), nullptr));
 }
 
+TEST_P(TaskRunnerImplTest, PostTaskAfterShutdown) {
+  TaskRunnerImpl task_runner{GetParam()};
+  task_runner.Shutdown();
+
+  EXPECT_FALSE(task_runner.PostTask([]() {}));
+}
+
+TEST_P(TaskRunnerImplTest, PostDelayedTaskAfterShutdown) {
+  TaskRunnerImpl task_runner{GetParam()};
+  task_runner.Shutdown();
+
+  EXPECT_FALSE(task_runner.PostDelayedTask(absl::Milliseconds(50), []() {}));
+}
+
 INSTANTIATE_TEST_SUITE_P(ParameterizedTaskRunnerImplTest, TaskRunnerImplTest,
                          ::testing::ValuesIn(kNumThreads));
 

internal/test/fake_task_runner.cc

@@ -32,7 +32,12 @@ namespace nearby {
 
 std::atomic_uint FakeTaskRunner::pending_tasks_count_ = 0;
 
-FakeTaskRunner::~FakeTaskRunner() { absl::MutexLock lock(&mutex_); }
+FakeTaskRunner::~FakeTaskRunner() { Shutdown(); }
+
+void FakeTaskRunner::Shutdown() {
+  absl::MutexLock lock(&mutex_);
+  task_executor_->Shutdown();
+}
 
 bool FakeTaskRunner::PostTask(absl::AnyInvocable<void()> task) {
   absl::MutexLock lock(&mutex_);
@@ -78,7 +83,6 @@ bool FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Duration timeout) {
     absl::SleepFor(absl::Milliseconds(50));
     --i;
   }
-
   return pending_tasks_count_ == 0;
 }
 

internal/test/fake_task_runner.h

@@ -47,6 +47,8 @@ class FakeTaskRunner : public TaskRunner {
                        absl::AnyInvocable<void()> task) override
       ABSL_LOCKS_EXCLUDED(mutex_);
 
+  void Shutdown() override ABSL_LOCKS_EXCLUDED(mutex_);
+
   // Wait for all thread completed.
   void Sync();