TaskScheduler: Set the IO allowed bit in TaskTracker::RunTask().

base/task_scheduler/task_tracker_unittest.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 <stdint.h>
 
 #include <memory>
 #include <utility>
@@ skipping 413 matching lines @@
 
   // Running the task should fail iff the task isn't allowed to use singletons.
   if (can_use_singletons) {
     EXPECT_TRUE(tracker.RunTask(std::move(task), SequenceToken::Create()));
   } else {
     EXPECT_DCHECK_DEATH(
         { tracker.RunTask(std::move(task), SequenceToken::Create()); });
   }
 }
 
+// Verify that AssertIOAllowed() succeeds only for a WithFileIO() task.
+TEST_P(TaskSchedulerTaskTrackerTest, IOAllowed) {
+  TaskTracker tracker;
+
+  // Unset the IO allowed bit. Expect TaskTracker to set it before running a
+  // task with the WithFileIO() trait.
+  ThreadRestrictions::SetIOAllowed(false);
+  auto task_with_file_io = MakeUnique<Task>(
+      FROM_HERE, Bind([]() {
+        // Shouldn't fail.
+        ThreadRestrictions::AssertIOAllowed();
+      }),
+      TaskTraits().WithFileIO().WithShutdownBehavior(GetParam()), TimeDelta());
+  EXPECT_TRUE(tracker.WillPostTask(task_with_file_io.get()));
+  tracker.RunTask(std::move(task_with_file_io), SequenceToken::Create());
+
+  // Set the IO allowed bit. Expect TaskTracker to unset it before running a
+  // task without the WithFileIO() trait.
+  ThreadRestrictions::SetIOAllowed(true);
+  auto task_without_file_io = MakeUnique<Task>(
+      FROM_HERE, Bind([]() {
+        EXPECT_DCHECK_DEATH({ ThreadRestrictions::AssertIOAllowed(); });
+      }),
+      TaskTraits().WithShutdownBehavior(GetParam()), TimeDelta());
+  EXPECT_TRUE(tracker.WillPostTask(task_without_file_io.get()));
+  tracker.RunTask(std::move(task_without_file_io), SequenceToken::Create());
+}
+
 static void RunTaskRunnerHandleVerificationTask(
     TaskTracker* tracker,
     std::unique_ptr<Task> verify_task) {
   // Pretend |verify_task| is posted to respect TaskTracker's contract.
   EXPECT_TRUE(tracker->WillPostTask(verify_task.get()));
 
   // Confirm that the test conditions are right (no TaskRunnerHandles set
   // already).
   EXPECT_FALSE(ThreadTaskRunnerHandle::IsSet());
   EXPECT_FALSE(SequencedTaskRunnerHandle::IsSet());
@@ skipping 355 matching lines @@
 
   // Unblock shutdown by running |block_shutdown_task|.
   EXPECT_TRUE(tracker_.RunTask(std::move(block_shutdown_task),
                                SequenceToken::Create()));
   EXPECT_EQ(kLoadTestNumIterations + 1, NumTasksExecuted());
   WAIT_FOR_ASYNC_SHUTDOWN_COMPLETED();
 }
 
 }  // namespace internal
 }  // namespace base