Bug 1682045 - Allow nsPresContext to store and release the last registered OneShotPostRefreshObserver r=smaug

OneShotPostRefreshObserver works as the caller registers it, and
let it deletes itself via the DidRefresh method. The issue is that
DidRefresh is not guaranteed to run, and it'll leak PresShell
if it doesn't run.

This patch allows nsPresContext to store and release the last
registered OneShotPostRefreshObserver, and converted the existing
usage of OneShotPostRefreshObserver to use that. So instead of asking
OneShotPostRefreshObserver to delete itself, we now ask nsPresContext
to release it.

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

dom/base/nsDOMWindowUtils.cpp

@@ -2797,12 +2797,16 @@ nsDOMWindowUtils::ZoomToFocusedInput() {
     }
   }
   if (waitForRefresh) {
-    waitForRefresh =
-        presShell->AddPostRefreshObserver(new OneShotPostRefreshObserver(
-            presShell, [widget = RefPtr<nsIWidget>(widget), presShellId, viewId,
-                        bounds, flags](PresShell*) {
-              widget->ZoomToRect(presShellId, viewId, bounds, flags);
-            }));
+    waitForRefresh = false;
+    if (nsPresContext* presContext = presShell->GetPresContext()) {
+      waitForRefresh = presContext->RegisterOneShotPostRefreshObserver(
+          new OneShotPostRefreshObserver(
+              presShell,
+              [widget = RefPtr<nsIWidget>(widget), presShellId, viewId, bounds,
+               flags](PresShell*, OneShotPostRefreshObserver*) {
+                widget->ZoomToRect(presShellId, viewId, bounds, flags);
+              }));
+    }
   }
   if (!waitForRefresh) {
     widget->ZoomToRect(presShellId, viewId, bounds, flags);

gfx/layers/apz/util/APZCCallbackHelper.cpp

@@ -682,17 +682,23 @@ static void SendLayersDependentApzcTargetConfirmation(
 DisplayportSetListener::DisplayportSetListener(
     nsIWidget* aWidget, PresShell* aPresShell, const uint64_t& aInputBlockId,
     nsTArray<ScrollableLayerGuid>&& aTargets)
-    : OneShotPostRefreshObserver(
-          aPresShell,
-          [this](PresShell* aPresShell) { OnPostRefresh(this, aPresShell); }),
+    : OneShotPostRefreshObserver(aPresShell),
       mWidget(aWidget),
       mInputBlockId(aInputBlockId),
-      mTargets(std::move(aTargets)) {}
+      mTargets(std::move(aTargets)) {
+  MOZ_ASSERT(!mAction, "Setting Action twice");
+  mAction = [](PresShell* aPresShell,
+               OneShotPostRefreshObserver* aThisObserver) {
+    OnPostRefresh(static_cast<DisplayportSetListener*>(aThisObserver),
+                  aPresShell);
+  };
+}
 
 DisplayportSetListener::~DisplayportSetListener() = default;
 
 bool DisplayportSetListener::Register() {
-  if (mPresShell->AddPostRefreshObserver(this)) {
+  if (nsPresContext* presContext = mPresShell->GetPresContext()) {
+    presContext->RegisterOneShotPostRefreshObserver(this);
     APZCCH_LOG("Successfully registered post-refresh observer\n");
     return true;
   }

layout/base/PresShell.cpp

@@ -1293,6 +1293,7 @@ void PresShell::Destroy() {
         }
       }
     }
+    mPresContext->ClearOneShotPostRefreshObservers();
   }
 
 #ifdef MOZ_REFLOW_PERF

layout/base/PresShell.h

@@ -1152,7 +1152,11 @@ class PresShell final : public nsStubDocumentObserver,
                              FlushType aFlushType);
 
   bool AddPostRefreshObserver(nsAPostRefreshObserver* aObserver);
+  bool AddPostRefreshObserver(mozilla::OneShotPostRefreshObserver* aObserver) =
+      delete;
   bool RemovePostRefreshObserver(nsAPostRefreshObserver* aObserver);
+  bool RemovePostRefreshObserver(
+      mozilla::OneShotPostRefreshObserver* aObserver) = delete;
 
   // Represents an update to the visual scroll offset that will be sent to APZ.
   // The update type is used to determine priority compared to other scroll

layout/base/nsPresContext.cpp

@@ -298,6 +298,7 @@ void nsPresContext::Destroy() {
                                    gExactCallbackPrefs, this);
 
   mRefreshDriver = nullptr;
+  MOZ_ASSERT(mOneShotPostRefreshObservers.IsEmpty());
 }
 
 nsPresContext::~nsPresContext() {
@@ -1507,6 +1508,33 @@ void nsPresContext::ContentLanguageChanged() {
                                RestyleHint::RecascadeSubtree());
 }
 
+bool nsPresContext::RegisterOneShotPostRefreshObserver(
+    mozilla::OneShotPostRefreshObserver* aObserver) {
+  RefreshDriver()->AddPostRefreshObserver(
+      static_cast<nsAPostRefreshObserver*>(aObserver));
+  mOneShotPostRefreshObservers.AppendElement(aObserver);
+  return true;
+}
+
+void nsPresContext::UnregisterOneShotPostRefreshObserver(
+    mozilla::OneShotPostRefreshObserver* aObserver) {
+  RefreshDriver()->RemovePostRefreshObserver(
+      static_cast<nsAPostRefreshObserver*>(aObserver));
+  DebugOnly<bool> removed =
+      mOneShotPostRefreshObservers.RemoveElement(aObserver);
+  MOZ_ASSERT(removed,
+             "OneShotPostRefreshObserver should be owned by PresContext");
+}
+
+void nsPresContext::ClearOneShotPostRefreshObservers() {
+  for (const auto& observer : mOneShotPostRefreshObservers) {
+    RefreshDriver()->RemovePostRefreshObserver(
+        static_cast<nsAPostRefreshObserver*>(observer));
+  }
+
+  mOneShotPostRefreshObservers.Clear();
+}
+
 void nsPresContext::RebuildAllStyleData(nsChangeHint aExtraHint,
                                         const RestyleHint& aRestyleHint) {
   if (!mPresShell) {

layout/base/nsPresContext.h

@@ -75,6 +75,7 @@ class EffectCompositor;
 class Encoding;
 class EventStateManager;
 class CounterStyleManager;
+class OneShotPostRefreshObserver;
 class PresShell;
 class RestyleManager;
 class ServoStyleSet;
@@ -511,6 +512,12 @@ class nsPresContext : public nsISupports, public mozilla::SupportsWeakPtr {
 
   mozilla::ScreenIntMargin GetSafeAreaInsets() const { return mSafeAreaInsets; }
 
+  bool RegisterOneShotPostRefreshObserver(
+      mozilla::OneShotPostRefreshObserver* aObserver);
+  void UnregisterOneShotPostRefreshObserver(
+      mozilla::OneShotPostRefreshObserver* aObserver);
+  void ClearOneShotPostRefreshObservers();
+
  protected:
   void UpdateEffectiveTextZoom();
 
@@ -1249,6 +1256,9 @@ class nsPresContext : public nsISupports, public mozilla::SupportsWeakPtr {
   // During page load we use slower frame rate.
   uint32_t mNextFrameRateMultiplier;
 
+  nsTArray<RefPtr<mozilla::OneShotPostRefreshObserver>>
+      mOneShotPostRefreshObservers;
+
   ScrollStyles mViewportScrollStyles;
 
   uint16_t mImageAnimationMode;

layout/base/nsRefreshDriver.h

@@ -127,7 +127,11 @@ class nsRefreshDriver final : public mozilla::layers::TransactionIdAllocator,
    * refresh driver ticks.
    */
   void AddPostRefreshObserver(nsAPostRefreshObserver* aObserver);
+  void AddPostRefreshObserver(mozilla::OneShotPostRefreshObserver* aObserver) =
+      delete;
   void RemovePostRefreshObserver(nsAPostRefreshObserver* aObserver);
+  void RemovePostRefreshObserver(
+      mozilla::OneShotPostRefreshObserver* aObserver) = delete;
 
   /**
    * Add/Remove imgIRequest versions of observers.

layout/base/nsRefreshObservers.cpp

@@ -13,6 +13,9 @@ OneShotPostRefreshObserver::OneShotPostRefreshObserver(PresShell* aPresShell,
                                                        Action&& aAction)
     : mPresShell(aPresShell), mAction(std::move(aAction)) {}
 
+OneShotPostRefreshObserver::OneShotPostRefreshObserver(PresShell* aPresShell)
+    : mPresShell(aPresShell) {}
+
 OneShotPostRefreshObserver::~OneShotPostRefreshObserver() = default;
 
 void OneShotPostRefreshObserver::DidRefresh() {
@@ -23,9 +26,12 @@ void OneShotPostRefreshObserver::DidRefresh() {
     return;
   }
 
-  mAction(mPresShell);
+  RefPtr<OneShotPostRefreshObserver> kungfuDeathGrip = this;
 
-  if (!mPresShell->RemovePostRefreshObserver(this)) {
+  mAction(mPresShell, this);
+
+  nsPresContext* presContext = mPresShell->GetPresContext();
+  if (!presContext) {
     MOZ_ASSERT_UNREACHABLE(
         "Unable to unregister post-refresh observer! Leaking it instead of "
         "leaving garbage registered");
@@ -34,8 +40,7 @@ void OneShotPostRefreshObserver::DidRefresh() {
     mAction = Action();
     return;
   }
-
-  delete this;
+  presContext->UnregisterOneShotPostRefreshObserver(this);
 }
 
 }  // namespace mozilla

layout/base/nsRefreshObservers.h

@@ -76,12 +76,15 @@ namespace mozilla {
  */
 class OneShotPostRefreshObserver : public nsAPostRefreshObserver {
  public:
-  using Action = std::function<void(mozilla::PresShell*)>;
+  using Action =
+      std::function<void(mozilla::PresShell*, OneShotPostRefreshObserver*)>;
+  NS_INLINE_DECL_REFCOUNTING(OneShotPostRefreshObserver)
   OneShotPostRefreshObserver(mozilla::PresShell* aPresShell, Action&& aAction);
-  virtual ~OneShotPostRefreshObserver();
+  explicit OneShotPostRefreshObserver(mozilla::PresShell* aPresShell);
   void DidRefresh() override;
 
  protected:
+  virtual ~OneShotPostRefreshObserver();
   RefPtr<mozilla::PresShell> mPresShell;
   Action mAction;
 };

layout/xul/nsSliderFrame.cpp

@@ -990,12 +990,15 @@ void nsSliderFrame::StartAPZDrag(WidgetGUIEvent* aEvent) {
   bool waitForRefresh = InputAPZContext::HavePendingLayerization();
   nsIWidget* widget = this->GetNearestWidget();
   if (waitForRefresh) {
-    waitForRefresh =
-        presShell->AddPostRefreshObserver(new OneShotPostRefreshObserver(
-            presShell, [widget = RefPtr<nsIWidget>(widget),
-                        dragMetrics](mozilla::PresShell*) {
-              widget->StartAsyncScrollbarDrag(dragMetrics);
-            }));
+    waitForRefresh = false;
+    if (nsPresContext* presContext = presShell->GetPresContext()) {
+      waitForRefresh = presContext->RegisterOneShotPostRefreshObserver(
+          new OneShotPostRefreshObserver(
+              presShell, [widget = RefPtr<nsIWidget>(widget), dragMetrics](
+                             mozilla::PresShell*, OneShotPostRefreshObserver*) {
+                widget->StartAsyncScrollbarDrag(dragMetrics);
+              }));
+    }
   }
   if (!waitForRefresh) {
     widget->StartAsyncScrollbarDrag(dragMetrics);