fix: track revalidate / cdn purge to ensure it finishes execution and is not suspended mid-execution

src/run/handlers/cache.cts

@@ -283,19 +283,33 @@ 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`)
+            }),
+          )
         }
       }
     })
   }
 
   // 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]
@@ -314,7 +328,7 @@ export class NetlifyCacheHandler implements CacheHandler {
       }),
     )
 
-    purgeCache({ tags }).catch((error) => {
+    await purgeCache({ tags }).catch((error) => {
       // TODO: add reporting here
       getLogger()
         .withError(error)

src/run/revalidate.ts

@@ -1,18 +1,34 @@
 import type { ServerResponse } from 'node:http'
+import { isPromise } from 'node:util/types'
+
+import type { NextApiResponse } from 'next'
 
 import type { RequestContext } from './handlers/request-context.cjs'
 
+type ResRevalidateMethod = NextApiResponse['revalidate']
+
+function isRevalidateMethod(
+  key: string,
+  nextResponseField: unknown,
+): nextResponseField is ResRevalidateMethod {
+  return key === 'revalidate' && typeof nextResponseField === 'function'
+}
+
 // Needing to proxy the response object to intercept the revalidate call for on-demand revalidation on page routes
 export const nextResponseProxy = (res: ServerResponse, requestContext: RequestContext) => {
   return new Proxy(res, {
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    get(target: any[string], key: string) {
-      const originalValue = target[key]
-      if (key === 'revalidate') {
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        return async function newRevalidate(...args: any[]) {
+    get(target: ServerResponse, key: string) {
+      const originalValue = Reflect.get(target, key)
+      if (isRevalidateMethod(key, originalValue)) {
+        return function newRevalidate(...args: Parameters<ResRevalidateMethod>) {
           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

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 })

tests/utils/fixture.ts

@@ -364,14 +364,24 @@ export async function invokeFunction(
       NETLIFY_BLOBS_CONTEXT: createBlobContext(ctx),
       ...(env || {}),
     }
+
+    const envVarsToRestore = {}
+
+    // We are not using lambda-local's environment variable setting because it cleans up
+    // environment vars to early (before stream is closed)
+    Object.keys(environment).forEach(function (key) {
+      if (typeof process.env[key] !== 'undefined') {
+        envVarsToRestore[key] = process.env[key]
+      }
+      process.env[key] = environment[key]
+    })
+
     const response = (await execute({
       event: {
         headers: headers || {},
         httpMethod: httpMethod || 'GET',
         rawUrl: new URL(url || '/', 'https://example.netlify').href,
       },
-      environment,
-      envdestroy: true,
       lambdaFunc: { handler },
       timeoutMs: 4_000,
     })) as LambdaResponse
@@ -386,6 +396,14 @@ export async function invokeFunction(
 
     const bodyBuffer = await streamToBuffer(response.body)
 
+    Object.keys(environment).forEach(function (key) {
+      if (typeof envVarsToRestore[key] !== 'undefined') {
+        process.env[key] = envVarsToRestore[key]
+      } else {
+        delete process.env[key]
+      }
+    })
+
     return {
       statusCode: response.statusCode,
       bodyBuffer,

tests/utils/sandbox-child.mjs

@@ -48,14 +48,23 @@ process.on('message', async (msg) => {
         ...(env || {}),
       }
 
+      const envVarsToRestore = {}
+
+      // We are not using lambda-local's environment variable setting because it cleans up
+      // environment vars to early (before stream is closed)
+      Object.keys(environment).forEach(function (key) {
+        if (typeof process.env[key] !== 'undefined') {
+          envVarsToRestore[key] = process.env[key]
+        }
+        process.env[key] = environment[key]
+      })
+
       const response = await execute({
         event: {
           headers: headers || {},
           httpMethod: httpMethod || 'GET',
           rawUrl: new URL(url || '/', 'https://example.netlify').href,
         },
-        environment,
-        envdestroy: true,
         lambdaFunc: { handler },
         timeoutMs: 4_000,
       })
@@ -70,6 +79,14 @@ process.on('message', async (msg) => {
 
       const bodyBuffer = await streamToBuffer(response.body)
 
+      Object.keys(environment).forEach(function (key) {
+        if (typeof envVarsToRestore[key] !== 'undefined') {
+          process.env[key] = envVarsToRestore[key]
+        } else {
+          delete process.env[key]
+        }
+      })
+
       const result = {
         statusCode: response.statusCode,
         bodyBuffer,