From 08d157399c09ac2cc4819b440e68ef03e52c8738 Mon Sep 17 00:00:00 2001
From: Janusz Sobczak <jsobczak@google.com>
Date: Thu, 6 Jul 2023 01:59:49 -0700
Subject: [PATCH] Fix TaskRunnerImpl::PostDelayedTask

The timers were destroyed from timer's callbacks. Moving the destruction
to a separate callback fixes the issue.

PiperOrigin-RevId: 545916450
---
 internal/platform/task_runner_impl.cc      | 17 +++--
 internal/platform/task_runner_impl_test.cc | 80 ++++++++--------------
 2 files changed, 35 insertions(+), 62 deletions(-)

diff --git a/internal/platform/task_runner_impl.cc b/internal/platform/task_runner_impl.cc
index 4f25101c..d7f2f49d 100644
--- a/internal/platform/task_runner_impl.cc
+++ b/internal/platform/task_runner_impl.cc
@@ -18,6 +18,7 @@
 #include <utility>
 
 #include "absl/functional/any_invocable.h"
+#include "absl/time/time.h"
 #include "internal/platform/implementation/crypto.h"
 #include "internal/platform/single_thread_executor.h"
 #include "internal/platform/timer_impl.h"
@@ -52,17 +53,15 @@ bool TaskRunnerImpl::PostDelayedTask(absl::Duration delay,
 
   absl::MutexLock lock(&mutex_);
   uint64_t id = GenerateId();
-
   std::unique_ptr<Timer> timer = std::make_unique<TimerImpl>();
-  if (timer->Start(delay / absl::Milliseconds(1), 0,
+  if (timer->Start(absl::ToInt64Milliseconds(delay), 0,
                    [this, id, task = std::move(task)]() mutable {
-                     if (task) {
-                       PostTask(std::move(task));
-                     }
-                     {
-                       absl::MutexLock lock(&mutex_);
-                       timers_map_.erase(id);
-                     }
+                     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)]() {});
                    })) {
     timers_map_.emplace(id, std::move(timer));
     return true;
diff --git a/internal/platform/task_runner_impl_test.cc b/internal/platform/task_runner_impl_test.cc
index ec247a5f..82cdaddd 100644
--- a/internal/platform/task_runner_impl_test.cc
+++ b/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 "internal/platform/count_down_latch.h"
 
 namespace nearby {
 namespace {
@@ -70,75 +71,48 @@ TEST_F(BaseTaskRunnerImplTest, PostSequenceTasks) {
   EXPECT_EQ(completed_tasks[1], "task2");
 }
 
-TEST_P(BaseTaskRunnerImplTest, DISABLED_PostDelayedTask) {
+TEST_P(BaseTaskRunnerImplTest, PostDelayedTask) {
   TaskRunnerImpl task_runner{GetParam()};
-  std::vector<std::string> completed_tasks;
-  absl::Notification notification;
+  std::atomic_bool first_task_started = false;
+  CountDownLatch latch(2);
 
   // Run the first task
-  task_runner.PostDelayedTask(absl::Milliseconds(50),
-                              [&completed_tasks, &notification]() {
-                                completed_tasks.push_back("task1");
-                                if (completed_tasks.size() == 2) {
-                                  notification.Notify();
-                                }
-                              });
-
-  // Run the second task
-  task_runner.PostTask([&completed_tasks, &notification]() {
-    completed_tasks.push_back("task2");
-    if (completed_tasks.size() == 2) {
-      notification.Notify();
-    }
+  task_runner.PostDelayedTask(absl::Milliseconds(50), [&]() {
+    first_task_started = true;
+    latch.CountDown();
   });
 
-  notification.WaitForNotificationWithTimeout(absl::Milliseconds(200));
-  ASSERT_EQ(completed_tasks.size(), 2u);
-  EXPECT_EQ(completed_tasks[0], "task2");
-  EXPECT_EQ(completed_tasks[1], "task1");
+  // Run the second task
+  task_runner.PostTask([&]() {
+    EXPECT_FALSE(first_task_started);
+    latch.CountDown();
+  });
+
+  latch.Await();
 }
 
-TEST_P(BaseTaskRunnerImplTest, DISABLED_PostTwoDelayedTask) {
+TEST_P(BaseTaskRunnerImplTest, PostTwoDelayedTask) {
   TaskRunnerImpl task_runner{GetParam()};
-  std::vector<std::string> completed_tasks;
-  absl::Notification notification;
+  std::atomic_bool first_task_started = false;
+  CountDownLatch latch(2);
 
   // Run the first task
-  task_runner.PostDelayedTask(absl::Milliseconds(100),
-                              [&completed_tasks, &notification]() {
-                                completed_tasks.push_back("task1");
-                                if (completed_tasks.size() == 2) {
-                                  notification.Notify();
-                                }
-                              });
+  task_runner.PostDelayedTask(absl::Milliseconds(100), [&]() {
+    first_task_started = true;
+    latch.CountDown();
+  });
 
   // Run the second task
-  task_runner.PostDelayedTask(absl::Milliseconds(50),
-                              [&completed_tasks, &notification]() {
-                                completed_tasks.push_back("task2");
-                                if (completed_tasks.size() == 2) {
-                                  notification.Notify();
-                                }
-                              });
+  task_runner.PostDelayedTask(absl::Milliseconds(50), [&]() {
+    EXPECT_FALSE(first_task_started);
+    latch.CountDown();
+  });
 
-  notification.WaitForNotificationWithTimeout(absl::Milliseconds(150));
-  ASSERT_EQ(completed_tasks.size(), 2u);
-  EXPECT_EQ(completed_tasks[0], "task2");
-  EXPECT_EQ(completed_tasks[1], "task1");
-
-  absl::Notification notification2;
-  task_runner.PostDelayedTask(absl::Milliseconds(100),
-                              [&completed_tasks, &notification2]() {
-                                completed_tasks.push_back("task3");
-                                notification2.Notify();
-                              });
-  notification2.WaitForNotificationWithTimeout(absl::Milliseconds(150));
-  ASSERT_EQ(completed_tasks.size(), 3u);
-  EXPECT_EQ(completed_tasks[2], "task3");
+  latch.Await();
 }
 
 TEST_F(BaseTaskRunnerImplTest, PostTasksOnRunnerWithOneThread) {
-  TaskRunnerImpl task_runner{10};
+  TaskRunnerImpl task_runner{1};
   std::atomic_int count = 0;
   absl::Notification notification;