Bug 1658571 - Add MainThreadWeakPtr and use it in PreloaderBase::RedirectSink. r=sg,smaug

Per matrix discussion, this is safer to uplift (instead of relying on
the RedirectSink being released on the main thread).

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

mfbt/WeakPtr.h

@@ -60,6 +60,12 @@
  *
  * Note that multiple base classes inheriting from SupportsWeakPtr is not
  * currently supported. We could support it if needed though.
+ *
+ * For Gecko-internal usage there is also MainThreadWeakPtr<T>, a version of
+ * WeakPtr that can be destroyed on any thread, but whose release gets proxied
+ * to the main thread. This is a similar API to nsMainThreadPtrHandle, but
+ * without keeping a strong reference to the main-thread object. Said WeakPtr
+ * can't be accessed from any other thread that isn't the main thread.
  */
 
 #ifndef mozilla_WeakPtr_h
@@ -78,6 +84,8 @@
 #if defined(MOZILLA_INTERNAL_API)
 // For thread safety checking.
 #  include "nsISupportsImpl.h"
+// For main thread destructor behavior.
+#  include "nsProxyRelease.h"
 #endif
 
 #if defined(MOZILLA_INTERNAL_API) && \
@@ -145,7 +153,19 @@
 
 namespace mozilla {
 
-template <typename T>
+namespace detail {
+
+enum class WeakPtrDestructorBehavior {
+  Normal,
+#ifdef MOZILLA_INTERNAL_API
+  ProxyToMainThread,
+#endif
+};
+
+}  // namespace detail
+
+template <typename T, detail::WeakPtrDestructorBehavior =
+                          detail::WeakPtrDestructorBehavior::Normal>
 class WeakPtr;
 class SupportsWeakPtr;
 
@@ -211,13 +231,13 @@ class SupportsWeakPtr {
     return mSelfReferencingWeakReference.get();
   }
 
-  template <typename U>
+  template <typename U, detail::WeakPtrDestructorBehavior>
   friend class WeakPtr;
 
   mutable RefPtr<WeakReference> mSelfReferencingWeakReference;
 };
 
-template <typename T>
+template <typename T, detail::WeakPtrDestructorBehavior Destruct>
 class WeakPtr {
   using WeakReference = detail::WeakReference;
 
@@ -264,7 +284,15 @@ class WeakPtr {
     return *this;
   }
 
-  MOZ_IMPLICIT WeakPtr(T* aOther) { *this = aOther; }
+  MOZ_IMPLICIT WeakPtr(T* aOther) {
+    *this = aOther;
+#ifdef MOZILLA_INTERNAL_API
+    if (Destruct == detail::WeakPtrDestructorBehavior::ProxyToMainThread) {
+      MOZ_ASSERT(NS_IsMainThread(),
+                 "MainThreadWeakPtr makes no sense on non-main threads");
+    }
+#endif
+  }
 
   // Ensure that mRef is dereferenceable in the uninitialized state.
   WeakPtr() : mRef(new WeakReference(nullptr)) {}
@@ -275,7 +303,15 @@ class WeakPtr {
   T& operator*() const { return *get(); }
   T* operator->() const MOZ_NO_ADDREF_RELEASE_ON_RETURN { return get(); }
 
-  ~WeakPtr() { MOZ_WEAKPTR_ASSERT_THREAD_SAFETY_DELEGATED(mRef); }
+#ifdef MOZILLA_INTERNAL_API
+  ~WeakPtr() {
+    if (Destruct == detail::WeakPtrDestructorBehavior::ProxyToMainThread) {
+      NS_ReleaseOnMainThread("WeakPtr::mRef", mRef.forget());
+    } else {
+      MOZ_WEAKPTR_ASSERT_THREAD_SAFETY_DELEGATED(mRef);
+    }
+  }
+#endif
 
  private:
   friend class SupportsWeakPtr;
@@ -285,6 +321,14 @@ class WeakPtr {
   RefPtr<WeakReference> mRef;
 };
 
+#ifdef MOZILLA_INTERNAL_API
+
+template <typename T>
+using MainThreadWeakPtr =
+    WeakPtr<T, detail::WeakPtrDestructorBehavior::ProxyToMainThread>;
+
+#endif
+
 #define NS_IMPL_CYCLE_COLLECTION_UNLINK_WEAK_PTR tmp->DetachWeakPtr();
 
 #define NS_IMPL_CYCLE_COLLECTION_WEAK_PTR(class_, ...) \

mfbt/tests/gtest/TestMainThreadWeakPtr.cpp

@@ -0,0 +1,43 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* 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 "gtest/gtest.h"
+
+#include "mozilla/WeakPtr.h"
+#include "mozilla/UniquePtr.h"
+#include <thread>
+
+using namespace mozilla;
+
+struct C : public SupportsWeakPtr {
+  int mNum{0};
+};
+
+struct HasWeakPtrToC {
+  HasWeakPtrToC(C* c) : mPtr(c) {}
+
+  MainThreadWeakPtr<C> mPtr;
+
+  ~HasWeakPtrToC() {
+    MOZ_RELEASE_ASSERT(!NS_IsMainThread(), "Should be released OMT");
+  }
+};
+
+TEST(MFBT_MainThreadWeakPtr, Basic)
+{
+  auto c = MakeUnique<C>();
+  MOZ_RELEASE_ASSERT(NS_IsMainThread());
+
+  auto weakRef = MakeUnique<HasWeakPtrToC>(c.get());
+
+  std::thread t([weakRef = std::move(weakRef)] {});
+
+  MOZ_RELEASE_ASSERT(!weakRef);
+  c = nullptr;
+
+  t.join();
+}
+

mfbt/tests/gtest/moz.build

@@ -14,6 +14,7 @@ UNIFIED_SOURCES += [
 SOURCES += [
     'TestAlgorithm.cpp',
     'TestInitializedOnce.cpp',
+    'TestMainThreadWeakPtr.cpp',
     'TestResultExtensions.cpp',
 ]
 

uriloader/preload/PreloaderBase.cpp

@@ -37,7 +37,7 @@ class PreloaderBase::RedirectSink final : public nsIInterfaceRequestor,
   RedirectSink(PreloaderBase* aPreloader, nsIInterfaceRequestor* aCallbacks);
 
  private:
-  WeakPtr<PreloaderBase> mPreloader;
+  MainThreadWeakPtr<PreloaderBase> mPreloader;
   nsCOMPtr<nsIInterfaceRequestor> mCallbacks;
   nsCOMPtr<nsIChannel> mRedirectChannel;
 };
@@ -46,11 +46,7 @@ PreloaderBase::RedirectSink::RedirectSink(PreloaderBase* aPreloader,
                                           nsIInterfaceRequestor* aCallbacks)
     : mPreloader(aPreloader), mCallbacks(aCallbacks) {}
 
-PreloaderBase::RedirectSink::~RedirectSink() {
-  MOZ_DIAGNOSTIC_ASSERT(NS_IsMainThread(),
-                        "Should figure out how to safely drop mPreloader "
-                        "otherwise");
-}
+PreloaderBase::RedirectSink::~RedirectSink() = default;
 
 NS_IMPL_ISUPPORTS(PreloaderBase::RedirectSink, nsIInterfaceRequestor,
                   nsIChannelEventSink, nsIRedirectResultListener)