TaskScheduler: Add TaskTraits::WithWait().

base/task_scheduler/task_tracker_unittest.cc

Index: base/task_scheduler/task_tracker_unittest.cc
diff --git a/base/task_scheduler/task_tracker_unittest.cc b/base/task_scheduler/task_tracker_unittest.cc
index 10e116c49d45b228520d80b23e5c12f81df019cd..d65ee3075b973d933527261c7c0caae25df5dba4 100644
--- a/base/task_scheduler/task_tracker_unittest.cc
+++ b/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()));
@@ -459,6 +459,59 @@ TEST_P(TaskSchedulerTaskTrackerTest, IOAllowed) {
   tracker.RunTask(std::move(task_without_file_io), SequenceToken::Create());
 }
 
+namespace {
+
+class WaitAllowedTestThread : public SimpleThread {
+ public:
+  WaitAllowedTestThread(TaskShutdownBehavior shutdown_behavior)
+      : SimpleThread("WaitAllowedTestThread"),
+        shutdown_behavior_(shutdown_behavior) {}
+
+ private:
+  void Run() override {
+    TaskTracker tracker;
+
+    // Waiting is allowed by default. Expect TaskTracker to disallow it before
+    // running a task without the WithWait() trait.
+    auto task_without_wait = MakeUnique<Task>(
+        FROM_HERE, Bind([]() {
+          EXPECT_DCHECK_DEATH({ ThreadRestrictions::AssertWaitAllowed(); });
+        }),
+        TaskTraits().WithShutdownBehavior(shutdown_behavior_), 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().WithShutdownBehavior(shutdown_behavior_),
+        TimeDelta());
+    EXPECT_TRUE(tracker.WillPostTask(task_with_wait.get()));
+    tracker.RunTask(std::move(task_with_wait), SequenceToken::Create());
+  }
+
+  const TaskShutdownBehavior shutdown_behavior_;
+
+  DISALLOW_COPY_AND_ASSIGN(WaitAllowedTestThread);
+};
+
+}  // namespace
+
+// Verify that AssertIOAllowed() succeeds for a WithWait() task.
+TEST_P(TaskSchedulerTaskTrackerTest, 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(GetParam());
+  wait_allowed_test_thread.Start();
+  wait_allowed_test_thread.Join();
+}
+
 static void RunTaskRunnerHandleVerificationTask(
     TaskTracker* tracker,
     std::unique_ptr<Task> verify_task) {