From c6a86997660be1b672982054013ad91449ad582c Mon Sep 17 00:00:00 2001
From: Michal Piechowiak <misiek.piechowiak@gmail.com>
Date: Fri, 14 Jun 2024 16:52:04 +0200
Subject: [PATCH] fix: track revalidate / cdn purge to ensure it finishes
 execution and is not suspended mid-execution

---
 src/run/handlers/cache.cts                    | 28 ++++++++++++++-----
 src/run/revalidate.ts                         |  9 +++++-
 .../page-router/pages/api/revalidate.js       |  5 +++-
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/run/handlers/cache.cts b/src/run/handlers/cache.cts
index ff3bd60db5..1ddbde5df0 100644
--- a/src/run/handlers/cache.cts
+++ b/src/run/handlers/cache.cts
@@ -244,12 +244,14 @@ export class NetlifyCacheHandler implements CacheHandler {
         if (requestContext?.didPagesRouterOnDemandRevalidate) {
           const tag = `_N_T_${key === '/index' ? '/' : key}`
           getLogger().debug(`Purging CDN cache for: [${tag}]`)
-          purgeCache({ tags: [tag] }).catch((error) => {
-            // TODO: add reporting here
-            getLogger()
-              .withError(error)
-              .error(`[NetlifyCacheHandler]: Purging the cache for tag ${tag} failed`)
-          })
+          requestContext.trackBackgroundWork(
+            purgeCache({ tags: [tag] }).catch((error) => {
+              // TODO: add reporting here
+              getLogger()
+                .withError(error)
+                .error(`[NetlifyCacheHandler]: Purging the cache for tag ${tag} failed`)
+            }),
+          )
         }
       }
     })
@@ -257,6 +259,18 @@ export class NetlifyCacheHandler implements CacheHandler {
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   async revalidateTag(tagOrTags: string | string[], ...args: any) {
+    const revalidateTagPromise = this.doRevalidateTag(tagOrTags, ...args)
+
+    const requestContext = getRequestContext()
+    if (requestContext) {
+      requestContext.trackBackgroundWork(revalidateTagPromise)
+    }
+
+    return revalidateTagPromise
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private async doRevalidateTag(tagOrTags: string | string[], ...args: any) {
     getLogger().withFields({ tagOrTags, args }).debug('NetlifyCacheHandler.revalidateTag')
 
     const tags = Array.isArray(tagOrTags) ? tagOrTags : [tagOrTags]
@@ -275,7 +289,7 @@ export class NetlifyCacheHandler implements CacheHandler {
       }),
     )
 
-    purgeCache({ tags }).catch((error) => {
+    await purgeCache({ tags }).catch((error) => {
       // TODO: add reporting here
       getLogger()
         .withError(error)
diff --git a/src/run/revalidate.ts b/src/run/revalidate.ts
index 6bc9928b7e..1a991d92d4 100644
--- a/src/run/revalidate.ts
+++ b/src/run/revalidate.ts
@@ -1,4 +1,5 @@
 import type { ServerResponse } from 'node:http'
+import { isPromise } from 'node:util/types'
 
 import type { RequestContext } from './handlers/request-context.cjs'
 
@@ -12,7 +13,13 @@ export const nextResponseProxy = (res: ServerResponse, requestContext: RequestCo
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         return async function newRevalidate(...args: any[]) {
           requestContext.didPagesRouterOnDemandRevalidate = true
-          return originalValue?.apply(target, args)
+
+          const result = originalValue?.apply(target, args)
+          if (result && isPromise(result)) {
+            requestContext.trackBackgroundWork(result)
+          }
+
+          return result
         }
       }
       return originalValue
diff --git a/tests/fixtures/page-router/pages/api/revalidate.js b/tests/fixtures/page-router/pages/api/revalidate.js
index 79c2428b51..8e4fd4d476 100644
--- a/tests/fixtures/page-router/pages/api/revalidate.js
+++ b/tests/fixtures/page-router/pages/api/revalidate.js
@@ -1,6 +1,9 @@
 export default async function handler(req, res) {
   try {
-    await res.revalidate('/static/revalidate-manual')
+    // res.revalidate returns a promise that can be awaited to wait for the revalidation to complete
+    // if user doesn't await it, we still want to ensure the revalidation is completed, so we internally track
+    // this as "background work" to ensure it completes before function suspends execution
+    res.revalidate('/static/revalidate-manual')
     return res.json({ code: 200, message: 'success' })
   } catch (err) {
     return res.status(500).send({ code: 500, message: err.message })