From a46f27f59a4ce8b1b61cd37eaacdb0b1937e5d74 Mon Sep 17 00:00:00 2001
From: JJ Kasper <jj@jjsweb.site>
Date: Tue, 7 Jan 2025 10:09:02 -0800
Subject: [PATCH] Use provided waitUntil for pending revalidates (#74164)
 (#74604)

This backports https://github.com/vercel/next.js/pull/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)
---
 .../next/src/server/app-render/app-render.tsx | 28 ++++++++++++++++---
 packages/next/src/server/base-server.ts       | 10 +++++++
 .../server/route-modules/app-route/module.ts  |  9 +++++-
 .../app-fetch-deduping.test.ts                | 11 ++++----
 4 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index 5cb7087f768e9..d2fa77e404844 100644
--- 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
index 1a0ab96b7d33e..367116abf0d87 100644
--- 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
index 0424c1f5c4f1b..0b8b4b4df1162 100644
--- 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
index d08294db501f3..3ea90dc18cdcd 100644
--- 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 {