SchedulerWorker Refcounting for Destruction in Production

base/task_scheduler/scheduler_worker.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/scheduler_worker.h"
 
 #include <stddef.h>
 
 #include <utility>
 
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
 #include "base/task_scheduler/task_tracker.h"
 
 #if defined(OS_MACOSX)
 #include "base/mac/scoped_nsautorelease_pool.h"
 #elif defined(OS_WIN)
 #include "base/win/scoped_com_initializer.h"
 #endif
 
 namespace base {
 namespace internal {
 
 class SchedulerWorker::Thread : public PlatformThread::Delegate {
  public:
   ~Thread() override = default;
 
-  static std::unique_ptr<Thread> Create(SchedulerWorker* outer) {
-    std::unique_ptr<Thread> thread(new Thread(outer));
+  static std::unique_ptr<Thread> Create(scoped_refptr<SchedulerWorker> outer) {
+    std::unique_ptr<Thread> thread(new Thread(std::move(outer)));
     thread->Initialize();
     if (thread->thread_handle_.is_null())
       return nullptr;
     return thread;
   }
 
   // PlatformThread::Delegate.
   void ThreadMain() override {
     // Set if this thread was detached.
     std::unique_ptr<Thread> detached_thread;
 
-    outer_->delegate_->OnMainEntry(outer_);
+    outer_->delegate_->OnMainEntry(outer_.get());
 
     // A SchedulerWorker starts out waiting for work.
     WaitForWork();
 
 #if defined(OS_WIN)
     std::unique_ptr<win::ScopedCOMInitializer> com_initializer;
     if (outer_->backward_compatibility_ ==
         SchedulerBackwardCompatibility::INIT_COM_STA) {
       com_initializer = MakeUnique<win::ScopedCOMInitializer>();
     }
 #endif
 
-    while (!outer_->task_tracker_->IsShutdownComplete() &&
-           !outer_->should_exit_for_testing_.IsSet()) {
+    while (!outer_->ShouldExit()) {
       DCHECK(outer_);
 
 #if defined(OS_MACOSX)
       mac::ScopedNSAutoreleasePool autorelease_pool;
 #endif
 
       UpdateThreadPriority(GetDesiredThreadPriority());
 
       // Get the sequence containing the next task to execute.
-      scoped_refptr<Sequence> sequence = outer_->delegate_->GetWork(outer_);
+      scoped_refptr<Sequence> sequence =
+          outer_->delegate_->GetWork(outer_.get());
       if (!sequence) {
-        if (outer_->delegate_->CanDetach(outer_)) {
-          detached_thread = outer_->Detach();
+        if (outer_->delegate_->CanDetach(outer_.get())) {
+          detached_thread = outer_->DetachThreadObject(DetachNotify::DELEGATE);
           if (detached_thread) {
             outer_ = nullptr;
             DCHECK_EQ(detached_thread.get(), this);
             PlatformThread::Detach(thread_handle_);
             break;
           }
         }
         WaitForWork();
         continue;
       }
@@ skipping 20 matching lines @@
       // sleep if WakeUp() is called while this SchedulerWorker is awake.
       wake_up_event_.Reset();
     }
 
     // If a wake up is pending and we successfully detached, somehow |outer_|
     // was able to signal us which means it probably thinks we're still alive.
     // This is bad as it will cause the WakeUp to no-op and |outer_| will be
     // stuck forever.
     DCHECK(!detached_thread || !IsWakeUpPending()) <<
         "This thread was detached and woken up at the same time.";
+
+    // This thread is generally responsible for cleaning itself up except when
+    // JoinForTesting() is called.
+    // We arrive here in the following cases:
+    // Thread Detachment Request:
+    //   * |detached_thread| will not be nullptr.
+    // ShouldExit() returns true:
+    //   * Shutdown: DetachThreadObject() returns the thread object.
+    //   * Cleanup: DetachThreadObject() returns the thread object.
+    //   * Join: DetachThreadObject() could return either the thread object or
+    //           nullptr. JoinForTesting() cleans up if we get nullptr.
+    if (!detached_thread)
+      detached_thread = outer_->DetachThreadObject(DetachNotify::SILENT);
   }
 
   void Join() { PlatformThread::Join(thread_handle_); }
 
   void WakeUp() { wake_up_event_.Signal(); }
 
   bool IsWakeUpPending() { return wake_up_event_.IsSignaled(); }
 
  private:
-  Thread(SchedulerWorker* outer)
-      : outer_(outer),
+  Thread(scoped_refptr<SchedulerWorker> outer)
+      : outer_(std::move(outer)),
         wake_up_event_(WaitableEvent::ResetPolicy::MANUAL,
                        WaitableEvent::InitialState::NOT_SIGNALED),
         current_thread_priority_(GetDesiredThreadPriority()) {
     DCHECK(outer_);
   }
 
   void Initialize() {
     constexpr size_t kDefaultStackSize = 0;
     PlatformThread::CreateWithPriority(kDefaultStackSize, this, &thread_handle_,
                                        current_thread_priority_);
@@ skipping 38 matching lines @@
   void UpdateThreadPriority(ThreadPriority desired_thread_priority) {
     if (desired_thread_priority == current_thread_priority_)
       return;
 
     PlatformThread::SetCurrentThreadPriority(desired_thread_priority);
     current_thread_priority_ = desired_thread_priority;
   }
 
   PlatformThreadHandle thread_handle_;
 
-  SchedulerWorker* outer_;
+  scoped_refptr<SchedulerWorker> outer_;
 
   // Event signaled to wake up this thread.
   WaitableEvent wake_up_event_;
 
   // Current priority of this thread. May be different from
   // |outer_->priority_hint_|.
   ThreadPriority current_thread_priority_;
 
   DISALLOW_COPY_AND_ASSIGN(Thread);
 };
 
-std::unique_ptr<SchedulerWorker> SchedulerWorker::Create(
+scoped_refptr<SchedulerWorker> SchedulerWorker::Create(
     ThreadPriority priority_hint,
     std::unique_ptr<Delegate> delegate,
     TaskTracker* task_tracker,
     InitialState initial_state,
     SchedulerBackwardCompatibility backward_compatibility) {
-  auto worker =
-      WrapUnique(new SchedulerWorker(priority_hint, std::move(delegate),
-                                     task_tracker, backward_compatibility));
+  scoped_refptr<SchedulerWorker> worker(
+      new SchedulerWorker(priority_hint, std::move(delegate), task_tracker,
+                          backward_compatibility));
   // Creation happens before any other thread can reference this one, so no
   // synchronization is necessary.
   if (initial_state == SchedulerWorker::InitialState::ALIVE) {
     worker->CreateThread();
     if (!worker->thread_) {
       return nullptr;
     }
   }
 
   return worker;
 }
 
-SchedulerWorker::~SchedulerWorker() {
-  // It is unexpected for |thread_| to be alive and for SchedulerWorker to
-  // destroy since SchedulerWorker owns the delegate needed by |thread_|.
-  // For testing, this generally means JoinForTesting was not called.
-  DCHECK(!thread_);
-}
-
 void SchedulerWorker::WakeUp() {
   AutoSchedulerLock auto_lock(thread_lock_);
 
-  DCHECK(!should_exit_for_testing_.IsSet());
+  DCHECK(!join_called_for_testing_.IsSet());
 
   if (!thread_)
     CreateThreadAssertSynchronized();
 
   if (thread_)
     thread_->WakeUp();
 }
 
 void SchedulerWorker::JoinForTesting() {
-  DCHECK(!should_exit_for_testing_.IsSet());
-  should_exit_for_testing_.Set();
+  DCHECK(!join_called_for_testing_.IsSet());
+  join_called_for_testing_.Set();
 
   std::unique_ptr<Thread> thread;
 
   {
     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.
+      // |join_called_for_testing_| is set and exit shortly after.
       thread_->WakeUp();
       thread = std::move(thread_);
     }
   }
 
   if (thread)
     thread->Join();
 }
 
 bool SchedulerWorker::ThreadAliveForTesting() const {
   AutoSchedulerLock auto_lock(thread_lock_);
   return !!thread_;
 }
 
+void SchedulerWorker::Cleanup() {
+  // |should_exit_| is synchronized with |thread_| for writes here so that we
+  // can maintain access to |thread_| for wakeup. Otherwise, the thread may take
+  // away |thread_| for destruction.
+  AutoSchedulerLock auto_lock(thread_lock_);
+  DCHECK(!should_exit_.IsSet());
+  if (thread_) {
+    should_exit_.Set();
+    thread_->WakeUp();
+  }
+}
+
 SchedulerWorker::SchedulerWorker(
     ThreadPriority priority_hint,
     std::unique_ptr<Delegate> delegate,
     TaskTracker* task_tracker,
     SchedulerBackwardCompatibility backward_compatibility)
     : priority_hint_(priority_hint),
       delegate_(std::move(delegate)),
       task_tracker_(task_tracker)
 #if defined(OS_WIN)
       ,
       backward_compatibility_(backward_compatibility)
 #endif
 {
   DCHECK(delegate_);
   DCHECK(task_tracker_);
 }
 
-std::unique_ptr<SchedulerWorker::Thread> SchedulerWorker::Detach() {
+SchedulerWorker::~SchedulerWorker() {
+  // It is unexpected for |thread_| to be alive and for SchedulerWorker to
+  // destroy since SchedulerWorker owns the delegate needed by |thread_|.
+  // For testing, this generally means JoinForTesting was not called.
+  DCHECK(!thread_);
+}
+
+std::unique_ptr<SchedulerWorker::Thread> SchedulerWorker::DetachThreadObject(
+    DetachNotify detach_notify) {
   AutoSchedulerLock auto_lock(thread_lock_);
 
   // Do not detach if the thread is being joined.
   if (!thread_) {
-    DCHECK(should_exit_for_testing_.IsSet());
+    DCHECK(join_called_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.
   if (thread_->IsWakeUpPending())
     return nullptr;
 
-  // Call OnDetach() within the scope of |thread_lock_| to prevent the delegate
-  // from being used concurrently from an old and a new thread.
-  delegate_->OnDetach();
+  if (detach_notify == DetachNotify::DELEGATE) {
+    // Call OnDetach() within the scope of |thread_lock_| to prevent the
+    // delegate from being used concurrently from an old and a new thread.
+    delegate_->OnDetach();
+  }
 
   return std::move(thread_);
 }
 
 void SchedulerWorker::CreateThread() {
-  thread_ = Thread::Create(this);
+  thread_ = Thread::Create(make_scoped_refptr(this));
 }
 
 void SchedulerWorker::CreateThreadAssertSynchronized() {
   thread_lock_.AssertAcquired();
   CreateThread();
 }
 
+bool SchedulerWorker::ShouldExit() {
+  return task_tracker_->IsShutdownComplete() ||
+         join_called_for_testing_.IsSet() || should_exit_.IsSet();
+}
+
 }  // namespace internal
 }  // namespace base