Bug 1747526 - Return a handle from asyncShutdown to observe and cancel thread shutdown progress, r=xpcom-reviewers,KrisWright

This change introduces a new interface, nsIThreadShutdown, which is used
to handle the relevant state for communicating thread shutdown state
between the joining and terminating threads. This type is now returned
from `nsIThread::AsyncShutdown` and can be used to register callbacks
for when thread shutdown is complete, as well as cancel shutdown
entirely, leading to the underlying PRThread never being joined using
PR_JoinThread. This leaking limitation may be avoidable if support for
detaching PRThreads is added to NSPR, or nsThread switches to a more
feature-complete threading API.

This patch also uses the new interface to rework nsThreadPool's Shutdown
and ShutdownWithTimeout methods to avoid poking at nsThread internals
and instead use the publicly facing methods. This allows us to start
async shutdown for all threads and spin the event loop until they all
complete, or a timeout timer fires.

Differential Revision: https://phabricator.services.mozilla.com/D136045

xpcom/threads/LazyIdleThread.cpp

@@ -473,6 +473,13 @@ LazyIdleThread::AsyncShutdown() {
   return NS_ERROR_NOT_IMPLEMENTED;
 }
 
+NS_IMETHODIMP
+LazyIdleThread::BeginShutdown(nsIThreadShutdown** aShutdown) {
+  ASSERT_OWNING_THREAD();
+  *aShutdown = nullptr;
+  return NS_ERROR_NOT_IMPLEMENTED;
+}
+
 NS_IMETHODIMP
 LazyIdleThread::Shutdown() {
   ASSERT_OWNING_THREAD();

xpcom/threads/moz.build

@@ -18,6 +18,7 @@ XPIDL_SOURCES += [
     "nsIThreadInternal.idl",
     "nsIThreadManager.idl",
     "nsIThreadPool.idl",
+    "nsIThreadShutdown.idl",
     "nsITimer.idl",
 ]
 

xpcom/threads/nsIThread.idl

@@ -5,9 +5,11 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 #include "nsISerialEventTarget.idl"
+#include "nsIThreadShutdown.idl"
 
 %{C++
 #include "mozilla/AlreadyAddRefed.h"
+
 namespace mozilla {
 class TimeStamp;
 class TimeDurationValueCalculator;
@@ -130,11 +132,17 @@ interface nsIThread : nsISerialEventTarget
    * @throws NS_ERROR_UNEXPECTED
    *   Indicates that this method was erroneously called when this thread was
    *   the current thread, that this thread was not created with a call to
-   *   nsIThreadManager::NewThread, or if this method was called more than once
-   *   on the thread object.
+   *   nsIThreadManager::NewNamedThread, or that this method was called more
+   *   than once on the thread object.
    */
   void asyncShutdown();
 
+  /**
+   * Like `asyncShutdown`, but also returns a nsIThreadShutdown instance to
+   * allow observing and controlling the thread's async shutdown progress.
+   */
+  nsIThreadShutdown beginShutdown();
+
   /**
    * Dispatch an event to a specified queue for the thread.  This function
    * may be called from any thread, and it may be called re-entrantly.

xpcom/threads/nsIThreadShutdown.idl

@@ -0,0 +1,57 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim:set ts=2 sw=2 sts=2 et cindent: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "nsISupports.idl"
+
+interface nsIRunnable;
+
+/**
+ * Handle for the ongoing shutdown progress of a given thread which can be used
+ * to observe and interrupt async shutdown progress. Methods on this interface
+ * may generally only be used on the thread which called
+ * `nsIThread::beginShutdown`.
+ */
+[scriptable, builtinclass, uuid(70a43748-6130-4ea6-a440-7c74e1b7dd7c)]
+interface nsIThreadShutdown : nsISupports
+{
+  /**
+   * Register a runnable to be executed when the thread has completed shutdown,
+   * or shutdown has been cancelled due to `stopWaitingAndLeakThread()`.
+   *
+   * If the thread has already completed or cancelled shutdown, the runnable
+   * may be executed synchronously.
+   *
+   * May only be called on the thread which invoked `nsIThread::beginShutdown`.
+   */
+  void onCompletion(in nsIRunnable aEvent);
+
+  /**
+   * Check if the target thread has completed shutdown.
+   *
+   * May only be accessed on the thread which called `nsIThread::beginShutdown`.
+   */
+  [infallible] readonly attribute boolean completed;
+
+  /**
+   * Give up on waiting for the shutting down thread to exit. Calling this
+   * method will allow the thread to continue running, no longer block shutdown,
+   * and the thread will never be joined or have its resources reclaimed.
+   *
+   * Completion callbacks attached to this `nsIThreadShutdown` may be executed
+   * during this call.
+   *
+   * This method should NOT be called except in exceptional circumstances during
+   * shutdown, as it will cause resources for the shutting down thread to be
+   * leaked.
+   *
+   * May only be called on the thread which called `nsIThread::beginShutdown`
+   *
+   * @throws NS_ERROR_NOT_AVAILABLE
+   *   Indicates that the target thread has already stopped running and a
+   *   request to be joined is already being dispatched to the waiting thread.
+   */
+  void stopWaitingAndLeakThread();
+};

xpcom/threads/nsThread.cpp

@@ -196,12 +196,6 @@ NS_IMPL_CI_INTERFACE_GETTER(nsThread, nsIThread, nsIThreadInternal,
 
 //-----------------------------------------------------------------------------
 
-bool nsThread::ShutdownContextsComp::Equals(
-    const ShutdownContexts::elem_type& a,
-    const ShutdownContexts::elem_type::Pointer b) const {
-  return a.get() == b;
-}
-
 // This event is responsible for notifying nsThread::Shutdown that it is time
 // to call PR_JoinThread. It implements nsICancelableRunnable so that it can
 // run on a DOM Worker thread (where all events must implement
@@ -220,7 +214,7 @@ class nsThreadShutdownAckEvent : public CancelableRunnable {
  private:
   virtual ~nsThreadShutdownAckEvent() = default;
 
-  NotNull<nsThreadShutdownContext*> mShutdownContext;
+  NotNull<RefPtr<nsThreadShutdownContext>> mShutdownContext;
 };
 
 // This event is responsible for setting mShutdownContext
@@ -232,6 +226,8 @@ class nsThreadShutdownEvent : public Runnable {
         mThread(aThr),
         mShutdownContext(aCtx) {}
   NS_IMETHOD Run() override {
+    // Creates a cycle between `mThread` and the shutdown context which will be
+    // broken when the thread exits.
     mThread->mShutdownContext = mShutdownContext;
     if (mThread->mEventTarget) {
       mThread->mEventTarget->NotifyShutdown();
@@ -242,7 +238,7 @@ class nsThreadShutdownEvent : public Runnable {
 
  private:
   NotNull<RefPtr<nsThread>> mThread;
-  NotNull<nsThreadShutdownContext*> mShutdownContext;
+  NotNull<RefPtr<nsThreadShutdownContext>> mShutdownContext;
 };
 
 //-----------------------------------------------------------------------------
@@ -422,36 +418,34 @@ void nsThread::ThreadFunc(void* aArg) {
     PROFILER_UNREGISTER_THREAD();
   }
 
-  // Dispatch shutdown ACK
-  NotNull<nsThreadShutdownContext*> context =
+  NotNull<RefPtr<nsThreadShutdownContext>> context =
       WrapNotNull(self->mShutdownContext);
+  self->mShutdownContext = nullptr;
   MOZ_ASSERT(context->mTerminatingThread == self);
-  nsCOMPtr<nsIRunnable> event =
-      do_QueryObject(new nsThreadShutdownAckEvent(context));
-  if (context->mIsMainThreadJoining) {
-    DebugOnly<nsresult> dispatch_ack_rv =
-        SchedulerGroup::Dispatch(TaskCategory::Other, event.forget());
-#ifdef DEBUG
-    // On the main thread, dispatch may fail if this thread is part of a
-    // `nsIThreadPool` which was shut down using `ShutdownWithTimeout`, and
-    // that shutdown attempt timed out. In that case, the main thread may have
-    // already completed thread shutdown before this dispatch attempt,
-    // allowing it to fail. At that point, it is impossible for us to join
-    // this thread anymore, so give up and warn instead.
-    if (NS_FAILED(dispatch_ack_rv)) {
-      NS_WARNING(
-          "Thread shudown ack dispatch failed, the main thread may no longer "
-          "be waiting.");
-    }
-#endif
-  } else {
+
+  // Take the joining thread from our shutdown context. This may have been
+  // cleared by the joining thread if it decided to cancel waiting on us, in
+  // which case we won't notify our caller, and leak.
+  RefPtr<nsThread> joiningThread;
+  {
+    auto lock = context->mJoiningThread.Lock();
+    joiningThread = lock->forget();
+  }
+  if (joiningThread) {
+    // Dispatch shutdown ACK
+    nsCOMPtr<nsIRunnable> event = new nsThreadShutdownAckEvent(context);
     nsresult dispatch_ack_rv =
-        context->mJoiningThread->Dispatch(event, NS_DISPATCH_NORMAL);
+        joiningThread->Dispatch(event, NS_DISPATCH_NORMAL);
+
     // We do not expect this to ever happen, but If we cannot dispatch
     // the ack event, someone probably blocks waiting on us and will
     // crash with a hang later anyways. The best we can do is to tell
     // the world what happened right here.
     MOZ_RELEASE_ASSERT(NS_SUCCEEDED(dispatch_ack_rv));
+  } else {
+    NS_WARNING(
+        "nsThread exiting after StopWaitingAndLeakThread was called, thread "
+        "resources will be leaked!");
   }
 
   // Release any observer of the thread here.
@@ -548,6 +542,7 @@ nsThread::nsThread(NotNull<SynchronizedEventQueue*> aQueue,
     : mEvents(aQueue.get()),
       mEventTarget(
           new ThreadEventTarget(mEvents.get(), aMainThread == MAIN_THREAD)),
+      mOutstandingShutdownContexts(0),
       mShutdownContext(nullptr),
       mScriptObserver(nullptr),
       mThreadName("<uninitialized>"),
@@ -572,6 +567,7 @@ nsThread::nsThread(NotNull<SynchronizedEventQueue*> aQueue,
 nsThread::nsThread()
     : mEvents(nullptr),
       mEventTarget(nullptr),
+      mOutstandingShutdownContexts(0),
       mShutdownContext(nullptr),
       mScriptObserver(nullptr),
       mThreadName("<uninitialized>"),
@@ -590,22 +586,10 @@ nsThread::nsThread()
 }
 
 nsThread::~nsThread() {
-  NS_ASSERTION(mRequestedShutdownContexts.IsEmpty(),
+  NS_ASSERTION(mOutstandingShutdownContexts == 0,
                "shouldn't be waiting on other threads to shutdown");
 
   MaybeRemoveFromThreadList();
-
-#ifdef DEBUG
-  // We deliberately leak these so they can be tracked by the leak checker.
-  // If you're having nsThreadShutdownContext leaks, you can set:
-  //   XPCOM_MEM_LOG_CLASSES=nsThreadShutdownContext
-  // during a test run and that will at least tell you what thread is
-  // requesting shutdown on another, which can be helpful for diagnosing
-  // the leak.
-  for (size_t i = 0; i < mRequestedShutdownContexts.Length(); ++i) {
-    Unused << mRequestedShutdownContexts[i].release();
-  }
-#endif
 }
 
 nsresult nsThread::Init(const nsACString& aName) {
@@ -778,36 +762,44 @@ NS_IMETHODIMP
 nsThread::AsyncShutdown() {
   LOG(("THRD(%p) async shutdown\n", this));
 
-  ShutdownInternal(/* aSync = */ false);
+  nsCOMPtr<nsIThreadShutdown> shutdown;
+  BeginShutdown(getter_AddRefs(shutdown));
   return NS_OK;
 }
 
-nsThreadShutdownContext* nsThread::ShutdownInternal(bool aSync) {
+NS_IMETHODIMP
+nsThread::BeginShutdown(nsIThreadShutdown** aShutdown) {
+  LOG(("THRD(%p) begin shutdown\n", this));
+
   MOZ_ASSERT(mEvents);
   MOZ_ASSERT(mEventTarget);
   MOZ_ASSERT(mThread != PR_GetCurrentThread());
   if (NS_WARN_IF(mThread == PR_GetCurrentThread())) {
-    return nullptr;
+    return NS_ERROR_UNEXPECTED;
   }
 
   // Prevent multiple calls to this method.
   if (!mShutdownRequired.compareExchange(true, false)) {
-    return nullptr;
+    return NS_ERROR_UNEXPECTED;
   }
   MOZ_ASSERT(mThread);
 
   MaybeRemoveFromThreadList();
 
-  NotNull<nsThread*> currentThread =
-      WrapNotNull(nsThreadManager::get().GetCurrentThread());
+  RefPtr<nsThread> currentThread = nsThreadManager::get().GetCurrentThread();
 
   MOZ_DIAGNOSTIC_ASSERT(currentThread->EventQueue(),
                         "Shutdown() may only be called from an XPCOM thread");
 
-  // Allocate a shutdown context and store a strong ref.
-  auto context =
-      new nsThreadShutdownContext(WrapNotNull(this), currentThread, aSync);
-  Unused << *currentThread->mRequestedShutdownContexts.EmplaceBack(context);
+  // Allocate a shutdown context, and record that we're waiting for it.
+  RefPtr<nsThreadShutdownContext> context =
+      new nsThreadShutdownContext(WrapNotNull(this), currentThread);
+
+  ++currentThread->mOutstandingShutdownContexts;
+  nsCOMPtr<nsIRunnable> clearOutstanding = NS_NewRunnableFunction(
+      "nsThread::ClearOutstandingShutdownContext",
+      [currentThread] { --currentThread->mOutstandingShutdownContexts; });
+  context->OnCompletion(clearOutstanding);
 
   // Set mShutdownContext and wake up the thread in case it is waiting for
   // events to process.
@@ -816,7 +808,7 @@ nsThreadShutdownContext* nsThread::ShutdownInternal(bool aSync) {
   if (!mEvents->PutEvent(event.forget(), EventQueuePriority::Normal)) {
     // We do not expect this to happen. Let's collect some diagnostics.
     nsAutoCString threadName;
-    currentThread->GetThreadName(threadName);
+    GetThreadName(threadName);
     MOZ_CRASH_UNSAFE_PRINTF("Attempt to shutdown an already dead thread: %s",
                             threadName.get());
   }
@@ -824,7 +816,8 @@ nsThreadShutdownContext* nsThread::ShutdownInternal(bool aSync) {
   // We could still end up with other events being added after the shutdown
   // task, but that's okay because we process pending events in ThreadFunc
   // after setting mShutdownContext just before exiting.
-  return context;
+  context.forget(aShutdown);
+  return NS_OK;
 }
 
 void nsThread::ShutdownComplete(NotNull<nsThreadShutdownContext*> aContext) {
@@ -834,13 +827,6 @@ void nsThread::ShutdownComplete(NotNull<nsThreadShutdownContext*> aContext) {
 
   MaybeRemoveFromThreadList();
 
-  if (aContext->mAwaitingShutdownAck) {
-    // We're in a synchronous shutdown, so tell whatever is up the stack that
-    // we're done and unwind the stack so it can call us again.
-    aContext->mAwaitingShutdownAck = false;
-    return;
-  }
-
   // Now, it should be safe to join without fear of dead-locking.
   PR_JoinThread(aContext->mTerminatingPRThread);
   MOZ_ASSERT(!mThread);
@@ -850,11 +836,7 @@ void nsThread::ShutdownComplete(NotNull<nsThreadShutdownContext*> aContext) {
   MOZ_ASSERT(!obs, "Should have been cleared at shutdown!");
 #endif
 
-  // Delete aContext.
-  // aContext might not be in mRequestedShutdownContexts if it belongs to a
-  // thread that was leaked by calling nsIThreadPool::ShutdownWithTimeout.
-  aContext->mJoiningThread->mRequestedShutdownContexts.RemoveElement(
-      aContext, ShutdownContextsComp{});
+  aContext->MarkCompleted();
 }
 
 void nsThread::WaitForAllAsynchronousShutdowns() {
@@ -862,32 +844,27 @@ void nsThread::WaitForAllAsynchronousShutdowns() {
   // has the template parameter we are providing here.
   SpinEventLoopUntil<ProcessFailureBehavior::IgnoreAndContinue>(
       "nsThread::WaitForAllAsynchronousShutdowns"_ns,
-      [&]() { return mRequestedShutdownContexts.IsEmpty(); }, this);
+      [&]() { return mOutstandingShutdownContexts == 0; }, this);
 }
 
 NS_IMETHODIMP
 nsThread::Shutdown() {
   LOG(("THRD(%p) sync shutdown\n", this));
 
-  nsThreadShutdownContext* maybeContext = ShutdownInternal(/* aSync = */ true);
-  if (!maybeContext) {
+  nsCOMPtr<nsIThreadShutdown> context;
+  nsresult rv = BeginShutdown(getter_AddRefs(context));
+  if (NS_FAILED(rv)) {
     return NS_OK;  // The thread has already shut down.
   }
 
-  NotNull<nsThreadShutdownContext*> context = WrapNotNull(maybeContext);
-
   // If we are going to hang here we want to see the thread's name
   nsAutoCString threadName;
   GetThreadName(threadName);
 
   // Process events on the current thread until we receive a shutdown ACK.
   // Allows waiting; ensure no locks are held that would deadlock us!
-  SpinEventLoopUntil(
-      "nsThread::Shutdown: "_ns + threadName,
-      [&, context]() { return !context->mAwaitingShutdownAck; },
-      context->mJoiningThread);
-
-  ShutdownComplete(context);
+  SpinEventLoopUntil("nsThread::Shutdown: "_ns + threadName,
+                     [&]() { return context->GetCompleted(); });
 
   return NS_OK;
 }
@@ -1025,7 +1002,6 @@ size_t nsThread::ShallowSizeOfIncludingThis(
   if (mShutdownContext) {
     n += aMallocSizeOf(mShutdownContext);
   }
-  n += mRequestedShutdownContexts.ShallowSizeOfExcludingThis(aMallocSizeOf);
   return aMallocSizeOf(this) + aMallocSizeOf(mThread) + n;
 }
 
@@ -1448,6 +1424,52 @@ nsLocalExecutionGuard::~nsLocalExecutionGuard() {
   mEventQueueStack.PopEventQueue(mLocalEventTarget);
 }
 
+NS_IMPL_ISUPPORTS(nsThreadShutdownContext, nsIThreadShutdown)
+
+NS_IMETHODIMP
+nsThreadShutdownContext::OnCompletion(nsIRunnable* aEvent) {
+  if (mCompleted) {
+    aEvent->Run();
+  } else {
+    mCompletionCallbacks.AppendElement(aEvent);
+  }
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+nsThreadShutdownContext::GetCompleted(bool* aCompleted) {
+  *aCompleted = mCompleted;
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+nsThreadShutdownContext::StopWaitingAndLeakThread() {
+  // Take the joining thread from `mJoiningThread` so that the terminating
+  // thread won't try to dispatch nsThreadShutdownAckEvent to us anymore.
+  RefPtr<nsThread> joiningThread;
+  {
+    auto lock = mJoiningThread.Lock();
+    joiningThread = lock->forget();
+  }
+  if (!joiningThread) {
+    // Shutdown is already being resolved, so there's nothing for us to do.
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  MOZ_DIAGNOSTIC_ASSERT(joiningThread->IsOnCurrentThread());
+  MarkCompleted();
+  return NS_OK;
+}
+
+void nsThreadShutdownContext::MarkCompleted() {
+  MOZ_ASSERT(!mCompleted);
+  mCompleted = true;
+  nsTArray<nsCOMPtr<nsIRunnable>> callbacks(std::move(mCompletionCallbacks));
+  for (auto& callback : callbacks) {
+    callback->Run();
+  }
+}
+
 namespace mozilla {
 PerformanceCounterState::Snapshot PerformanceCounterState::RunnableWillRun(
     PerformanceCounter* aCounter, TimeStamp aNow, bool aIsIdleRunnable) {

xpcom/threads/nsThread.h

@@ -47,6 +47,7 @@ using mozilla::NotNull;
 class nsIRunnable;
 class nsLocalExecutionRecord;
 class nsThreadEnumerator;
+class nsThreadShutdownContext;
 
 // See https://www.w3.org/TR/longtasks
 #define LONGTASK_BUSY_WINDOW_MS 50
@@ -219,7 +220,7 @@ class nsThread : public nsIThreadInternal,
 
   uint32_t RecursionDepth() const;
 
-  void ShutdownComplete(NotNull<struct nsThreadShutdownContext*> aContext);
+  void ShutdownComplete(NotNull<nsThreadShutdownContext*> aContext);
 
   void WaitForAllAsynchronousShutdowns();
 
@@ -292,7 +293,7 @@ class nsThread : public nsIThreadInternal,
     return already_AddRefed<nsIThreadObserver>(obs);
   }
 
-  struct nsThreadShutdownContext* ShutdownInternal(bool aSync);
+  already_AddRefed<nsThreadShutdownContext> ShutdownInternal(bool aSync);
 
   friend class nsThreadManager;
   friend class nsThreadPool;
@@ -313,19 +314,11 @@ class nsThread : public nsIThreadInternal,
   RefPtr<mozilla::SynchronizedEventQueue> mEvents;
   RefPtr<mozilla::ThreadEventTarget> mEventTarget;
 
-  // The shutdown contexts for any other threads we've asked to shut down.
-  using ShutdownContexts =
-      nsTArray<mozilla::UniquePtr<struct nsThreadShutdownContext>>;
-
-  // Helper for finding a ShutdownContext in the contexts array.
-  struct ShutdownContextsComp {
-    bool Equals(const ShutdownContexts::elem_type& a,
-                const ShutdownContexts::elem_type::Pointer b) const;
-  };
-
-  ShutdownContexts mRequestedShutdownContexts;
+  // The number of outstanding nsThreadShutdownContext started by this thread.
+  // The thread will not be allowed to exit until this number reaches 0.
+  uint32_t mOutstandingShutdownContexts;
   // The shutdown context for ourselves.
-  struct nsThreadShutdownContext* mShutdownContext;
+  RefPtr<nsThreadShutdownContext> mShutdownContext;
 
   mozilla::CycleCollectedJSContext* mScriptObserver;
 
@@ -370,26 +363,40 @@ class nsThread : public nsIThreadInternal,
   mozilla::SimpleTaskQueue mDirectTasks;
 };
 
-struct nsThreadShutdownContext {
+class nsThreadShutdownContext final : public nsIThreadShutdown {
+ public:
+  NS_DECL_THREADSAFE_ISUPPORTS
+  NS_DECL_NSITHREADSHUTDOWN
+
+ private:
+  friend class nsThread;
+  friend class nsThreadShutdownEvent;
+  friend class nsThreadShutdownAckEvent;
+
   nsThreadShutdownContext(NotNull<nsThread*> aTerminatingThread,
-                          NotNull<nsThread*> aJoiningThread,
-                          bool aAwaitingShutdownAck)
+                          nsThread* aJoiningThread)
       : mTerminatingThread(aTerminatingThread),
         mTerminatingPRThread(aTerminatingThread->GetPRThread()),
-        mJoiningThread(aJoiningThread),
-        mAwaitingShutdownAck(aAwaitingShutdownAck),
-        mIsMainThreadJoining(NS_IsMainThread()) {
-    MOZ_COUNT_CTOR(nsThreadShutdownContext);
-  }
-  MOZ_COUNTED_DTOR(nsThreadShutdownContext)
+        mJoiningThread(aJoiningThread,
+                       "nsThreadShutdownContext::mJoiningThread") {}
 
-  // NB: This will be the last reference.
-  NotNull<RefPtr<nsThread>> mTerminatingThread;
+  ~nsThreadShutdownContext() = default;
+
+  // Must be called on the joining thread.
+  void MarkCompleted();
+
+  // NB: This may be the last reference.
+  NotNull<RefPtr<nsThread>> const mTerminatingThread;
   PRThread* const mTerminatingPRThread;
-  NotNull<nsThread*> MOZ_UNSAFE_REF(
-      "Thread manager is holding reference to joining thread") mJoiningThread;
-  bool mAwaitingShutdownAck;
-  bool mIsMainThreadJoining;
+
+  // May only be accessed on the joining thread.
+  bool mCompleted = false;
+  nsTArray<nsCOMPtr<nsIRunnable>> mCompletionCallbacks;
+
+  // The thread waiting for this thread to shut down. Will either be cleared by
+  // the joining thread if `StopWaitingAndLeakThread` is called or by the
+  // terminating thread upon exiting and notifying the joining thread.
+  mozilla::DataMutex<RefPtr<nsThread>> mJoiningThread;
 };
 
 // This RAII class controls the duration of the associated nsThread's local

xpcom/threads/nsThreadPool.cpp

@@ -11,6 +11,7 @@
 #include "nsThreadManager.h"
 #include "nsThread.h"
 #include "nsMemory.h"
+#include "nsThreadUtils.h"
 #include "prinrval.h"
 #include "mozilla/Logging.h"
 #include "mozilla/ProfilerLabels.h"
@@ -381,61 +382,10 @@ nsThreadPool::IsOnCurrentThread(bool* aResult) {
 }
 
 NS_IMETHODIMP
-nsThreadPool::Shutdown() {
-  nsCOMArray<nsIThread> threads;
-  nsCOMPtr<nsIThreadPoolListener> listener;
-  {
-    MutexAutoLock lock(mMutex);
-    if (mShutdown) {
-      return NS_ERROR_ILLEGAL_DURING_SHUTDOWN;
-    }
-    mShutdown = true;
-    mEventsAvailable.NotifyAll();
-
-    threads.AppendObjects(mThreads);
-    mThreads.Clear();
-
-    // Swap in a null listener so that we release the listener at the end of
-    // this method. The listener will be kept alive as long as the other threads
-    // that were created when it was set.
-    mListener.swap(listener);
-  }
-
-  // It's important that we shutdown the threads while outside the event queue
-  // monitor.  Otherwise, we could end up dead-locking.
-
-  for (int32_t i = 0; i < threads.Count(); ++i) {
-    threads[i]->Shutdown();
-  }
-
-  return NS_OK;
-}
-
-template <typename Pred>
-static void SpinMTEventLoopUntil(Pred&& aPredicate, nsIThread* aThread,
-                                 TimeDuration aTimeout) {
-  MOZ_ASSERT(NS_IsMainThread(), "Must be run on the main thread");
-
-  // From a latency perspective, spinning the event loop is like leaving script
-  // and returning to the event loop. Tell the watchdog we stopped running
-  // script (until we return).
-  mozilla::Maybe<xpc::AutoScriptActivity> asa;
-  asa.emplace(false);
-
-  TimeStamp deadline = TimeStamp::Now() + aTimeout;
-  while (!aPredicate() && TimeStamp::Now() < deadline) {
-    if (!NS_ProcessNextEvent(aThread, false)) {
-      PR_Sleep(PR_MillisecondsToInterval(1));
-    }
-  }
-}
+nsThreadPool::Shutdown() { return ShutdownWithTimeout(-1); }
 
 NS_IMETHODIMP
 nsThreadPool::ShutdownWithTimeout(int32_t aTimeoutMs) {
-  if (!NS_IsMainThread()) {
-    return NS_ERROR_NOT_AVAILABLE;
-  }
-
   nsCOMArray<nsIThread> threads;
   nsCOMPtr<nsIThreadPoolListener> listener;
   {
@@ -455,56 +405,47 @@ nsThreadPool::ShutdownWithTimeout(int32_t aTimeoutMs) {
     mListener.swap(listener);
   }
 
-  // IMPORTANT! Never dereference these pointers, as the objects may go away at
-  // any time. We just use the pointers values for comparison, to check if the
-  // thread has been shut down or not.
-  nsTArray<nsThreadShutdownContext*> contexts;
-
-  // It's important that we shutdown the threads while outside the event queue
-  // monitor.  Otherwise, we could end up dead-locking.
+  nsTArray<nsCOMPtr<nsIThreadShutdown>> contexts;
   for (int32_t i = 0; i < threads.Count(); ++i) {
-    // Shutdown async
-    nsThreadShutdownContext* maybeContext =
-        static_cast<nsThread*>(threads[i])->ShutdownInternal(false);
-    contexts.AppendElement(maybeContext);
-  }
-
-  NotNull<nsThread*> currentThread =
-      WrapNotNull(nsThreadManager::get().GetCurrentThread());
-
-  // We spin the event loop until all of the threads in the thread pool
-  // have shut down, or the timeout expires.
-  SpinMTEventLoopUntil(
-      [&]() {
-        for (nsIThread* thread : threads) {
-          if (static_cast<nsThread*>(thread)->mThread) {
-            return false;
-          }
-        }
-        return true;
-      },
-      currentThread, TimeDuration::FromMilliseconds(aTimeoutMs));
-
-  // For any threads that have not shutdown yet, we need to remove them from
-  // mRequestedShutdownContexts so the thread manager does not wait for them
-  // at shutdown.
-  static const nsThread::ShutdownContextsComp comparator{};
-  for (int32_t i = 0; i < threads.Count(); ++i) {
-    nsThread* thread = static_cast<nsThread*>(threads[i]);
-    // If mThread is not null on the thread it means that it hasn't shutdown
-    // context[i] corresponds to thread[i]
-    if (thread->mThread && contexts[i]) {
-      auto index = currentThread->mRequestedShutdownContexts.IndexOf(
-          contexts[i], 0, comparator);
-      if (index != nsThread::ShutdownContexts::NoIndex) {
-        // We must leak the shutdown context just in case the leaked thread
-        // does get unstuck and completes before the main thread is done.
-        Unused << currentThread->mRequestedShutdownContexts[index].release();
-        currentThread->mRequestedShutdownContexts.RemoveElementAt(index);
-      }
+    nsCOMPtr<nsIThreadShutdown> context;
+    if (NS_SUCCEEDED(threads[i]->BeginShutdown(getter_AddRefs(context)))) {
+      contexts.AppendElement(std::move(context));
     }
   }
 
+  // Start a timer which will stop waiting & leak the thread, forcing
+  // onCompletion to be called when it expires.
+  nsCOMPtr<nsITimer> timer;
+  if (aTimeoutMs >= 0) {
+    NS_NewTimerWithCallback(
+        getter_AddRefs(timer),
+        [&](nsITimer*) {
+          for (auto& context : contexts) {
+            context->StopWaitingAndLeakThread();
+          }
+        },
+        aTimeoutMs, nsITimer::TYPE_ONE_SHOT,
+        "nsThreadPool::ShutdownWithTimeout");
+  }
+
+  // Start a counter and register a callback to decrement outstandingThreads
+  // when the threads finish exiting. We'll spin an event loop until
+  // outstandingThreads reaches 0.
+  uint32_t outstandingThreads = contexts.Length();
+  RefPtr onCompletion = NS_NewCancelableRunnableFunction(
+      "nsThreadPool thread completion", [&] { --outstandingThreads; });
+  for (auto& context : contexts) {
+    context->OnCompletion(onCompletion);
+  }
+
+  mozilla::SpinEventLoopUntil("nsThreadPool::ShutdownWithTimeout"_ns,
+                              [&] { return outstandingThreads == 0; });
+
+  if (timer) {
+    timer->Cancel();
+  }
+  onCompletion->Cancel();
+
   return NS_OK;
 }