From af4484e3c7f30b87efd19615b01ef418cafa09dc Mon Sep 17 00:00:00 2001
From: Tianyu Yao <skyyao@fb.com>
Date: Tue, 23 May 2023 18:23:01 -0700
Subject: [PATCH] Always trigger componentWillUnmount in StrictMode

---
 .../src/ReactFiberClassComponent.js           |  37 +++---
 .../src/ReactFiberCommitWork.js               |  10 +-
 .../src/__tests__/StrictEffectsMode-test.js   | 116 ++++++++++++++++++
 3 files changed, 138 insertions(+), 25 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberClassComponent.js b/packages/react-reconciler/src/ReactFiberClassComponent.js
index b071808a9117d..c816fb0ba21ea 100644
--- a/packages/react-reconciler/src/ReactFiberClassComponent.js
+++ b/packages/react-reconciler/src/ReactFiberClassComponent.js
@@ -10,7 +10,6 @@
 import type {Fiber} from './ReactInternalTypes';
 import type {Lanes} from './ReactFiberLane';
 import type {UpdateQueue} from './ReactFiberClassUpdateQueue';
-import type {Flags} from './ReactFiberFlags';
 
 import {
   LayoutStatic,
@@ -897,11 +896,10 @@ function mountClassInstance(
   }
 
   if (typeof instance.componentDidMount === 'function') {
-    let fiberFlags: Flags = Update | LayoutStatic;
-    if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
-      fiberFlags |= MountLayoutDev;
-    }
-    workInProgress.flags |= fiberFlags;
+    workInProgress.flags |= Update | LayoutStatic;
+  }
+  if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
+    workInProgress.flags |= MountLayoutDev;
   }
 }
 
@@ -971,11 +969,10 @@ function resumeMountClassInstance(
     // If an update was already in progress, we should schedule an Update
     // effect even though we're bailing out, so that cWU/cDU are called.
     if (typeof instance.componentDidMount === 'function') {
-      let fiberFlags: Flags = Update | LayoutStatic;
-      if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
-        fiberFlags |= MountLayoutDev;
-      }
-      workInProgress.flags |= fiberFlags;
+      workInProgress.flags |= Update | LayoutStatic;
+    }
+    if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
+      workInProgress.flags |= MountLayoutDev;
     }
     return false;
   }
@@ -1018,21 +1015,19 @@ function resumeMountClassInstance(
       }
     }
     if (typeof instance.componentDidMount === 'function') {
-      let fiberFlags: Flags = Update | LayoutStatic;
-      if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
-        fiberFlags |= MountLayoutDev;
-      }
-      workInProgress.flags |= fiberFlags;
+      workInProgress.flags |= Update | LayoutStatic;
+    }
+    if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
+      workInProgress.flags |= MountLayoutDev;
     }
   } else {
     // If an update was already in progress, we should schedule an Update
     // effect even though we're bailing out, so that cWU/cDU are called.
     if (typeof instance.componentDidMount === 'function') {
-      let fiberFlags: Flags = Update | LayoutStatic;
-      if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
-        fiberFlags |= MountLayoutDev;
-      }
-      workInProgress.flags |= fiberFlags;
+      workInProgress.flags |= Update | LayoutStatic;
+    }
+    if (__DEV__ && (workInProgress.mode & StrictEffectsMode) !== NoMode) {
+      workInProgress.flags |= MountLayoutDev;
     }
 
     // If shouldComponentUpdate returned false, we should still update the
diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.js b/packages/react-reconciler/src/ReactFiberCommitWork.js
index 1960a4707e244..89c160df460a2 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.js
+++ b/packages/react-reconciler/src/ReactFiberCommitWork.js
@@ -4572,10 +4572,12 @@ function invokeLayoutEffectMountInDEV(fiber: Fiber): void {
       }
       case ClassComponent: {
         const instance = fiber.stateNode;
-        try {
-          instance.componentDidMount();
-        } catch (error) {
-          captureCommitPhaseError(fiber, fiber.return, error);
+        if (typeof instance.componentDidMount === 'function') {
+          try {
+            instance.componentDidMount();
+          } catch (error) {
+            captureCommitPhaseError(fiber, fiber.return, error);
+          }
         }
         break;
       }
diff --git a/packages/react-reconciler/src/__tests__/StrictEffectsMode-test.js b/packages/react-reconciler/src/__tests__/StrictEffectsMode-test.js
index d7e5e29b0e7d6..f7c30616c98c1 100644
--- a/packages/react-reconciler/src/__tests__/StrictEffectsMode-test.js
+++ b/packages/react-reconciler/src/__tests__/StrictEffectsMode-test.js
@@ -349,6 +349,47 @@ describe('StrictEffectsMode', () => {
     assertLog(['componentWillUnmount']);
   });
 
+  it('invokes componentWillUnmount for class components without componentDidMount', async () => {
+    class App extends React.PureComponent {
+      componentDidUpdate() {
+        Scheduler.log('componentDidUpdate');
+      }
+
+      componentWillUnmount() {
+        Scheduler.log('componentWillUnmount');
+      }
+
+      render() {
+        return this.props.text;
+      }
+    }
+
+    let renderer;
+    await act(() => {
+      renderer = ReactTestRenderer.create(<App text={'mount'} />, {
+        unstable_isConcurrent: true,
+      });
+    });
+
+    if (supportsDoubleInvokeEffects()) {
+      assertLog(['componentWillUnmount']);
+    } else {
+      assertLog([]);
+    }
+
+    await act(() => {
+      renderer.update(<App text={'update'} />);
+    });
+
+    assertLog(['componentDidUpdate']);
+
+    await act(() => {
+      renderer.unmount();
+    });
+
+    assertLog(['componentWillUnmount']);
+  });
+
   it('should not double invoke class lifecycles in legacy mode', async () => {
     class App extends React.PureComponent {
       componentDidMount() {
@@ -590,4 +631,79 @@ describe('StrictEffectsMode', () => {
       'useEffect unmount',
     ]);
   });
+
+  it('classes without componentDidMount and functions are double invoked together correctly', async () => {
+    class ClassChild extends React.PureComponent {
+      componentWillUnmount() {
+        Scheduler.log('componentWillUnmount');
+      }
+
+      render() {
+        return this.props.text;
+      }
+    }
+
+    function FunctionChild({text}) {
+      React.useEffect(() => {
+        Scheduler.log('useEffect mount');
+        return () => Scheduler.log('useEffect unmount');
+      });
+      React.useLayoutEffect(() => {
+        Scheduler.log('useLayoutEffect mount');
+        return () => Scheduler.log('useLayoutEffect unmount');
+      });
+      return text;
+    }
+
+    function App({text}) {
+      return (
+        <>
+          <ClassChild text={text} />
+          <FunctionChild text={text} />
+        </>
+      );
+    }
+
+    let renderer;
+    await act(() => {
+      renderer = ReactTestRenderer.create(<App text={'mount'} />, {
+        unstable_isConcurrent: true,
+      });
+    });
+
+    if (supportsDoubleInvokeEffects()) {
+      assertLog([
+        'useLayoutEffect mount',
+        'useEffect mount',
+        'componentWillUnmount',
+        'useLayoutEffect unmount',
+        'useEffect unmount',
+        'useLayoutEffect mount',
+        'useEffect mount',
+      ]);
+    } else {
+      assertLog(['useLayoutEffect mount', 'useEffect mount']);
+    }
+
+    await act(() => {
+      renderer.update(<App text={'mount'} />);
+    });
+
+    assertLog([
+      'useLayoutEffect unmount',
+      'useLayoutEffect mount',
+      'useEffect unmount',
+      'useEffect mount',
+    ]);
+
+    await act(() => {
+      renderer.unmount();
+    });
+
+    assertLog([
+      'componentWillUnmount',
+      'useLayoutEffect unmount',
+      'useEffect unmount',
+    ]);
+  });
 });