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..b1809fd26cac8b418cadd4a8fc35f8a88860aa0c 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,7 +37,67 @@ void RecordNumBlockShutdownTasksPostedDuringShutdown(
 
 }  // namespace
 
-TaskTracker::TaskTracker() = default;
+class TaskTracker::State {
+ public:
+  State() = default;
+
+  // Sets a flag indicating that shutdown has started.
+  void StartShutdown() {
+    subtle::Barrier_AtomicIncrement(&bits_, kShutdownHasStartedMask);

[robliao, 2016/05/27 22:16:14]

If this gets called twice, doesn't that mean the LSB is now 0?

[fdoray, 2016/05/30 15:48:04]

Yes. StartShutdown() can't be called twice. Added a DCHECK.

+  }
+
+  // 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(

[robliao, 2016/05/27 22:16:14]

Is this safe without a barrier?
If the atomic increment in IncrementNumTasks gets reordered after
DecrementNumTasks, we just told someone that the num of tasks blocking shutdown
hit 0.

[fdoray, 2016/05/30 15:48:04]

I believe a barrier "specifies how regular, non-atomic memory accesses are to be
ordered around an atomic operation"
http://en.cppreference.com/w/cpp/atomic/memory_order The atomic increment cannot
be reordered after the atomic decrement.

[robliao, 2016/06/01 18:32:26]

Our usage is not unlike shared_ptr refcounts. From the doc:
Typical use for relaxed memory ordering is incrementing counters, such as the
reference counters of std::shared_ptr, since this only requires atomicity, but
not ordering or synchronization (note that decrementing the shared_ptr counters
requires acquire-release synchronization with the destructor)

[fdoray, 2016/06/06 16:43:35]

Are you suggesting that I use Barrier_AtomicIncrement instead of
NoBarrier_AtomicIncrement here? That would add a memory barrier which I don't
think is useful since after calling DecrementNumTasksBlockingShutdown(), we
always call OnBlockingShutdownTasksComplete() which acquires a lock (implicit
memory barrier).

[robliao, 2016/06/06 19:16:30]

Yup, because it seems like we'll have a race without the barrier.

We decrement the count. An increment gets reordered before we make the call to
OnBlockingShutdownTasksComplete. The count is now non-zero when we make the
call.

[fdoray, 2016/06/16 13:19:37]

Do you worry about increment on thread A occurring between decrement and call to
OnBlockingShutdownTasksComplete() on thread B? This increment must have been
caused by one of these:

- A BLOCK_SHUTDOWN task was posted from a CONTINUE_ON_SHUTDOWN task or from a
thread that isn't joined before the TaskScheduler is shut down [1]: We should
disallow that.
- A SKIP_ON_SHUTDOWN task is about to be scheduled (BeforeRunTask): The counter
will be decremented immediately after because the shutdown bit is set
(https://codereview.chromium.org/2019763002/diff/60001/base/task_scheduler/tas...).
It doesn't matter if OnBlockingShutdownTasksComplete() is called on thread B
when the counter hasn't been decremented yet because no task is really blocking
shutdown.

[1] The task can't have been posted from a SKIP_ON_SHUTDOWN or BLOCK_SHUTDOWN
task. Otherwise, the counter wouldn't have become zero when it was decremented
on thread B.

[robliao, 2016/06/16 20:40:45]

I'm considering the atomics in isolation to the rest of the system. Relaxing the
memory consistency requirement based off of existing knowledge on how the rest
of the system is currently calling the state makes this less resilient towards
changes in calling pattern. It would be difficult for us to trace bugs through
here if we ran OnBlockingShutdownTasksComplete with a non-zero value in
num_tasks_blocking_shutdown.

[fdoray, 2016/06/17 20:28:49]

I'm not against adding a memory barrier... but it won't prevent another thread
from incrementing the counter (can happen in BeforeRunTask with
shutdown_behavior == SKIP_ON_SHUTDOWN) between a counter decrement and a call to
OnBlockingShutdownTasksComplete().

[robliao, 2016/06/18 00:10:27]

Oddly enough, I think the same thing can happen in the lock world.

An unlucky thread needs to increment the count and waits on the currently
running thread that happens to decrement the count to zero.

The old behavior is to DCHECK and then disallow execution. 

We can plausibly do the same here. The increment "fails" and the task is
rejected.

+        &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() {
@@ -45,18 +106,24 @@ void TaskTracker::Shutdown() {
   // This method should only be called once.
   DCHECK(!shutdown_completed_);
   DCHECK(!shutdown_cv_);
+  DCHECK(!num_block_shutdown_tasks_posted_during_shutdown_);
+
+  state_->StartShutdown();
+
+  // From now, it a thread causes the number of tasks blocking shutdown to
+  // become zero, it will call OnNumTasksBlockingShutdownIsZero().
 
   shutdown_cv_ = lock_.CreateConditionVariable();
 
   // Wait until the number of tasks blocking shutdown is zero.
-  while (num_tasks_blocking_shutdown_ != 0)
+  while (state_->TasksAreBlockingShutdown())
     shutdown_cv_->Wait();
 
   shutdown_cv_.reset();
   shutdown_completed_ = true;
 
-  // Record the TaskScheduler.BlockShutdownTasksPostedDuringShutdown if less
-  // than |kMaxBlockShutdownTasksPostedDuringShutdown| BLOCK_SHUTDOWN tasks were
+  // 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_ <
@@ -122,23 +189,18 @@ bool TaskTracker::IsShuttingDownForTesting() const {
 }
 
 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(lock_);
+
+      // A BLOCK_SHUTDOWN task posted after shutdown has completed is an
+      // ordering bug. This DCHECK aims to catch those early.
+      DCHECK(!shutdown_completed_);
 
-    if (shutdown_cv_) {
       ++num_block_shutdown_tasks_posted_during_shutdown_;
 
       if (num_block_shutdown_tasks_posted_during_shutdown_ ==
@@ -152,45 +214,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_->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() || !shutdown_completed());
 
-  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)
+          OnNumTasksBlockingShutdownIsZero();
 
-    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 +269,29 @@ 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)
+      OnNumTasksBlockingShutdownIsZero();
   }
 }
 
+void TaskTracker::OnNumTasksBlockingShutdownIsZero() {
+  AutoSchedulerLock auto_lock(lock_);
+
+  // This method can only be called after shutdown has started.
+  DCHECK(state_->ShutdownHasStarted());
+
+  // The condition below must be true because |lock_| is held:
+  // - Between the time |state_|->StartShutdown() is called and the time
+  //   |shutdown_cv_| is instantiated.
+  // - Between the time |shutdown_cv_| is deleted and the time
+  //   |shutdown_completed_| is set.
+  DCHECK(!!shutdown_cv_ || shutdown_completed_);
+
+  if (!!shutdown_cv_)
+    shutdown_cv_->Signal();
+}
+
 }  // namespace internal
 }  // namespace base