Fixed data race in FakeTaskRunner

PiperOrigin-RevId: 528627577

internal/test/BUILD

@@ -1,3 +1,17 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 licenses(["notice"])
 
 cc_library(
@@ -50,6 +64,7 @@ cc_test(
         "//internal/platform/implementation:types",
         "//internal/platform/implementation/g3",
         "@com_github_protobuf_matchers//protobuf-matchers",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
     ],

internal/test/fake_task_runner.cc

@@ -1,4 +1,4 @@
-// Copyright 2022 Google LLC
+// Copyright 2022-2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,21 +14,36 @@
 
 #include "internal/test/fake_task_runner.h"
 
-#include <functional>
+#include <atomic>
 #include <future>  // NOLINT
 #include <memory>
 #include <utility>
 #include <vector>
 
+#include "absl/synchronization/mutex.h"
+#include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
+#include "internal/test/fake_timer.h"
 
 namespace nearby {
 
-std::atomic_uint FakeTaskRunner::running_thread_count_ = 0;
+std::atomic_uint FakeTaskRunner::total_running_thread_count_ = 0;
+
+FakeTaskRunner::~FakeTaskRunner() {
+  absl::MutexLock lock(&mutex_);
+  CleanThreads();
+}
 
 bool FakeTaskRunner::PostTask(absl::AnyInvocable<void()> task) {
-  if (mode_ == Mode::kNoPending) {
-    run(std::move(task));
+  absl::MutexLock lock(&mutex_);
+  if (mode_ == Mode::kActive) {
+    if (running_thread_count_ >= count_) {
+      queued_tasks_.push_back(std::move(task));
+      return true;
+    }
+
+    ++running_thread_count_;
+    Run(std::move(task));
     return true;
   }
   pending_tasks_.push_back(std::move(task));
@@ -37,56 +52,80 @@ bool FakeTaskRunner::PostTask(absl::AnyInvocable<void()> task) {
 
 bool FakeTaskRunner::PostDelayedTask(absl::Duration delay,
                                      absl::AnyInvocable<void()> task) {
+  absl::MutexLock lock(&mutex_);
   std::unique_ptr<Timer> timer = std::make_unique<FakeTimer>(clock_);
   Timer* timer_ptr = timer.get();
   uint32_t id = GenerateId();
-  pending_delayed_tasks_.emplace(id, std::move(timer));
+  queued_delayed_tasks_.emplace(id, std::move(timer));
   timer_ptr->Start(delay / absl::Milliseconds(1), 0,
                    [this, task = std::move(task), id]() mutable {
                      PostTask(std::move(task));
-                     completed_delayed_tasks_.push_back(id);
+                     {
+                       absl::MutexLock lock(&mutex_);
+                       completed_delayed_tasks_.push_back(id);
+                     }
                    });
   return true;
 }
 
-void FakeTaskRunner::RunNextTask() {
-  if (pending_tasks_.empty()) {
-    return;
-  }
+void FakeTaskRunner::SetMode(Mode mode) {
+  absl::MutexLock lock(&mutex_);
+  mode_ = mode;
+}
 
-  run(std::move(pending_tasks_.front()));
-  pending_tasks_.erase(pending_tasks_.begin());
+FakeTaskRunner::Mode FakeTaskRunner::GetMode() const {
+  absl::MutexLock lock(&mutex_);
+  return mode_;
+}
+
+void FakeTaskRunner::RunNextPendingTask() {
+  absl::MutexLock lock(&mutex_);
+  InternalRunNextPendingTask();
 }
 
 void FakeTaskRunner::RunAllPendingTasks() {
+  absl::MutexLock lock(&mutex_);
   while (!pending_tasks_.empty()) {
-    RunNextTask();
+    InternalRunNextPendingTask();
   }
 }
 
-const std::vector<absl::AnyInvocable<void()>>& FakeTaskRunner::GetPendingTasks()
-    const {
+void FakeTaskRunner::Sync() {
+  absl::Notification notification;
+  PostTask([&] { notification.Notify(); });
+  notification.WaitForNotification();
+}
+
+const std::vector<absl::AnyInvocable<void()>>&
+FakeTaskRunner::GetAllPendingTasks() const {
+  absl::MutexLock lock(&mutex_);
   return pending_tasks_;
 }
 
 const absl::flat_hash_map<uint32_t, std::unique_ptr<Timer>>&
-FakeTaskRunner::GetPendingDelayedTask() {
+FakeTaskRunner::GetAllDelayedTasks() {
+  absl::MutexLock lock(&mutex_);
   if (!completed_delayed_tasks_.empty()) {
     for (uint32_t id : completed_delayed_tasks_) {
-      pending_delayed_tasks_.erase(id);
+      queued_delayed_tasks_.erase(id);
     }
   }
-  return pending_delayed_tasks_;
+  return queued_delayed_tasks_;
+}
+
+int FakeTaskRunner::GetConcurrentCount() const {
+  absl::MutexLock lock(&mutex_);
+  return count_;
 }
 
 bool FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Duration timeout) {
   int i = (timeout / absl::Milliseconds(1)) / 50;
-  while (running_thread_count_ != 0 && i > 0) {
+  while (total_running_thread_count_ != 0 && i > 0) {
     absl::SleepFor(absl::Milliseconds(50));
     --i;
   }
 
-  return running_thread_count_ == 0;
+  return total_running_thread_count_ == 0;
 }
 
 uint32_t FakeTaskRunner::GenerateId() {
@@ -107,17 +146,39 @@ void FakeTaskRunner::CleanThreads() {
   }
 }
 
-void FakeTaskRunner::run(absl::AnyInvocable<void()> task) {
-  absl::MutexLock lock(&mutex_);
+void FakeTaskRunner::Run(absl::AnyInvocable<void()> task) {
   CleanThreads();
-  ++running_thread_count_;
+  ++total_running_thread_count_;
   // Run the task in a new thread, to simulate the real environment.
   std::future<void> thread =
       std::async(std::launch::async, [&, task = std::move(task)]() mutable {
         task();
-        --running_thread_count_;
+        RunNextQueueTask();
+        --total_running_thread_count_;
       });
   threads_.push_back(std::move(thread));
 }
 
+void FakeTaskRunner::InternalRunNextPendingTask() {
+  if (pending_tasks_.empty()) {
+    return;
+  }
+
+  Run(std::move(pending_tasks_.front()));
+  pending_tasks_.erase(pending_tasks_.begin());
+}
+
+void FakeTaskRunner::RunNextQueueTask() {
+  absl::MutexLock lock(&mutex_);
+  --running_thread_count_;
+  if (queued_tasks_.empty()) {
+    return;
+  }
+
+  auto task = std::move(queued_tasks_.front());
+  queued_tasks_.erase(queued_tasks_.begin());
+  ++running_thread_count_;
+  Run(std::move(task));
+}
+
 }  // namespace nearby

internal/test/fake_task_runner.h

@@ -1,4 +1,4 @@
-// Copyright 2022 Google LLC
+// Copyright 2022-2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -19,8 +19,6 @@
 #include <cstdint>
 #include <future>  //NOLINT
 #include <memory>
-#include <thread>  //NOLINT
-#include <utility>
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
@@ -34,52 +32,68 @@ namespace nearby {
 
 class FakeTaskRunner : public TaskRunner {
  public:
-  enum class Mode { kNoPending, kPending };
+  enum class Mode { kActive, kPending };
 
   FakeTaskRunner(FakeClock* clock, uint32_t count)
       : clock_(clock), count_(count) {}
-  ~FakeTaskRunner() override = default;
+  ~FakeTaskRunner() override ABSL_LOCKS_EXCLUDED(mutex_);
 
-  bool PostTask(absl::AnyInvocable<void()> task) override;
+  bool PostTask(absl::AnyInvocable<void()> task) override
+      ABSL_LOCKS_EXCLUDED(mutex_);
 
   // No matter the mode is pending or not, always put the task in timer control.
   // Caller can move forward time to trigger it.
   bool PostDelayedTask(absl::Duration delay,
-                       absl::AnyInvocable<void()> task) override;
+                       absl::AnyInvocable<void()> task) override
+      ABSL_LOCKS_EXCLUDED(mutex_);
 
   // Mocked methods.
-  void SetMode(Mode mode) { mode_ = mode; }
-  Mode GetMode() const { return mode_; }
+  void SetMode(Mode mode) ABSL_LOCKS_EXCLUDED(mutex_);
+  Mode GetMode() const ABSL_LOCKS_EXCLUDED(mutex_);
 
-  void RunNextTask();
-  void RunAllPendingTasks();
+  void RunNextPendingTask() ABSL_LOCKS_EXCLUDED(mutex_);
+  void RunAllPendingTasks() ABSL_LOCKS_EXCLUDED(mutex_);
+  void Sync();
 
-  const std::vector<absl::AnyInvocable<void()>>& GetPendingTasks() const;
+  const std::vector<absl::AnyInvocable<void()>>& GetAllPendingTasks() const
+      ABSL_LOCKS_EXCLUDED(mutex_);
   const absl::flat_hash_map<uint32_t, std::unique_ptr<Timer>>&
-  GetPendingDelayedTask();
+  GetAllDelayedTasks() ABSL_LOCKS_EXCLUDED(mutex_);
 
-  int GetConcurrentCount() const { return count_; }
+  int GetConcurrentCount() const ABSL_LOCKS_EXCLUDED(mutex_);
 
-  // In some testcases, we needs to make sure all running tasks completion
+  // In some test cases, we needs to make sure all running tasks completion
   // before go to next task. This method can be used for the purpose.
   static bool WaitForRunningTasksWithTimeout(absl::Duration timeout);
+  static int GetTotalRunningThreadCount() {
+    return total_running_thread_count_;
+  }
 
  private:
   uint32_t GenerateId();
   void CleanThreads() ABSL_SHARED_LOCKS_REQUIRED(mutex_);
-  void run(absl::AnyInvocable<void()> task) ABSL_LOCKS_EXCLUDED(mutex_);
+  void Run(absl::AnyInvocable<void()> task) ABSL_SHARED_LOCKS_REQUIRED(mutex_);
+  void InternalRunNextPendingTask() ABSL_SHARED_LOCKS_REQUIRED(mutex_);
+  void RunNextQueueTask() ABSL_LOCKS_EXCLUDED(mutex_);
 
-  Mode mode_ = Mode::kNoPending;
+  mutable absl::Mutex mutex_;
+  mutable absl::Mutex thread_mutex_;
+  Mode mode_ ABSL_GUARDED_BY(mutex_) = Mode::kActive;
   std::atomic_uint current_id_ = 0;
   FakeClock* clock_ = nullptr;
-  uint32_t count_;
-  std::vector<absl::AnyInvocable<void()>> pending_tasks_;
-  std::vector<uint32_t> completed_delayed_tasks_;
-  absl::flat_hash_map<uint32_t, std::unique_ptr<Timer>> pending_delayed_tasks_;
-  absl::Mutex mutex_;
-  std::vector<std::future<void>> threads_ ABSL_GUARDED_BY(mutex_);
+  uint32_t count_ ABSL_GUARDED_BY(mutex_);
 
-  static std::atomic_uint running_thread_count_;
+  // Used for pending mode
+  std::vector<absl::AnyInvocable<void()>> pending_tasks_
+      ABSL_GUARDED_BY(mutex_);
+  std::vector<absl::AnyInvocable<void()>> queued_tasks_ ABSL_GUARDED_BY(mutex_);
+  std::vector<uint32_t> completed_delayed_tasks_ ABSL_GUARDED_BY(mutex_);
+  absl::flat_hash_map<uint32_t, std::unique_ptr<Timer>> queued_delayed_tasks_
+      ABSL_GUARDED_BY(mutex_);
+  std::vector<std::future<void>> threads_ ABSL_GUARDED_BY(mutex_);
+  int running_thread_count_ ABSL_GUARDED_BY(mutex_) = 0;
+
+  static std::atomic_uint total_running_thread_count_;
 };
 
 }  // namespace nearby

internal/test/fake_task_runner_test.cc

@@ -1,4 +1,4 @@
-// Copyright 2022 Google LLC
+// Copyright 2022-2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,7 +14,11 @@
 
 #include "internal/test/fake_task_runner.h"
 
+#include <list>
+
 #include "gtest/gtest.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "internal/test/fake_clock.h"
 
@@ -28,7 +32,7 @@ TEST(FakeTaskRunner, PostTask) {
   task_runner.PostTask([&count] { ++count; });
   ASSERT_TRUE(
       FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Milliseconds(100)));
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 0);
   EXPECT_EQ(count, 1);
 }
 
@@ -37,14 +41,14 @@ TEST(FakeTaskRunner, PostDelayedTask) {
   int count = 0;
   FakeTaskRunner task_runner{&clock, 1};
   task_runner.PostDelayedTask(absl::Seconds(10), [&count] { ++count; });
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 0);
-  EXPECT_EQ(task_runner.GetPendingDelayedTask().size(), 1);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllDelayedTasks().size(), 1);
   EXPECT_EQ(count, 0);
   clock.FastForward(absl::Seconds(10));
   ASSERT_TRUE(
       FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Milliseconds(100)));
   EXPECT_EQ(count, 1);
-  EXPECT_EQ(task_runner.GetPendingDelayedTask().size(), 0);
+  EXPECT_EQ(task_runner.GetAllDelayedTasks().size(), 0);
 }
 
 TEST(FakeTaskRunner, PostTasksInPendingMode) {
@@ -54,11 +58,11 @@ TEST(FakeTaskRunner, PostTasksInPendingMode) {
   EXPECT_EQ(task_runner.GetConcurrentCount(), 1);
   task_runner.PostTask([]() {});
   task_runner.PostTask([]() {});
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 2);
-  task_runner.RunNextTask();
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 1);
-  task_runner.RunNextTask();
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 2);
+  task_runner.RunNextPendingTask();
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 1);
+  task_runner.RunNextPendingTask();
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 0);
 }
 
 TEST(FakeTaskRunner, RunAllPostedTasksInPendingMode) {
@@ -67,11 +71,11 @@ TEST(FakeTaskRunner, RunAllPostedTasksInPendingMode) {
   task_runner.SetMode(FakeTaskRunner::Mode::kPending);
   task_runner.PostTask([]() {});
   task_runner.PostTask([]() {});
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 2);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 2);
   task_runner.RunAllPendingTasks();
   ASSERT_TRUE(
       FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Milliseconds(100)));
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 0);
 }
 
 TEST(FakeTaskRunner, PostDelayedTaskInPendingMode) {
@@ -80,17 +84,36 @@ TEST(FakeTaskRunner, PostDelayedTaskInPendingMode) {
   FakeTaskRunner task_runner{&clock, 1};
   task_runner.SetMode(FakeTaskRunner::Mode::kPending);
   task_runner.PostDelayedTask(absl::Seconds(1), [&called]() { called = true; });
-  EXPECT_EQ(task_runner.GetPendingDelayedTask().size(), 1);
+  EXPECT_EQ(task_runner.GetAllDelayedTasks().size(), 1);
   clock.FastForward(absl::Seconds(1));
-  EXPECT_EQ(task_runner.GetPendingDelayedTask().size(), 0);
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 1);
+  EXPECT_EQ(task_runner.GetAllDelayedTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 1);
   task_runner.RunAllPendingTasks();
   ASSERT_TRUE(
       FakeTaskRunner::WaitForRunningTasksWithTimeout(absl::Milliseconds(100)));
-  EXPECT_EQ(task_runner.GetPendingTasks().size(), 0);
+  EXPECT_EQ(task_runner.GetAllPendingTasks().size(), 0);
   EXPECT_TRUE(called);
 }
 
+TEST(FakeTaskRunner, PostTasksRunInSequence) {
+  std::list<int> result;
+  absl::Mutex mutex;
+  FakeClock clock;
+  FakeTaskRunner task_runner{&clock, 1};
+  for (int i = 0; i < 100; ++i) {
+    task_runner.PostTask([&, i]() {
+      absl::MutexLock lock(&mutex);
+      absl::SleepFor(absl::Milliseconds(40));
+      result.push_back(i);
+    });
+  }
+  task_runner.Sync();
+  for (int i = 0; i < 100; ++i) {
+    EXPECT_EQ(result.front(), i);
+    result.pop_front();
+  }
+}
+
 TEST(FakeTaskRunner, PostDelayedTaskInDelayedTask) {
   FakeClock clock;
   int called_count = 0;