Fixed remaining tearing cases

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberHooks.js

@@ -76,7 +76,6 @@ import {
 } from './SchedulerWithReactIntegration';
 import {
   getFirstPendingExpirationTime,
-  getLastPendingExpirationTime,
   getWorkInProgressVersion,
   markSourceAsDirty,
   setPendingExpirationTime,
@@ -1026,13 +1025,22 @@ function useMutableSource<Source, Snapshot>(
         if (!is(snapshot, maybeNewSnapshot)) {
           setSnapshot(maybeNewSnapshot);
 
+          const currentTime = requestCurrentTimeForUpdate();
+          const suspenseConfig = requestCurrentSuspenseConfig();
+          const expirationTime = computeExpirationForFiber(
+            currentTime,
+            fiber,
+            suspenseConfig,
+          );
+          setPendingExpirationTime(root, expirationTime);
+
           // If the source mutated between render and now,
           // there may be state updates already scheduled from the old getSnapshot.
           // Those updates should not commit without this value.
           // There is no mechanism currently to associate these updates though,
           // so for now we fall back to synchronously flushing all pending updates.
           // TODO: Improve this later.
-          markRootExpiredAtTime(root, getLastPendingExpirationTime(root));
+          markRootExpiredAtTime(root, getFirstPendingExpirationTime(root));
         }
       }
     }
@@ -1088,10 +1096,19 @@ function useMutableSource<Source, Snapshot>(
       if (!is(snapshot, maybeNewSnapshot)) {
         setSnapshot(maybeNewSnapshot);
 
+        const currentTime = requestCurrentTimeForUpdate();
+        const suspenseConfig = requestCurrentSuspenseConfig();
+        const expirationTime = computeExpirationForFiber(
+          currentTime,
+          fiber,
+          suspenseConfig,
+        );
+        setPendingExpirationTime(root, expirationTime);
+
         // We missed a mutation before committing.
         // It's possible that other components using this source also have pending updates scheduled.
         // In that case, we should ensure they all commit together.
-        markRootExpiredAtTime(root, getLastPendingExpirationTime(root));
+        markRootExpiredAtTime(root, getFirstPendingExpirationTime(root));
       }
     }
 

packages/react-reconciler/src/ReactMutableSource.js

@@ -30,11 +30,14 @@ export function clearPendingUpdates(
   root: FiberRoot,
   expirationTime: ExpirationTime,
 ): void {
-  if (root.mutableSourceFirstPendingUpdateTime <= expirationTime) {
+  if (expirationTime <= root.mutableSourceLastPendingUpdateTime) {
+    // All updates for this source have been processed.
     root.mutableSourceFirstPendingUpdateTime = NoWork;
-  }
-  if (root.mutableSourceLastPendingUpdateTime <= expirationTime) {
     root.mutableSourceLastPendingUpdateTime = NoWork;
+  } else if (expirationTime <= root.mutableSourceFirstPendingUpdateTime) {
+    // The highest priority pending updates have been processed,
+    // but we don't how many updates remain between the current and lowest priority.
+    root.mutableSourceFirstPendingUpdateTime = expirationTime - 1;
   }
 }
 
@@ -53,18 +56,14 @@ export function setPendingExpirationTime(
   // Because we only track one pending update time per root,
   // track the lowest priority update.
   // It's inclusive of all other pending updates.
-  root.mutableSourceLastPendingUpdateTime = Math.max(
-    root.mutableSourceLastPendingUpdateTime,
-    expirationTime,
-  );
+  if (expirationTime > root.mutableSourceLastPendingUpdateTime) {
+    root.mutableSourceLastPendingUpdateTime = expirationTime;
+  }
 
   if (root.mutableSourceFirstPendingUpdateTime === NoWork) {
     root.mutableSourceFirstPendingUpdateTime = expirationTime;
-  } else {
-    root.mutableSourceFirstPendingUpdateTime = Math.min(
-      root.mutableSourceFirstPendingUpdateTime,
-      expirationTime,
-    );
+  } else if (expirationTime < root.mutableSourceFirstPendingUpdateTime) {
+    root.mutableSourceFirstPendingUpdateTime = expirationTime;
   }
 }
 

packages/react-reconciler/src/__tests__/useMutableSource-test.internal.js

@@ -1385,9 +1385,7 @@ describe('useMutableSource', () => {
         });
         expect(Scheduler).toFlushUntilNextPaint([]);
 
-        // TODO: This test currently tears.
-        // Fill in with correct values once the entanglement issue has been fixed.
-        expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a0');
+        expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
       });
 
       expect(root.getChildrenAsJSX()).toEqual('first: a1, second: a1');
@@ -1481,20 +1479,13 @@ describe('useMutableSource', () => {
         );
       });
 
-      // TODO: This test currently tears.
-      // Fill in with correct values once the entanglement issue has been fixed.
-      // Here's the current behavior:
       expect(Scheduler).toHaveYielded([
         // The partial render completes
         'Child: 2',
         'Commit: 2, 2',
 
         // Then we start rendering the low priority mutation
         'Parent: 3',
-        // But the child never received a mutation event, because it hadn't
-        // mounted yet. So the render tears.
-        'Child: 2',
-        'Commit: Oops, tearing!',
 
         // Eventually the child corrects itself, because of the check that
         // occurs when re-subscribing.