Update useOnMountEffect to pass isMountedRef to callback (#2413)

## Summary:
Without this, consumers of useOnMountEffect have to keep track of whether the hosting component is mounted or not when running async code within the effect.

Issue: None

## Test plan:
- let CI run automated checks

Author: kevinb-khan

Reviewers: marcysutton

Required Reviewers:

Approved By: marcysutton

Checks: ✅ Chromatic - Get results on regular PRs (ubuntu-latest, 20.x), ✅ Lint / Lint (ubuntu-latest, 20.x), ✅ Test / Test (ubuntu-latest, 20.x, 2/2), ✅ Test / Test (ubuntu-latest, 20.x, 1/2), ✅ Check build sizes (ubuntu-latest, 20.x), ✅ Chromatic - Build on regular PRs / chromatic (ubuntu-latest, 20.x), ⏭️  Chromatic - Skip on Release PR (changesets), ✅ Publish npm snapshot (ubuntu-latest, 20.x), ✅ Prime node_modules cache for primary configuration (ubuntu-latest, 20.x), ✅ Check for .changeset entries for all changed files (ubuntu-latest, 20.x), ✅ gerald, ⏭️  dependabot

Pull Request URL: #2413

Author: kevinb-khan

.changeset/tall-shoes-enjoy.md

@@ -0,0 +1,5 @@
+---
+"@khanacademy/wonder-blocks-core": minor
+---
+
+Update useOnMountEffect to pass isMountedRef to callback

packages/wonder-blocks-core/src/hooks/__tests__/use-on-mount-effect.test.ts

@@ -1,4 +1,5 @@
-import {renderHook} from "@testing-library/react";
+import * as React from "react";
+import {renderHook, waitFor} from "@testing-library/react";
 
 import {useOnMountEffect} from "../use-on-mount-effect";
 
@@ -27,4 +28,90 @@ describe("#useOnMountEffect", () => {
         // Assert
         expect(cleanup).toHaveBeenCalled();
     });
+
+    it("should pass {current: true} as isMountedRef on initial render", () => {
+        // Arrange
+        const callback = jest.fn();
+
+        // Act
+        renderHook(() => useOnMountEffect(callback));
+
+        // Assert
+        expect(callback).toHaveBeenCalledWith({current: true});
+    });
+
+    it("should pass {current: false} as isMountedRef after being unmounted", () => {
+        // Arrange
+        const callback = jest.fn();
+
+        // Act
+        const {unmount} = renderHook(() => useOnMountEffect(callback));
+        unmount();
+
+        // Assert
+        expect(callback).toHaveBeenCalledWith({current: false});
+    });
+
+    describe("async", () => {
+        it("should pass {current: true} while mounted", async () => {
+            // Arrange
+            const wait = async (duration: number) => {
+                return new Promise((resolve) => {
+                    setTimeout(resolve, duration);
+                });
+            };
+            let foo = false;
+            const callback = (
+                isMountedRef: React.MutableRefObject<boolean>,
+            ) => {
+                const action = async () => {
+                    await wait(100);
+                    if (isMountedRef.current) {
+                        foo = true;
+                    }
+                };
+
+                action();
+            };
+
+            // Act
+            renderHook(() => useOnMountEffect(callback));
+
+            // Assert
+            await waitFor(() => {
+                expect(foo).toEqual(true);
+            });
+        });
+
+        it("should pass {current: false} after being unmounted", async () => {
+            // Arrange
+            const wait = async (duration: number) => {
+                return new Promise((resolve) => {
+                    setTimeout(resolve, duration);
+                });
+            };
+            let foo = false;
+            const callback = (
+                isMountedRef: React.MutableRefObject<boolean>,
+            ) => {
+                const action = async () => {
+                    await wait(100);
+                    if (isMountedRef.current) {
+                        foo = true;
+                    }
+                };
+
+                action();
+            };
+
+            // Act
+            const {unmount} = renderHook(() => useOnMountEffect(callback));
+            unmount();
+
+            // Assert
+            await waitFor(() => {
+                expect(foo).toEqual(false);
+            });
+        });
+    });
 });

packages/wonder-blocks-core/src/hooks/use-on-mount-effect.ts

@@ -5,7 +5,7 @@ import * as React from "react";
  *
  * If the callback returns a cleanup function, it will be called when the component is unmounted.
  *
- * @param {() => (void | (() => void))} callback function that forces the component to update.
+ * @param {(isMountedRef: React.MutableRefObject<boolean>) => void | (() => void)} callback function that forces the component to update.
  *
  * The following code snippets are equivalent
  * ```
@@ -27,8 +27,42 @@ import * as React from "react";
  * }, []);
  *
  * If you only need to do something on mount, don't return a cleanup function from `callback`.
+ *
+ * If your callback is async, use the `isMountedRef` ref that's passed to the callback to ensure
+ * that the component using `useOnMountEffect` hasn't been unmounted, e.g.
+ *
+ * ```
+ * const MyComponent = () => {
+ *    const [foo, setFoo] = React.useState("");
+ *    useOnMountEffect((isMountedRef) => {
+ *        const action = async () => {
+ *            const res = await fetch("/foo");
+ *            const text = res.text();
+ *            if (isMountedRef.current) {
+ *                setFoo(text);
+ *            }
+ *        }
+ *
+ *        action();
+ *    });
+ *
+ *    return foo !== "" ? <h1>Loading...</h1> : <h1>{foo}</h1>;
+ * }
  */
-export const useOnMountEffect = (callback: () => void | (() => void)): void => {
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    React.useEffect(callback, []);
+export const useOnMountEffect = (
+    callback: (
+        isMountedRef: React.MutableRefObject<boolean>,
+    ) => void | (() => void),
+): void => {
+    const isMountedRef = React.useRef(true);
+
+    React.useEffect(() => {
+        const cleanup = callback(isMountedRef);
+
+        return () => {
+            cleanup?.();
+            isMountedRef.current = false;
+        };
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
 };