address review feedback, simplify

Author: phryneas

src/renderStream/__tests__/passiveEffects.test.tsx

@@ -73,7 +73,7 @@ describe('passive effects', () => {
     },
   )
 
-  test('takeRender waits for effects even when the commit is already recorded', async () => {
+  test('takeRender waits for effects even when onRender has already run', async () => {
     let commitRecorded: () => void = () => {}
     const stream = createRenderStream<{value: number}>({
       onRender: () => commitRecorded(),
@@ -97,7 +97,6 @@ describe('passive effects', () => {
     })
     update!(1)
     await committed
-    expect(stream.renders).toHaveLength(2)
     expect((await stream.takeRender()).snapshot).toEqual({value: 1})
     expect(effectValue).toBe(1)
   })
@@ -222,6 +221,7 @@ describe('passive effects', () => {
     for (const value of [0, 1, 2]) {
       // eslint-disable-next-line no-await-in-loop
       const render = await stream.takeRender()
+      expect(render.count).toBe(value + 1)
       expect(render.snapshot).toEqual({value})
       expect(render.withinDOM().getByTestId('value').textContent).toBe(
         String(value),

src/renderStream/__tests__/skippedRenders.test.tsx

@@ -0,0 +1,68 @@
+import {expect, test} from '@jest/globals'
+import {
+  createRenderStream,
+  useTrackRenders,
+  WaitForRenderTimeoutError,
+} from '@testing-library/react-render-stream'
+import * as React from 'react'
+
+test.each(['render', 'rerender'] as const)(
+  '%s completes a skipped commit without notifying snapshot consumers',
+  async method => {
+    const stream = createRenderStream({skipNonTrackingRenders: true})
+    let subscribed = false
+    function Tracked() {
+      useTrackRenders()
+      React.useEffect(() => {
+        subscribed = true
+        return () => {
+          subscribed = false
+        }
+      }, [])
+      return <span>tracked</span>
+    }
+    const container = document.createElement('div')
+    const {rerender} = await stream.render(<Tracked />, {container})
+    await stream.takeRender()
+    const render = (ui: React.ReactNode) =>
+      method === 'render' ? stream.render(ui, {container}) : rerender(ui)
+    let notified = false
+    const next = stream.takeRender().then(result => {
+      notified = true
+      return result
+    })
+
+    await render(null)
+    expect(container.innerHTML).toBe('')
+    expect(subscribed).toBe(false)
+    expect(stream.totalRenderCount()).toBe(1)
+    expect(notified).toBe(false)
+
+    await render(<Tracked />)
+    expect(subscribed).toBe(true)
+    expect((await next).count).toBe(2)
+  },
+)
+
+test('initial and overlapping skipped renders finish while snapshot waits still time out', async () => {
+  const stream = createRenderStream({skipNonTrackingRenders: true})
+  let effectValue = -1
+  function Untracked({value}: {value: number}) {
+    React.useEffect(() => {
+      effectValue = value
+    }, [value])
+    return null
+  }
+  const {rerender} = await stream.render(<Untracked value={0} />)
+  expect(effectValue).toBe(0)
+  const timeout = expect(
+    stream.takeRender({timeout: 100}),
+  ).rejects.toBeInstanceOf(WaitForRenderTimeoutError)
+  await Promise.all([
+    rerender(<Untracked value={1} />),
+    rerender(<Untracked value={2} />),
+  ])
+  expect(effectValue).toBe(2)
+  expect(stream.totalRenderCount()).toBe(0)
+  await timeout
+})

src/renderStream/createRenderStream.tsx

@@ -11,8 +11,6 @@ import {type RenderStreamContextValue} from './context.js'
 import {RenderStreamContextProvider} from './context.js'
 import {syncQueries, type Queries, type SyncQueries} from './syncQueries.js'
 
-let nextProfilerId = 0
-
 export type ValidSnapshot =
   // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
   void | (object & {/* not a function */ call?: never})
@@ -140,19 +138,20 @@ export function createRenderStream<
   // created, which is important for the `markAssertable` function
   const stream = {} as any as RenderStreamWithRenderFn<Snapshot, Q>
 
+  type Commit = PromiseWithResolvers<
+    (typeof stream.renders)[number] | undefined
+  > & {
+    finalize?: () => (typeof stream.renders)[number]
+  }
   let nextRender: Promise<Render<Snapshot, Q>> | undefined,
-    resolveNextRender: ((render: Render<Snapshot, Q>) => void) | undefined,
-    rejectNextRender: ((error: unknown) => void) | undefined,
-    nextRenderTimeout: ReturnType<typeof setTimeout> | undefined
-  function resetNextRender() {
-    clearTimeout(nextRenderTimeout)
-    nextRenderTimeout = undefined
-    nextRender = undefined
-    resolveNextRender = undefined
-    rejectNextRender = undefined
+    upcomingCommit: Commit | undefined
+  const pendingCommits: Array<{commitTime: number; commit: Commit}> = []
+
+  function nextCommit(): Commit {
+    return (upcomingCommit ??=
+      Promise.withResolvers<Awaited<Commit['promise']>>())
   }
-  let effectsFlushedCount = 0
-  const pendingCommits = new Map<string, number>()
+
   const snapshotRef = {current: initialSnapshot}
   const replaceSnapshot: ReplaceSnapshot<Snapshot> = snap => {
     if (typeof snap === 'function') {
@@ -193,6 +192,16 @@ export function createRenderStream<
     startTime,
     commitTime,
   ) => {
+    const commit = nextCommit()
+    upcomingCommit = undefined
+    pendingCommits.push({commitTime, commit})
+    if (/^(16|17)\./.test(React.version)) {
+      // React 16/17 have no onPostCommit hook. Let the commit finish scheduling
+      // passive effects before yielding a task for them to run.
+      void Promise.resolve().then(() => {
+        setTimeout(() => profilerOnPostCommit(id, phase, 0, commitTime), 0)
+      })
+    }
     if (
       skipNonTrackingRenders &&
       renderStreamContext.renderedComponents.length === 0
@@ -207,7 +216,10 @@ export function createRenderStream<
       baseDuration,
       startTime,
       commitTime,
-      count: stream.renders.length + 1,
+      count:
+        stream.renders.length +
+        pendingCommits.filter(pending => pending.commit.finalize).length +
+        1,
     }
     try {
       /*
@@ -236,51 +248,43 @@ export function createRenderStream<
         queries,
       )
       renderStreamContext.renderedComponents = []
-      stream.renders.push(render)
+      commit.finalize = () => {
+        stream.renders.push(render)
+        return render
+      }
     } catch (error) {
-      stream.renders.push({
-        phase: 'snapshotError',
-        count: stream.renders.length,
-        error,
-      })
-    } finally {
-      pendingCommits.set(`${id}:${commitTime}`, stream.renders.length)
-      if (/^(16|17)\./.test(React.version)) {
-        // React 16/17 have no onPostCommit hook. Let the commit finish scheduling
-        // passive effects before yielding a task for them to run.
-        void Promise.resolve().then(() => {
-          setTimeout(() => profilerOnPostCommit(id, phase, 0, commitTime), 0)
-        })
+      commit.finalize = () => {
+        const err = {
+          phase: 'snapshotError' as const,
+          count: stream.renders.length,
+          error,
+        }
+        stream.renders.push(err)
+        return err
       }
     }
   }
 
   function profilerOnPostCommit(
-    id: string,
+    _id: string,
     _phase: string,
     _passiveEffectDuration: number,
     commitTime: number,
   ) {
-    const key = `${id}:${commitTime}`
-    const renderCount = pendingCommits.get(key)
-    if (renderCount === undefined) return
-    pendingCommits.delete(key)
-    const render = stream.renders[effectsFlushedCount]
-    effectsFlushedCount = renderCount
-    if (render.phase === 'snapshotError') {
-      rejectNextRender?.(render.error)
-    } else {
-      resolveNextRender?.(render)
-    }
-    resetNextRender()
+    const index = pendingCommits.findIndex(
+      pending => pending.commitTime === commitTime,
+    )
+    if (index === -1) return
+    const {commit} = pendingCommits.splice(index, 1)[0]
+    commit.resolve(commit.finalize?.())
   }
 
   let iteratorPosition = 0
-  function Wrapper({children, id}: {children: React.ReactNode; id: string}) {
+  function Wrapper({children}: {children: React.ReactNode}) {
     return (
       <RenderStreamContextProvider value={renderStreamContext}>
         <React.Profiler
-          id={id}
+          id="test"
           onRender={profilerOnRender}
           // Supported in React 18+ development/profiling builds, but untyped.
           {...{onPostCommit: profilerOnPostCommit}}
@@ -295,33 +299,24 @@ export function createRenderStream<
     ui: React.ReactNode,
     options?: RenderOptions<any, any, any>,
   ) => {
-    const id = `render-stream-${nextProfilerId++}`
-    const previousRenderCount = stream.renders.length
+    const commit = nextCommit()
     const ret = await renderWithoutAct(ui, {
       ...options,
       wrapper: props => {
         const ParentWrapper = options?.wrapper ?? React.Fragment
         return (
           <ParentWrapper>
-            <Wrapper id={id}>{props.children}</Wrapper>
+            <Wrapper>{props.children}</Wrapper>
           </ParentWrapper>
         )
       },
     })
-    // eslint-disable-next-line no-unmodified-loop-condition -- updated by onPostCommit
-    while (effectsFlushedCount <= previousRenderCount) {
-      // eslint-disable-next-line no-await-in-loop
-      await stream.waitForNextRender()
-    }
+    await waitForCommit(commit)
     const origRerender = ret.rerender
     ret.rerender = async function rerender(rerenderUi: React.ReactNode) {
-      const rerenderPosition = stream.renders.length
+      const rerenderCommit = nextCommit()
       await origRerender(rerenderUi)
-      // eslint-disable-next-line no-unmodified-loop-condition -- updated by onPostCommit
-      while (effectsFlushedCount <= rerenderPosition) {
-        // eslint-disable-next-line no-await-in-loop
-        await stream.waitForNextRender()
-      }
+      await waitForCommit(rerenderCommit)
     }
     return ret
   }) as unknown as RenderWithoutActAsync // TODO
@@ -337,7 +332,7 @@ export function createRenderStream<
       return stream.renders.length
     },
     async peekRender(options: NextRenderOptions = {}) {
-      if (iteratorPosition < effectsFlushedCount) {
+      if (iteratorPosition < stream.renders.length) {
         const peekedRender = stream.renders[iteratorPosition]
 
         if (peekedRender.phase === 'snapshotError') {
@@ -394,26 +389,59 @@ export function createRenderStream<
     },
     waitForNextRender({timeout = 1000}: NextRenderOptions = {}) {
       if (!nextRender) {
+        const position = stream.renders.length
+        const wait = {active: true}
+        let timer: ReturnType<typeof setTimeout> | undefined
         nextRender = Promise.race<Render<Snapshot, Q>>([
-          new Promise<Render<Snapshot, Q>>((resolve, reject) => {
-            resolveNextRender = resolve
-            rejectNextRender = reject
-          }),
+          (async () => {
+            while (wait.active && stream.renders.length === position) {
+              // A skipped commit completes without publishing a snapshot.
+              // eslint-disable-next-line no-await-in-loop
+              await (pendingCommits[0]?.commit ?? nextCommit()).promise
+            }
+            if (!wait.active) throw new WaitForRenderTimeoutError()
+            const next = stream.renders[position]
+            if (next.phase === 'snapshotError') throw next.error
+            return next
+          })(),
           new Promise<Render<Snapshot, Q>>((_, reject) => {
-            nextRenderTimeout = setTimeout(() => {
+            timer = setTimeout(() => {
               const error = new WaitForRenderTimeoutError()
               Error.captureStackTrace(error, stream.waitForNextRender)
               reject(error)
-              resetNextRender()
             }, timeout)
           }),
-        ])
+        ]).finally(() => {
+          wait.active = false
+          clearTimeout(timer)
+          nextRender = undefined
+        })
       }
       return nextRender
     },
     render,
   })
   return stream
+
+  async function waitForCommit(commit: Commit) {
+    let timer: ReturnType<typeof setTimeout> | undefined
+    try {
+      return await Promise.race([
+        commit.promise.then(result => {
+          if (result?.phase === 'snapshotError') throw result.error
+          return result
+        }),
+        new Promise<never>((_, reject) => {
+          timer = setTimeout(
+            () => reject(new WaitForRenderTimeoutError()),
+            1000,
+          )
+        }),
+      ])
+    } finally {
+      clearTimeout(timer)
+    }
+  }
 }
 
 function rethrowWithCapturedStackTrace(constructorOpt: Function | undefined) {