Bug 1738106 - Part 5: Track TaskQueue lifetimes to avoid leaking TaskQueues, r=xpcom-reviewers,mccr8

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

xpcom/threads/TaskQueue.cpp

@@ -11,14 +11,56 @@
 #include "nsIEventTarget.h"
 #include "nsITargetShutdownTask.h"
 #include "nsThreadUtils.h"
+#include "nsQueryObject.h"
 
 namespace mozilla {
 
+// Handle for a TaskQueue being tracked by a TaskQueueTracker. When created,
+// it is registered with the TaskQueueTracker, and when destroyed it is
+// unregistered. Holds a threadsafe weak reference to the TaskQueue.
+class TaskQueueTrackerEntry final
+    : private LinkedListElement<TaskQueueTrackerEntry> {
+ public:
+  TaskQueueTrackerEntry(TaskQueueTracker* aTracker,
+                        const RefPtr<TaskQueue>& aQueue)
+      : mTracker(aTracker), mQueue(aQueue) {
+    MutexAutoLock lock(mTracker->mMutex);
+    mTracker->mEntries.insertFront(this);
+  }
+  ~TaskQueueTrackerEntry() {
+    MutexAutoLock lock(mTracker->mMutex);
+    removeFrom(mTracker->mEntries);
+  }
+
+  TaskQueueTrackerEntry(const TaskQueueTrackerEntry&) = delete;
+  TaskQueueTrackerEntry(TaskQueueTrackerEntry&&) = delete;
+  TaskQueueTrackerEntry& operator=(const TaskQueueTrackerEntry&) = delete;
+  TaskQueueTrackerEntry& operator=(TaskQueueTrackerEntry&&) = delete;
+
+  RefPtr<TaskQueue> GetQueue() const { return RefPtr<TaskQueue>(mQueue); }
+
+ private:
+  friend class LinkedList<TaskQueueTrackerEntry>;
+  friend class LinkedListElement<TaskQueueTrackerEntry>;
+
+  const RefPtr<TaskQueueTracker> mTracker;
+  const ThreadSafeWeakPtr<TaskQueue> mQueue;
+};
+
 RefPtr<TaskQueue> TaskQueue::Create(already_AddRefed<nsIEventTarget> aTarget,
                                     const char* aName,
                                     bool aSupportsTailDispatch) {
+  nsCOMPtr<nsIEventTarget> target(std::move(aTarget));
   RefPtr<TaskQueue> queue =
-      new TaskQueue(std::move(aTarget), aName, aSupportsTailDispatch);
+      new TaskQueue(do_AddRef(target), aName, aSupportsTailDispatch);
+
+  // If |target| is a TaskQueueTracker, register this TaskQueue with it. It will
+  // be unregistered when the TaskQueue is destroyed or shut down.
+  if (RefPtr<TaskQueueTracker> tracker = do_QueryObject(target)) {
+    MonitorAutoLock lock(queue->mQueueMonitor);
+    queue->mTrackerEntry = MakeUnique<TaskQueueTrackerEntry>(tracker, queue);
+  }
+
   return queue;
 }
 
@@ -166,6 +208,16 @@ RefPtr<ShutdownPromise> TaskQueue::BeginShutdown() {
   return p;
 }
 
+void TaskQueue::MaybeResolveShutdown() {
+  mQueueMonitor.AssertCurrentThreadOwns();
+  if (mIsShutdown && !mIsRunning) {
+    mShutdownPromise.ResolveIfExists(true, __func__);
+    // Disconnect from our target as we won't try to dispatch any more events.
+    mTrackerEntry = nullptr;
+    mTarget = nullptr;
+  }
+}
+
 bool TaskQueue::IsEmpty() {
   MonitorAutoLock mon(mQueueMonitor);
   return mTasks.IsEmpty();
@@ -279,4 +331,17 @@ NS_IMETHODIMP TaskQueue::HaveDirectTasks(bool* aValue) {
   return NS_OK;
 }
 
+nsTArray<RefPtr<TaskQueue>> TaskQueueTracker::GetAllTrackedTaskQueues() {
+  MutexAutoLock lock(mMutex);
+  nsTArray<RefPtr<TaskQueue>> queues;
+  for (auto* entry : mEntries) {
+    if (auto queue = entry->GetQueue()) {
+      queues.AppendElement(queue);
+    }
+  }
+  return queues;
+}
+
+TaskQueueTracker::~TaskQueueTracker() = default;
+
 }  // namespace mozilla

xpcom/threads/TaskQueue.h

@@ -22,6 +22,8 @@ namespace mozilla {
 
 typedef MozPromise<bool, bool, false> ShutdownPromise;
 
+class TaskQueueTrackerEntry;
+
 // Abstracts executing runnables in order on an arbitrary event target. The
 // runnables dispatched to the TaskQueue will be executed in the order in which
 // they're received, and are guaranteed to not be executed concurrently.
@@ -142,16 +144,14 @@ class TaskQueue final : public AbstractThread,
   nsresult DispatchLocked(nsCOMPtr<nsIRunnable>& aRunnable, uint32_t aFlags,
                           DispatchReason aReason = NormalDispatch);
 
-  void MaybeResolveShutdown() {
-    mQueueMonitor.AssertCurrentThreadOwns();
-    if (mIsShutdown && !mIsRunning) {
-      mShutdownPromise.ResolveIfExists(true, __func__);
-      mTarget = nullptr;
-    }
-  }
+  void MaybeResolveShutdown();
 
   nsCOMPtr<nsIEventTarget> mTarget GUARDED_BY(mQueueMonitor);
 
+  // Handle for this TaskQueue being registered with our target if it implements
+  // TaskQueueTracker.
+  UniquePtr<TaskQueueTrackerEntry> mTrackerEntry GUARDED_BY(mQueueMonitor);
+
   // Monitor that protects the queue, mIsRunning, mIsShutdown and
   // mShutdownTasks;
   Monitor mQueueMonitor;
@@ -240,6 +240,41 @@ class TaskQueue final : public AbstractThread,
   };
 };
 
+#define MOZILLA_TASKQUEUETRACKER_IID                 \
+  {                                                  \
+    0x765c4b56, 0xd5f6, 0x4a9f, {                    \
+      0x91, 0xcf, 0x51, 0x47, 0xb3, 0xc1, 0x7e, 0xa6 \
+    }                                                \
+  }
+
+// XPCOM "interface" which may be implemented by nsIEventTarget implementations
+// which want to keep track of what TaskQueue instances are currently targeting
+// them. This may be used to asynchronously shutdown TaskQueues targeting a
+// threadpool or other event target before the threadpool goes away.
+//
+// This explicitly TaskQueue-aware tracker is used instead of
+// `nsITargetShutdownTask` as the operations required to shut down a TaskQueue
+// are asynchronous, which is not a requirement of that interface.
+class TaskQueueTracker : public nsISupports {
+ public:
+  NS_DECLARE_STATIC_IID_ACCESSOR(MOZILLA_TASKQUEUETRACKER_IID)
+
+  // Get a strong reference to every TaskQueue currently tracked by this
+  // TaskQueueTracker. May be called from any thraed.
+  nsTArray<RefPtr<TaskQueue>> GetAllTrackedTaskQueues();
+
+ protected:
+  virtual ~TaskQueueTracker();
+
+ private:
+  friend class TaskQueueTrackerEntry;
+
+  Mutex mMutex{"TaskQueueTracker"};
+  LinkedList<TaskQueueTrackerEntry> mEntries GUARDED_BY(mMutex);
+};
+
+NS_DEFINE_STATIC_IID_ACCESSOR(TaskQueueTracker, MOZILLA_TASKQUEUETRACKER_IID)
+
 }  // namespace mozilla
 
 #endif  // TaskQueue_h_

xpcom/threads/nsThreadManager.cpp

@@ -43,12 +43,13 @@ static MOZ_THREAD_LOCAL(bool) sTLSIsMainThread;
 
 bool NS_IsMainThreadTLSInitialized() { return sTLSIsMainThread.initialized(); }
 
-class BackgroundEventTarget final : public nsIEventTarget {
+class BackgroundEventTarget final : public nsIEventTarget,
+                                    public TaskQueueTracker {
  public:
   NS_DECL_THREADSAFE_ISUPPORTS
   NS_DECL_NSIEVENTTARGET_FULL
 
-  BackgroundEventTarget();
+  BackgroundEventTarget() = default;
 
   nsresult Init();
 
@@ -63,15 +64,9 @@ class BackgroundEventTarget final : public nsIEventTarget {
 
   nsCOMPtr<nsIThreadPool> mPool;
   nsCOMPtr<nsIThreadPool> mIOPool;
-
-  Mutex mMutex;
-  nsTArray<RefPtr<TaskQueue>> mTaskQueues GUARDED_BY(mMutex);
 };
 
-NS_IMPL_ISUPPORTS(BackgroundEventTarget, nsIEventTarget)
-
-BackgroundEventTarget::BackgroundEventTarget()
-    : mMutex("BackgroundEventTarget::mMutex") {}
+NS_IMPL_ISUPPORTS(BackgroundEventTarget, nsIEventTarget, TaskQueueTracker)
 
 nsresult BackgroundEventTarget::Init() {
   nsCOMPtr<nsIThreadPool> pool(new nsThreadPool());
@@ -192,8 +187,8 @@ BackgroundEventTarget::UnregisterShutdownTask(nsITargetShutdownTask* aTask) {
 
 void BackgroundEventTarget::BeginShutdown(
     nsTArray<RefPtr<ShutdownPromise>>& promises) {
-  MutexAutoLock lock(mMutex);
-  for (auto& queue : mTaskQueues) {
+  auto queues = GetAllTrackedTaskQueues();
+  for (auto& queue : queues) {
     promises.AppendElement(queue->BeginShutdown());
   }
 }
@@ -205,12 +200,7 @@ void BackgroundEventTarget::FinishShutdown() {
 
 already_AddRefed<nsISerialEventTarget>
 BackgroundEventTarget::CreateBackgroundTaskQueue(const char* aName) {
-  MutexAutoLock lock(mMutex);
-
-  RefPtr<TaskQueue> queue = TaskQueue::Create(do_AddRef(this), aName);
-  mTaskQueues.AppendElement(queue);
-
-  return queue.forget();
+  return TaskQueue::Create(do_AddRef(this), aName).forget();
 }
 
 extern "C" {