Improve thread-safety in TaskRunnerImpl

PiperOrigin-RevId: 571981130

internal/platform/task_runner_impl.cc

@@ -14,30 +14,41 @@
 
 #include "internal/platform/task_runner_impl.h"
 
+#include <cstdint>
 #include <memory>
 #include <utility>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/time/time.h"
 #include "internal/platform/implementation/crypto.h"
+#include "internal/platform/multi_thread_executor.h"
 #include "internal/platform/single_thread_executor.h"
+#include "internal/platform/timer.h"
 #include "internal/platform/timer_impl.h"
 
 namespace nearby {
 
 TaskRunnerImpl::TaskRunnerImpl(uint32_t runner_count) {
   if (runner_count == 1) {
-    executor_ = std::make_unique<::nearby::SingleThreadExecutor>();
+    executor_ = std::make_unique<SingleThreadExecutor>();
   } else {
-    executor_ = std::make_unique<::nearby::MultiThreadExecutor>(runner_count);
+    executor_ = std::make_unique<MultiThreadExecutor>(runner_count);
   }
 }
 
 TaskRunnerImpl::~TaskRunnerImpl() {
+  absl::flat_hash_map<uint64_t, std::unique_ptr<Timer>> timers;
   {
     absl::MutexLock lock(&mutex_);
-    timers_map_.clear();
+    closed_ = true;
+    timers = std::move(timers_map_);
   }
+  for (auto& timer : timers) {
+    timer.second->Stop();
+  }
+  // We expect that all timers are stopped, no new timers will be added, and the
+  // timer callbacks are not running.
   executor_->Shutdown();
 }
 
@@ -58,14 +69,20 @@ bool TaskRunnerImpl::PostDelayedTask(absl::Duration delay,
   }
 
   absl::MutexLock lock(&mutex_);
+  if (closed_) {
+    return false;
+  }
   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;
+                     }
                      PostTask(std::move(task));
                      // We can't destroy the timer directly from the timer
                      // callback.
-                     absl::MutexLock lock(&mutex_);
                      auto timer = timers_map_.extract(id);
                      PostTask([timer = std::move(timer)]() {});
                    })) {