Use provided waitUntil for pending revalidates (#74164) (#74604)

This backports #74164 to leverage built-in waitUntil if available instead of using the approach that keeps the stream open until the waitUntil promise resolves. x-ref: [slack thread](https://vercel.slack.com/archives/C02K2HCH5V4/p1736211642221149?thread_ts=1734707275.666089&cid=C02K2HCH5V4)
vercel · Jan 7, 2025 · a46f27f · a46f27f
1 parent bc6acf6
commit a46f27f
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 11 deletions.
diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
@@ -1276,13 +1276,23 @@ async function renderToHTMLOrFlightImpl(
       workStore.pendingRevalidateWrites ||
       workStore.revalidatedTags
     ) {
-      options.waitUntil = Promise.all([
+      const pendingPromise = Promise.all([
         workStore.incrementalCache?.revalidateTag(
           workStore.revalidatedTags || []
         ),
         ...Object.values(workStore.pendingRevalidates || {}),
         ...(workStore.pendingRevalidateWrites || []),
-      ])
+      ]).finally(() => {
+        if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+          console.log('pending revalidates promise finished for:', url)
+        }
+      })
+
+      if (renderOpts.waitUntil) {
+        renderOpts.waitUntil(pendingPromise)
+      } else {
+        options.waitUntil = pendingPromise
+      }
     }
 
     if (response.collectedTags) {
@@ -1433,13 +1443,23 @@ async function renderToHTMLOrFlightImpl(
       workStore.pendingRevalidateWrites ||
       workStore.revalidatedTags
     ) {
-      options.waitUntil = Promise.all([
+      const pendingPromise = Promise.all([
         workStore.incrementalCache?.revalidateTag(
           workStore.revalidatedTags || []
         ),
         ...Object.values(workStore.pendingRevalidates || {}),
         ...(workStore.pendingRevalidateWrites || []),
-      ])
+      ]).finally(() => {
+        if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+          console.log('pending revalidates promise finished for:', url)
+        }
+      })
+
+      if (renderOpts.waitUntil) {
+        renderOpts.waitUntil(pendingPromise)
+      } else {
+        options.waitUntil = pendingPromise
+      }
     }
 
     // Create the new render result for the response.

diff --git a/packages/next/src/server/base-server.ts b/packages/next/src/server/base-server.ts
@@ -2541,6 +2541,16 @@ export default abstract class Server<
 
               return cacheEntry
             }
+            let pendingWaitUntil = context.renderOpts.pendingWaitUntil
+
+            // Attempt using provided waitUntil if available
+            // if it's not we fallback to sendResponse's handling
+            if (pendingWaitUntil) {
+              if (context.renderOpts.waitUntil) {
+                context.renderOpts.waitUntil(pendingWaitUntil)
+                pendingWaitUntil = undefined
+              }
+            }
 
             // Send the response now that we have copied it into the cache.
             await sendResponse(

diff --git a/packages/next/src/server/route-modules/app-route/module.ts b/packages/next/src/server/route-modules/app-route/module.ts
@@ -597,7 +597,14 @@ export class AppRouteRouteModule extends RouteModule<
         workStore.revalidatedTags || []
       ),
       ...Object.values(workStore.pendingRevalidates || {}),
-    ])
+    ]).finally(() => {
+      if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+        console.log(
+          'pending revalidates promise finished for:',
+          requestStore.url
+        )
+      }
+    })
 
     if (prerenderStore) {
       context.renderOpts.collectedTags = prerenderStore.tags?.join(',')

diff --git a/test/e2e/app-dir/app-fetch-deduping/app-fetch-deduping.test.ts b/test/e2e/app-dir/app-fetch-deduping/app-fetch-deduping.test.ts
@@ -1,4 +1,4 @@
-import { findPort, waitFor } from 'next-test-utils'
+import { findPort, retry } from 'next-test-utils'
 import http from 'http'
 import url from 'url'
 import { outdent } from 'outdent'
@@ -112,11 +112,10 @@ describe('app-fetch-deduping', () => {
         expect(invocation(next.cliOutput)).toBe(1)
 
         // wait for the revalidation to finish
-        await waitFor(revalidate * 1000 + 1000)
-
-        await next.render('/test')
-
-        expect(invocation(next.cliOutput)).toBe(2)
+        await retry(async () => {
+          await next.render('/test')
+          expect(invocation(next.cliOutput)).toBe(2)
+        }, 10_000)
       })
     })
   } else {