TaskScheduler: Atomic operations in TaskTracker

base/task_scheduler/task_tracker_unittest.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/task_scheduler/task_tracker.h"
 
 #include <memory>
+#include <vector>
 
 #include "base/bind.h"
 #include "base/logging.h"
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
 #include "base/sequenced_task_runner.h"
 #include "base/single_thread_task_runner.h"
 #include "base/synchronization/waitable_event.h"
+#include "base/task_scheduler/scheduler_lock.h"
 #include "base/task_scheduler/task.h"
 #include "base/task_scheduler/task_traits.h"
 #include "base/task_scheduler/test_utils.h"
 #include "base/test/test_simple_task_runner.h"
 #include "base/threading/platform_thread.h"
 #include "base/threading/sequenced_task_runner_handle.h"
 #include "base/threading/simple_thread.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace base {
 namespace internal {
 
 namespace {
 
+constexpr size_t kLoadTestNumIterations = 100;
+
 // Calls TaskTracker::Shutdown() asynchronously.
 class ThreadCallingShutdown : public SimpleThread {
  public:
   explicit ThreadCallingShutdown(TaskTracker* tracker)
       : SimpleThread("ThreadCallingShutdown"),
         tracker_(tracker),
         has_returned_(true, false) {}
 
   // Returns true once the async call to Shutdown() has returned.
   bool has_returned() { return has_returned_.IsSignaled(); }
 
  private:
   void Run() override {
     tracker_->Shutdown();
     has_returned_.Signal();
   }
 
   TaskTracker* const tracker_;
   WaitableEvent has_returned_;
 
   DISALLOW_COPY_AND_ASSIGN(ThreadCallingShutdown);
 };
 
-// Runs a task asynchronously.
-class ThreadRunningTask : public SimpleThread {
+class ThreadPostingAndRunningTask : public SimpleThread {
  public:
-  ThreadRunningTask(TaskTracker* tracker, const Task* task)
-      : SimpleThread("ThreadRunningTask"), tracker_(tracker), task_(task) {}
+  enum Action {
+    WILL_POST = 0x1,
+    RUN = 0x2,
+    WILL_POST_AND_RUN = WILL_POST | RUN,
+  };
+
+  ThreadPostingAndRunningTask(TaskTracker* tracker,
+                              const Task* task,
+                              Action action,
+                              bool expect_post_succeeds)
+      : SimpleThread("ThreadPostingAndRunningTask"),
+        tracker_(tracker),
+        task_(task),
+        action_(action),
+        expect_post_succeeds_(expect_post_succeeds) {}
 
  private:
-  void Run() override { tracker_->RunTask(task_); }
+  void Run() override {
+    bool post_succeeded = true;
+    if (action_ & WILL_POST) {
+      post_succeeded = tracker_->WillPostTask(task_);
+      EXPECT_EQ(expect_post_succeeds_, post_succeeded);
+    }
+    if (post_succeeded && action_ & RUN)
+      tracker_->RunTask(task_);
+  }
 
   TaskTracker* const tracker_;
   const Task* const task_;
+  const Action action_;
+  const bool expect_post_succeeds_;
 
-  DISALLOW_COPY_AND_ASSIGN(ThreadRunningTask);
+  DISALLOW_COPY_AND_ASSIGN(ThreadPostingAndRunningTask);
+  ;
 };
 
 class ScopedSetSingletonAllowed {
  public:
   ScopedSetSingletonAllowed(bool singleton_allowed)
       : previous_value_(
             ThreadRestrictions::SetSingletonAllowed(singleton_allowed)) {}
   ~ScopedSetSingletonAllowed() {
     ThreadRestrictions::SetSingletonAllowed(previous_value_);
   }
@@ skipping 35 matching lines @@
     EXPECT_TRUE(tracker_.shutdown_completed());
   }
 
   void VerifyAsyncShutdownInProgress() {
     ASSERT_TRUE(thread_calling_shutdown_);
     EXPECT_FALSE(thread_calling_shutdown_->has_returned());
     EXPECT_FALSE(tracker_.shutdown_completed());
     EXPECT_TRUE(tracker_.IsShuttingDownForTesting());
   }
 
+  size_t NumTasksExecuted() {
+    AutoSchedulerLock auto_lock(lock_);
+    return num_tasks_executed_;
+  }
+
   TaskTracker tracker_;
-  size_t num_tasks_executed_ = 0;
 
  private:
-  void RunTaskCallback() { ++num_tasks_executed_; }
+  void RunTaskCallback() {
+    AutoSchedulerLock auto_lock(lock_);
+    ++num_tasks_executed_;
+  }
 
   std::unique_ptr<ThreadCallingShutdown> thread_calling_shutdown_;
 
+  // Synchronizes accesses to |num_tasks_executed_|.
+  SchedulerLock lock_;
+
+  size_t num_tasks_executed_ = 0;
+
   DISALLOW_COPY_AND_ASSIGN(TaskSchedulerTaskTrackerTest);
 };
 
 #define WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED() \
   do {                                      \
     SCOPED_TRACE("");                       \
     WaitForAsyncShutdownCompleted();        \
   } while (false)
 
 #define VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS() \
   do {                                      \
     SCOPED_TRACE("");                       \
     VerifyAsyncShutdownInProgress();        \
   } while (false)
 
 }  // namespace
 
 TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunBeforeShutdown) {
   std::unique_ptr<Task> task(CreateTask(GetParam()));
 
   // Inform |task_tracker_| that |task| will be posted.
   EXPECT_TRUE(tracker_.WillPostTask(task.get()));
 
   // Run the task.
-  EXPECT_EQ(0U, num_tasks_executed_);
+  EXPECT_EQ(0U, NumTasksExecuted());
   tracker_.RunTask(task.get());
-  EXPECT_EQ(1U, num_tasks_executed_);
+  EXPECT_EQ(1U, NumTasksExecuted());
 
   // Shutdown() shouldn't block.
   tracker_.Shutdown();
 }
 
 TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunLongTaskBeforeShutdown) {
   // Create a task that will block until |event| is signaled.
   WaitableEvent event(false, false);
   std::unique_ptr<Task> blocked_task(
       new Task(FROM_HERE, Bind(&WaitableEvent::Wait, Unretained(&event)),
                TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta()));
 
   // Inform |task_tracker_| that |blocked_task| will be posted.
   EXPECT_TRUE(tracker_.WillPostTask(blocked_task.get()));
 
   // Run the task asynchronouly.
-  ThreadRunningTask thread_running_task(&tracker_, blocked_task.get());
+  ThreadPostingAndRunningTask thread_running_task(
+      &tracker_, blocked_task.get(), ThreadPostingAndRunningTask::RUN, false);
   thread_running_task.Start();
 
   // Initiate shutdown while the task is running.
   CallShutdownAsync();
 
   if (GetParam() == TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN) {
     // Shutdown should complete even with a CONTINUE_ON_SHUTDOWN in progress.
     WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
   } else {
     // Shutdown should block with any non CONTINUE_ON_SHUTDOWN task in progress.
@@ skipping 19 matching lines @@
   std::unique_ptr<Task> block_shutdown_task(
       CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
   EXPECT_TRUE(tracker_.WillPostTask(block_shutdown_task.get()));
 
   // Call Shutdown() asynchronously.
   CallShutdownAsync();
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
   // Try to run |task|. It should only run it it's BLOCK_SHUTDOWN. Otherwise it
   // should be discarded.
-  EXPECT_EQ(0U, num_tasks_executed_);
+  EXPECT_EQ(0U, NumTasksExecuted());
   tracker_.RunTask(task.get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 1U : 0U,
-            num_tasks_executed_);
+            NumTasksExecuted());
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
   // Unblock shutdown by running the remaining BLOCK_SHUTDOWN task.
   tracker_.RunTask(block_shutdown_task.get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 2U : 1U,
-            num_tasks_executed_);
+            NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 }
 
 TEST_P(TaskSchedulerTaskTrackerTest, WillPostBeforeShutdownRunAfterShutdown) {
   // Inform |task_tracker_| that a task will be posted.
   std::unique_ptr<Task> task(CreateTask(GetParam()));
   EXPECT_TRUE(tracker_.WillPostTask(task.get()));
 
   // Call Shutdown() asynchronously.
   CallShutdownAsync();
-  EXPECT_EQ(0U, num_tasks_executed_);
+  EXPECT_EQ(0U, NumTasksExecuted());
 
   if (GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN) {
     VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
     // Run the task to unblock shutdown.
     tracker_.RunTask(task.get());
-    EXPECT_EQ(1U, num_tasks_executed_);
+    EXPECT_EQ(1U, NumTasksExecuted());
     WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 
     // It is not possible to test running a BLOCK_SHUTDOWN task posted before
     // shutdown after shutdown because Shutdown() won't return if there are
     // pending BLOCK_SHUTDOWN tasks.
   } else {
     WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 
     // The task shouldn't be allowed to run after shutdown.
     tracker_.RunTask(task.get());
-    EXPECT_EQ(0U, num_tasks_executed_);
+    EXPECT_EQ(0U, NumTasksExecuted());
   }
 }
 
 TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunDuringShutdown) {
   // Inform |task_tracker_| that a BLOCK_SHUTDOWN task will be posted just to
   // block shutdown.
   std::unique_ptr<Task> block_shutdown_task(
       CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
   EXPECT_TRUE(tracker_.WillPostTask(block_shutdown_task.get()));
 
   // Call Shutdown() asynchronously.
   CallShutdownAsync();
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
 
   if (GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN) {
     // Inform |task_tracker_| that a BLOCK_SHUTDOWN task will be posted.
     std::unique_ptr<Task> task(CreateTask(GetParam()));
     EXPECT_TRUE(tracker_.WillPostTask(task.get()));
 
     // Run the BLOCK_SHUTDOWN task.
-    EXPECT_EQ(0U, num_tasks_executed_);
+    EXPECT_EQ(0U, NumTasksExecuted());
     tracker_.RunTask(task.get());
-    EXPECT_EQ(1U, num_tasks_executed_);
+    EXPECT_EQ(1U, NumTasksExecuted());
   } else {
     // It shouldn't be allowed to post a non BLOCK_SHUTDOWN task.
     std::unique_ptr<Task> task(CreateTask(GetParam()));
     EXPECT_FALSE(tracker_.WillPostTask(task.get()));
 
     // Don't try to run the task, because it wasn't allowed to be posted.
   }
 
   // Unblock shutdown by running |block_shutdown_task|.
   VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
   tracker_.RunTask(block_shutdown_task.get());
   EXPECT_EQ(GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN ? 2U : 1U,
-            num_tasks_executed_);
+            NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 }
 
 TEST_P(TaskSchedulerTaskTrackerTest, WillPostAfterShutdown) {
   tracker_.Shutdown();
 
   std::unique_ptr<Task> task(CreateTask(GetParam()));
 
   // |task_tracker_| shouldn't allow a task to be posted after shutdown.
   if (GetParam() == TaskShutdownBehavior::BLOCK_SHUTDOWN) {
@@ skipping 113 matching lines @@
     ::testing::Values(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
 INSTANTIATE_TEST_CASE_P(
     SkipOnShutdown,
     TaskSchedulerTaskTrackerTest,
     ::testing::Values(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
 INSTANTIATE_TEST_CASE_P(
     BlockShutdown,
     TaskSchedulerTaskTrackerTest,
     ::testing::Values(TaskShutdownBehavior::BLOCK_SHUTDOWN));
 
+TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunBeforeShutdown) {
+  // Post and run tasks asynchronously.
+  std::vector<std::unique_ptr<Task>> tasks;
+  std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> threads;
+
+  for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
+    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, true)));
+    threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, true)));
+    threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, true)));
+    threads.back()->Start();
+  }
+
+  for (const auto& thread : threads)
+    thread->Join();
+
+  // Expect all tasks to be executed.
+  EXPECT_EQ(kLoadTestNumIterations * 3, NumTasksExecuted());
+
+  // Should return immediately because no tasks are blocking shutdown.
+  tracker_.Shutdown();
+}
+
+TEST_F(TaskSchedulerTaskTrackerTest,
+       LoadWillPostBeforeShutdownAndRunDuringShutdown) {
+  // Post tasks asynchronously.
+  std::vector<std::unique_ptr<Task>> tasks;
+  std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> post_threads;
+
+  for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
+    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
+    post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(), ThreadPostingAndRunningTask::WILL_POST,
+        true)));
+    post_threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
+    post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(), ThreadPostingAndRunningTask::WILL_POST,
+        true)));
+    post_threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
+    post_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(), ThreadPostingAndRunningTask::WILL_POST,
+        true)));
+    post_threads.back()->Start();
+  }
+
+  for (const auto& thread : post_threads)
+    thread->Join();
+
+  // Call Shutdown() asynchronously.
+  CallShutdownAsync();
+
+  // Run tasks asynchronously.
+  std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> run_threads;
+
+  for (const auto& task : tasks) {
+    run_threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, task.get(), ThreadPostingAndRunningTask::RUN, false)));
+    run_threads.back()->Start();
+  }
+
+  for (const auto& thread : run_threads)
+    thread->Join();
+
+  WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
+
+  // Expect BLOCK_SHUTDOWN tasks to have been executed.
+  EXPECT_EQ(kLoadTestNumIterations, NumTasksExecuted());
+}
+
+TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunDuringShutdown) {
+  // Inform |task_tracker_| that a BLOCK_SHUTDOWN task will be posted just to
+  // block shutdown.
+  std::unique_ptr<Task> block_shutdown_task(
+      CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
+  EXPECT_TRUE(tracker_.WillPostTask(block_shutdown_task.get()));
+
+  // Call Shutdown() asynchronously.
+  CallShutdownAsync();
+
+  // Post and run tasks asynchronously.
+  std::vector<std::unique_ptr<Task>> tasks;
+  std::vector<std::unique_ptr<ThreadPostingAndRunningTask>> threads;
+
+  for (size_t i = 0; i < kLoadTestNumIterations; ++i) {
+    tasks.push_back(CreateTask(TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, false)));
+    threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::SKIP_ON_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, false)));
+    threads.back()->Start();
+
+    tasks.push_back(CreateTask(TaskShutdownBehavior::BLOCK_SHUTDOWN));
+    threads.push_back(WrapUnique(new ThreadPostingAndRunningTask(
+        &tracker_, tasks.back().get(),
+        ThreadPostingAndRunningTask::WILL_POST_AND_RUN, true)));
+    threads.back()->Start();
+  }
+
+  for (const auto& thread : threads)
+    thread->Join();
+
+  // Expect BLOCK_SHUTDOWN tasks to have been executed.
+  EXPECT_EQ(kLoadTestNumIterations, NumTasksExecuted());
+
+  // Shutdown() shouldn't return before |block_shutdown_task| is executed.
+  VERIFY_ASYNC_SHUTDOWN_IN_PROGRESS();
+
+  // Unblock shutdown by running |block_shutdown_task|.
+  tracker_.RunTask(block_shutdown_task.get());
+  EXPECT_EQ(kLoadTestNumIterations + 1, NumTasksExecuted());
+  WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
+}
+
 }  // namespace internal
 }  // namespace base