Add a new optional AbortController parameter to defer

packages/router/__tests__/router-test.ts

@@ -12597,6 +12597,130 @@ describe("a router", () => {
       expect(shouldRevalidateSpy).not.toHaveBeenCalled();
     });
 
+    it("should allow user to provide an AbortController for defer cancellations after navigation completes", async () => {
+      let t = setup({
+        routes: [
+          {
+            id: "index",
+            index: true,
+            loader: true,
+          },
+          {
+            id: "lazy",
+            path: "lazy",
+            loader: true,
+          },
+        ],
+        hydrationData: { loaderData: { index: "INDEX" } },
+      });
+
+      let A = await t.navigate("/lazy");
+      let dfd = createDeferred();
+      let deferController = new AbortController();
+      await A.loaders.lazy.resolve(
+        defer({ lazy: dfd.promise }, null, deferController)
+      );
+
+      // Interrupt pending deferred's from /lazy navigation
+      let B = await t.navigate("/");
+      expect(A.loaders.lazy.signal.aborted).toBe(false);
+      expect(deferController.signal.aborted).toBe(true);
+
+      await B.loaders.index.resolve("INDEX*");
+      expect(t.router.state).toMatchObject({
+        location: { pathname: "/" },
+        navigation: IDLE_NAVIGATION,
+      });
+    });
+
+    it("should proxy navigation AbortController through to defer AbortController when request aborts prior to defer call", async () => {
+      let t = setup({
+        routes: [
+          {
+            id: "index",
+            index: true,
+            loader: true,
+          },
+          {
+            id: "lazy",
+            path: "lazy",
+            loader: true,
+          },
+        ],
+        hydrationData: { loaderData: { index: "INDEX" } },
+      });
+
+      // Route to lazy
+      let A = await t.navigate("/lazy");
+
+      // Navigate away from /lazy before it's loader returns the defer()
+      let B = await t.navigate("/");
+      expect(A.loaders.lazy.signal.aborted).toBe(true);
+
+      let dfd = createDeferred();
+      let deferController = new AbortController();
+      await A.loaders.lazy.resolve(
+        defer({ lazy: dfd.promise }, null, deferController)
+      );
+      expect(deferController.signal.aborted).toBe(true);
+
+      await B.loaders.index.resolve("INDEX*");
+      expect(t.router.state).toMatchObject({
+        location: { pathname: "/" },
+        navigation: IDLE_NAVIGATION,
+      });
+    });
+
+    it("should proxy navigation AbortController through to defer AbortController when request aborts after defer call", async () => {
+      let t = setup({
+        routes: [
+          {
+            id: "index",
+            index: true,
+            loader: true,
+          },
+          {
+            id: "lazy",
+            path: "lazy",
+            loader: true,
+            children: [
+              {
+                id: "child",
+                path: "child",
+                loader: true,
+              },
+            ],
+          },
+        ],
+        hydrationData: { loaderData: { index: "INDEX" } },
+      });
+
+      // Route to lazy
+      let A = await t.navigate("/lazy/child");
+
+      let dfd = createDeferred();
+      let deferController = new AbortController();
+      await A.loaders.lazy.resolve(
+        defer({ lazy: dfd.promise }, null, deferController)
+      );
+      expect(deferController.signal.aborted).toBe(false);
+
+      // Navigate away from /lazy after it's loader returns the defer() but
+      // before /child's loader finishes so the request is aborted which gets
+      // proxied through to the lazy loaders already resolved defer() instance
+      let B = await t.navigate("/");
+      expect(A.loaders.lazy.signal.aborted).toBe(true);
+      expect(deferController.signal.aborted).toBe(true);
+
+      await A.loaders.child.resolve("CHILD");
+
+      await B.loaders.index.resolve("INDEX*");
+      expect(t.router.state).toMatchObject({
+        location: { pathname: "/" },
+        navigation: IDLE_NAVIGATION,
+      });
+    });
+
     it("does not put resolved deferred's back into a loading state during revalidation", async () => {
       let shouldRevalidateSpy = jest.fn(() => false);
       let t = setup({

packages/router/router.ts

@@ -9,6 +9,7 @@ import {
 } from "./history";
 import type {
   ActionFunction,
+  ActionFunctionArgs,
   AgnosticDataRouteMatch,
   AgnosticDataRouteObject,
   AgnosticRouteMatch,
@@ -23,6 +24,7 @@ import type {
   HTMLFormMethod,
   ImmutableRouteKey,
   LoaderFunction,
+  LoaderFunctionArgs,
   MapRoutePropertiesFunction,
   MutationFormMethod,
   RedirectResult,
@@ -3627,8 +3629,25 @@ async function callLoaderOrAction(
     let abortPromise = new Promise((_, r) => (reject = r));
     onReject = () => reject();
     request.signal.addEventListener("abort", onReject);
+
+    let runHandlerAndMaybeAbortDeferreds = async (
+      handler: ActionFunction | LoaderFunction,
+      args: LoaderFunctionArgs | ActionFunctionArgs
+    ) => {
+      // Since we Promise.race the handler against the request signal, if the
+      // request is aborted we never get access to the returned value from the
+      // handler.  If we get a defer() instance _after_ we've already aborted
+      // the request, we should proxy that cancellation along so the defer()
+      // controller also aborts.
+      let result = await handler(args);
+      if (args.request.signal.aborted && isDeferredData(result)) {
+        result.cancel();
+      }
+      return result;
+    };
+
     return Promise.race([
-      handler({
+      runHandlerAndMaybeAbortDeferreds(handler, {
         request,
         params: match.params,
         context: opts.requestContext,
@@ -3792,6 +3811,17 @@ async function callLoaderOrAction(
   }
 
   if (isDeferredData(result)) {
+    if (request.signal.aborted) {
+      // Don't think this is technically possible since we race the loader
+      // against the request signal and we would short circuit via the error path
+      // above on interruption. Included to be safe though :)
+      result.cancel();
+    } else {
+      // Abort the defer() instance if we abort while waiting on other loaders
+      // from this navigation
+      let deferResult = result;
+      request.signal.addEventListener("abort", () => deferResult.cancel());
+    }
     return {
       type: ResultType.deferred,
       deferredData: result,

packages/router/utils.ts

@@ -1276,7 +1276,11 @@ export class DeferredData {
   init?: ResponseInit;
   deferredKeys: string[] = [];
 
-  constructor(data: Record<string, unknown>, responseInit?: ResponseInit) {
+  constructor(
+    data: Record<string, unknown>,
+    responseInit: ResponseInit,
+    controller: AbortController
+  ) {
     invariant(
       data && typeof data === "object" && !Array.isArray(data),
       "defer() only accepts plain objects"
@@ -1286,7 +1290,7 @@ export class DeferredData {
     // cancellation
     let reject: (e: AbortedDeferredError) => void;
     this.abortPromise = new Promise((_, r) => (reject = r));
-    this.controller = new AbortController();
+    this.controller = controller;
     let onAbort = () =>
       reject(new AbortedDeferredError("Deferred data aborted"));
     this.unlistenAbortSignal = () =>
@@ -1455,13 +1459,18 @@ function unwrapTrackedPromise(value: any) {
 
 export type DeferFunction = (
   data: Record<string, unknown>,
-  init?: number | ResponseInit
+  init?: number | ResponseInit | null,
+  controller?: AbortController
 ) => DeferredData;
 
-export const defer: DeferFunction = (data, init = {}) => {
-  let responseInit = typeof init === "number" ? { status: init } : init;
-
-  return new DeferredData(data, responseInit);
+export const defer: DeferFunction = (data, init, controller) => {
+  let responseInit =
+    typeof init === "number" ? { status: init } : init != null ? init : {};
+  return new DeferredData(
+    data,
+    responseInit,
+    controller || new AbortController()
+  );
 };
 
 export type RedirectFunction = (