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..a9fdd627ef11abb9013eae8eeffc01abadc5fd29 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,125 @@ 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::NoBarrier_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;
+    return num_tasks_blocking_shutdown != 0;
+  }
+
+  // Returns true if shutdown has started.
+  bool ShutdownHasStarted() const {

[danakj, 2016/06/30 22:25:36]

nit: i prefer names that start with an action, in the cases of getters,
indicating a query: HasShutdownStarted, AreTasksBlockingShutdown.

[fdoray, 2016/07/04 20:53:11]

Done.

+    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);
+
+#if DCHECK_IS_ON()
+    // Verify that there is no overflow.

[danakj, 2016/06/30 22:25:36]

Signed integer overflow is undefined is it not? I think you'll need to instead
DCHECK on the iteration before where adding anything else *would* be overflow.
But is DCHECK enough? What prevents it from happening, should we reject tasks if
it would overflow?

[fdoray, 2016/07/04 20:53:11]

I moved the DCHECK before the increment. Note that it is racy.

This counter will overflow if there are 2^30 *pending* tasks. I assume that if
this happens, something is wrong. Getting a crash report (CHECK instead of
DCHECK) might be a better solution than silently rejecting tasks.

+    const auto num_tasks_blocking_shutdown =
+        new_bits >> kShutdownHasStartedMask;

[danakj, 2016/06/30 22:25:36]

You mean / Increment?

[fdoray, 2016/07/04 20:53:11]

Gab had a similar comment
https://codereview.chromium.org/2019763002/diff/80001/base/task_scheduler/tas...

It is now new_bits >> kNumTasksBlockingShutdownBitOffset.

[danakj, 2016/07/06 18:51:48]

Are you sure you want to >> a signed integer? Does DCHECK_GE even make sense
after you did a right shift? I assumed you were doing division to preserve the
sign bit?

+    DCHECK_GE(num_tasks_blocking_shutdown, 0);
+#endif
+
+    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 >> kShutdownHasStartedMask;

[danakj, 2016/06/30 22:25:36]

you mean / Increment?

[fdoray, 2016/07/04 20:53:11]

ditto

+    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 =
+      1 << kShutdownHasStartedMask;

[danakj, 2016/06/30 22:25:36]

I don't think this is quite right logically, if mask was 11b you'd want to shift
by 2, not by 3. Maybe just write 2 (or mask+1 if you need to do this mathy) with
a comment explaining it's the number of bits in the mask?

[fdoray, 2016/07/04 20:53:11]

ditto

+
+  // The LSB indicates whether shutdown has started. The other bits count the
+  // number of tasks blocking shutdown.
+  //
+  // This atomic variable is always read/written without a barrier. This is
+  // correct because it encapsulates all shutdown state. Atomic behavior would

[gab, 2016/06/29 21:38:47]

s/Atomic behavior/Barrier semantics/ ?

[fdoray, 2016/07/04 20:53:11]

Done.

+  // need to be reassessed if another variable was added.
+  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_);
+
+    shutdown_event_.reset(
+        new WaitableEvent(WaitableEvent::ResetPolicy::MANUAL,
+                          WaitableEvent::InitialState::NOT_SIGNALED));
 
-  // This method should only be called once.
-  DCHECK(!shutdown_completed_);
-  DCHECK(!shutdown_cv_);
+    const bool tasks_are_blocking_shutdown = state_->StartShutdown();
 
-  shutdown_cv_ = lock_.CreateConditionVariable();
+    // From now, if a thread causes the number of tasks blocking shutdown to

[danakj, 2016/06/30 22:25:36]

What happens if a BLOCK_SHUTDOWN task is posted right here?

It sees shutdown has started, takes the lock waits for this block to exit, then
increments num_block_shutdown_tasks_posted_during_shutdown_ and returns *true*?

I think that's different than before right? Then we have a task in queue after
Shutdown is done.

[fdoray, 2016/07/04 20:53:11]

You mean if a BLOCK_SHUTDOWN task is posted right here when the number of tasks
blocking shutdown is zero? Line 238 will DCHECK.

Posting a BLOCK_SHUTDOWN task when shutdown is complete isn't supported. I would
argue that shutdown is complete when Shutdown() has entered and the number of
tasks blocking shutdown is zero.

If we disallowed posting BLOCK_SHUTDOWN tasks from CONTINUE_ON_SHUTDOWN tasks
and from external threads that aren't joined before the task scheduler is shut
down, this couldn't happen.

[danakj, 2016/07/06 18:50:24]

The DCHECK is that the event is not Signaled. It is not Signalled yet here, that
happens just below. There is some time between the StartShutdown() call the
Signal of the shutdown event, which is where BLOCK_SHUTDOWN can be posted
(shutdown isn't done), but it will see that shutdown has started and assumes we
will Wait here, which we wouldn't.
But the code is using the Signal to determine it right?

[fdoray, 2016/07/06 19:19:40]

Thread A: state_->StartShutdown();
Thread B: Starts posting a BLOCK_SHUTDOWN task. BeforePostTask blocks when it
tries to acquire |shutdown_lock_| because it is still held by thread A.
Thread A: Signals the shutdown event to indicate that shutdown is complete.
Releases the lock.
Thread B: Acquires the lock. DCHECK(!shutdown_event_->IsSignaled()) fails
because it is not valid to post a task when shutdown is complete (and shutdown
is complete as soon as StartShutdown() has been called and num tasks blocking
shutdown is zero).

[danakj, 2016/07/07 22:59:10]

Ah I see. The lock makes this work. Ok thanks for the explanation. Please
consider adding some comments in places to explain this, like on the lock grabs
and stuff.

+    // become zero, it will call OnBlockingShutdownTasksComplete().
 
-  // Wait until the number of tasks blocking shutdown is zero.
-  while (num_tasks_blocking_shutdown_ != 0)
-    shutdown_cv_->Wait();
+    if (!tasks_are_blocking_shutdown) {
+      shutdown_event_->Signal();
+      return;
+    }
+  }
 
-  shutdown_cv_.reset();
-  shutdown_completed_ = true;
+  // 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 +209,30 @@ void TaskTracker::RunTask(const Task* task) {
   AfterRunTask(shutdown_behavior);
 }
 
+bool TaskTracker::ShutdownCompleted() const {

[danakj, 2016/06/30 22:25:36]

same naming nit: IsShutdownComplete?

[fdoray, 2016/07/04 20:53:11]

Done.

+  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) {

[danakj, 2016/06/30 22:25:36]

If shutdown started, it may also have finished by now, how do we handle that?

[fdoray, 2016/07/04 20:53:11]

If shutdown is complete, line 238 will DCHECK:
DCHECK(!shutdown_event_->IsSignaled());

+      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 +246,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();

[danakj, 2016/06/30 22:25:36]

What if shutdown started between this and the check above?

[fdoray, 2016/07/04 20:53:11]

Which check above? The rest of this method is for BLOCK_SHUTDOWN tasks and this
line is for non-BLOCK_SHUTDOWN tasks.

 }
 
 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 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_->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.
+        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 +301,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_->ShutdownHasStarted());
+  DCHECK(shutdown_event_);
+
+  shutdown_event_->Signal();
+}
+
 }  // namespace internal
 }  // namespace base