TaskScheduler: Atomic operations in TaskTracker

base/task_scheduler/task_tracker.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 <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"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/trace_event/trace_event.h"
 
 namespace base {
 namespace internal {
 
 namespace {
@@ skipping 10 matching lines @@
 
 void RecordNumBlockShutdownTasksPostedDuringShutdown(
     HistogramBase::Sample value) {
   UMA_HISTOGRAM_CUSTOM_COUNTS(
       "TaskScheduler.BlockShutdownTasksPostedDuringShutdown", value, 1,
       kMaxBlockShutdownTasksPostedDuringShutdown, 50);
 }
 
 }  // 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 >> kNumTasksBlockingShutdownBitOffset;
+    return num_tasks_blocking_shutdown != 0;
+  }
+
+  // Returns true if shutdown has started.
+  bool HasShutdownStarted() const {
+    return subtle::NoBarrier_Load(&bits_) & kShutdownHasStartedMask;
+  }
+
+  // Returns true if there are tasks blocking shutdown.
+  bool AreTasksBlockingShutdown() const {
+    const auto num_tasks_blocking_shutdown =
+        subtle::NoBarrier_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::NoBarrier_Load(&bits_) >> kNumTasksBlockingShutdownBitOffset;
+    DCHECK_LT(num_tasks_blocking_shutdown,

[danakj, 2016/07/08 21:21:29]

So I'm still a bit unsure why this is a DCHECK instead of handling it. What
makes you confident this won't actually happen?

[fdoray, 2016/07/18 19:01:15]

Having 2^30 pending tasks requires more than 2^36 bytes of RAM (sizeof(Task) ==
120). That means that a 32-bit Chrome will get an OOM error before we reach 2^30
pending tasks. We don't handle that gracefully. Similarly, I don't think we
should handle incrementing the number of tasks blocking shutdown above 2^30
(maybe we should even remove the DCHECK?).

Note that |num_tasks_blocking_shutdown| is the number of tasks that are
*currently* blocking shutdown, not the number of tasks that have blocked
shutdown over the lifetime of the process.

+              std::numeric_limits<subtle::Atomic32>::max() -
+                  kNumTasksBlockingShutdownIncrement);
+#endif
+
+    const auto new_bits = subtle::NoBarrier_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::NoBarrier_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.
+  //
+  // This atomic variable is always read/written without a barrier. This is
+  // correct because it encapsulates all shutdown state. Barrier semantics would

[danakj, 2016/07/08 21:21:29]

I spent some time learning about atomics yesterday, and I don't understand this
comment about not needing barriers. Can you explain this further?

AFAIK using atomic increment etc ensures that the operation happens atomically,
so that no one would ever see the value half-changed. But it does not guarantee
that other threads will see the new value at that time, they may never see it.
But the logic in this patch depends on changes to the bits_ value being visible
on other threads immediately.

[gab, 2016/07/13 18:49:10]

Quoting Francois from way further up in this review:
"""
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.
"""

So what we understand is that atomic operations themselves are always atomic
(and see each other in memory). Adding a barrier merely makes sure that
everything else around it has also been synchronized. In this case since we
merely care about that specific value and nothing else around it: NoBarrier is
appropriate.

For example: atomic_ref_count.h uses NoBarrier on increment but Barrier on
decrement. It does so not to make sure multiple threads don't hit 0 but to make
sure that the thread that hits zero inserts a barrier so that it sees the most
recent version of surrounding non-atomic state (e.g. to ensure a clean deletion
or whatever else the zero'ed ref triggers). In our case we don't need barriers
either way as no non-atomic state depends on being synced based on calls to
TaskTracker::State.

[danakj, 2016/07/13 22:39:17]

That's not my understanding at all. The atomic operation means that you'll never
see the operation half-complete. It does not promise that other threads will
ever see the new value vs the old value.

Here's what I read the other day on this:
http://preshing.com/20130618/atomic-vs-non-atomic-operations/

Our base atomics almost always use std::memory_order_relaxed, which means that
"Relaxed operation: there are no synchronization or ordering constraints, only
atomicity is required of this operation." it only guarantees the memory
modified/read by the atomic operation is seen entirely before or entirely after
the operation. But provides no guarantees about what value is seen otherwise.

Without the barrier thread A may change the value, and thread B may never read
the new value (but the atomics will guarantee it doesn't read a half-changed
value).
I think the problem is where we have situtations like:

Thread A: Posts a task, TaskTracker::BeforePostTask does
IncrementNumTasksBlockingShutdown().
Thread B: Runs a task, DCHECKs AreTasksBlockingShutdown(), what guantees it sees
it? Maybe in this case there's a barrier somewhere due to the fact the task has
made it to the other thread. But are we guaranteed that the atomic op made it
too, since we used relaxed?

Or

Thread A: TaskTracker::BeforePostTask() does
IncrementNumTasksBlockingShutdown(), sees false for shutdown started.
Thread B: TaskTracker::Shutdown() sets the shutdown flag.
Thread A: checks HasShutdownStarted(), this check may not see the shutdown state
changed. Is that okay? Maybe that should just be returning the already queried
bool then anyways? Shutdown could start just after it checks the flag too.

[fdoray, 2016/07/18 19:01:15]

According to your link, atomic operations without barriers offer no ordering
guarantees. Knowing that, I agree that we should use barriers every time we
read/write |bits_|.

On 2016/07/13 22:39:17, danakj wrote:
This should no longer be an issue with Barrier_AtomicIncrement and Acquire_Load
(since IncrementNumTasksBlockingShutdown runs before AreTasksBlockingShutdown
and a memory barrier is used, the DCHECK should never fail).
We never call both IncrementNumTasksBlockingShutdown() and HasShutdownStarted()
in BeforePostTask() (there is a code path for BLOCK_SHUTDOWN tasks and a code
path for other tasks).

[gab, 2016/07/19 16:22:20]

Now that we've changed to use barriers, it would be worth it to redo the
benchmarks in the CL description to get another column for barriers versus no
barriers.

Great reads I found to educate myself further:

http://preshing.com/20130618/atomic-vs-non-atomic-operations/
http://preshing.com/20120710/memory-barriers-are-like-source-control-operations/
http://preshing.com/20120913/acquire-and-release-semantics/
http://preshing.com/20120612/an-introduction-to-lock-free-programming/
http://preshing.com/20130922/acquire-and-release-fences/

Given all of this I think the barriers are correct for the consistency of this
API alone (although as was mentioned in practice there are other things ensuring
consistencies of the actual calls but this API can't know that).

So long as the benchmarks w/ barriers show that it doesn't matter much I'm happy
with this -- otherwise we may want to go a notch further.

+  // 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 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) {
+      // 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();
+
+  {
+    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_);
+    }
   }
 }
 
 bool TaskTracker::WillPostTask(const Task* task) {
   DCHECK(task);
 
   if (!BeforePostTask(task->traits.shutdown_behavior()))
     return false;
 
   debug::TaskAnnotator task_annotator;
@@ skipping 33 matching lines @@
 
     TRACE_TASK_EXECUTION(kRunFunctionName, *task);
 
     debug::TaskAnnotator task_annotator;
     task_annotator.RunTask(kQueueFunctionName, *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_cv_) {
+    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());
+
       ++num_block_shutdown_tasks_posted_during_shutdown_;
 
       if (num_block_shutdown_tasks_posted_during_shutdown_ ==
           kMaxBlockShutdownTasksPostedDuringShutdown) {
         // Record the TaskScheduler.BlockShutdownTasksPostedDuringShutdown
         // histogram as soon as its upper bound is hit. That way, a value will
         // be recorded even if an infinite number of BLOCK_SHUTDOWN tasks are
         // posted, preventing shutdown to complete.
         RecordNumBlockShutdownTasksPostedDuringShutdown(
             num_block_shutdown_tasks_posted_during_shutdown_);
       }
     }
 
-    // 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_);
+  switch (shutdown_behavior) {
+    case TaskShutdownBehavior::BLOCK_SHUTDOWN: {
+      // The number of tasks blocking shutdown has been incremented when the
+      // task was posted.
+      DCHECK(state_->AreTasksBlockingShutdown());
 
-  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);
+      // 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());
 
-    // A worker might extract a non BLOCK_SHUTDOWN task from a PriorityQueue
-    // after shutdown. It shouldn't be allowed to run it.
-    return false;
-  }
+      return true;
+    }
 
-  switch (shutdown_behavior) {
-    case TaskShutdownBehavior::BLOCK_SHUTDOWN:
-      DCHECK_GT(num_tasks_blocking_shutdown_, 0U);
+    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();
+
+        return false;
+      }
+
       return true;
+    }
 
-    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();
   return false;
 }
 
 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