Add TaskScheduler::JoinForTesting().

base/task_scheduler/scheduler_worker.cc

Index: base/task_scheduler/scheduler_worker.cc
diff --git a/base/task_scheduler/scheduler_worker.cc b/base/task_scheduler/scheduler_worker.cc
index 970a7d4cf908a9b4aa8249f988e1ebeaed1e41ba..14f35658efca904a105190ad3c51970e52fa71c9 100644
--- a/base/task_scheduler/scheduler_worker.cc
+++ b/base/task_scheduler/scheduler_worker.cc
@@ -216,6 +216,9 @@ SchedulerWorker::~SchedulerWorker() {
 
 void SchedulerWorker::WakeUp() {
   AutoSchedulerLock auto_lock(thread_lock_);
+
+  DCHECK(!should_exit_for_testing_.IsSet());
+
   if (!thread_)
     CreateThreadAssertSynchronized();
 
@@ -227,17 +230,21 @@ void SchedulerWorker::JoinForTesting() {
   DCHECK(!should_exit_for_testing_.IsSet());
   should_exit_for_testing_.Set();
 
-  WakeUp();
+  std::unique_ptr<Thread> thread;
 
-  // Normally holding a lock and joining is dangerous. However, since this is
-  // only for testing, we're okay since the only scenario that could impact this
-  // is a call to Detach, which is disallowed by having the delegate always
-  // return false for the CanDetach call.
-  AutoSchedulerLock auto_lock(thread_lock_);
-  if (thread_)
-    thread_->Join();
+  {
+    AutoSchedulerLock auto_lock(thread_lock_);
+
+    if (thread_) {
+      // Make sure the thread is awake. It will see that
+      // |should_exit_for_testing_| is set and exit shortly after.
+      thread_->WakeUp();
+      thread = std::move(thread_);
+    }
+  }
 
-  thread_.reset();
+  if (thread)
+    thread->Join();
 }
 
 bool SchedulerWorker::ThreadAliveForTesting() const {
@@ -256,8 +263,14 @@ SchedulerWorker::SchedulerWorker(ThreadPriority priority_hint,
 }
 
 std::unique_ptr<SchedulerWorker::Thread> SchedulerWorker::Detach() {
-  DCHECK(!should_exit_for_testing_.IsSet()) << "Worker was already joined";
   AutoSchedulerLock auto_lock(thread_lock_);
+
+  // Do not detach if the thread is being joined.
+  if (!thread_) {
+    DCHECK(should_exit_for_testing_.IsSet());
+    return nullptr;
+  }
+
   // If a wakeup is pending, then a WakeUp() came in while we were deciding to
   // detach. This means we can't go away anymore since we would break the
   // guarantee that we call GetWork() after a successful wakeup.