fix ISR not working reliably

conico974 · khuezy · commit f5ee4da1ce4a · 2023-09-08T08:42:28.000-07:00
diff --git a/packages/open-next/src/adapters/http/responseStreaming.ts b/packages/open-next/src/adapters/http/responseStreaming.ts
@@ -9,6 +9,7 @@ export class StreamingServerResponse extends http.ServerResponse {
   [HEADERS]: Record<string, string> = {};
   responseStream: ResponseStream;
   fixHeaders: (headers: Record<string, string>) => void;
+  onEnd: (headers: Record<string, string>) => Promise<void>;
   private _wroteHeader = false;
   private _hasWritten = false;
 
@@ -54,8 +55,12 @@ export class StreamingServerResponse extends http.ServerResponse {
         statusCode: statusCode as number,
         headers: this[HEADERS],
       });
-      this.internalWrite(prelude);
-      this.internalWrite(new Uint8Array(8));
+      process.nextTick(() => {
+        this.responseStream.write(prelude);
+      });
+      process.nextTick(() => {
+        this.responseStream.write(new Uint8Array(8));
+      });
       // this.responseStream = awslambda.HttpResponseStream.from(
       //   this.responseStream,
       //   {
@@ -96,9 +101,10 @@ export class StreamingServerResponse extends http.ServerResponse {
     }
 
     process.nextTick(() => {
-      this.responseStream.end(() => {
+      this.responseStream.end(async () => {
         // The callback seems necessary here
         debug("stream end", chunk);
+        await this.onEnd(this[HEADERS]);
       });
     });
     // debug("stream end", chunk);
@@ -116,13 +122,15 @@ export class StreamingServerResponse extends http.ServerResponse {
     { method, headers }: { method?: string; headers?: Record<string, string> },
     responseStream: ResponseStream,
     fixHeaders: (headers: Record<string, string>) => void,
+    onEnd: (headers: Record<string, string>) => Promise<void>,
   ) {
     //@ts-ignore
     super({ method });
 
     this[HEADERS] = headers || {};
 
     this.fixHeaders = fixHeaders;
+    this.onEnd = onEnd;
     this.responseStream = responseStream;
 
     this.useChunkedEncodingByDefault = false;
diff --git a/packages/open-next/src/adapters/plugins/routing/default.replacement.ts b/packages/open-next/src/adapters/plugins/routing/default.replacement.ts
@@ -19,6 +19,7 @@ import { loadBuildId, loadConfigHeaders, loadRoutesManifest } from "../../util";
 import {
   addOpenNextHeader,
   fixCacheHeaderForHtmlPages,
+  fixISRHeaders,
   fixSWRCacheHeader,
   revalidateIfRequired,
 } from "./util";
@@ -129,6 +130,7 @@ export async function postProcessResponse({
     fixCacheHeaderForHtmlPages(internalEvent.rawPath, headers);
     fixSWRCacheHeader(headers);
     addOpenNextHeader(headers);
+    fixISRHeaders(headers);
 
     await revalidateIfRequired(
       internalEvent.headers.host,
diff --git a/packages/open-next/src/adapters/plugins/routing/default.ts b/packages/open-next/src/adapters/plugins/routing/default.ts
@@ -11,6 +11,7 @@ import { ServerResponse } from "../../http/response";
 import {
   addOpenNextHeader,
   fixCacheHeaderForHtmlPages,
+  fixISRHeaders,
   fixSWRCacheHeader,
   revalidateIfRequired,
 } from "./util";
@@ -58,6 +59,7 @@ export async function postProcessResponse({
     fixCacheHeaderForHtmlPages(internalEvent.rawPath, headers);
     fixSWRCacheHeader(headers);
     addOpenNextHeader(headers);
+    fixISRHeaders(headers);
 
     await revalidateIfRequired(
       internalEvent.headers.host,
diff --git a/packages/open-next/src/adapters/plugins/routing/util.ts b/packages/open-next/src/adapters/plugins/routing/util.ts
@@ -73,25 +73,33 @@ export function addOpenNextHeader(headers: Record<string, string | undefined>) {
   headers["X-OpenNext"] = process.env.OPEN_NEXT_VERSION;
 }
 
-export async function revalidateIfRequired(
-  host: string,
-  rawPath: string,
-  headers: Record<string, string | undefined>,
-  req: IncomingMessage,
-) {
+export function fixISRHeaders(headers: Record<string, string | undefined>) {
+  if (headers["x-nextjs-cache"] === "REVALIDATED") {
+    headers["cache-control"] =
+      "private, no-cache, no-store, max-age=0, must-revalidate";
+    return;
+  }
   if (headers["x-nextjs-cache"] !== "STALE") return;
 
   // If the cache is stale, we revalidate in the background
   // In order for CloudFront SWR to work, we set the stale-while-revalidate value to 2 seconds
   // This will cause CloudFront to cache the stale data for a short period of time while we revalidate in the background
   // Once the revalidation is complete, CloudFront will serve the fresh data
   headers["cache-control"] = "s-maxage=2, stale-while-revalidate=2592000";
+}
 
+export async function revalidateIfRequired(
+  host: string,
+  rawPath: string,
+  headers: Record<string, string | undefined>,
+  req?: IncomingMessage,
+) {
+  if (headers["x-nextjs-cache"] !== "STALE") return;
   // If the URL is rewritten, revalidation needs to be done on the rewritten URL.
   // - Link to Next.js doc: https://nextjs.org/docs/pages/building-your-application/data-fetching/incremental-static-regeneration#on-demand-revalidation
   // - Link to NextInternalRequestMeta: https://github.com/vercel/next.js/blob/57ab2818b93627e91c937a130fb56a36c41629c3/packages/next/src/server/request-meta.ts#L11
   // @ts-ignore
-  const internalMeta = req[Symbol.for("NextInternalRequestMeta")];
+  const internalMeta = req?.[Symbol.for("NextInternalRequestMeta")];
 
   // When using Pages Router, two requests will be received:
   // 1. one for the page: /foo
diff --git a/packages/open-next/src/adapters/plugins/streaming.replacement.ts b/packages/open-next/src/adapters/plugins/streaming.replacement.ts
@@ -15,6 +15,7 @@ import { processInternalEvent } from "./routing/default.js";
 import {
   addOpenNextHeader,
   fixCacheHeaderForHtmlPages,
+  fixISRHeaders,
   fixSWRCacheHeader,
   revalidateIfRequired,
 } from "./routing/util";
@@ -51,10 +52,19 @@ export const lambdaHandler = awslambda.streamifyResponse(async function (
       { method, headers },
       responseStream,
       // We need to fix the cache header before sending any response
-      async (headers) => {
+      (headers) => {
         fixCacheHeaderForHtmlPages(internalEvent.rawPath, headers);
         fixSWRCacheHeader(headers);
         addOpenNextHeader(headers);
+        fixISRHeaders(headers);
+      },
+      // This run in the callback of the response stream end
+      async (headers) => {
+        await revalidateIfRequired(
+          internalEvent.headers.host,
+          internalEvent.rawPath,
+          headers,
+        );
       },
     );
 
@@ -86,13 +96,6 @@ export const lambdaHandler = awslambda.streamifyResponse(async function (
 
     //@ts-expect-error - processRequest is already defined in serverHandler.ts
     await processRequest(req, res, overwrittenInternalEvent, isExternalRewrite);
-
-    await revalidateIfRequired(
-      internalEvent.headers.host,
-      internalEvent.rawPath,
-      res.headers,
-      req,
-    );
   }
 });
 //#endOverride
