Bug 1716481 - Improve ManagedPostRefreshObserver. r=smaug

Make it less sketchy.

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

dom/base/nsDOMWindowUtils.cpp

@@ -3078,10 +3078,11 @@ nsDOMWindowUtils::ZoomToFocusedInput() {
   if (waitForRefresh) {
     waitForRefresh = false;
     if (nsPresContext* presContext = presShell->GetPresContext()) {
-      waitForRefresh = presContext->RegisterManagedPostRefreshObserver(
+      waitForRefresh = true;
+      presContext->RegisterManagedPostRefreshObserver(
           new ManagedPostRefreshObserver(
-              presShell, [widget = RefPtr<nsIWidget>(widget), presShellId,
-                          viewId, bounds, flags](bool aWasCanceled) {
+              presContext, [widget = RefPtr<nsIWidget>(widget), presShellId,
+                            viewId, bounds, flags](bool aWasCanceled) {
                 if (!aWasCanceled) {
                   widget->ZoomToRect(presShellId, viewId, bounds, flags);
                 }

dom/base/nsGlobalWindowInner.cpp

@@ -7071,8 +7071,9 @@ bool nsGlobalWindowInner::TryToObserveRefresh() {
     return false;
   }
 
+  mObservingRefresh = true;
   auto observer = MakeRefPtr<ManagedPostRefreshObserver>(
-      pc->PresShell(), [win = RefPtr{this}](bool aWasCanceled) {
+      pc, [win = RefPtr{this}](bool aWasCanceled) {
         if (win->MaybeCallDocumentFlushedResolvers(
                 /* aUntilExhaustion = */ aWasCanceled)) {
           return ManagedPostRefreshObserver::Unregister::No;
@@ -7080,7 +7081,7 @@ bool nsGlobalWindowInner::TryToObserveRefresh() {
         win->mObservingRefresh = false;
         return ManagedPostRefreshObserver::Unregister::Yes;
       });
-  mObservingRefresh = pc->RegisterManagedPostRefreshObserver(observer.get());
+  pc->RegisterManagedPostRefreshObserver(observer.get());
   return mObservingRefresh;
 }
 

dom/ipc/BrowserChild.cpp

@@ -1754,7 +1754,7 @@ void BrowserChild::HandleRealMouseButtonEvent(const WidgetMouseEvent& aEvent,
   // APZ before the SetTargetAPZC message. This ensures the drag input block
   // gets the drag metrics before handling the input events.
   if (postLayerization) {
-    postLayerization->TryRegister();
+    postLayerization->Register();
   }
 }
 
@@ -1800,7 +1800,7 @@ void BrowserChild::DispatchWheelEvent(const WidgetWheelEvent& aEvent,
         APZCCallbackHelper::SendSetTargetAPZCNotification(
             mPuppetWidget, document, aEvent, aGuid.mLayersId, aInputBlockId);
     if (postLayerization) {
-      postLayerization->TryRegister();
+      postLayerization->Register();
     }
   }
 
@@ -1904,7 +1904,7 @@ mozilla::ipc::IPCResult BrowserChild::RecvRealTouchEvent(
             mPuppetWidget, document, localEvent, aGuid.mLayersId,
             aInputBlockId);
     if (postLayerization) {
-      postLayerization->TryRegister();
+      postLayerization->Register();
     }
   }
 

dom/performance/PerformanceMainThread.cpp

@@ -221,17 +221,18 @@ void PerformanceMainThread::InsertEventTimingEntry(
   // run any JS between the `mark paint timing` step and the
   // `pending Event Timing entries` step. So mixing the order
   // here is fine.
-  mHasQueuedRefreshdriverObserver =
-      presContext->RegisterManagedPostRefreshObserver(
-          new ManagedPostRefreshObserver(
-              presShell, [performance = RefPtr<PerformanceMainThread>(this)](
-                             bool aWasCanceled) {
-                if (!aWasCanceled) {
-                  // XXX Should we do this even if canceled?
-                  performance->DispatchPendingEventTimingEntries();
-                }
-                return ManagedPostRefreshObserver::Unregister::Yes;
-              }));
+  mHasQueuedRefreshdriverObserver = true;
+  presContext->RegisterManagedPostRefreshObserver(
+      new ManagedPostRefreshObserver(
+          presContext, [performance = RefPtr<PerformanceMainThread>(this)](
+                           bool aWasCanceled) {
+            if (!aWasCanceled) {
+              // XXX Should we do this even if canceled?
+              performance->DispatchPendingEventTimingEntries();
+            }
+            performance->mHasQueuedRefreshdriverObserver = false;
+            return ManagedPostRefreshObserver::Unregister::Yes;
+          }));
 }
 
 void PerformanceMainThread::BufferEventTimingEntryIfNeeded(
@@ -285,7 +286,6 @@ void PerformanceMainThread::DispatchPendingEventTimingEntries() {
       }
     }
   }
-  mHasQueuedRefreshdriverObserver = false;
 }
 
 // To be removed once bug 1124165 lands

gfx/layers/apz/util/APZCCallbackHelper.cpp

@@ -681,9 +681,14 @@ static bool PrepareForSetTargetAPZCNotification(
 }
 
 static void SendLayersDependentApzcTargetConfirmation(
-    PresShell* aPresShell, uint64_t aInputBlockId,
+    nsPresContext* aPresContext, uint64_t aInputBlockId,
     nsTArray<ScrollableLayerGuid>&& aTargets) {
-  LayerManager* lm = aPresShell->GetLayerManager();
+  PresShell* ps = aPresContext->GetPresShell();
+  if (!ps) {
+    return;
+  }
+
+  LayerManager* lm = ps->GetLayerManager();
   if (!lm) {
     return;
   }
@@ -711,9 +716,9 @@ static void SendLayersDependentApzcTargetConfirmation(
 }  // namespace
 
 DisplayportSetListener::DisplayportSetListener(
-    nsIWidget* aWidget, PresShell* aPresShell, const uint64_t& aInputBlockId,
-    nsTArray<ScrollableLayerGuid>&& aTargets)
-    : ManagedPostRefreshObserver(aPresShell),
+    nsIWidget* aWidget, nsPresContext* aPresContext,
+    const uint64_t& aInputBlockId, nsTArray<ScrollableLayerGuid>&& aTargets)
+    : ManagedPostRefreshObserver(aPresContext),
       mWidget(aWidget),
       mInputBlockId(aInputBlockId),
       mTargets(std::move(aTargets)) {
@@ -726,23 +731,15 @@ DisplayportSetListener::DisplayportSetListener(
 
 DisplayportSetListener::~DisplayportSetListener() = default;
 
-void DisplayportSetListener::TryRegister() {
-  if (nsPresContext* presContext = mPresShell->GetPresContext()) {
-    if (presContext->RegisterManagedPostRefreshObserver(this)) {
-      APZCCH_LOG("Successfully registered post-refresh observer\n");
-      return;
-    }
-  }
-  // In case of failure just send the notification right away
-  APZCCH_LOG("Sending target APZCs for input block %" PRIu64 "\n",
-             mInputBlockId);
-  mWidget->SetConfirmedTargetAPZC(mInputBlockId, mTargets);
+void DisplayportSetListener::Register() {
+  APZCCH_LOG("DisplayportSetListener::Register\n");
+  mPresContext->RegisterManagedPostRefreshObserver(this);
 }
 
 void DisplayportSetListener::OnPostRefresh() {
   APZCCH_LOG("Got refresh, sending target APZCs for input block %" PRIu64 "\n",
              mInputBlockId);
-  SendLayersDependentApzcTargetConfirmation(mPresShell, mInputBlockId,
+  SendLayersDependentApzcTargetConfirmation(mPresContext, mInputBlockId,
                                             std::move(mTargets));
 }
 
@@ -793,7 +790,8 @@ APZCCallbackHelper::SendSetTargetAPZCNotification(nsIWidget* aWidget,
               "At least one target got a new displayport, need to wait for "
               "refresh\n");
           return MakeAndAddRef<DisplayportSetListener>(
-              aWidget, presShell, aInputBlockId, std::move(targets));
+              aWidget, presShell->GetPresContext(), aInputBlockId,
+              std::move(targets));
         }
         APZCCH_LOG("Sending target APZCs for input block %" PRIu64 "\n",
                    aInputBlockId);

gfx/layers/apz/util/APZCCallbackHelper.h

@@ -18,6 +18,7 @@
 class nsIContent;
 class nsIScrollableFrame;
 class nsIWidget;
+class nsPresContext;
 template <class T>
 struct already_AddRefed;
 template <class T>
@@ -38,11 +39,11 @@ typedef std::function<void(uint64_t, const nsTArray<TouchBehaviorFlags>&)>
 /* Refer to documentation on SendSetTargetAPZCNotification for this class */
 class DisplayportSetListener : public ManagedPostRefreshObserver {
  public:
-  DisplayportSetListener(nsIWidget* aWidget, PresShell* aPresShell,
+  DisplayportSetListener(nsIWidget* aWidget, nsPresContext*,
                          const uint64_t& aInputBlockId,
                          nsTArray<ScrollableLayerGuid>&& aTargets);
   virtual ~DisplayportSetListener();
-  void TryRegister();
+  void Register();
 
  private:
   RefPtr<nsIWidget> mWidget;

layout/base/nsPresContext.cpp

@@ -1499,16 +1499,23 @@ void nsPresContext::ContentLanguageChanged() {
                                RestyleHint::RecascadeSubtree());
 }
 
-bool nsPresContext::RegisterManagedPostRefreshObserver(
-    mozilla::ManagedPostRefreshObserver* aObserver) {
+void nsPresContext::RegisterManagedPostRefreshObserver(
+    ManagedPostRefreshObserver* aObserver) {
+  if (MOZ_UNLIKELY(!mPresShell)) {
+    // If we're detached from our pres shell already, refuse to keep observer
+    // around, as that'd create a cycle.
+    RefPtr<ManagedPostRefreshObserver> obs = aObserver;
+    obs->Cancel();
+    return;
+  }
+
   RefreshDriver()->AddPostRefreshObserver(
       static_cast<nsAPostRefreshObserver*>(aObserver));
   mManagedPostRefreshObservers.AppendElement(aObserver);
-  return true;
 }
 
 void nsPresContext::UnregisterManagedPostRefreshObserver(
-    mozilla::ManagedPostRefreshObserver* aObserver) {
+    ManagedPostRefreshObserver* aObserver) {
   RefreshDriver()->RemovePostRefreshObserver(
       static_cast<nsAPostRefreshObserver*>(aObserver));
   DebugOnly<bool> removed =

layout/base/nsPresContext.h

@@ -510,7 +510,7 @@ class nsPresContext : public nsISupports, public mozilla::SupportsWeakPtr {
 
   mozilla::ScreenIntMargin GetSafeAreaInsets() const { return mSafeAreaInsets; }
 
-  bool RegisterManagedPostRefreshObserver(mozilla::ManagedPostRefreshObserver*);
+  void RegisterManagedPostRefreshObserver(mozilla::ManagedPostRefreshObserver*);
   void UnregisterManagedPostRefreshObserver(
       mozilla::ManagedPostRefreshObserver*);
   void CancelManagedPostRefreshObservers();

layout/base/nsRefreshDriver.cpp

@@ -1317,12 +1317,15 @@ void nsRefreshDriver::DispatchVisualViewportScrollEvents() {
 
 void nsRefreshDriver::AddPostRefreshObserver(
     nsAPostRefreshObserver* aObserver) {
+  MOZ_DIAGNOSTIC_ASSERT(!mPostRefreshObservers.Contains(aObserver));
   mPostRefreshObservers.AppendElement(aObserver);
 }
 
 void nsRefreshDriver::RemovePostRefreshObserver(
     nsAPostRefreshObserver* aObserver) {
-  mPostRefreshObservers.RemoveElement(aObserver);
+  bool removed = mPostRefreshObservers.RemoveElement(aObserver);
+  MOZ_DIAGNOSTIC_ASSERT(removed);
+  Unused << removed;
 }
 
 bool nsRefreshDriver::AddImageRequest(imgIRequest* aRequest) {

layout/base/nsRefreshObservers.cpp

@@ -5,45 +5,39 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 #include "nsRefreshObservers.h"
-#include "PresShell.h"
+#include "nsPresContext.h"
 
 namespace mozilla {
 
-ManagedPostRefreshObserver::ManagedPostRefreshObserver(PresShell* aPresShell,
+ManagedPostRefreshObserver::ManagedPostRefreshObserver(nsPresContext* aPc,
                                                        Action&& aAction)
-    : mPresShell(aPresShell), mAction(std::move(aAction)) {}
+    : mPresContext(aPc), mAction(std::move(aAction)) {}
 
-ManagedPostRefreshObserver::ManagedPostRefreshObserver(PresShell* aPresShell)
-    : mPresShell(aPresShell) {}
+ManagedPostRefreshObserver::ManagedPostRefreshObserver(nsPresContext* aPc)
+    : mPresContext(aPc) {}
 
 ManagedPostRefreshObserver::~ManagedPostRefreshObserver() = default;
 
 void ManagedPostRefreshObserver::Cancel() {
+  // Caller holds a strong reference, so no need to reference stuff from here.
   mAction(true);
   mAction = nullptr;
-  mPresShell = nullptr;
+  mPresContext = nullptr;
 }
 
 void ManagedPostRefreshObserver::DidRefresh() {
-  if (!mPresShell) {
-    MOZ_ASSERT_UNREACHABLE(
-        "Post-refresh observer fired again after failed attempt at "
-        "unregistering it");
-    return;
-  }
+  RefPtr<ManagedPostRefreshObserver> thisObject = this;
+
   Unregister unregister = mAction(false);
-  if (bool(unregister)) {
-    nsPresContext* presContext = mPresShell->GetPresContext();
-    if (!presContext) {
-      MOZ_ASSERT_UNREACHABLE(
-          "Unable to unregister post-refresh observer! Leaking it instead of "
-          "leaving garbage registered");
-      // Graceful handling, just in case...
-      mPresShell = nullptr;
+  if (unregister == Unregister::Yes) {
+    if (RefPtr<nsPresContext> pc = std::move(mPresContext)) {
+      // In theory mAction could've ended up in `Cancel` being called. In which
+      // case we're already unregistered so no need to do anything.
       mAction = nullptr;
-      return;
+      pc->UnregisterManagedPostRefreshObserver(this);
+    } else {
+      MOZ_DIAGNOSTIC_ASSERT(!mAction);
     }
-    presContext->UnregisterManagedPostRefreshObserver(this);
   }
 }
 

layout/base/nsRefreshObservers.h

@@ -18,6 +18,8 @@
 #include "mozilla/TimeStamp.h"
 #include "nsISupports.h"
 
+class nsPresContext;
+
 namespace mozilla {
 class AnimationEventDispatcher;
 class PendingFullscreenEvent;
@@ -80,14 +82,14 @@ class ManagedPostRefreshObserver : public nsAPostRefreshObserver {
   enum class Unregister : bool { No, Yes };
   using Action = std::function<Unregister(bool aWasCanceled)>;
   NS_INLINE_DECL_REFCOUNTING(ManagedPostRefreshObserver)
-  ManagedPostRefreshObserver(mozilla::PresShell*, Action&&);
-  explicit ManagedPostRefreshObserver(mozilla::PresShell*);
+  ManagedPostRefreshObserver(nsPresContext*, Action&&);
+  explicit ManagedPostRefreshObserver(nsPresContext*);
   void DidRefresh() override;
   void Cancel();
 
  protected:
   virtual ~ManagedPostRefreshObserver();
-  RefPtr<mozilla::PresShell> mPresShell;
+  RefPtr<nsPresContext> mPresContext;
   Action mAction;
 };
 

layout/xul/nsSliderFrame.cpp

@@ -1140,15 +1140,16 @@ void nsSliderFrame::StartAPZDrag(WidgetGUIEvent* aEvent) {
   if (waitForRefresh) {
     waitForRefresh = false;
     if (nsPresContext* presContext = presShell->GetPresContext()) {
-      waitForRefresh = presContext->RegisterManagedPostRefreshObserver(
+      presContext->RegisterManagedPostRefreshObserver(
           new ManagedPostRefreshObserver(
-              presShell, [widget = RefPtr<nsIWidget>(widget),
-                          dragMetrics](bool aWasCanceled) {
+              presContext, [widget = RefPtr<nsIWidget>(widget),
+                            dragMetrics](bool aWasCanceled) {
                 if (!aWasCanceled) {
                   widget->StartAsyncScrollbarDrag(dragMetrics);
                 }
                 return ManagedPostRefreshObserver::Unregister::Yes;
               }));
+      waitForRefresh = true;
     }
   }
   if (!waitForRefresh) {

widget/nsBaseWidget.cpp

@@ -1094,7 +1094,7 @@ nsEventStatus nsBaseWidget::ProcessUntransformedAPZEvent(
       mAPZEventState->ProcessMouseEvent(*mouseEvent, inputBlockId);
     }
     if (postLayerization) {
-      postLayerization->TryRegister();
+      postLayerization->Register();
     }
   }