From 281032ee7835d7c857c1a7279fde01e0d9b61b2b Mon Sep 17 00:00:00 2001
From: Janusz Sobczak <jsobczak@google.com>
Date: Tue, 18 Jul 2023 14:54:14 -0700
Subject: [PATCH] Fix flakiness in TaskRunnerImpl

PiperOrigin-RevId: 549119512
---
 internal/platform/task_runner_impl.cc      |  8 +++-
 internal/platform/task_runner_impl_test.cc | 47 ++++++----------------
 2 files changed, 19 insertions(+), 36 deletions(-)

diff --git a/internal/platform/task_runner_impl.cc b/internal/platform/task_runner_impl.cc
index d7f2f49d..4892ae76 100644
--- a/internal/platform/task_runner_impl.cc
+++ b/internal/platform/task_runner_impl.cc
@@ -33,7 +33,13 @@ TaskRunnerImpl::TaskRunnerImpl(uint32_t runner_count) {
   }
 }
 
-TaskRunnerImpl::~TaskRunnerImpl() = default;
+TaskRunnerImpl::~TaskRunnerImpl() {
+  {
+    absl::MutexLock lock(&mutex_);
+    timers_map_.clear();
+  }
+  executor_->Shutdown();
+}
 
 bool TaskRunnerImpl::PostTask(absl::AnyInvocable<void()> task) {
   if (task) {
diff --git a/internal/platform/task_runner_impl_test.cc b/internal/platform/task_runner_impl_test.cc
index 82cdaddd..f7323ef7 100644
--- a/internal/platform/task_runner_impl_test.cc
+++ b/internal/platform/task_runner_impl_test.cc
@@ -91,19 +91,19 @@ TEST_P(BaseTaskRunnerImplTest, PostDelayedTask) {
   latch.Await();
 }
 
-TEST_P(BaseTaskRunnerImplTest, PostTwoDelayedTask) {
+TEST_P(BaseTaskRunnerImplTest, PostTwoDelayedTasks) {
   TaskRunnerImpl task_runner{GetParam()};
   std::atomic_bool first_task_started = false;
   CountDownLatch latch(2);
 
   // Run the first task
-  task_runner.PostDelayedTask(absl::Milliseconds(100), [&]() {
+  task_runner.PostDelayedTask(absl::Milliseconds(500), [&]() {
     first_task_started = true;
     latch.CountDown();
   });
 
   // Run the second task
-  task_runner.PostDelayedTask(absl::Milliseconds(50), [&]() {
+  task_runner.PostDelayedTask(absl::Milliseconds(1), [&]() {
     EXPECT_FALSE(first_task_started);
     latch.CountDown();
   });
@@ -111,42 +111,19 @@ TEST_P(BaseTaskRunnerImplTest, PostTwoDelayedTask) {
   latch.Await();
 }
 
-TEST_F(BaseTaskRunnerImplTest, PostTasksOnRunnerWithOneThread) {
-  TaskRunnerImpl task_runner{1};
-  std::atomic_int count = 0;
-  absl::Notification notification;
+TEST_P(BaseTaskRunnerImplTest, PostMultipleTasks) {
+  TaskRunnerImpl task_runner(GetParam());
+  constexpr int kNumTasks = 10;
+  CountDownLatch latch(kNumTasks);
 
-  for (int i = 0; i < 10; i++) {
-    task_runner.PostTask([&count, &notification]() {
-      absl::SleepFor(absl::Milliseconds(100));
-      count++;
-      if (count == 10) {
-        notification.Notify();
-      }
+  for (int i = 0; i < kNumTasks; i++) {
+    task_runner.PostTask([&]() {
+      absl::SleepFor(absl::Milliseconds(10));
+      latch.CountDown();
     });
   }
 
-  notification.WaitForNotificationWithTimeout(absl::Milliseconds(1900));
-  EXPECT_EQ(count, 10);
-}
-
-TEST_F(BaseTaskRunnerImplTest, PostTasksOnRunnerWithMultipleThreads) {
-  TaskRunnerImpl task_runner{10};
-  std::atomic_int count = 0;
-  absl::Notification notification;
-
-  for (int i = 0; i < 10; i++) {
-    task_runner.PostTask([&count, &notification]() {
-      absl::SleepFor(absl::Milliseconds(100));
-      count++;
-      if (count == 10) {
-        notification.Notify();
-      }
-    });
-  }
-
-  notification.WaitForNotificationWithTimeout(absl::Milliseconds(190));
-  EXPECT_EQ(count, 10);
+  EXPECT_TRUE(latch.Await());
 }
 
 TEST_P(BaseTaskRunnerImplTest, PostEmptyTask) {