Revert of Revert "Allow SequencedTaskRunnerHandle::Get() while running unsequenced tasks."

base/threading/sequenced_worker_pool.cc

Index: base/threading/sequenced_worker_pool.cc
diff --git a/base/threading/sequenced_worker_pool.cc b/base/threading/sequenced_worker_pool.cc
index d8d10f34e9f2998d79ba601a46fa6dfe0a4fdab4..0f2eac4cc45fb1751f5d882175e248e2691b7054 100644
--- a/base/threading/sequenced_worker_pool.cc
+++ b/base/threading/sequenced_worker_pool.cc
@@ -325,6 +325,11 @@
   bool RunsTasksOnCurrentThread() const;
 
   bool IsRunningSequenceOnCurrentThread(SequenceToken sequence_token) const;
+
+  bool IsRunningSequence(SequenceToken sequence_token) const;
+
+  void SetRunningTaskInfoForCurrentThread(SequenceToken sequence_token,
+                                          WorkerShutdown shutdown_behavior);
 
   void CleanupForTesting();
 
@@ -701,6 +706,28 @@
          sequence_token.Equals(found->second->task_sequence_token());
 }
 
+bool SequencedWorkerPool::Inner::IsRunningSequence(
+    SequenceToken sequence_token) const {
+  DCHECK(sequence_token.IsValid());
+  AutoLock lock(lock_);
+  return !IsSequenceTokenRunnable(sequence_token.id_);
+}
+
+void SequencedWorkerPool::Inner::SetRunningTaskInfoForCurrentThread(
+    SequenceToken sequence_token,
+    WorkerShutdown shutdown_behavior) {
+  AutoLock lock(lock_);
+  ThreadMap::const_iterator found = threads_.find(PlatformThread::CurrentId());
+  DCHECK(found != threads_.end());
+  DCHECK(found->second->is_processing_task());
+  DCHECK(!found->second->task_sequence_token().IsValid());
+  found->second->set_running_task_info(sequence_token, shutdown_behavior);
+
+  // Mark the sequence token as in use.
+  bool success = current_sequences_.insert(sequence_token.id_).second;
+  DCHECK(success);
+}
+
 // See https://code.google.com/p/chromium/issues/detail?id=168415
 void SequencedWorkerPool::Inner::CleanupForTesting() {
   DCHECK(!RunsTasksOnCurrentThread());
@@ -825,6 +852,11 @@
 
           tracked_objects::ThreadData::TallyRunOnNamedThreadIfTracking(
               task, stopwatch);
+
+          // Update the sequence token in case it has been set from within the
+          // task, so it can be removed from the set of currently running
+          // sequences in DidRunWorkerTask() below.
+          task.sequence_token_id = this_worker->task_sequence_token().id_;
 
           // Make sure our task is erased outside the lock for the
           // same reason we do this with delete_these_oustide_lock.
@@ -1209,6 +1241,33 @@
   return worker->worker_pool();
 }
 
+// static
+scoped_refptr<SequencedTaskRunner>
+SequencedWorkerPool::GetSequencedTaskRunnerForCurrentThread() {
+  Worker* worker = Worker::GetForCurrentThread();
+
+  // If there is no worker, this thread is not a worker thread. Otherwise, it is
+  // currently running a task (sequenced or unsequenced).
+  if (!worker)
+    return nullptr;
+
+  scoped_refptr<SequencedWorkerPool> pool = worker->worker_pool();
+  SequenceToken sequence_token = worker->task_sequence_token();
+  WorkerShutdown shutdown_behavior = worker->task_shutdown_behavior();
+  if (!sequence_token.IsValid()) {
+    // Create a new sequence token and bind this thread to it, to make sure that
+    // a task posted to the SequencedTaskRunner we are going to return is not
+    // immediately going to run on a different thread.
+    sequence_token = Inner::GetSequenceToken();
+    pool->inner_->SetRunningTaskInfoForCurrentThread(sequence_token,
+                                                     shutdown_behavior);
+  }
+
+  DCHECK(pool->IsRunningSequenceOnCurrentThread(sequence_token));
+  return new SequencedWorkerPoolSequencedTaskRunner(
+      std::move(pool), sequence_token, shutdown_behavior);
+}
+
 SequencedWorkerPool::SequencedWorkerPool(size_t max_threads,
                                          const std::string& thread_name_prefix,
                                          base::TaskPriority task_priority)
@@ -1353,6 +1412,11 @@
   return inner_->IsRunningSequenceOnCurrentThread(sequence_token);
 }
 
+bool SequencedWorkerPool::IsRunningSequence(
+    SequenceToken sequence_token) const {
+  return inner_->IsRunningSequence(sequence_token);
+}
+
 void SequencedWorkerPool::FlushForTesting() {
   inner_->CleanupForTesting();
 }