Allow BLOCK_SHUTDOWN tasks during shutdown with SequencedWorkerPool redirected to TaskScheduler.

base/threading/sequenced_worker_pool_unittest.cc

 // Copyright (c) 2012 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/threading/sequenced_worker_pool.h"
 
 #include <stddef.h>
 
 #include <algorithm>
 #include <memory>
@@ skipping 578 matching lines @@
       SequencedWorkerPool::SKIP_ON_SHUTDOWN));
   EXPECT_FALSE(pool()->PostWorkerTaskWithShutdownBehavior(
       FROM_HERE,
       base::Bind(&TestTracker::FastTask, tracker(), 102),
       SequencedWorkerPool::BLOCK_SHUTDOWN));
 
   ASSERT_EQ(old_has_work_call_count, has_work_call_count());
 }
 
 TEST_P(SequencedWorkerPoolTest, AllowsAfterShutdown) {
-  // As tested by TaskSchedulerTaskTrackerTest.WillPostAndRunDuringShutdown,
-  // TaskScheduler allows tasks to be posted during shutdown. However, since it
-  // doesn't provide a way to run a callback from inside its Shutdown() method,
-  // it would be hard to make this test work with redirection enabled.
-  if (RedirectedToTaskScheduler())
-    return;
-
   // Test that <n> new blocking tasks are allowed provided they're posted
   // by a running tasks.
   EnsureAllWorkersCreated();
   ThreadBlocker blocker;
 
   // Start tasks to take all the threads and block them.
   const int kNumBlockTasks = static_cast<int>(kNumWorkerThreads);
   for (int i = 0; i < kNumBlockTasks; ++i) {
     EXPECT_TRUE(pool()->PostWorkerTask(
         FROM_HERE,
         base::Bind(&TestTracker::BlockTask, tracker(), i, &blocker)));
   }
   tracker()->WaitUntilTasksBlocked(kNumWorkerThreads);
 
   // Queue up shutdown blocking tasks behind those which will attempt to post
-  // additional tasks when run, PostAdditionalTasks attemtps to post 3
+  // additional tasks when run, PostAdditionalTasks attempts to post 3
   // new FastTasks, one for each shutdown_behavior.
   const int kNumQueuedTasks = static_cast<int>(kNumWorkerThreads);
   for (int i = 0; i < kNumQueuedTasks; ++i) {
     EXPECT_TRUE(pool()->PostWorkerTaskWithShutdownBehavior(
         FROM_HERE, base::Bind(&TestTracker::PostAdditionalTasks, tracker(), i,
                               base::RetainedRef(pool()), false),
         SequencedWorkerPool::BLOCK_SHUTDOWN));
   }
 
-  // Setup to open the floodgates from within Shutdown().
-  SetWillWaitForShutdownCallback(
-      base::Bind(&EnsureTasksToCompleteCountAndUnblock,
-                 scoped_refptr<TestTracker>(tracker()),
-                 0, &blocker, kNumBlockTasks));
+  // Half the additional blocking tasks will be allowed to run.
+  constexpr int kNumNewBlockingTasksToAllow = kNumWorkerThreads / 2;
 
-  // Allow half of the additional blocking tasks thru.
-  const int kNumNewBlockingTasksToAllow = kNumWorkerThreads / 2;
-  pool()->Shutdown(kNumNewBlockingTasksToAllow);
+  if (RedirectedToTaskScheduler()) {
+    // When redirection to TaskScheduler is enabled,
+    // SequencedWorkerPool::Shutdown() sets the number of additional
+    // BLOCK_SHUTDOWN tasks that can be posted and returns without waiting for
+    // pending BLOCK_SHUTDOWN tasks to complete their execution.
+    pool()->Shutdown(kNumNewBlockingTasksToAllow);
+
+    // Unblock tasks.
+    EnsureTasksToCompleteCountAndUnblock(tracker(), 0, &blocker,
+                                         kNumBlockTasks);
+
+    // TaskScheduler::Shutdown() waits for pending BLOCK_SHUTDOWN tasks to
+    // complete their execution.
+    TaskScheduler::GetInstance()->Shutdown();
+  } else {
+    // Once shutdown starts, unblock tasks.
+    SetWillWaitForShutdownCallback(base::Bind(
+        &EnsureTasksToCompleteCountAndUnblock,
+        scoped_refptr<TestTracker>(tracker()), 0, &blocker, kNumBlockTasks));
+
+    // Set the number of additional BLOCK_SHUTDOWN tasks that can be posted and
+    // wait for pending BLOCK_SHUTDOWN tasks to complete their execution.
+    pool()->Shutdown(kNumNewBlockingTasksToAllow);
+  }
 
   // Ensure that the correct number of tasks actually got run.
   tracker()->WaitUntilTasksComplete(static_cast<size_t>(
       kNumBlockTasks + kNumQueuedTasks + kNumNewBlockingTasksToAllow));
 
   // Clean up the task IDs we added and go home.
   tracker()->ClearCompleteSequence();
 }
 
 // Tests that blocking tasks can still be posted during shutdown, as long as
 // the task is not being posted within the context of a running task.
 TEST_P(SequencedWorkerPoolTest,
        AllowsBlockingTasksDuringShutdownOutsideOfRunningTask) {
-  // As tested by TaskSchedulerTaskTrackerTest.WillPostAndRunDuringShutdown,
-  // TaskScheduler allows tasks to be posted during shutdown. However, since it
-  // doesn't provide a way to run a callback from inside its Shutdown() method,
-  // it would be hard to make this test work with redirection enabled.
-  if (RedirectedToTaskScheduler())
-    return;
-
   EnsureAllWorkersCreated();
   ThreadBlocker blocker;
 
   // Start tasks to take all the threads and block them.
   const int kNumBlockTasks = static_cast<int>(kNumWorkerThreads);
   for (int i = 0; i < kNumBlockTasks; ++i) {
     EXPECT_TRUE(pool()->PostWorkerTask(
         FROM_HERE,
         base::Bind(&TestTracker::BlockTask, tracker(), i, &blocker)));
   }
   tracker()->WaitUntilTasksBlocked(kNumWorkerThreads);
 
-  // Setup to open the floodgates from within Shutdown().
-  SetWillWaitForShutdownCallback(
-      base::Bind(&TestTracker::PostBlockingTaskThenUnblockThreads,
-                 scoped_refptr<TestTracker>(tracker()), pool(), &blocker,
-                 kNumWorkerThreads));
-  pool()->Shutdown(kNumWorkerThreads + 1);
+  constexpr int kNumNewBlockingTasksToAllow = 1;
+
+  if (RedirectedToTaskScheduler()) {
+    // When redirection to TaskScheduler is enabled,
+    // SequencedWorkerPool::Shutdown() sets the number of additional
+    // BLOCK_SHUTDOWN tasks that can be posted and returns without waiting for
+    // pending BLOCK_SHUTDOWN tasks to complete their execution.
+    pool()->Shutdown(kNumNewBlockingTasksToAllow);
+
+    // Post a blocking task and unblock tasks.
+    tracker()->PostBlockingTaskThenUnblockThreads(pool(), &blocker,
+                                                  kNumWorkerThreads);
+
+    // TaskScheduler::Shutdown() waits for pending BLOCK_SHUTDOWN tasks to
+    // complete their execution.
+    TaskScheduler::GetInstance()->Shutdown();
+  } else {
+    // Once shutdown starts, post a blocking task and unblock tasks.
+    SetWillWaitForShutdownCallback(
+        base::Bind(&TestTracker::PostBlockingTaskThenUnblockThreads,
+                   scoped_refptr<TestTracker>(tracker()), pool(), &blocker,
+                   kNumWorkerThreads));
+
+    // Set the number of additional BLOCK_SHUTDOWN tasks that can be posted and
+    // wait for pending BLOCK_SHUTDOWN tasks to complete their execution.
+    pool()->Shutdown(kNumNewBlockingTasksToAllow);
+  }
 
   // Ensure that the correct number of tasks actually got run.
-  tracker()->WaitUntilTasksComplete(static_cast<size_t>(kNumWorkerThreads + 1));
+  tracker()->WaitUntilTasksComplete(
+      static_cast<size_t>(kNumWorkerThreads + kNumNewBlockingTasksToAllow));
   tracker()->ClearCompleteSequence();
 }
 
 // Tests that unrun tasks are discarded properly according to their shutdown
 // mode.
 TEST_P(SequencedWorkerPoolTest, DiscardOnShutdown) {
   // As tested by
   // TaskSchedulerTaskTrackerTest.WillPostBeforeShutdownRunDuringShutdown, on
   // shutdown, the TaskScheduler discards SKIP_ON_SHUTDOWN and
   // CONTINUE_ON_SHUTDOWN tasks and runs BLOCK_SHUTDOWN tasks. However, since it
@@ skipping 514 matching lines @@
     SequencedWorkerPoolSequencedTaskRunner, SequencedTaskRunnerTest,
     SequencedWorkerPoolSequencedTaskRunnerTestDelegate);
 INSTANTIATE_TYPED_TEST_CASE_P(
     SequencedWorkerPoolSequencedTaskRunner,
     SequencedTaskRunnerDelayedTest,
     SequencedWorkerPoolSequencedTaskRunnerTestDelegate);
 
 }  // namespace
 
 }  // namespace base