From 4e80990827e790fe37e4dba7d7443a7574029bfd Mon Sep 17 00:00:00 2001
From: JJ Kasper <jj@jjsweb.site>
Date: Thu, 19 Dec 2024 21:41:47 -0600
Subject: [PATCH] Use provided waitUntil for pending revalidates (#74164)

Currently we are using `pendingWaitUntil` on `renderOpts` to handle
pending revalidations (fetch updates and revalidate tag calls). This
uses our old strategy of `waitUntil` inside of `sendResponse` which just
keeps the stream open until the promise resolves. This isn't ideal if we
have a proper waitUntil strategy provided so this updates to ensure we
use that instead if available.

Also adds some debug logs so we can track when this pending revalidates
promise is resolved.
---
 .../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 1a6510248aa14..6c2f0dd08b92f 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -1298,13 +1298,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) {
@@ -1455,13 +1465,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 564ffc68f4aec..bcd09130c97b5 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 23f1f79679453..3816331c89763 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 {