Fix Overlay2 storing stale onClose callback

packages/core/src/components/overlay2/overlay2.tsx

@@ -195,17 +195,6 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
         [getThisOverlayAndDescendants, id, onClose],
     );
 
-    // Important: clean up old document-level event listeners if their memoized values change (this is rare, but
-    // may happen, for example, if a user forgets to use `React.useCallback` in the `props.onClose` value).
-    // Otherwise, we will lose the reference to those values and create a memory leak since we won't be able
-    // to successfully detach them inside overlayWillClose.
-    React.useEffect(() => {
-        document.removeEventListener("mousedown", handleDocumentMousedown);
-    }, [handleDocumentMousedown]);
-    React.useEffect(() => {
-        document.removeEventListener("focus", handleDocumentFocus, /* useCapture */ true);
-    }, [handleDocumentFocus]);
-
     // send this instance's imperative handle to the the forwarded ref as well as our local ref
     const ref = React.useMemo(() => mergeRefs(forwardedRef, instance), [forwardedRef]);
     React.useImperativeHandle(
@@ -264,28 +253,8 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
             bringFocusInsideOverlay();
         }
 
-        if (enforceFocus) {
-            // Focus events do not bubble, but setting useCapture allows us to listen in and execute
-            // our handler before all others
-            document.addEventListener("focus", handleDocumentFocus, /* useCapture */ true);
-        }
-
-        if (canOutsideClickClose && !hasBackdrop) {
-            document.addEventListener("mousedown", handleDocumentMousedown);
-        }
-
         setRef(lastActiveElementBeforeOpened, getActiveElement(getRef(containerElement)));
-    }, [
-        autoFocus,
-        bringFocusInsideOverlay,
-        canOutsideClickClose,
-        enforceFocus,
-        getLastOpened,
-        handleDocumentMousedown,
-        handleDocumentFocus,
-        hasBackdrop,
-        openOverlay,
-    ]);
+    }, [autoFocus, bringFocusInsideOverlay, getLastOpened, openOverlay]);
 
     const overlayWillClose = React.useCallback(() => {
         document.removeEventListener("focus", handleDocumentFocus, /* useCapture */ true);
@@ -307,6 +276,8 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
             }
         }
     }, [closeOverlay, getLastOpened, handleDocumentFocus, handleDocumentMousedown, id]);
+    const mostRecetOverlayWillClose = React.useRef(overlayWillClose);
+    mostRecetOverlayWillClose.current = overlayWillClose;
 
     const prevIsOpen = usePrevious(isOpen) ?? false;
     React.useEffect(() => {
@@ -325,14 +296,41 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
         }
     }, [isOpen, overlayWillOpen, overlayWillClose, prevIsOpen]);
 
-    // run once on unmount
+    // Important: clean up old document-level event listeners if their memoized values change (this is rare, but
+    // may happen, for example, if a user forgets to use `React.useCallback` in the `props.onClose` value).
+    // Otherwise, we will lose the reference to those values and create a memory leak since we won't be able
+    // to successfully detach them inside overlayWillClose.
     React.useEffect(() => {
+        if (!isOpen || !(canOutsideClickClose && !hasBackdrop)) {
+            return;
+        }
+
+        document.addEventListener("mousedown", handleDocumentMousedown);
+
         return () => {
-            // we need to run cleanup code to remove some event handlers from the overlay element
-            overlayWillClose();
+            document.removeEventListener("mousedown", handleDocumentMousedown);
+        };
+    }, [handleDocumentMousedown, isOpen, canOutsideClickClose, hasBackdrop]);
+    React.useEffect(() => {
+        if (!isOpen || !enforceFocus) {
+            return;
+        }
+
+        // Focus events do not bubble, but setting useCapture allows us to listen in and execute
+        // our handler before all others
+        document.addEventListener("focus", handleDocumentFocus, /* useCapture */ true);
+
+        return () => {
+            document.removeEventListener("focus", handleDocumentFocus, /* useCapture */ true);
         };
+    }, [handleDocumentFocus, enforceFocus, isOpen]);
 
-        // eslint-disable-next-line react-hooks/exhaustive-deps
+    React.useEffect(() => {
+        // run cleanup code once on unmount, ensuring we call the most recent overlayWillClose callback
+        // by storing in a ref and keeping up to date
+        return () => {
+            mostRecetOverlayWillClose.current();
+        };
     }, []);
 
     const handleTransitionExited = React.useCallback(

packages/core/test/overlay2/overlay2Tests.tsx

@@ -184,8 +184,13 @@ describe("<Overlay2>", () => {
                     {createOverlayContents()}
                 </OverlayWrapper>,
             );
+
+            // ensure onClose can be updated
+            const onClose2 = spy();
+            overlay.setProps({ onClose: onClose2 });
             overlay.find(BACKDROP_SELECTOR).simulate("mousedown");
-            assert.isTrue(onClose.calledOnce);
+            assert.isTrue(onClose.notCalled);
+            assert.isTrue(onClose2.calledOnce);
         });
 
         it("not invoked on backdrop mousedown when canOutsideClickClose=false", () => {
@@ -201,15 +206,18 @@ describe("<Overlay2>", () => {
 
         it("invoked on document mousedown when hasBackdrop=false", () => {
             const onClose = spy();
-            // mounting cuz we need document events + lifecycle
-            mountWrapper(
-                <OverlayWrapper hasBackdrop={false} isOpen={true} onClose={onClose} usePortal={false}>
+            const overlay = mountWrapper(
+                <OverlayWrapper isOpen={true} onClose={onClose} usePortal={false} hasBackdrop={false}>
                     {createOverlayContents()}
                 </OverlayWrapper>,
             );
 
+            // ensure onClose can be updated
+            const onClose2 = spy();
+            overlay.setProps({ onClose: onClose2 });
             dispatchMouseEvent(document.documentElement, "mousedown");
-            assert.isTrue(onClose.calledOnce);
+            assert.isTrue(onClose.notCalled);
+            assert.isTrue(onClose2.calledOnce);
         });
 
         it("not invoked on document mousedown when hasBackdrop=false and canOutsideClickClose=false", () => {
@@ -249,13 +257,18 @@ describe("<Overlay2>", () => {
 
         it("invoked on escape key", () => {
             const onClose = spy();
-            mountWrapper(
+            const overlay = mountWrapper(
                 <OverlayWrapper isOpen={true} onClose={onClose} usePortal={false}>
                     {createOverlayContents()}
                 </OverlayWrapper>,
             );
+
+            // ensure onClose can be updated
+            const onClose2 = spy();
+            overlay.setProps({ onClose: onClose2 });
             wrapper.simulate("keydown", { key: "Escape" });
-            assert.isTrue(onClose.calledOnce);
+            assert.isTrue(onClose.notCalled);
+            assert.isTrue(onClose2.calledOnce);
         });
 
         it("not invoked on escape key when canEscapeKeyClose=false", () => {