From cbb705cc07d86c1396724d003f8bc63a729adab4 Mon Sep 17 00:00:00 2001
From: Adi Dahiya <adahiya@palantir.com>
Date: Tue, 30 Jan 2024 12:08:53 -0500
Subject: [PATCH] [core] fix(ContextMenuPopover): clean up document handlers on
 close (#6685)

---
 .../core/src/common/utils/mountOptions.ts     | 54 +++++++++++++++++++
 .../context-menu/context-menu-popover.md      | 27 ++++++----
 .../components/context-menu/context-menu.md   | 14 +++--
 .../context-menu/contextMenuSingleton.tsx     | 30 ++++++++---
 packages/core/src/components/dialog/dialog.md |  5 +-
 packages/core/src/components/drawer/drawer.md |  4 +-
 packages/core/src/components/index.ts         |  3 +-
 .../core/src/components/overlay2/overlay2.tsx | 49 +++++------------
 .../components/overlay2/overlayInstance.ts    | 49 +++++++++++++++++
 .../src/components/toast/overlayToaster.tsx   | 21 +-------
 .../src/context/overlays/overlaysProvider.tsx |  2 +-
 .../hooks/overlays/useLegacyOverlayStack.ts   | 10 ++--
 .../src/hooks/overlays/useOverlayStack.ts     | 22 +++++---
 .../core/test/hooks/useOverlayStackTests.tsx  | 29 ++--------
 .../examples/core-examples/drawerExample.tsx  |  7 ++-
 15 files changed, 200 insertions(+), 126 deletions(-)
 create mode 100644 packages/core/src/common/utils/mountOptions.ts
 create mode 100644 packages/core/src/components/overlay2/overlayInstance.ts

diff --git a/packages/core/src/common/utils/mountOptions.ts b/packages/core/src/common/utils/mountOptions.ts
new file mode 100644
index 0000000000..5798069b2f
--- /dev/null
+++ b/packages/core/src/common/utils/mountOptions.ts
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2024 Palantir Technologies, Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type * as React from "react";
+
+/** Identical to `import("react-dom").Container`, copied here to avoid an unncessary type dependency. */
+type Container = Element | Document | DocumentFragment;
+
+/**
+ * Generic options interface for Blueprint APIs which imperatively mount a React component to the
+ * DOM using `"react-dom"`: `OverlayToaster.create`, `showContextMenu`, etc.
+ *
+ * The `domRenderer` currently defaults to React 16's `ReactDOM.render()`; a future version of Blueprint
+ * will default to using React 18's `createRoot()` instead, but it's possible to configure this
+ * function to use the newer API by overriding the default.
+ */
+export interface DOMMountOptions<P> {
+    /**
+     * A new DOM element will be created and appended to this container.
+     *
+     * @default document.body
+     */
+    container?: HTMLElement;
+
+    /**
+     * A function to render the React component onto a newly created DOM element.
+     *
+     * @default ReactDOM.render
+     */
+    domRenderer?: (
+        element: React.ReactElement<P>,
+        container: Container | null,
+    ) => React.Component<P, any> | Element | void;
+
+    /**
+     * A function to unmount the React component from its DOM element.
+     *
+     * @default ReactDOM.unmountComponentAtNode
+     */
+    domUnmounter?: (container: Element | DocumentFragment) => void;
+}
diff --git a/packages/core/src/components/context-menu/context-menu-popover.md b/packages/core/src/components/context-menu/context-menu-popover.md
index 25982db14f..a180ae0401 100644
--- a/packages/core/src/components/context-menu/context-menu-popover.md
+++ b/packages/core/src/components/context-menu/context-menu-popover.md
@@ -1,20 +1,21 @@
-@# ContextMenuPopover
+@# Context Menu Popover
 
 <div class="@ns-callout @ns-intent-primary @ns-icon-info-sign @ns-callout-has-body-content">
     <h5 class="@ns-heading">
 
-Consider [ContextMenu](#core/components/context-menu) first
+Consider [Context Menu](#core/components/context-menu) first
 
 </h5>
 
 The APIs described on this page are lower-level and have some limitations compared to
-[ContextMenu](#core/components/context-menu), so you should try that API _first_ to see if it addresses your use case.
+[Context Menu](#core/components/context-menu), so you should try that API _first_ to see if it addresses your use case.
 
 </div>
 
-__ContextMenuPopover__ is a lower-level API for [ContextMenu](#core/components/context-menu) which does not hook up any
-interaction handlers for you and simply renders an opinionated [Popover](#core/components/popover) at a particular
-target offset on the page through a [Portal](#core/components/portal).
+**Context Menu Popover** is a lower-level API for [**Context Menu**](#core/components/context-menu) which does
+not hook up any interaction handlers for you and simply renders an opinionated
+[**Popover**](#core/components/popover) at a particular target offset on the page through a
+[**Portal**](#core/components/portal).
 
 @reactExample ContextMenuPopoverExample
 
@@ -27,17 +28,21 @@ component which requires its `isOpen` and `targetOffset` props to be defined.
 
 @## Imperative API
 
-Two functions are provided as an imperative API for showing and hiding a singleton ContextMenuPopover on the page:
+Two functions are provided as an imperative API for showing and hiding a singleton **Context Menu Popover** on
+the page:
 
 ```ts
-export function showContextMenu(props: ContextMenuPopoverProps): void;
-export function hideContextMenu(): void;
+export function showContextMenu(
+    props: ContextMenuPopoverProps,
+    options?: DOMMountOptions<ContextMenuPopoverProps>,
+): void;
+export function hideContextMenu(options?: DOMMountOptions<ContextMenuPopoverProps>): void;
 ```
 
 These are useful in some cases when working with imperative code that does not follow typical React paradigms.
 Note that these functions come with come caveats, and thus they should be used with caution:
 
 -   they rely on global state stored in Blueprint library code.
--   they create a new React DOM tree via `ReactDOM.render()`, which means they do not preserve any existing React
-    context from the calling code.
+-   they create a new React DOM tree via `ReactDOM.render()` (or `ReactDOM.createRoot()` if you override the
+    default renderer via `options`), which means they do not preserve any existing React context from the calling code.
 -   they do _not_ automatically detect dark theme, so you must manualy set the `{ isDarkTheme: true }` property
diff --git a/packages/core/src/components/context-menu/context-menu.md b/packages/core/src/components/context-menu/context-menu.md
index 445b648057..466759e977 100644
--- a/packages/core/src/components/context-menu/context-menu.md
+++ b/packages/core/src/components/context-menu/context-menu.md
@@ -1,8 +1,8 @@
-@# ContextMenu
+@# Context Menu
 
-Context menus present the user with a list of actions when right-clicking on a target element.
-They essentially generate an opinionated Popover instance configured with the appropriate
-interaction handlers.
+**Context menus** present the user with a list of actions when right-clicking on a target element.
+They essentially generate an opinionated [**Popover**](#core/components/popover) instance configured
+with the appropriate interaction handlers.
 
 @reactExample ContextMenuExample
 
@@ -24,9 +24,7 @@ export default function ContextMenuExample() {
                 </Menu>
             }
         >
-            <div className="my-context-menu-target">
-                Right click me!
-            </div>
+            <div className="my-context-menu-target">Right click me!</div>
         </ContextMenu>
     );
 }
@@ -72,7 +70,7 @@ export default function AdvancedContextMenuExample() {
                 </div>
             )}
         </ContextMenu>
-    )
+    );
 }
 ```
 
diff --git a/packages/core/src/components/context-menu/contextMenuSingleton.tsx b/packages/core/src/components/context-menu/contextMenuSingleton.tsx
index 3fb704eac2..fa250a2503 100644
--- a/packages/core/src/components/context-menu/contextMenuSingleton.tsx
+++ b/packages/core/src/components/context-menu/contextMenuSingleton.tsx
@@ -18,6 +18,8 @@ import * as React from "react";
 import * as ReactDOM from "react-dom";
 
 import { Classes } from "../../common";
+import type { DOMMountOptions } from "../../common/utils/mountOptions";
+import { OverlaysProvider } from "../../context/overlays/overlaysProvider";
 
 import { ContextMenuPopover, type ContextMenuPopoverProps } from "./contextMenuPopover";
 
@@ -42,19 +44,33 @@ let contextMenuElement: HTMLElement | undefined;
  *
  * @see https://blueprintjs.com/docs/#core/components/context-menu-popover.imperative-api
  */
-export function showContextMenu(props: Omit<ContextMenuPopoverProps, "isOpen">) {
+export function showContextMenu(
+    props: Omit<ContextMenuPopoverProps, "isOpen">,
+    options: DOMMountOptions<ContextMenuPopoverProps> = {},
+) {
+    const {
+        container = document.body,
+        domRenderer = ReactDOM.render,
+        domUnmounter = ReactDOM.unmountComponentAtNode,
+    } = options;
+
     if (contextMenuElement === undefined) {
         contextMenuElement = document.createElement("div");
         contextMenuElement.classList.add(Classes.CONTEXT_MENU);
-        document.body.appendChild(contextMenuElement);
+        container.appendChild(contextMenuElement);
     } else {
         // N.B. It's important to unmount previous instances of the ContextMenuPopover rendered by this function.
         // Otherwise, React will detect no change in props sent to the already-mounted component, and therefore
         // do nothing after the first call to this function, leading to bugs like https://github.com/palantir/blueprint/issues/5949
-        ReactDOM.unmountComponentAtNode(contextMenuElement);
+        domUnmounter(contextMenuElement);
     }
 
-    ReactDOM.render(<UncontrolledContextMenuPopover {...props} />, contextMenuElement);
+    domRenderer(
+        <OverlaysProvider>
+            <UncontrolledContextMenuPopover {...props} />
+        </OverlaysProvider>,
+        contextMenuElement,
+    );
 }
 
 /**
@@ -64,9 +80,11 @@ export function showContextMenu(props: Omit<ContextMenuPopoverProps, "isOpen">)
  *
  * @see https://blueprintjs.com/docs/#core/components/context-menu-popover.imperative-api
  */
-export function hideContextMenu() {
+export function hideContextMenu(options: DOMMountOptions<ContextMenuPopoverProps> = {}) {
+    const { domUnmounter = ReactDOM.unmountComponentAtNode } = options;
+
     if (contextMenuElement !== undefined) {
-        ReactDOM.unmountComponentAtNode(contextMenuElement);
+        domUnmounter(contextMenuElement);
         contextMenuElement = undefined;
     }
 }
diff --git a/packages/core/src/components/dialog/dialog.md b/packages/core/src/components/dialog/dialog.md
index 664254c3a0..5387241cd2 100644
--- a/packages/core/src/components/dialog/dialog.md
+++ b/packages/core/src/components/dialog/dialog.md
@@ -1,6 +1,7 @@
-@# Dialogs
+@# Dialog
 
-**Dialog** presents content overlaid over other parts of the UI.
+The **Dialog** component presents content overlaid over other parts of the UI via
+[**Overlay2**](#core/components/overlay2).
 
 <div class="@ns-callout @ns-intent-primary @ns-icon-info-sign @ns-callout-has-body-content">
     <h5 class="@ns-heading">Terminology note</h5>
diff --git a/packages/core/src/components/drawer/drawer.md b/packages/core/src/components/drawer/drawer.md
index bd99c09d4a..d232dad9a4 100644
--- a/packages/core/src/components/drawer/drawer.md
+++ b/packages/core/src/components/drawer/drawer.md
@@ -1,7 +1,7 @@
 @# Drawer
 
-**Drawers** overlay content over existing parts of the UI and are anchored to the edge of the screen. It is built using
-the lower-level [**Overlay**](#core/components/overlay) component.
+**Drawers** overlay content over existing parts of the UI and are anchored to the edge of the screen.
+It is built using the lower-level [**Overlay2**](#core/components/overlay2) component.
 
 @reactExample DrawerExample
 
diff --git a/packages/core/src/components/index.ts b/packages/core/src/components/index.ts
index 7b14377c41..d46b8f038e 100644
--- a/packages/core/src/components/index.ts
+++ b/packages/core/src/components/index.ts
@@ -73,7 +73,8 @@ export { OverflowList, type OverflowListProps } from "./overflow-list/overflowLi
 // eslint-disable-next-line deprecation/deprecation
 export { Overlay } from "./overlay/overlay";
 export type { OverlayLifecycleProps, OverlayProps, OverlayableProps } from "./overlay/overlayProps";
-export { Overlay2, type Overlay2Props, type OverlayInstance } from "./overlay2/overlay2";
+export { Overlay2, type Overlay2Props } from "./overlay2/overlay2";
+export type { OverlayInstance } from "./overlay2/overlayInstance";
 export { Text, type TextProps } from "./text/text";
 // eslint-disable-next-line deprecation/deprecation
 export { PanelStack, type PanelStackProps } from "./panel-stack/panelStack";
diff --git a/packages/core/src/components/overlay2/overlay2.tsx b/packages/core/src/components/overlay2/overlay2.tsx
index 46efc7c159..8d9789daf9 100644
--- a/packages/core/src/components/overlay2/overlay2.tsx
+++ b/packages/core/src/components/overlay2/overlay2.tsx
@@ -42,28 +42,7 @@ import type { OverlayProps } from "../overlay/overlayProps";
 import { getKeyboardFocusableElements } from "../overlay/overlayUtils";
 import { Portal } from "../portal/portal";
 
-/**
- * Public instance properties & methods for an overlay in the current overlay stack.
- */
-export interface OverlayInstance {
-    /** Bring document focus inside this overlay element. */
-    bringFocusInsideOverlay: () => void;
-
-    /** Reference to the overlay container element which may or may not be in a Portal. */
-    containerElement: React.RefObject<HTMLDivElement>;
-
-    /** Document "focus" event handler which needs to be attached & detached appropriately. */
-    handleDocumentFocus: (e: FocusEvent) => void;
-
-    /** Document "mousedown" event handler which needs to be attached & detached appropriately. */
-    handleDocumentMousedown?: (e: MouseEvent) => void;
-
-    /** Unique ID for this overlay which helps to identify it across prop changes. */
-    id: string;
-
-    /** Subset of props necessary for some overlay stack focus management logic. */
-    props: Pick<OverlayProps, "autoFocus" | "enforceFocus" | "usePortal" | "hasBackdrop">;
-}
+import type { OverlayInstance } from "./overlayInstance";
 
 export interface Overlay2Props extends OverlayProps, React.RefAttributes<OverlayInstance> {
     /**
@@ -194,15 +173,11 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
     // N.B. this listener is only kept attached when `isOpen={true}` and `canOutsideClickClose={true}`
     const handleDocumentMousedown = React.useCallback(
         (e: MouseEvent) => {
-            if (instance.current == null) {
-                return;
-            }
-
             // get the actual target even in the Shadow DOM
             // see https://github.com/palantir/blueprint/issues/4220
             const eventTarget = (e.composed ? e.composedPath()[0] : e.target) as HTMLElement;
 
-            const thisOverlayAndDescendants = getThisOverlayAndDescendants(instance.current);
+            const thisOverlayAndDescendants = getThisOverlayAndDescendants(id);
             const isClickInThisOverlayOrDescendant = thisOverlayAndDescendants.some(
                 ({ containerElement: containerRef }) => {
                     // `elem` is the container of backdrop & content, so clicking directly on that container
@@ -217,7 +192,7 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
                 onClose?.(e as any);
             }
         },
-        [getThisOverlayAndDescendants, onClose],
+        [getThisOverlayAndDescendants, id, onClose],
     );
 
     // Important: clean up old document-level event listeners if their memoized values change (this is rare, but
@@ -279,7 +254,7 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
         }
 
         const lastOpenedOverlay = getLastOpened();
-        if (lastOpenedOverlay !== undefined) {
+        if (lastOpenedOverlay?.handleDocumentFocus !== undefined) {
             document.removeEventListener("focus", lastOpenedOverlay.handleDocumentFocus, /* useCapture */ true);
         }
         openOverlay(instance.current);
@@ -313,25 +288,25 @@ export const Overlay2 = React.forwardRef<OverlayInstance, Overlay2Props>((props,
     ]);
 
     const overlayWillClose = React.useCallback(() => {
-        if (instance.current == null) {
-            return;
-        }
-
         document.removeEventListener("focus", handleDocumentFocus, /* useCapture */ true);
         document.removeEventListener("mousedown", handleDocumentMousedown);
 
-        closeOverlay(instance.current);
+        // N.B. `instance.current` may be null at this point if we are cleaning up an open overlay during the unmount phase
+        // (this is common, for example, with context menu's singleton `showContextMenu` / `hideContextMenu` imperative APIs).
+        closeOverlay(id);
         const lastOpenedOverlay = getLastOpened();
         if (lastOpenedOverlay !== undefined) {
             // Only bring focus back to last overlay if it had autoFocus _and_ enforceFocus enabled.
             // If `autoFocus={false}`, it's likely that the overlay never received focus in the first place,
             // so it would be surprising for us to send it there. See https://github.com/palantir/blueprint/issues/4921
             if (lastOpenedOverlay.props.autoFocus && lastOpenedOverlay.props.enforceFocus) {
-                lastOpenedOverlay.bringFocusInsideOverlay();
-                document.addEventListener("focus", lastOpenedOverlay.handleDocumentFocus, /* useCapture */ true);
+                lastOpenedOverlay.bringFocusInsideOverlay?.();
+                if (lastOpenedOverlay.handleDocumentFocus !== undefined) {
+                    document.addEventListener("focus", lastOpenedOverlay.handleDocumentFocus, /* useCapture */ true);
+                }
             }
         }
-    }, [closeOverlay, getLastOpened, handleDocumentFocus, handleDocumentMousedown]);
+    }, [closeOverlay, getLastOpened, handleDocumentFocus, handleDocumentMousedown, id]);
 
     const prevIsOpen = usePrevious(isOpen) ?? false;
     React.useEffect(() => {
diff --git a/packages/core/src/components/overlay2/overlayInstance.ts b/packages/core/src/components/overlay2/overlayInstance.ts
new file mode 100644
index 0000000000..da3ca90abc
--- /dev/null
+++ b/packages/core/src/components/overlay2/overlayInstance.ts
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2024 Palantir Technologies, Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type { OverlayProps } from "../overlay/overlayProps";
+
+/**
+ * Public instance properties & methods for an overlay in the current overlay stack.
+ */
+export interface OverlayInstance {
+    /**
+     * Bring document focus inside this overlay element.
+     * This should be defined if `props.enforceFocus={true}` or `props.autoFocus={true}`.
+     */
+    bringFocusInsideOverlay?: () => void;
+
+    /** Reference to the overlay container element which may or may not be in a Portal. */
+    containerElement: React.RefObject<HTMLDivElement>;
+
+    /**
+     * Document "focus" event handler which needs to be attached & detached appropriately.
+     * This should be defined if `props.enforceFocus={true}`.
+     */
+    handleDocumentFocus?: (e: FocusEvent) => void;
+
+    /**
+     * Document "mousedown" event handler which needs to be attached & detached appropriately.
+     * This should be defined if `props.canOutsideClickClose={true}` and `props.hasBackdrop={false}`.
+     */
+    handleDocumentMousedown?: (e: MouseEvent) => void;
+
+    /** Unique ID for this overlay which helps to identify it across prop changes. */
+    id: string;
+
+    /** Subset of props necessary for some overlay stack focus management logic. */
+    props: Pick<OverlayProps, "autoFocus" | "enforceFocus" | "usePortal" | "hasBackdrop">;
+}
diff --git a/packages/core/src/components/toast/overlayToaster.tsx b/packages/core/src/components/toast/overlayToaster.tsx
index 42e046db25..508d368b06 100644
--- a/packages/core/src/components/toast/overlayToaster.tsx
+++ b/packages/core/src/components/toast/overlayToaster.tsx
@@ -27,6 +27,7 @@ import {
 } from "../../common/errors";
 import { DISPLAYNAME_PREFIX } from "../../common/props";
 import { isElementOfType, isNodeEnv } from "../../common/utils";
+import type { DOMMountOptions } from "../../common/utils/mountOptions";
 import { Overlay2 } from "../overlay2/overlay2";
 
 import type { OverlayToasterProps } from "./overlayToasterProps";
@@ -40,25 +41,7 @@ export interface OverlayToasterState {
     toastRefs: Record<string, React.RefObject<HTMLElement>>;
 }
 
-export interface OverlayToasterCreateOptions {
-    /**
-     * A new DOM element will be created to render the OverlayToaster component
-     * and appended to this container.
-     *
-     * @default document.body
-     */
-    container?: HTMLElement;
-
-    /**
-     * A function to render the OverlayToaster React component onto a newly
-     * created DOM element.
-     *
-     * Defaults to `ReactDOM.render`. A future version of Blueprint will default
-     * to using React 18's createRoot API, but it's possible to configure this
-     * function to use createRoot on earlier Blueprint versions.
-     */
-    domRenderer?: (toaster: React.ReactElement<OverlayToasterProps>, containerElement: HTMLElement) => void;
-}
+export type OverlayToasterCreateOptions = DOMMountOptions<OverlayToasterProps>;
 
 /**
  * OverlayToaster component.
diff --git a/packages/core/src/context/overlays/overlaysProvider.tsx b/packages/core/src/context/overlays/overlaysProvider.tsx
index b9613b083f..e32f90929a 100644
--- a/packages/core/src/context/overlays/overlaysProvider.tsx
+++ b/packages/core/src/context/overlays/overlaysProvider.tsx
@@ -16,7 +16,7 @@
 
 import * as React from "react";
 
-import type { OverlayInstance } from "../../components/overlay2/overlay2";
+import type { OverlayInstance } from "../../components/overlay2/overlayInstance";
 
 // N.B. using a mutable ref for the stack is much easier to work with in the world of hooks and FCs.
 // This matches the mutable global behavior of the old Overlay implementation in Blueprint v5. An alternative
diff --git a/packages/core/src/hooks/overlays/useLegacyOverlayStack.ts b/packages/core/src/hooks/overlays/useLegacyOverlayStack.ts
index 2e410fe7e8..c4740af17b 100644
--- a/packages/core/src/hooks/overlays/useLegacyOverlayStack.ts
+++ b/packages/core/src/hooks/overlays/useLegacyOverlayStack.ts
@@ -65,8 +65,8 @@ export function useLegacyOverlayStack(): UseOverlayStackReturnValue {
     const getLastOpened = React.useCallback(() => stack.at(-1), [stack]);
 
     const getThisOverlayAndDescendants = React.useCallback(
-        (overlay: OverlayInstance) => {
-            const stackIndex = stack.findIndex(o => o.id === overlay.id);
+        (id: string) => {
+            const stackIndex = stack.findIndex(o => o.id === id);
             return stack.slice(stackIndex);
         },
         [stack],
@@ -85,12 +85,12 @@ export function useLegacyOverlayStack(): UseOverlayStackReturnValue {
     }, []);
 
     const closeOverlay = React.useCallback(
-        (overlay: OverlayInstance) => {
+        (id: string) => {
             const otherOverlaysWithBackdrop = stack.filter(
-                o => o.props.usePortal && o.props.hasBackdrop && o.id !== overlay.id,
+                o => o.props.usePortal && o.props.hasBackdrop && o.id !== id,
             );
 
-            const index = globalStack.findIndex(o => o.id === overlay.id);
+            const index = globalStack.findIndex(o => o.id === id);
             if (index > -1) {
                 globalStack.splice(index, 1);
             }
diff --git a/packages/core/src/hooks/overlays/useOverlayStack.ts b/packages/core/src/hooks/overlays/useOverlayStack.ts
index 378cbb1699..73406dd38c 100644
--- a/packages/core/src/hooks/overlays/useOverlayStack.ts
+++ b/packages/core/src/hooks/overlays/useOverlayStack.ts
@@ -27,8 +27,14 @@ import { useLegacyOverlayStack } from "./useLegacyOverlayStack";
 export interface UseOverlayStackReturnValue {
     /**
      * Removes an existing overlay off the stack.
+     *
+     * N.B. This method accepts an id instead of an overlay instance because the latter may be
+     * null when an overlay is unmounting, and we may stil have some cleanup to do at that time.
+     * Also, this method is not idempotent: if the overlay is not found on the stack, nothing happens.
+     *
+     * @param id identifier of the overlay to be closed
      */
-    closeOverlay: (overlay: OverlayInstance) => void;
+    closeOverlay: (id: string) => void;
 
     /**
      * @returns the last opened overlay on the stack
@@ -36,10 +42,10 @@ export interface UseOverlayStackReturnValue {
     getLastOpened: () => OverlayInstance | undefined;
 
     /**
-     * @param overlay current overlay
+     * @param id current overlay identifier
      * @returns a list of the current overlay and all overlays which are descendants of it.
      */
-    getThisOverlayAndDescendants: (overlay: OverlayInstance) => OverlayInstance[];
+    getThisOverlayAndDescendants: (id: string) => OverlayInstance[];
 
     /**
      * Pushes a new overlay onto the stack.
@@ -68,8 +74,8 @@ export function useOverlayStack(): UseOverlayStackReturnValue {
     }, [stack]);
 
     const getThisOverlayAndDescendants = React.useCallback(
-        (overlay: OverlayInstance) => {
-            const index = stack.current.findIndex(o => o.id === overlay.id);
+        (id: string) => {
+            const index = stack.current.findIndex(o => o.id === id);
             if (index === -1) {
                 return [];
             }
@@ -94,12 +100,12 @@ export function useOverlayStack(): UseOverlayStackReturnValue {
     );
 
     const closeOverlay = React.useCallback(
-        (overlay: OverlayInstance) => {
+        (id: string) => {
             const otherOverlaysWithBackdrop = stack.current.filter(
-                o => o.props.usePortal && o.props.hasBackdrop && o.id !== overlay.id,
+                o => o.props.usePortal && o.props.hasBackdrop && o.id !== id,
             );
 
-            const index = stack.current.findIndex(o => o.id === overlay.id);
+            const index = stack.current.findIndex(o => o.id === id);
             if (index > -1) {
                 stack.current.splice(index, 1);
             }
diff --git a/packages/core/test/hooks/useOverlayStackTests.tsx b/packages/core/test/hooks/useOverlayStackTests.tsx
index bcca0ed57b..1361694343 100644
--- a/packages/core/test/hooks/useOverlayStackTests.tsx
+++ b/packages/core/test/hooks/useOverlayStackTests.tsx
@@ -21,7 +21,7 @@ import { useUID } from "react-uid";
 import { spy } from "sinon";
 
 import type { OverlayProps } from "../../src/components/overlay/overlayProps";
-import type { OverlayInstance } from "../../src/components/overlay2/overlay2";
+import type { OverlayInstance } from "../../src/components/overlay2/overlayInstance";
 import { OverlaysProvider } from "../../src/context";
 import { useOverlayStack, usePrevious } from "../../src/hooks";
 import { modifyGlobalStack } from "../../src/hooks/overlays/useLegacyOverlayStack";
@@ -43,20 +43,10 @@ const TestComponentWithoutProvider: React.FC<TestComponentProps> = ({
 }) => {
     const { openOverlay, getLastOpened, closeOverlay } = useOverlayStack();
 
-    const bringFocusInsideOverlay = React.useCallback(() => {
-        // unimplemented since it's not tested in this suite
-    }, []);
-
-    const handleDocumentFocus = React.useCallback((_e: FocusEvent) => {
-        // unimplemented since it's not tested in this suite
-    }, []);
-
     const id = useUID();
     const instance = React.useMemo<OverlayInstance>(
         () => ({
-            bringFocusInsideOverlay,
             containerElement,
-            handleDocumentFocus,
             id,
             props: {
                 autoFocus,
@@ -65,16 +55,7 @@ const TestComponentWithoutProvider: React.FC<TestComponentProps> = ({
                 usePortal,
             },
         }),
-        [
-            autoFocus,
-            bringFocusInsideOverlay,
-            containerElement,
-            enforceFocus,
-            handleDocumentFocus,
-            hasBackdrop,
-            id,
-            usePortal,
-        ],
+        [autoFocus, containerElement, enforceFocus, hasBackdrop, id, usePortal],
     );
 
     const prevIsOpen = usePrevious(isOpen) ?? false;
@@ -86,15 +67,15 @@ const TestComponentWithoutProvider: React.FC<TestComponentProps> = ({
 
         if (prevIsOpen && !isOpen) {
             // just closed
-            closeOverlay(instance);
+            closeOverlay(id);
         }
-    }, [isOpen, openOverlay, closeOverlay, prevIsOpen, instance]);
+    }, [isOpen, openOverlay, closeOverlay, prevIsOpen, instance, id]);
 
     // run once on unmount
     React.useEffect(() => {
         return () => {
             if (isOpen) {
-                closeOverlay(instance);
+                closeOverlay(id);
             }
         };
         // eslint-disable-next-line react-hooks/exhaustive-deps
diff --git a/packages/docs-app/src/examples/core-examples/drawerExample.tsx b/packages/docs-app/src/examples/core-examples/drawerExample.tsx
index bd6a7e49c1..54f4ca2e7e 100644
--- a/packages/docs-app/src/examples/core-examples/drawerExample.tsx
+++ b/packages/docs-app/src/examples/core-examples/drawerExample.tsx
@@ -79,6 +79,8 @@ export class DrawerExample extends React.PureComponent<ExampleProps<BlueprintExa
     private handleSizeChange = handleStringChange(size => this.setState({ size }));
 
     public render() {
+        const { size, ...drawerProps } = this.state;
+
         return (
             <Example options={this.renderOptions()} {...this.props}>
                 <Button onClick={this.handleOpen}>Show Drawer</Button>
@@ -87,7 +89,8 @@ export class DrawerExample extends React.PureComponent<ExampleProps<BlueprintExa
                     icon="info-sign"
                     onClose={this.handleClose}
                     title="Palantir Foundry"
-                    {...this.state}
+                    size={size === "default" ? undefined : size}
+                    {...drawerProps}
                 >
                     <div className={Classes.DRAWER_BODY}>
                         {/* HACKHACK: strange use of unrelated dialog class, should be refactored */}
@@ -183,7 +186,7 @@ export class DrawerExample extends React.PureComponent<ExampleProps<BlueprintExa
 }
 
 const SIZES: Array<string | OptionProps> = [
-    { label: "Default", value: undefined },
+    { label: "Default", value: "default" },
     { label: "Small", value: DrawerSize.SMALL },
     { label: "Standard", value: DrawerSize.STANDARD },
     { label: "Large", value: DrawerSize.LARGE },