TaskScheduler: Atomic operations in TaskTracker

base/task_scheduler/task_tracker.cc

Index: base/task_scheduler/task_tracker.cc
diff --git a/base/task_scheduler/task_tracker.cc b/base/task_scheduler/task_tracker.cc
index 4a272cd54a97044a3be127ac92f96cebecff17d9..8817baa04701088ece63b91005c68fe19ddca6e3 100644
--- a/base/task_scheduler/task_tracker.cc
+++ b/base/task_scheduler/task_tracker.cc
@@ -4,8 +4,12 @@
 
 #include "base/task_scheduler/task_tracker.h"
 
+#include <limits>
+
+#include "base/atomicops.h"
 #include "base/callback.h"
 #include "base/debug/task_annotator.h"
+#include "base/logging.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/threading/sequenced_task_runner_handle.h"
 #include "base/threading/thread_restrictions.h"
@@ -36,33 +40,127 @@ void RecordNumBlockShutdownTasksPostedDuringShutdown(
 
 }  // namespace
 
-TaskTracker::TaskTracker() = default;
+class TaskTracker::State {
+ public:
+  State() = default;
+
+  // Sets a flag indicating that shutdown has started. Returns true if there are
+  // tasks blocking shutdown. Can only be called once.
+  bool StartShutdown() {
+    const auto new_value =
+        subtle::Barrier_AtomicIncrement(&bits_, kShutdownHasStartedMask);
+
+    // Check that the "shutdown has started" bit isn't zero. This would happen
+    // if it was incremented twice.
+    DCHECK(new_value & kShutdownHasStartedMask);
+
+    const auto num_tasks_blocking_shutdown =
+        new_value >> kNumTasksBlockingShutdownBitOffset;
+    return num_tasks_blocking_shutdown != 0;
+  }
+
+  // Returns true if shutdown has started.
+  bool HasShutdownStarted() const {
+    return subtle::Acquire_Load(&bits_) & kShutdownHasStartedMask;
+  }
+
+  // Returns true if there are tasks blocking shutdown.
+  bool AreTasksBlockingShutdown() const {
+    const auto num_tasks_blocking_shutdown =
+        subtle::Acquire_Load(&bits_) >> kNumTasksBlockingShutdownBitOffset;
+    DCHECK_GE(num_tasks_blocking_shutdown, 0);
+    return num_tasks_blocking_shutdown != 0;
+  }
+
+  // Increments the number of tasks blocking shutdown. Returns true if shutdown
+  // has started.
+  bool IncrementNumTasksBlockingShutdown() {
+#if DCHECK_IS_ON()
+    // Verify that no overflow will occur.
+    const auto num_tasks_blocking_shutdown =
+        subtle::Acquire_Load(&bits_) >> kNumTasksBlockingShutdownBitOffset;
+    DCHECK_LT(num_tasks_blocking_shutdown,
+              std::numeric_limits<subtle::Atomic32>::max() -
+                  kNumTasksBlockingShutdownIncrement);
+#endif
+
+    const auto new_bits = subtle::Barrier_AtomicIncrement(
+        &bits_, kNumTasksBlockingShutdownIncrement);
+    return new_bits & kShutdownHasStartedMask;
+  }
+
+  // Decrements the number of tasks blocking shutdown. Returns true if shutdown
+  // has started and the number of tasks blocking shutdown becomes zero.
+  bool DecrementNumTasksBlockingShutdown() {
+    const auto new_bits = subtle::Barrier_AtomicIncrement(
+        &bits_, -kNumTasksBlockingShutdownIncrement);
+    const bool shutdown_has_started = new_bits & kShutdownHasStartedMask;
+    const auto num_tasks_blocking_shutdown =
+        new_bits >> kNumTasksBlockingShutdownBitOffset;
+    DCHECK_GE(num_tasks_blocking_shutdown, 0);
+    return shutdown_has_started && num_tasks_blocking_shutdown == 0;
+  }
+
+ private:
+  static constexpr subtle::Atomic32 kShutdownHasStartedMask = 1;
+  static constexpr subtle::Atomic32 kNumTasksBlockingShutdownBitOffset = 1;
+  static constexpr subtle::Atomic32 kNumTasksBlockingShutdownIncrement =
+      1 << kNumTasksBlockingShutdownBitOffset;
+
+  // The LSB indicates whether shutdown has started. The other bits count the
+  // number of tasks blocking shutdown.
+  subtle::Atomic32 bits_ = 0;
+
+  DISALLOW_COPY_AND_ASSIGN(State);
+};
+
+TaskTracker::TaskTracker() : state_(new State) {}
 TaskTracker::~TaskTracker() = default;
 
 void TaskTracker::Shutdown() {
-  AutoSchedulerLock auto_lock(lock_);
+  {
+    AutoSchedulerLock auto_lock(shutdown_lock_);
+
+    // This method can only be called once.
+    DCHECK(!shutdown_event_);
+    DCHECK(!num_block_shutdown_tasks_posted_during_shutdown_);
 
-  // This method should only be called once.
-  DCHECK(!shutdown_completed_);
-  DCHECK(!shutdown_cv_);
+    shutdown_event_.reset(
+        new WaitableEvent(WaitableEvent::ResetPolicy::MANUAL,
+                          WaitableEvent::InitialState::NOT_SIGNALED));
 
-  shutdown_cv_ = lock_.CreateConditionVariable();
+    const bool tasks_are_blocking_shutdown = state_->StartShutdown();
 
-  // Wait until the number of tasks blocking shutdown is zero.
-  while (num_tasks_blocking_shutdown_ != 0)
-    shutdown_cv_->Wait();
+    // From now, if a thread causes the number of tasks blocking shutdown to
+    // become zero, it will call OnBlockingShutdownTasksComplete().
 
-  shutdown_cv_.reset();
-  shutdown_completed_ = true;
+    if (!tasks_are_blocking_shutdown) {
+      // If another thread posts a BLOCK_SHUTDOWN task at this moment, it will
+      // block until this method releases |shutdown_lock_|. Then, it will fail
+      // DCHECK(!shutdown_event_->IsSignaled()). This is the desired behavior
+      // because posting a BLOCK_SHUTDOWN task when TaskTracker::Shutdown() has
+      // started and no tasks are blocking shutdown isn't allowed.
+      shutdown_event_->Signal();
+      return;
+    }
+  }
+
+  // It is safe to access |shutdown_event_| without holding |lock_| because the
+  // pointer never changes after being set above.
+  shutdown_event_->Wait();
 
-  // Record the TaskScheduler.BlockShutdownTasksPostedDuringShutdown if less
-  // than |kMaxBlockShutdownTasksPostedDuringShutdown| BLOCK_SHUTDOWN tasks were
-  // posted during shutdown. Otherwise, the histogram has already been recorded
-  // in BeforePostTask().
-  if (num_block_shutdown_tasks_posted_during_shutdown_ <
-      kMaxBlockShutdownTasksPostedDuringShutdown) {
-    RecordNumBlockShutdownTasksPostedDuringShutdown(
-        num_block_shutdown_tasks_posted_during_shutdown_);
+  {
+    AutoSchedulerLock auto_lock(shutdown_lock_);
+
+    // Record TaskScheduler.BlockShutdownTasksPostedDuringShutdown if less than
+    // |kMaxBlockShutdownTasksPostedDuringShutdown| BLOCK_SHUTDOWN tasks were
+    // posted during shutdown. Otherwise, the histogram has already been
+    // recorded in BeforePostTask().
+    if (num_block_shutdown_tasks_posted_during_shutdown_ <
+        kMaxBlockShutdownTasksPostedDuringShutdown) {
+      RecordNumBlockShutdownTasksPostedDuringShutdown(
+          num_block_shutdown_tasks_posted_during_shutdown_);
+    }
   }
 }
 
@@ -116,29 +214,30 @@ void TaskTracker::RunTask(const Task* task) {
   AfterRunTask(shutdown_behavior);
 }
 
+bool TaskTracker::IsShutdownComplete() const {
+  AutoSchedulerLock auto_lock(shutdown_lock_);
+  return shutdown_event_ && shutdown_event_->IsSignaled();
+}
+
 bool TaskTracker::IsShuttingDownForTesting() const {
-  AutoSchedulerLock auto_lock(lock_);
-  return !!shutdown_cv_;
+  AutoSchedulerLock auto_lock(shutdown_lock_);
+  return shutdown_event_ && !shutdown_event_->IsSignaled();
 }
 
 bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
-  AutoSchedulerLock auto_lock(lock_);
-
-  if (shutdown_completed_) {
-    // A BLOCK_SHUTDOWN task posted after shutdown has completed is an ordering
-    // bug. This DCHECK aims to catch those early.
-    DCHECK_NE(shutdown_behavior, TaskShutdownBehavior::BLOCK_SHUTDOWN);
-
-    // No task is allowed to be posted after shutdown.
-    return false;
-  }
-
   if (shutdown_behavior == TaskShutdownBehavior::BLOCK_SHUTDOWN) {
     // BLOCK_SHUTDOWN tasks block shutdown between the moment they are posted
     // and the moment they complete their execution.
-    ++num_tasks_blocking_shutdown_;
+    const bool shutdown_started = state_->IncrementNumTasksBlockingShutdown();
+
+    if (shutdown_started) {
+      AutoSchedulerLock auto_lock(shutdown_lock_);
+
+      // A BLOCK_SHUTDOWN task posted after shutdown has completed is an
+      // ordering bug. This aims to catch those early.
+      DCHECK(shutdown_event_);
+      DCHECK(!shutdown_event_->IsSignaled());
 
-    if (shutdown_cv_) {
       ++num_block_shutdown_tasks_posted_during_shutdown_;
 
       if (num_block_shutdown_tasks_posted_during_shutdown_ ==
@@ -152,45 +251,52 @@ bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
       }
     }
 
-    // A BLOCK_SHUTDOWN task is allowed to be posted iff shutdown hasn't
-    // completed.
     return true;
   }
 
   // A non BLOCK_SHUTDOWN task is allowed to be posted iff shutdown hasn't
   // started.
-  return !shutdown_cv_;
+  return !state_->HasShutdownStarted();
 }
 
 bool TaskTracker::BeforeRunTask(TaskShutdownBehavior shutdown_behavior) {
-  AutoSchedulerLock auto_lock(lock_);
-
-  if (shutdown_completed_) {
-    // Trying to run a BLOCK_SHUTDOWN task after shutdown has completed is
-    // unexpected as it either shouldn't have been posted if shutdown completed
-    // or should be blocking shutdown if it was posted before it did.
-    DCHECK_NE(shutdown_behavior, TaskShutdownBehavior::BLOCK_SHUTDOWN);
+  switch (shutdown_behavior) {
+    case TaskShutdownBehavior::BLOCK_SHUTDOWN: {
+      // The number of tasks blocking shutdown has been incremented when the
+      // task was posted.
+      DCHECK(state_->AreTasksBlockingShutdown());
 
-    // A worker might extract a non BLOCK_SHUTDOWN task from a PriorityQueue
-    // after shutdown. It shouldn't be allowed to run it.
-    return false;
-  }
+      // Trying to run a BLOCK_SHUTDOWN task after shutdown has completed is
+      // unexpected as it either shouldn't have been posted if shutdown
+      // completed or should be blocking shutdown if it was posted before it
+      // did.
+      DCHECK(!state_->HasShutdownStarted() || !IsShutdownComplete());
 
-  switch (shutdown_behavior) {
-    case TaskShutdownBehavior::BLOCK_SHUTDOWN:
-      DCHECK_GT(num_tasks_blocking_shutdown_, 0U);
       return true;
+    }
+
+    case TaskShutdownBehavior::SKIP_ON_SHUTDOWN: {
+      // SKIP_ON_SHUTDOWN tasks block shutdown while they are running.
+      const bool shutdown_started = state_->IncrementNumTasksBlockingShutdown();
+
+      if (shutdown_started) {
+        // The SKIP_ON_SHUTDOWN task isn't allowed to run during shutdown.
+        // Decrement the number of tasks blocking shutdown that was wrongly
+        // incremented.
+        const bool shutdown_started_and_no_tasks_block_shutdown =
+            state_->DecrementNumTasksBlockingShutdown();
+        if (shutdown_started_and_no_tasks_block_shutdown)
+          OnBlockingShutdownTasksComplete();
 
-    case TaskShutdownBehavior::SKIP_ON_SHUTDOWN:
-      if (shutdown_cv_)
         return false;
+      }
 
-      // SKIP_ON_SHUTDOWN tasks block shutdown while they are running.
-      ++num_tasks_blocking_shutdown_;
       return true;
+    }
 
-    case TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN:
-      return !shutdown_cv_;
+    case TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN: {
+      return !state_->HasShutdownStarted();
+    }
   }
 
   NOTREACHED();
@@ -200,13 +306,22 @@ bool TaskTracker::BeforeRunTask(TaskShutdownBehavior shutdown_behavior) {
 void TaskTracker::AfterRunTask(TaskShutdownBehavior shutdown_behavior) {
   if (shutdown_behavior == TaskShutdownBehavior::BLOCK_SHUTDOWN ||
       shutdown_behavior == TaskShutdownBehavior::SKIP_ON_SHUTDOWN) {
-    AutoSchedulerLock auto_lock(lock_);
-    DCHECK_GT(num_tasks_blocking_shutdown_, 0U);
-    --num_tasks_blocking_shutdown_;
-    if (num_tasks_blocking_shutdown_ == 0 && shutdown_cv_)
-      shutdown_cv_->Signal();
+    const bool shutdown_started_and_no_tasks_block_shutdown =
+        state_->DecrementNumTasksBlockingShutdown();
+    if (shutdown_started_and_no_tasks_block_shutdown)
+      OnBlockingShutdownTasksComplete();
   }
 }
 
+void TaskTracker::OnBlockingShutdownTasksComplete() {
+  AutoSchedulerLock auto_lock(shutdown_lock_);
+
+  // This method can only be called after shutdown has started.
+  DCHECK(state_->HasShutdownStarted());
+  DCHECK(shutdown_event_);
+
+  shutdown_event_->Signal();
+}
+
 }  // namespace internal
 }  // namespace base