From 5f96520f4b0cb548ea816181dbda02ad99b3f73f Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Mon, 29 Apr 2019 13:51:49 -0700
Subject: [PATCH] Fixed potential interaction tracing leak in Suspense
 thennable memoization

Audited the other places we call unstable_wrap() in React DOM and verified that they didn't have this similar problem.
---
 .../src/ReactFiberCommitWork.js               |  6 +-
 .../__tests__/ReactSuspense-test.internal.js  | 75 +++++++++++++++++++
 .../__tests__/ReactProfiler-test.internal.js  |  6 +-
 3 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberCommitWork.js b/packages/react-reconciler/src/ReactFiberCommitWork.js
index 5aaf20195b191..b8043c2b8297b 100644
--- a/packages/react-reconciler/src/ReactFiberCommitWork.js
+++ b/packages/react-reconciler/src/ReactFiberCommitWork.js
@@ -1317,10 +1317,10 @@ function commitSuspenseComponent(finishedWork: Fiber) {
     thenables.forEach(thenable => {
       // Memoize using the boundary fiber to prevent redundant listeners.
       let retry = resolveRetryThenable.bind(null, finishedWork, thenable);
-      if (enableSchedulerTracing) {
-        retry = Schedule_tracing_wrap(retry);
-      }
       if (!retryCache.has(thenable)) {
+        if (enableSchedulerTracing) {
+          retry = Schedule_tracing_wrap(retry);
+        }
         retryCache.add(thenable);
         thenable.then(retry, retry);
       }
diff --git a/packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js b/packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js
index 18d477d839860..d7a3b894290dc 100644
--- a/packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js
@@ -2,6 +2,7 @@ let React;
 let ReactTestRenderer;
 let ReactFeatureFlags;
 let Scheduler;
+let SchedulerTracing;
 let ReactCache;
 let Suspense;
 let act;
@@ -17,10 +18,12 @@ describe('ReactSuspense', () => {
     ReactFeatureFlags = require('shared/ReactFeatureFlags');
     ReactFeatureFlags.debugRenderPhaseSideEffectsForStrictMode = false;
     ReactFeatureFlags.replayFailedUnitOfWorkWithInvokeGuardedCallback = false;
+    ReactFeatureFlags.enableSchedulerTracing = true;
     React = require('react');
     ReactTestRenderer = require('react-test-renderer');
     act = ReactTestRenderer.act;
     Scheduler = require('scheduler');
+    SchedulerTracing = require('scheduler/tracing');
     ReactCache = require('react-cache');
 
     Suspense = React.Suspense;
@@ -914,6 +917,78 @@ describe('ReactSuspense', () => {
       ]);
     });
 
+    it('should call onInteractionScheduledWorkCompleted after suspending', done => {
+      const subscriber = {
+        onInteractionScheduledWorkCompleted: jest.fn(),
+        onInteractionTraced: jest.fn(),
+        onWorkCanceled: jest.fn(),
+        onWorkScheduled: jest.fn(),
+        onWorkStarted: jest.fn(),
+        onWorkStopped: jest.fn(),
+      };
+      SchedulerTracing.unstable_subscribe(subscriber);
+      SchedulerTracing.unstable_trace('test', performance.now(), () => {
+        function App() {
+          return (
+            <React.Suspense fallback={<Text text="Loading..." />}>
+              <AsyncText text="A" ms={1000} />
+              <AsyncText text="B" ms={2000} />
+              <AsyncText text="C" ms={3000} />
+            </React.Suspense>
+          );
+        }
+
+        const root = ReactTestRenderer.create(null);
+        root.update(<App />);
+
+        expect(Scheduler).toHaveYielded([
+          'Suspend! [A]',
+          'Suspend! [B]',
+          'Suspend! [C]',
+          'Loading...',
+        ]);
+
+        // Resolve A
+        jest.advanceTimersByTime(1000);
+
+        expect(Scheduler).toHaveYielded(['Promise resolved [A]']);
+        expect(Scheduler).toFlushExpired([
+          'A',
+          // The promises for B and C have now been thrown twice
+          'Suspend! [B]',
+          'Suspend! [C]',
+        ]);
+
+        // Resolve B
+        jest.advanceTimersByTime(1000);
+
+        expect(Scheduler).toHaveYielded(['Promise resolved [B]']);
+        expect(Scheduler).toFlushExpired([
+          // Even though the promise for B was thrown twice, we should only
+          // re-render once.
+          'B',
+          // The promise for C has now been thrown three times
+          'Suspend! [C]',
+        ]);
+
+        // Resolve C
+        jest.advanceTimersByTime(1000);
+
+        expect(Scheduler).toHaveYielded(['Promise resolved [C]']);
+        expect(Scheduler).toFlushExpired([
+          // Even though the promise for C was thrown three times, we should only
+          // re-render once.
+          'C',
+        ]);
+
+        done();
+      });
+
+      expect(
+        subscriber.onInteractionScheduledWorkCompleted,
+      ).toHaveBeenCalledTimes(1);
+    });
+
     it('#14162', () => {
       const {lazy} = React;
 
diff --git a/packages/react/src/__tests__/ReactProfiler-test.internal.js b/packages/react/src/__tests__/ReactProfiler-test.internal.js
index e894f2c3173e4..508d7814c821a 100644
--- a/packages/react/src/__tests__/ReactProfiler-test.internal.js
+++ b/packages/react/src/__tests__/ReactProfiler-test.internal.js
@@ -2607,7 +2607,11 @@ describe('Profiler', () => {
         ]);
         onRender.mockClear();
 
-        expect(onInteractionScheduledWorkCompleted).not.toHaveBeenCalled();
+        expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(1);
+        expect(
+          onInteractionScheduledWorkCompleted.mock.calls[0][0],
+        ).toMatchInteraction(highPriUpdateInteraction);
+        onInteractionScheduledWorkCompleted.mockClear();
 
         Scheduler.advanceTime(100);
         jest.advanceTimersByTime(100);