feat(cache): cache event handler stream response

nitrojs · Dec 12, 2024 · 789bb93 · 789bb93
1 parent f982f99
commit 789bb93
Show file tree

Hide file tree

Showing 6 changed files with 92 additions and 1 deletion.
diff --git a/src/runtime/internal/cache.ts b/src/runtime/internal/cache.ts
@@ -5,6 +5,7 @@ import {
   fetchWithEvent,
   handleCacheHeaders,
   isEvent,
+  isStream,
   splitCookiesString,
 } from "h3";
 import type { EventHandlerRequest, EventHandlerResponse, H3Event } from "h3";
@@ -67,6 +68,8 @@ export function defineCachedFunction<T, ArgsT extends unknown[] = any[]>(
           useNitroApp().captureError(error, { event, tags: ["cache"] });
         })) as unknown) || {};
 
+    entry.key = cacheKey;
+
     // https://github.com/nitrojs/nitro/issues/2160
     if (typeof entry !== "object") {
       entry = {};
@@ -272,6 +275,32 @@ export function defineCachedEventHandler<
       }
       return true;
     },
+    // transform(entry) {
+    //   if (!entry.value?.body) return entry.value;
+    //
+    //   if (!(entry.value.body instanceof ReadableStream)) return entry.value;
+    //
+    //   const [stream1, stream2] = entry.value.body.tee();
+    //
+    //   const td = new TextDecoder();
+    //   let res = "";
+    //   stream2
+    //     .pipeTo(
+    //       new WritableStream({
+    //         write: (chunk) => {
+    //           res += td.decode(chunk);
+    //         },
+    //       })
+    //     )
+    //     .then(() => {
+    //       const en = entry as any;
+    //       en.value.body = res;
+    //       useStorage().setItem(entry.key!, entry);
+    //     });
+    //
+    //   entry.value.body = stream1 as any;
+    //   return entry.value;
+    // },
     group: opts.group || "nitro/handlers",
     integrity: opts.integrity || hash([handler, opts]),
   };
@@ -406,11 +435,26 @@ export function defineCachedEventHandler<
         headers["cache-control"] = cacheControl.join(", ");
       }
 
+      let cachedBody = body as any;
+      // When handler response is a stream, we cache the result of this stream
+      if (body instanceof ReadableStream) {
+        const td = new TextDecoder();
+        let buffer = "";
+        await body.pipeTo(
+          new WritableStream({
+            write(chunk) {
+              buffer += td.decode(chunk);
+            },
+          })
+        );
+        cachedBody = buffer;
+      }
+
       // Create cache entry for response
       const cacheEntry: ResponseCacheEntry<Response> = {
         code: event.node.res.statusCode,
         headers,
-        body,
+        body: cachedBody,
       };
 
       return cacheEntry;

diff --git a/test/fixture/nitro.config.ts b/test/fixture/nitro.config.ts
@@ -91,6 +91,7 @@ export default defineNitroConfig({
     "/rules/_/cached/noncached": { cache: false, swr: false, isr: false },
     "/rules/_/cached/**": { swr: true },
     "/api/proxy/**": { proxy: "/api/echo" },
+    "/stream-cached": { swr: true },
   },
   prerender: {
     crawlLinks: true,

diff --git a/test/fixture/routes/stream-cached.ts b/test/fixture/routes/stream-cached.ts
@@ -0,0 +1,13 @@
+export default eventHandler(() => {
+  const encoder = new TextEncoder();
+  const stream = new ReadableStream({
+    start(controller) {
+      controller.enqueue(encoder.encode("nitro"));
+      controller.enqueue(encoder.encode("is"));
+      controller.enqueue(encoder.encode("awesome,"));
+      controller.enqueue(encoder.encode(Date.now().toString()));
+      controller.close();
+    },
+  });
+  return stream;
+});
diff --git a/test/presets/netlify-legacy.test.ts b/test/presets/netlify-legacy.test.ts
@@ -71,6 +71,7 @@ describe("nitro:preset:netlify-legacy", async () => {
         /rules/isr-ttl/*	/.netlify/builders/server 200
         /rules/isr/*	/.netlify/builders/server 200
         /rules/dynamic	/.netlify/functions/server 200
+        /stream-cached	/.netlify/builders/server 200
         /build/* /build/:splat 200
         /* /.netlify/functions/server 200"
       `);

diff --git a/test/presets/vercel.test.ts b/test/presets/vercel.test.ts
@@ -142,6 +142,10 @@ describe("nitro:preset:vercel", async () => {
                 "dest": "/__nitro--rules-swr-ttl?url=$url",
                 "src": "(?<url>/rules/swr-ttl/.*)",
               },
+              {
+                "dest": "/stream-cached?url=$url",
+                "src": "/stream-cached",
+              },
               {
                 "dest": "/__nitro",
                 "src": "/(.*)",

diff --git a/test/tests.ts b/test/tests.ts
@@ -729,6 +729,34 @@ export function testNitro(
         }
       }
     );
+
+    it.skipIf(
+      ctx.isIsolated ||
+        (isWindows && ctx.preset === "nitro-dev") ||
+        ctx.isLambda
+    )("should cache stream result", async () => {
+      const { data } = await callHandler({
+        url: "/stream-cached",
+      });
+
+      const [str, timestamp] = data.split(",") as string[];
+
+      expect(str).toBe(
+        ctx.isLambda ? btoa("nitroisawesome") : "nitroisawesome"
+      );
+
+      const calls = await Promise.all([
+        callHandler({ url: "/stream-cached" }),
+        callHandler({ url: "/stream-cached" }),
+        callHandler({ url: "/stream-cached" }),
+      ]);
+
+      for (const call of calls) {
+        expect(call.data).toBe(
+          ctx.isLambda ? btoa(`${str},${timestamp}`) : `${str},${timestamp}`
+        );
+      }
+    });
   });
 
   describe("scanned files", () => {