TaskScheduler: Add TaskTraits::WithWait().

WithWait() tasks are allowed to wait on other things than file I/O. In
particular, they may wait on a WaitableEvent or a ConditionVariable, join a
thread or a process, or make a blocking system call that doesn't involve
interactions with the file system.

In upcoming CLs, this trait will be used in scheduling decisions (more
details on the bug).

BUG=669247

Review-Url: https://codereview.chromium.org/2531663003
Cr-Commit-Position: refs/heads/master@{#435054}

base/task_scheduler/scheduler_worker_pool_impl_unittest.cc

@@ -115,13 +115,15 @@ class TaskSchedulerWorkerPoolImplTest
 scoped_refptr<TaskRunner> CreateTaskRunnerWithExecutionMode(
     SchedulerWorkerPoolImpl* worker_pool,
     test::ExecutionMode execution_mode) {
+  // Allow tasks posted to the returned TaskRunner to wait on a WaitableEvent.
+  const TaskTraits traits = TaskTraits().WithWait();
   switch (execution_mode) {
     case test::ExecutionMode::PARALLEL:
-      return worker_pool->CreateTaskRunnerWithTraits(TaskTraits());
+      return worker_pool->CreateTaskRunnerWithTraits(traits);
     case test::ExecutionMode::SEQUENCED:
-      return worker_pool->CreateSequencedTaskRunnerWithTraits(TaskTraits());
+      return worker_pool->CreateSequencedTaskRunnerWithTraits(traits);
     case test::ExecutionMode::SINGLE_THREADED:
-      return worker_pool->CreateSingleThreadTaskRunnerWithTraits(TaskTraits());
+      return worker_pool->CreateSingleThreadTaskRunnerWithTraits(traits);
   }
   ADD_FAILURE() << "Unknown ExecutionMode";
   return nullptr;
@@ -603,7 +605,7 @@ TEST_F(TaskSchedulerWorkerPoolCheckTlsReuse, CheckDetachedThreads) {
   std::vector<std::unique_ptr<test::TestTaskFactory>> factories;
   for (size_t i = 0; i < kNumWorkersInWorkerPool; ++i) {
     factories.push_back(MakeUnique<test::TestTaskFactory>(
-        worker_pool_->CreateTaskRunnerWithTraits(TaskTraits()),
+        worker_pool_->CreateTaskRunnerWithTraits(TaskTraits().WithWait()),
         test::ExecutionMode::PARALLEL));
     ASSERT_TRUE(factories.back()->PostTask(
         PostNestedTask::NO,
@@ -674,8 +676,8 @@ TEST_F(TaskSchedulerWorkerPoolHistogramTest, NumTasksBetweenWaits) {
   WaitableEvent event(WaitableEvent::ResetPolicy::MANUAL,
                       WaitableEvent::InitialState::NOT_SIGNALED);
   InitializeWorkerPool(TimeDelta::Max(), kNumWorkersInWorkerPool);
-  auto task_runner =
-      worker_pool_->CreateSequencedTaskRunnerWithTraits(TaskTraits());
+  auto task_runner = worker_pool_->CreateSequencedTaskRunnerWithTraits(
+      TaskTraits().WithWait());
 
   // Post a task.
   task_runner->PostTask(FROM_HERE,
@@ -718,7 +720,8 @@ TEST_F(TaskSchedulerWorkerPoolHistogramTest, NumTasksBetweenWaitsWithDetach) {
   WaitableEvent tasks_can_exit_event(WaitableEvent::ResetPolicy::MANUAL,
                                      WaitableEvent::InitialState::NOT_SIGNALED);
   InitializeWorkerPool(kReclaimTimeForDetachTests, kNumWorkersInWorkerPool);
-  auto task_runner = worker_pool_->CreateTaskRunnerWithTraits(TaskTraits());
+  auto task_runner =
+      worker_pool_->CreateTaskRunnerWithTraits(TaskTraits().WithWait());
 
   // Post tasks to saturate the pool.
   std::vector<std::unique_ptr<WaitableEvent>> task_started_events;

base/task_scheduler/task_tracker.cc

@@ -226,6 +226,9 @@ bool TaskTracker::RunTask(std::unique_ptr<Task> task,
         task->traits.shutdown_behavior() !=
         TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN);
 
+    const bool previous_wait_allowed =
+        ThreadRestrictions::SetWaitAllowed(task->traits.with_wait());
+
     {
       // Set up SequenceToken as expected for the scope of the task.
       ScopedSetSequenceTokenForCurrentThread
@@ -261,6 +264,8 @@ bool TaskTracker::RunTask(std::unique_ptr<Task> task,
       PerformRunTask(std::move(task));
     }
 
+    ThreadRestrictions::SetWaitAllowed(previous_wait_allowed);
+
     AfterRunTask(shutdown_behavior);
   }
 

base/task_scheduler/task_tracker_unittest.cc

@@ -264,7 +264,7 @@ TEST_P(TaskSchedulerTaskTrackerTest, WillPostAndRunLongTaskBeforeShutdown) {
                       WaitableEvent::InitialState::NOT_SIGNALED);
   auto blocked_task = base::MakeUnique<Task>(
       FROM_HERE, Bind(&WaitableEvent::Wait, Unretained(&event)),
-      TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta());
+      TaskTraits().WithWait().WithShutdownBehavior(GetParam()), TimeDelta());
 
   // Inform |task_tracker_| that |blocked_task| will be posted.
   EXPECT_TRUE(tracker_.WillPostTask(blocked_task.get()));
@@ -804,5 +804,54 @@ TEST_F(TaskSchedulerTaskTrackerTest, LoadWillPostAndRunDuringShutdown) {
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 }
 
+namespace {
+
+class WaitAllowedTestThread : public SimpleThread {
+ public:
+  WaitAllowedTestThread() : SimpleThread("WaitAllowedTestThread") {}
+
+ private:
+  void Run() override {
+    TaskTracker tracker;
+
+    // Waiting is allowed by default. Expect TaskTracker to disallow it before
+    // running a task without the WithWait() trait.
+    ThreadRestrictions::AssertWaitAllowed();
+    auto task_without_wait = MakeUnique<Task>(
+        FROM_HERE, Bind([]() {
+          EXPECT_DCHECK_DEATH({ ThreadRestrictions::AssertWaitAllowed(); });
+        }),
+        TaskTraits(), TimeDelta());
+    EXPECT_TRUE(tracker.WillPostTask(task_without_wait.get()));
+    tracker.RunTask(std::move(task_without_wait), SequenceToken::Create());
+
+    // Disallow waiting. Expect TaskTracker to allow it before running a task
+    // with the WithWait() trait.
+    ThreadRestrictions::DisallowWaiting();
+    auto task_with_wait =
+        MakeUnique<Task>(FROM_HERE, Bind([]() {
+                           // Shouldn't fail.
+                           ThreadRestrictions::AssertWaitAllowed();
+                         }),
+                         TaskTraits().WithWait(), TimeDelta());
+    EXPECT_TRUE(tracker.WillPostTask(task_with_wait.get()));
+    tracker.RunTask(std::move(task_with_wait), SequenceToken::Create());
+  }
+
+  DISALLOW_COPY_AND_ASSIGN(WaitAllowedTestThread);
+};
+
+}  // namespace
+
+// Verify that AssertIOAllowed() succeeds for a WithWait() task.
+TEST(TaskSchedulerTaskTrackerWaitAllowedTest, WaitAllowed) {
+  // Run the test on the separate thread since it is not possible to reset the
+  // "wait allowed" bit of a thread without being a friend of
+  // ThreadRestrictions.
+  WaitAllowedTestThread wait_allowed_test_thread;
+  wait_allowed_test_thread.Start();
+  wait_allowed_test_thread.Join();
+}
+
 }  // namespace internal
 }  // namespace base

base/task_scheduler/task_traits.cc

@@ -17,6 +17,7 @@ namespace base {
 // request defaults if the behavior is critical to the task.
 TaskTraits::TaskTraits()
     : with_file_io_(false),
+      with_wait_(false),
       priority_(TaskPriority::BACKGROUND),
       shutdown_behavior_(TaskShutdownBehavior::SKIP_ON_SHUTDOWN) {}
 
@@ -27,6 +28,11 @@ TaskTraits& TaskTraits::WithFileIO() {
   return *this;
 }
 
+TaskTraits& TaskTraits::WithWait() {
+  with_wait_ = true;
+  return *this;
+}
+
 TaskTraits& TaskTraits::WithPriority(TaskPriority priority) {
   priority_ = priority;
   return *this;

base/task_scheduler/task_traits.h

@@ -89,18 +89,28 @@ class BASE_EXPORT TaskTraits {
   TaskTraits& operator=(const TaskTraits& other) = default;
   ~TaskTraits();
 
-  // Allows tasks with these traits to do file I/O.
+  // Allows tasks with these traits to wait on synchronous file I/O.
   TaskTraits& WithFileIO();
 
+  // Allows tasks with these traits to wait on things other than file I/O. In
+  // particular, they may wait on a WaitableEvent or a ConditionVariable, join a
+  // thread or a process, or make a blocking system call that doesn't involve
+  // interactions with the file system.
+  TaskTraits& WithWait();
+
   // Applies |priority| to tasks with these traits.
   TaskTraits& WithPriority(TaskPriority priority);
 
   // Applies |shutdown_behavior| to tasks with these traits.
   TaskTraits& WithShutdownBehavior(TaskShutdownBehavior shutdown_behavior);
 
-  // Returns true if file I/O is allowed by these traits.
+  // Returns true if waiting on synchronous file I/O is allowed by these traits.
   bool with_file_io() const { return with_file_io_; }
 
+  // Returns true if waiting on things other than file I/O is allowed by these
+  // traits.
+  bool with_wait() const { return with_wait_; }
+
   // Returns the priority of tasks with these traits.
   TaskPriority priority() const { return priority_; }
 
@@ -109,6 +119,7 @@ class BASE_EXPORT TaskTraits {
 
  private:
   bool with_file_io_;
+  bool with_wait_;
   TaskPriority priority_;
   TaskShutdownBehavior shutdown_behavior_;
 };

base/threading/sequenced_worker_pool.cc

@@ -808,6 +808,7 @@ bool SequencedWorkerPool::Inner::PostTaskToTaskScheduler(
       static_cast<TaskShutdownBehavior>(sequenced.shutdown_behavior);
   const TaskTraits traits = TaskTraits()
                                 .WithFileIO()
+                                .WithWait()
                                 .WithPriority(task_priority_)
                                 .WithShutdownBehavior(task_shutdown_behavior);
   return GetTaskSchedulerTaskRunner(sequenced.sequence_token_id, traits)
@@ -858,7 +859,7 @@ bool SequencedWorkerPool::Inner::RunsTasksOnCurrentThread() const {
   if (g_all_pools_state == AllPoolsState::REDIRECTED_TO_TASK_SCHEDULER) {
     if (!runs_tasks_on_verifier_) {
       runs_tasks_on_verifier_ = CreateTaskRunnerWithTraits(
-          TaskTraits().WithFileIO().WithPriority(task_priority_));
+          TaskTraits().WithFileIO().WithWait().WithPriority(task_priority_));
     }
     return runs_tasks_on_verifier_->RunsTasksOnCurrentThread();
   } else {