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

pieh · web-flow · commit 5a0fec519ac4 · 2024-06-26T16:09:19.000+02:00
* fix: track revalidate / cdn purge to ensure it finishes execution and is not suspended mid-execution

* test: don't destroy lambda env vars until response stream finished

* fix: remove 'any' types from next's response proxy

* test: add a case for not awaited res.revalidate
diff --git a/src/run/handlers/cache.cts b/src/run/handlers/cache.cts
@@ -326,19 +326,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]
@@ -357,7 +371,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
@@ -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
diff --git a/tests/e2e/page-router.test.ts b/tests/e2e/page-router.test.ts
@@ -50,25 +50,35 @@ export async function check(
 
 test.describe('Simple Page Router (no basePath, no i18n)', () => {
   test.describe('On-demand revalidate works correctly', () => {
-    for (const { label, prerendered, pagePath, expectedH1Content } of [
+    for (const { label, prerendered, pagePath, revalidateApiBasePath, expectedH1Content } of [
       {
-        label: 'prerendered page with static path',
+        label: 'prerendered page with static path and awaited res.revalidate()',
         prerendered: true,
         pagePath: '/static/revalidate-manual',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Show #71',
       },
       {
-        label: 'prerendered page with dynamic path',
+        label: 'prerendered page with dynamic path and awaited res.revalidate()',
         prerendered: true,
         pagePath: '/products/prerendered',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Product prerendered',
       },
       {
-        label: 'not prerendered page with dynamic path',
+        label: 'not prerendered page with dynamic path and awaited res.revalidate()',
         prerendered: false,
         pagePath: '/products/not-prerendered',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Product not-prerendered',
       },
+      {
+        label: 'not prerendered page with dynamic path and not awaited res.revalidate()',
+        prerendered: false,
+        pagePath: '/products/not-prerendered-and-not-awaited-revalidation',
+        revalidateApiBasePath: '/api/revalidate-no-await',
+        expectedH1Content: 'Product not-prerendered-and-not-awaited-revalidation',
+      },
     ]) {
       test(label, async ({ page, pollUntilHeadersMatch, pageRouter }) => {
         // in case there is retry or some other test did hit that path before
@@ -192,7 +202,7 @@ test.describe('Simple Page Router (no basePath, no i18n)', () => {
         expect(data2?.pageProps?.time).toBe(date1)
 
         const revalidate = await page.goto(
-          new URL(`/api/revalidate?path=${pagePath}`, pageRouter.url).href,
+          new URL(`${revalidateApiBasePath}?path=${pagePath}`, pageRouter.url).href,
         )
         expect(revalidate?.status()).toBe(200)
 
@@ -411,25 +421,35 @@ test.describe('Simple Page Router (no basePath, no i18n)', () => {
 
 test.describe('Page Router with basePath and i18n', () => {
   test.describe('Static revalidate works correctly', () => {
-    for (const { label, prerendered, pagePath, expectedH1Content } of [
+    for (const { label, prerendered, pagePath, revalidateApiBasePath, expectedH1Content } of [
       {
-        label: 'prerendered page with static path',
+        label: 'prerendered page with static path and awaited res.revalidate()',
         prerendered: true,
         pagePath: '/static/revalidate-manual',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Show #71',
       },
       {
-        label: 'prerendered page with dynamic path',
+        label: 'prerendered page with dynamic path and awaited res.revalidate()',
         prerendered: true,
         pagePath: '/products/prerendered',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Product prerendered',
       },
       {
-        label: 'not prerendered page with dynamic path',
+        label: 'not prerendered page with dynamic path and awaited res.revalidate()',
         prerendered: false,
         pagePath: '/products/not-prerendered',
+        revalidateApiBasePath: '/api/revalidate',
         expectedH1Content: 'Product not-prerendered',
       },
+      {
+        label: 'not prerendered page with dynamic path and not awaited res.revalidate()',
+        prerendered: false,
+        pagePath: '/products/not-prerendered-and-not-awaited-revalidation',
+        revalidateApiBasePath: '/api/revalidate-no-await',
+        expectedH1Content: 'Product not-prerendered-and-not-awaited-revalidation',
+      },
     ]) {
       test.describe(label, () => {
         test(`default locale`, async ({ page, pollUntilHeadersMatch, pageRouterBasePathI18n }) => {
@@ -622,7 +642,10 @@ test.describe('Page Router with basePath and i18n', () => {
 
           // revalidate implicit locale path
           const revalidateImplicit = await page.goto(
-            new URL(`/base/path/api/revalidate?path=${pagePath}`, pageRouterBasePathI18n.url).href,
+            new URL(
+              `/base/path${revalidateApiBasePath}?path=${pagePath}`,
+              pageRouterBasePathI18n.url,
+            ).href,
           )
           expect(revalidateImplicit?.status()).toBe(200)
 
@@ -713,8 +736,10 @@ test.describe('Page Router with basePath and i18n', () => {
 
           // revalidate implicit locale path
           const revalidateExplicit = await page.goto(
-            new URL(`/base/path/api/revalidate?path=/en${pagePath}`, pageRouterBasePathI18n.url)
-              .href,
+            new URL(
+              `/base/path${revalidateApiBasePath}?path=/en${pagePath}`,
+              pageRouterBasePathI18n.url,
+            ).href,
           )
           expect(revalidateExplicit?.status()).toBe(200)
 
@@ -934,8 +959,10 @@ test.describe('Page Router with basePath and i18n', () => {
           expect(data2?.pageProps?.time).toBe(date1)
 
           const revalidate = await page.goto(
-            new URL(`/base/path/api/revalidate?path=/de${pagePath}`, pageRouterBasePathI18n.url)
-              .href,
+            new URL(
+              `/base/path${revalidateApiBasePath}?path=/de${pagePath}`,
+              pageRouterBasePathI18n.url,
+            ).href,
           )
           expect(revalidate?.status()).toBe(200)
 
diff --git a/tests/fixtures/page-router-base-path-i18n/pages/api/revalidate-no-await.js b/tests/fixtures/page-router-base-path-i18n/pages/api/revalidate-no-await.js
@@ -0,0 +1,12 @@
+export default async function handler(req, res) {
+  try {
+    const pathToPurge = req.query.path ?? '/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(pathToPurge)
+    return res.json({ code: 200, message: 'success' })
+  } catch (err) {
+    return res.status(500).send({ code: 500, message: err.message })
+  }
+}
diff --git a/tests/fixtures/page-router/pages/api/revalidate-no-await.js b/tests/fixtures/page-router/pages/api/revalidate-no-await.js
@@ -0,0 +1,12 @@
+export default async function handler(req, res) {
+  try {
+    const pathToPurge = req.query.path ?? '/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(pathToPurge)
+    return res.json({ code: 200, message: 'success' })
+  } catch (err) {
+    return res.status(500).send({ code: 500, message: err.message })
+  }
+}
diff --git a/tests/utils/fixture.ts b/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,
diff --git a/tests/utils/sandbox-child.mjs b/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,
