Add per-request context, limit subrequests to 50, closes #117

Author: mrbbot

packages/cache/src/cache.ts

@@ -11,6 +11,7 @@ import {
   Clock,
   Storage,
   defaultClock,
+  getRequestContext,
   millisToSeconds,
   waitForOpenInputGate,
   waitForOpenOutputGate,
@@ -151,6 +152,7 @@ export class Cache implements CacheInterface {
     req: RequestInfo,
     res: BaseResponse | Response
   ): Promise<undefined> {
+    getRequestContext()?.incrementSubrequests();
     req = normaliseRequest(req);
 
     if (res instanceof Response && res.webSocket) {
@@ -192,6 +194,7 @@ export class Cache implements CacheInterface {
     req: RequestInfo,
     options?: CacheMatchOptions
   ): Promise<Response | undefined> {
+    getRequestContext()?.incrementSubrequests();
     req = normaliseRequest(req);
     // Cloudflare only caches GET requests
     if (req.method !== "GET" && !options?.ignoreMethod) return;
@@ -232,6 +235,7 @@ export class Cache implements CacheInterface {
     req: RequestInfo,
     options?: CacheMatchOptions
   ): Promise<boolean> {
+    getRequestContext()?.incrementSubrequests();
     req = normaliseRequest(req);
     // Cloudflare only caches GET requests
     if (req.method !== "GET" && !options?.ignoreMethod) return false;

packages/cache/test/cache.spec.ts

@@ -2,7 +2,7 @@ import assert from "assert";
 import { URL } from "url";
 import { Cache, CacheError, CachedMeta } from "@miniflare/cache";
 import { Request, RequestInitCfProperties, Response } from "@miniflare/core";
-import { Storage } from "@miniflare/shared";
+import { RequestContext, Storage } from "@miniflare/shared";
 import {
   getObjectProperties,
   utf8Decode,
@@ -191,6 +191,12 @@ test("Cache: put respects cf cacheTtlByStatus", async (t) => {
   t.is(await cache.match("http://localhost/302"), undefined);
 });
 
+test("Cache: put increments subrequest count", async (t) => {
+  const { cache } = t.context;
+  const ctx = new RequestContext();
+  await ctx.runWith(() => cache.put("http://localhost:8787/", testResponse()));
+  t.is(ctx.subrequests, 1);
+});
 test("Cache: put waits for output gate to open before storing", (t) => {
   const { cache } = t.context;
   return waitsForOutputGate(
@@ -243,6 +249,12 @@ test("Cache: only matches non-GET requests when ignoring method", async (t) => {
   t.not(await cache.match(req, { ignoreMethod: true }), undefined);
 });
 
+test("Cache: match increments subrequest count", async (t) => {
+  const { cache } = t.context;
+  const ctx = new RequestContext();
+  await ctx.runWith(() => cache.match("http://localhost:8787/"));
+  t.is(ctx.subrequests, 1);
+});
 test("Cache: match MISS waits for input gate to open before returning", async (t) => {
   const { cache } = t.context;
   await waitsForInputGate(t, () => cache.match("http://localhost:8787/"));
@@ -297,6 +309,12 @@ test("Cache: only deletes non-GET requests when ignoring method", async (t) => {
   t.true(await cache.delete(req, { ignoreMethod: true }));
 });
 
+test("Cache: delete increments subrequest count", async (t) => {
+  const { cache } = t.context;
+  const ctx = new RequestContext();
+  await ctx.runWith(() => cache.delete("http://localhost:8787/"));
+  t.is(ctx.subrequests, 1);
+});
 test("Cache: delete waits for output gate to open before deleting", async (t) => {
   const { cache } = t.context;
   await cache.put("http://localhost:8787/", testResponse());

packages/core/src/index.ts

@@ -14,6 +14,7 @@ import {
   PluginOptions,
   PluginOptionsUnion,
   PluginSignatures,
+  RequestContext,
   ScriptBlueprint,
   ScriptRunner,
   ScriptRunnerResult,
@@ -918,9 +919,12 @@ export class MiniflareCore<
     }
 
     const globalScope = this.#globalScope;
-    return globalScope![kDispatchFetch]<WaitUntil>(
-      withImmutableHeaders(request),
-      !!upstreamURL // only proxy if upstream URL set
+    // Each fetch gets its own context (e.g. 50 subrequests)
+    return new RequestContext().runWith(() =>
+      globalScope![kDispatchFetch]<WaitUntil>(
+        withImmutableHeaders(request),
+        !!upstreamURL // only proxy if upstream URL set
+      )
     );
   }
 
@@ -930,7 +934,10 @@ export class MiniflareCore<
   ): Promise<WaitUntil> {
     await this.#initPromise;
     const globalScope = this.#globalScope;
-    return globalScope![kDispatchScheduled]<WaitUntil>(scheduledTime, cron);
+    // Each fetch gets its own context (e.g. 50 subrequests)
+    return new RequestContext().runWith(() =>
+      globalScope![kDispatchScheduled]<WaitUntil>(scheduledTime, cron)
+    );
   }
 
   async dispose(): Promise<void> {

packages/core/src/standards/http.ts

@@ -13,6 +13,7 @@ import { URL } from "url";
 import {
   Compatibility,
   Log,
+  getRequestContext,
   nonCircularClone,
   waitForOpenInputGate,
   waitForOpenOutputGate,
@@ -591,14 +592,31 @@ export function withWaitUntil<WaitUntil extends any[]>(
   return resWaitUntil;
 }
 
+/** @internal */
+export function _getURLList(res: BaseResponse): URL[] | undefined {
+  // Extract the internal urlList property on Responses. It doesn't matter
+  // too much if the internal representation changes in the future: this code
+  // shouldn't throw. Currently we use this to count the number of redirects,
+  // and increment the subrequest count accordingly.
+  for (const symbol of Object.getOwnPropertySymbols(res)) {
+    if (symbol.description === "state") {
+      // @ts-expect-error symbol properties are not included type definitions
+      return res[symbol].urlList;
+    }
+  }
+}
+
 export async function fetch(
   input: RequestInfo,
   init?: RequestInit
 ): Promise<Response> {
-  // TODO (someday): support cache using fetch:
+  // TODO (someday): support cache using fetch (could add to request context?):
   //  https://developers.cloudflare.com/workers/learning/how-the-cache-works#fetch
   //  https://developers.cloudflare.com/workers/examples/cache-using-fetch
 
+  const ctx = getRequestContext();
+  ctx?.incrementSubrequests();
+
   await waitForOpenOutputGate();
 
   // Don't pass our strange hybrid Request to undici
@@ -616,6 +634,16 @@ export async function fetch(
 
   const baseRes = await baseFetch(req);
 
+  // Increment the subrequest count by the number of redirects
+  // TODO (someday): technically we should check the subrequest count before
+  //  each redirect, so requests don't actually get sent to the server if the
+  //  subrequest count exceeds the limit
+  if (baseRes.redirected && ctx) {
+    const urlList = _getURLList(baseRes);
+    // Last url is final destination, so subtract 1 for redirect count
+    if (urlList) ctx.incrementSubrequests(urlList.length - 1);
+  }
+
   // Convert the response to our hybrid Response
   const res = new Response(baseRes.body, baseRes);
 

packages/core/src/standards/index.ts

@@ -13,6 +13,7 @@ export {
   withImmutableHeaders,
   Response,
   withWaitUntil,
+  _getURLList,
   fetch,
   _urlFromRequestInput,
   _buildUnknownProtocolWarning,

packages/core/test/index.spec.ts

@@ -2,6 +2,7 @@ import assert from "assert";
 import fs from "fs/promises";
 import path from "path";
 import { setTimeout } from "timers/promises";
+import { CachePlugin } from "@miniflare/cache";
 import {
   BindingsPlugin,
   BuildPlugin,
@@ -25,6 +26,7 @@ import {
   Options,
   Storage,
   TypedEventListener,
+  getRequestContext,
 } from "@miniflare/shared";
 import {
   LogEntry,
@@ -44,6 +46,7 @@ import {
 import test, { Macro } from "ava";
 import { Request as BaseRequest, File, FormData } from "undici";
 
+const log = new NoOpLog();
 // Only use this shared storage factory when the test doesn't care about storage
 const storageFactory = new MemoryStorageFactory();
 const scriptRunner = new VMScriptRunner();
@@ -1085,6 +1088,37 @@ test("MiniflareCore: dispatchFetch: Request parse files in FormData as File obje
   t.is(await file.text(), "test");
   t.is(file.name, "test.txt");
 });
+test("MiniflareCore: dispatchFetch: creates new request context", async (t) => {
+  const mf = useMiniflareWithHandler(
+    { BindingsPlugin, CachePlugin },
+    {
+      globals: {
+        assertSubrequests(expected: number) {
+          t.is(getRequestContext()?.subrequests, expected);
+        },
+      },
+    },
+    async (globals, req) => {
+      globals.assertSubrequests(0);
+      await globals.caches.default.match("http://localhost/");
+      globals.assertSubrequests(1);
+
+      const n = parseInt(new globals.URL(req.url).searchParams.get("n"));
+      await Promise.all(
+        Array.from(Array(n)).map(() =>
+          globals.caches.default.match("http://localhost/")
+        )
+      );
+      return new globals.Response("body");
+    }
+  );
+  await t.throwsAsync(mf.dispatchFetch("http://localhost/?n=50"), {
+    instanceOf: Error,
+    message: /^Too many subrequests/,
+  });
+  const res = await mf.dispatchFetch("http://localhost/?n=1");
+  t.is(await res.text(), "body");
+});
 
 test("MiniflareCore: dispatchScheduled: dispatches scheduled event", async (t) => {
   const mf = useMiniflare(
@@ -1101,6 +1135,40 @@ test("MiniflareCore: dispatchScheduled: dispatches scheduled event", async (t) =
   const res = await mf.dispatchScheduled(1000, "30 * * * *");
   t.deepEqual(res, [1000, "30 * * * *"]);
 });
+test("MiniflareCore: dispatchScheduled: creates new request context", async (t) => {
+  const mf = new MiniflareCore(
+    { CorePlugin, BindingsPlugin, CachePlugin },
+    { log, storageFactory, scriptRunner },
+    {
+      globals: {
+        assertSubrequests(expected: number) {
+          t.is(getRequestContext()?.subrequests, expected);
+        },
+      },
+      modules: true,
+      script: `export default {
+        async scheduled(controller) {
+          assertSubrequests(0);
+          await caches.default.match("http://localhost/");
+          assertSubrequests(1);
+          
+          await Promise.all(
+            Array.from(Array(controller.scheduledTime)).map(() =>
+              caches.default.match("http://localhost/")
+            )
+          );
+          return true;
+        }
+      }`,
+    }
+  );
+  await t.throwsAsync(mf.dispatchScheduled(50), {
+    instanceOf: Error,
+    message: /^Too many subrequests/,
+  });
+  const waitUntil = await mf.dispatchScheduled(1);
+  t.true(waitUntil[0]);
+});
 
 test("MiniflareCore: dispose: runs dispose for all plugins", async (t) => {
   const log = new TestLog();

packages/core/test/standards/http.spec.ts

@@ -1,4 +1,5 @@
 import assert from "assert";
+import http from "http";
 import { text } from "stream/consumers";
 import { ReadableStream, TransformStream, WritableStream } from "stream/web";
 import { URL } from "url";
@@ -7,6 +8,7 @@ import {
   IncomingRequestCfProperties,
   Request,
   Response,
+  _getURLList,
   _isByteStream,
   createCompatFetch,
   fetch,
@@ -16,7 +18,13 @@ import {
   withStringFormDataFiles,
   withWaitUntil,
 } from "@miniflare/core";
-import { Compatibility, InputGate, LogLevel, NoOpLog } from "@miniflare/shared";
+import {
+  Compatibility,
+  InputGate,
+  LogLevel,
+  NoOpLog,
+  RequestContext,
+} from "@miniflare/shared";
 import {
   TestLog,
   triggerPromise,
@@ -35,6 +43,7 @@ import {
   File,
   FormData,
   Headers,
+  fetch as baseFetch,
 } from "undici";
 
 // @ts-expect-error filling out all properties is annoying
@@ -804,12 +813,51 @@ test("withWaitUntil: adds wait until to (Base)Response", async (t) => {
   t.is(await res.waitUntil(), baseWaitUntil);
 });
 
+function redirectingServerListener(
+  req: http.IncomingMessage,
+  res: http.ServerResponse
+) {
+  const { searchParams } = new URL(req.url ?? "", "http://localhost");
+  const n = parseInt(searchParams.get("n") ?? "0");
+  if (n > 0) {
+    res.writeHead(302, { Location: `/?n=${n - 1}` });
+  } else {
+    res.writeHead(200);
+  }
+  res.end();
+}
+test("_getURLList: extracts URL list from Response", async (t) => {
+  const upstream = (await useServer(t, redirectingServerListener)).http;
+  const url = new URL("/?n=3", upstream);
+  const res = await baseFetch(url);
+  const urlList = _getURLList(res);
+  t.deepEqual(urlList?.map(String), [
+    `${upstream.origin}/?n=3`,
+    `${upstream.origin}/?n=2`,
+    `${upstream.origin}/?n=1`,
+    `${upstream.origin}/?n=0`,
+  ]);
+});
+
 test("fetch: can fetch from existing Request", async (t) => {
   const upstream = (await useServer(t, (req, res) => res.end("upstream"))).http;
   const req = new Request(upstream);
   const res = await fetch(req);
   t.is(await res.text(), "upstream");
 });
+test("fetch: increments subrequest count", async (t) => {
+  const upstream = (await useServer(t, (req, res) => res.end("upstream"))).http;
+  const ctx = new RequestContext();
+  await ctx.runWith(() => fetch(upstream));
+  t.is(ctx.subrequests, 1);
+});
+test("fetch: increments subrequest count for each redirect", async (t) => {
+  const upstream = (await useServer(t, redirectingServerListener)).http;
+  const url = new URL("/?n=3", upstream);
+  const ctx = new RequestContext();
+  await ctx.runWith(() => fetch(url));
+  t.is(ctx.subrequests, 4);
+});
 test("fetch: waits for output gate to open before fetching", async (t) => {
   let fetched = false;
   const upstream = (

packages/durable-objects/package.json

@@ -41,6 +41,8 @@
     "undici": "^4.11.1"
   },
   "devDependencies": {
+    "@miniflare/cache": "2.0.0-rc.3",
+    "@miniflare/runner-vm": "2.0.0-rc.3",
     "@miniflare/shared-test": "2.0.0-rc.3"
   }
 }