fix: track background work (#407)

* tmp

* restore previous dynamic page with fetch-cache and add separate ~isr page

* await background tasks

* optimistically prepare to use context.waitUntil

* remove dev logs

* test: allow using different test site id

* revert timing changes in test fixture

* remove test page as behavior is background revalidation behavior is already tested in both integration and e2e tests

* add e2e for time-based revalidation

* test: update assertions

* chore: format with prettier

* update test

* tmp: don't await background promise

* Revert "tmp: don't await background promise"

This reverts commit d91cb976f5330b1cc9bc14902515c2cdf173a0d2.

* chore: update todo with reference to upstream pull request

* test: revert not strictly related e2e test setup change

---------

Co-authored-by: pieh <[email protected]>

Author: pieh

src/run/handlers/request-context.cts

@@ -10,13 +10,29 @@ export type RequestContext = {
   didPagesRouterOnDemandRevalidate?: boolean
   serverTiming?: string
   routeHandlerRevalidate?: NetlifyCachedRouteValue['revalidate']
+  /**
+   * Track promise running in the background and need to be waited for
+   */
+  trackBackgroundWork: (promise: Promise<unknown>) => void
+  /**
+   * Promise that need to be executed even if response was already sent
+   */
+  backgroundWorkPromise: Promise<unknown>
 }
 
 type RequestContextAsyncLocalStorage = AsyncLocalStorage<RequestContext>
 
 export function createRequestContext(debug = false): RequestContext {
+  const backgroundWorkPromises: Promise<unknown>[] = []
+
   return {
     debug,
+    trackBackgroundWork: (promise) => {
+      backgroundWorkPromises.push(promise)
+    },
+    get backgroundWorkPromise() {
+      return Promise.allSettled(backgroundWorkPromises)
+    },
   }
 }
 

src/run/handlers/server.ts

@@ -1,6 +1,7 @@
 import type { OutgoingHttpHeaders } from 'http'
 
 import { toComputeResponse, toReqRes, ComputeJsOutgoingMessage } from '@fastly/http-compute-js'
+import { Context } from '@netlify/functions'
 import type { NextConfigComplete } from 'next/dist/server/config-shared.js'
 import type { WorkerRequestHandler } from 'next/dist/server/lib/types.js'
 
@@ -45,7 +46,13 @@ const disableFaultyTransferEncodingHandling = (res: ComputeJsOutgoingMessage) =>
   }
 }
 
-export default async (request: Request) => {
+// TODO: remove once https://github.com/netlify/serverless-functions-api/pull/219
+// is released and public types are updated
+interface FutureContext extends Context {
+  waitUntil?: (promise: Promise<unknown>) => void
+}
+
+export default async (request: Request, context: FutureContext) => {
   const tracer = getTracer()
 
   if (!nextHandler) {
@@ -133,11 +140,19 @@ export default async (request: Request) => {
       return new Response(body || null, response)
     }
 
+    if (context.waitUntil) {
+      context.waitUntil(requestContext.backgroundWorkPromise)
+    }
+
     const keepOpenUntilNextFullyRendered = new TransformStream({
-      flush() {
-        // it's important to keep the stream open until the next handler has finished,
-        // or otherwise the cache revalidates might not go through
-        return nextHandlerPromise
+      async flush() {
+        // it's important to keep the stream open until the next handler has finished
+        await nextHandlerPromise
+        if (!context.waitUntil) {
+          // if waitUntil is not available, we have to keep response stream open until background promises are resolved
+          // to ensure that all background work executes
+          await requestContext.backgroundWorkPromise
+        }
       },
     })
 

src/run/next.cts

@@ -12,6 +12,33 @@ console.time('import next server')
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const { getRequestHandlers } = require('next/dist/server/lib/start-server.js')
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const ResponseCache = require('next/dist/server/response-cache/index.js').default
+
+// Next.js standalone doesn't expose background work promises (such as generating fresh response
+// while stale one is being served) that we could use so we regrettably have to use hacks to
+// gain access to them so that we can explicitly track them to ensure they finish before function
+// execution stops
+const originalGet = ResponseCache.prototype.get
+ResponseCache.prototype.get = function get(...getArgs: unknown[]) {
+  if (!this.didAddBackgroundWorkTracking) {
+    const originalBatcherBatch = this.batcher.batch
+    this.batcher.batch = async (key: string, fn: (...args: unknown[]) => unknown) => {
+      const trackedFn = async (...workFnArgs: unknown[]) => {
+        const workPromise = fn(...workFnArgs)
+        const requestContext = getRequestContext()
+        if (requestContext && workPromise instanceof Promise) {
+          requestContext.trackBackgroundWork(workPromise)
+        }
+        return await workPromise
+      }
+
+      return originalBatcherBatch.call(this.batcher, key, trackedFn)
+    }
+    this.didAddBackgroundWorkTracking = true
+  }
+  return originalGet.apply(this, getArgs)
+}
 
 console.timeEnd('import next server')
 

tests/e2e/page-router.test.ts

@@ -49,7 +49,11 @@ export async function check(
 }
 
 test.describe('Simple Page Router (no basePath, no i18n)', () => {
-  test('Static revalidate works correctly', async ({ page, pollUntilHeadersMatch, pageRouter }) => {
+  test('On-demand revalidate works correctly', async ({
+    page,
+    pollUntilHeadersMatch,
+    pageRouter,
+  }) => {
     // in case there is retry or some other test did hit that path before
     // we want to make sure that cdn cache is not warmed up
     const purgeCdnCache = await page.goto(
@@ -230,6 +234,70 @@ test.describe('Simple Page Router (no basePath, no i18n)', () => {
     expect(data3?.pageProps?.time).toBe(date3)
   })
 
+  test('Time based revalidate works correctly', async ({
+    page,
+    pollUntilHeadersMatch,
+    pageRouter,
+  }) => {
+    // in case there is retry or some other test did hit that path before
+    // we want to make sure that cdn cache is not warmed up
+    const purgeCdnCache = await page.goto(
+      new URL('/api/purge-cdn?path=/static/revalidate-slow-data', pageRouter.url).href,
+    )
+    expect(purgeCdnCache?.status()).toBe(200)
+
+    // wait a bit until cdn cache purge propagates and make sure page gets stale (revalidate 10)
+    await page.waitForTimeout(10_000)
+
+    const beforeFetch = new Date().toISOString()
+
+    const response1 = await pollUntilHeadersMatch(
+      new URL('static/revalidate-slow-data', pageRouter.url).href,
+      {
+        headersToMatch: {
+          // either first time hitting this route or we invalidated
+          // just CDN node in earlier step
+          // we will invoke function and see Next cache hit status \
+          // in the response because it was prerendered at build time
+          // or regenerated in previous attempt to run this test
+          'cache-status': [/"Netlify Edge"; fwd=(miss|stale)/m, /"Next.js"; hit/m],
+        },
+        headersNotMatchedMessage:
+          'First request to tested page (html) should be a miss or stale on the Edge and stale in Next.js',
+      },
+    )
+    expect(response1?.status()).toBe(200)
+    const date1 = (await page.textContent('[data-testid="date-now"]')) ?? ''
+
+    // ensure response was produced before invocation (served from cache)
+    expect(date1.localeCompare(beforeFetch)).toBeLessThan(0)
+
+    // wait a bit to ensure background work has a chance to finish
+    // (page is fresh for 10 seconds and it should take at least 5 seconds to regenerate, so we should wait at least more than 15 seconds)
+    await page.waitForTimeout(20_000)
+
+    const response2 = await pollUntilHeadersMatch(
+      new URL('static/revalidate-slow-data', pageRouter.url).href,
+      {
+        headersToMatch: {
+          // either first time hitting this route or we invalidated
+          // just CDN node in earlier step
+          // we will invoke function and see Next cache hit status \
+          // in the response because it was prerendered at build time
+          // or regenerated in previous attempt to run this test
+          'cache-status': [/"Netlify Edge"; fwd=(miss|stale)/m, /"Next.js"; hit;/m],
+        },
+        headersNotMatchedMessage:
+          'Second request to tested page (html) should be a miss or stale on the Edge and hit or stale in Next.js',
+      },
+    )
+    expect(response2?.status()).toBe(200)
+    const date2 = (await page.textContent('[data-testid="date-now"]')) ?? ''
+
+    // ensure response was produced after initial invocation
+    expect(beforeFetch.localeCompare(date2)).toBeLessThan(0)
+  })
+
   test('should serve 404 page when requesting non existing page (no matching route)', async ({
     page,
     pageRouter,

tests/fixtures/advanced-api-routes/tsconfig.json

@@ -1,10 +1,6 @@
 {
   "compilerOptions": {
-    "lib": [
-      "dom",
-      "dom.iterable",
-      "esnext"
-    ],
+    "lib": ["dom", "dom.iterable", "esnext"],
     "allowJs": true,
     "skipLibCheck": true,
     "strict": false,
@@ -17,12 +13,6 @@
     "isolatedModules": true,
     "jsx": "preserve"
   },
-  "include": [
-    "next-env.d.ts",
-    "**/*.ts",
-    "**/*.tsx"
-  ],
-  "exclude": [
-    "node_modules"
-  ]
+  "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx"],
+  "exclude": ["node_modules"]
 }

tests/fixtures/page-router/pages/static/revalidate-slow-data.js

@@ -0,0 +1,34 @@
+const Show = ({ show, time }) => (
+  <div>
+    <p>
+      This page uses getStaticProps() to pre-fetch a TV show at
+      <span data-testid="date-now">{time}</span>
+    </p>
+    <hr />
+    <h1>Show #{show.id}</h1>
+    <p>{show.name}</p>
+  </div>
+)
+
+export async function getStaticProps(context) {
+  const start = Date.now()
+
+  const res = await fetch(`https://tvproxy.netlify.app/shows/71`)
+  const data = await res.json()
+
+  const fetchDuration = Date.now() - start
+  if (fetchDuration < 5000) {
+    // simulate slow (~5s) response
+    await new Promise((resolve) => setTimeout(resolve, 5000 - fetchDuration))
+  }
+
+  return {
+    props: {
+      show: data,
+      time: new Date().toISOString(),
+    },
+    revalidate: +process.env.REVALIDATE_SECONDS || 10, // In seconds
+  }
+}
+
+export default Show

tests/integration/cache-handler.test.ts

@@ -43,6 +43,7 @@ describe('page router', () => {
       '/static/revalidate-automatic',
       '/static/revalidate-manual',
       '/static/revalidate-slow',
+      '/static/revalidate-slow-data',
       '404.html',
       '500.html',
       'static/fully-static.html',

tests/integration/static.test.ts

@@ -38,6 +38,7 @@ test<FixtureTestContext>('requesting a non existing page route that needs to be
     '/static/revalidate-automatic',
     '/static/revalidate-manual',
     '/static/revalidate-slow',
+    '/static/revalidate-slow-data',
     '404.html',
     '500.html',
     'static/fully-static.html',