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 9e85b5cfc17b7b58c890e964294fbfc0f3680e72..d964d1c5f715036cfa8d42697cd5dd7dc3e5c0cf 100644
--- a/base/task_scheduler/task_tracker.cc
+++ b/base/task_scheduler/task_tracker.cc
@@ -4,6 +4,7 @@
 
 #include "base/task_scheduler/task_tracker.h"
 
+#include "base/atomicops.h"
 #include "base/callback.h"
 #include "base/debug/task_annotator.h"
 #include "base/metrics/histogram_macros.h"
@@ -36,33 +37,118 @@ 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 / kNumTasksBlockingShutdownIncrement;

[gab, 2016/06/20 16:46:20]

I think it would be simpler to have the shutdown bit in the MSB. With constexpr
the mask should be computable in kShutdownHasStartedMask at compile time as
something like 0x1 << 31.

Then you don't need weird divisions and +2/-2 business.

[fdoray, 2016/06/20 19:31:35]

I don't know how I would set the shutdown bit it it was the MSB because
base/atomicops.h ...
- ... doesn't have an atomic "and" (can't do atomic &= 2^31)
- ... only works with signed integers (can't do atomic += 2^31)

+    return num_tasks_blocking_shutdown != 0;
+  }
+
+  // Returns true if shutdown has started.
+  bool ShutdownHasStarted() const {
+    return subtle::NoBarrier_Load(&bits_) & kShutdownHasStartedMask;
+  }
+
+  // Returns true if there are tasks blocking shutdown.
+  bool TasksAreBlockingShutdown() const {
+    const auto num_tasks_blocking_shutdown =
+        subtle::NoBarrier_Load(&bits_) / kNumTasksBlockingShutdownIncrement;
+    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() {
+    const auto new_bits = subtle::NoBarrier_AtomicIncrement(
+        &bits_, kNumTasksBlockingShutdownIncrement);
+
+    // Verify that there is no overflow.
+    const auto num_tasks_blocking_shutdown =
+        new_bits / kNumTasksBlockingShutdownIncrement;
+    DCHECK_GE(num_tasks_blocking_shutdown, 0);
+
+    const bool shutdown_has_started = new_bits & kShutdownHasStartedMask;
+    return shutdown_has_started;
+  }
+
+  // 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::NoBarrier_AtomicIncrement(
+        &bits_, -kNumTasksBlockingShutdownIncrement);
+    const bool shutdown_has_started = new_bits & kShutdownHasStartedMask;
+    const auto num_tasks_blocking_shutdown =
+        new_bits / kNumTasksBlockingShutdownIncrement;
+    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 kNumTasksBlockingShutdownIncrement = 2;
+
+  // 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(lock_);
 
-  // This method should only be called once.
-  DCHECK(!shutdown_completed_);
-  DCHECK(!shutdown_cv_);
+    // This method can only be called once.
+    DCHECK(!shutdown_event_);
+    DCHECK(!num_block_shutdown_tasks_posted_during_shutdown_);
 
-  shutdown_cv_ = lock_.CreateConditionVariable();
+    shutdown_event_.reset(
+        new WaitableEvent(WaitableEvent::ResetPolicy::MANUAL,
+                          WaitableEvent::InitialState::NOT_SIGNALED));
 
-  // Wait until the number of tasks blocking shutdown is zero.
-  while (num_tasks_blocking_shutdown_ != 0)
-    shutdown_cv_->Wait();
+    const bool tasks_are_blocking_shutdown = state_->StartShutdown();
 
-  shutdown_cv_.reset();
-  shutdown_completed_ = true;
+    // From now, if a thread causes the number of tasks blocking shutdown to
+    // become zero, it will call OnBlockingShutdownTasksComplete().
 
-  // 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_);
+    if (!tasks_are_blocking_shutdown) {
+      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();
+
+  {
+    AutoSchedulerLock auto_lock(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 +202,30 @@ void TaskTracker::RunTask(const Task* task) {
   AfterRunTask(shutdown_behavior);
 }
 
-bool TaskTracker::IsShuttingDownForTesting() const {
+bool TaskTracker::ShutdownCompleted() const {
   AutoSchedulerLock auto_lock(lock_);
-  return !!shutdown_cv_;
+  return shutdown_event_ && shutdown_event_->IsSignaled();
 }
 
-bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
+bool TaskTracker::IsShuttingDownForTesting() const {
   AutoSchedulerLock auto_lock(lock_);
+  return !!shutdown_event_ && !shutdown_event_->IsSignaled();

[gab, 2016/06/20 16:46:20]

Don't think you need !! because the && explicitly turns this into a boolean
expression already.

[fdoray, 2016/06/20 19:31:35]

Done.

+}
 
-  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;
-  }
-
+bool TaskTracker::BeforePostTask(TaskShutdownBehavior shutdown_behavior) {
   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(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 +239,57 @@ 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_->ShutdownHasStarted();
 }
 
 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_->TasksAreBlockingShutdown());
 
-    // A WorkerThread 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_->ShutdownHasStarted() || !ShutdownCompleted());
 
-  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.
+        //
+        // Note: Reading the "shutdown initiated" bit and then incrementing the
+        // number of tasks blocking shutdown if it isn't set (without re-
+        // checking it afterwards) wouldn't work because the bit can be set
+        // between the two operations.

[gab, 2016/06/20 16:46:20]

As mentioned in my CL description comments I don't think this note is required.

[fdoray, 2016/06/20 19:31:35]

Done.

+        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_->ShutdownHasStarted();
+    }
   }
 
   NOTREACHED();
@@ -200,13 +299,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(lock_);
+
+  // This method can only be called after shutdown has started.
+  DCHECK(state_->ShutdownHasStarted());
+  DCHECK(shutdown_event_);
+
+  shutdown_event_->Signal();
+}
+
 }  // namespace internal
 }  // namespace base