Cancel idle RPC readable streams when their consumers stop

.changeset/idle-readable-cancellation.md

@@ -0,0 +1,5 @@
+---
+"@iterate-com/capnweb": patch
+---
+
+Propagate cancellation of idle RPC readable streams immediately, and release stream failure callbacks when their imports are disposed.

__tests__/index.test.ts

@@ -2165,6 +2165,17 @@ describe("error serialization", () => {
 });
 
 describe("onRpcBroken", () => {
+  it("releases broken callbacks when their last stub is disposed", async () => {
+    const harness = new TestHarness(new TestTarget());
+    const stub = await harness.stub.makeCounter(0);
+    const errors: unknown[] = [];
+    stub.onRpcBroken(error => errors.push(error));
+    stub[Symbol.dispose]();
+    harness.clientTransport.forceReceiveError(new Error("disconnected"));
+    await pumpMicrotasks();
+    expect(errors).toEqual([]);
+  });
+
   it("signals when the connection is lost", async () => {
     class TestBroken extends RpcTarget {
       getValue() { return 42; }
@@ -3119,6 +3130,44 @@ describe("transport encoding levels", () => {
 });
 
 describe("ReadableStream over RPC", () => {
+  it("cancels an idle source when its RPC session disconnects", async () => {
+    const cancelled = Promise.withResolvers<unknown>();
+    class StreamProvider extends RpcTarget {
+      getStream() {
+        return new ReadableStream<string>({
+          start(controller) { controller.enqueue("first"); },
+          cancel(reason) { cancelled.resolve(reason); },
+        });
+      }
+    }
+    // The normal harness disposer expects live connections; this test ends both sessions.
+    const harness = new TestHarness(new StreamProvider());
+    const stream = await harness.stub.getStream();
+    const reader = stream.getReader();
+    expect(await reader.read()).toEqual({ done: false, value: "first" });
+    harness.serverTransport.forceReceiveError(new Error("test disconnect"));
+    expect(await cancelled.promise).toEqual(new Error("test disconnect"));
+    reader.releaseLock();
+  });
+
+  it("cancels an idle remote source without waiting for another chunk", async () => {
+    let cancelled = Promise.withResolvers<unknown>();
+    class StreamProvider extends RpcTarget {
+      getStream() {
+        return new ReadableStream<string>({
+          start(controller) { controller.enqueue("first"); },
+          cancel(reason) { cancelled.resolve(reason); },
+        });
+      }
+    }
+    await using harness = new TestHarness(new StreamProvider());
+    const stream = await harness.stub.getStream();
+    const reader = stream.getReader();
+    expect(await reader.read()).toEqual({ done: false, value: "first" });
+    await reader.cancel("finished");
+    expect(await cancelled.promise).toBe("finished");
+  });
+
   it("can send a ReadableStream and read all chunks", async () => {
     let stream = new ReadableStream<string>({
       start(controller) {
@@ -3530,6 +3579,29 @@ describe("ReadableStream over RPC", () => {
 // =======================================================================================
 
 describe("Fetch API types over RPC", () => {
+  it.each([
+    { state: "consumed", viaCallback: false },
+    { state: "locked", viaCallback: false },
+    { state: "consumed", viaCallback: true },
+    { state: "locked", viaCallback: true },
+  ])("rejects a $state Response (callback=$viaCallback) without breaking its session", async ({ state, viaCallback }) => {
+    const getResponse = async () => {
+      const response = new Response("test-body");
+      if (state === "consumed") await response.text();
+      else response.body!.getReader();
+      return response;
+    };
+    class ResponseProvider extends RpcTarget {
+      getResponse() { return getResponse(); }
+      call(callback: () => Promise<Response>) { return callback(); }
+      ping() { return "healthy"; }
+    }
+    await using harness = new TestHarness(new ResponseProvider());
+    using call = viaCallback ? harness.stub.call(getResponse) : harness.stub.getResponse();
+    await expect(call).rejects.toThrow();
+    expect(await harness.stub.ping()).toBe("healthy");
+  });
+
   it("can send Headers over RPC", async () => {
     class HeaderServer extends RpcTarget {
       getHeaders() {

src/rpc.ts

@@ -255,6 +255,10 @@ class ImportTableEntry {
     if (this.resolution) {
       this.resolution.dispose();
     } else {
+      if (this.onBrokenRegistrations) {
+        for (const index of this.onBrokenRegistrations) delete this.session.onBrokenCallbacks[index];
+        this.onBrokenRegistrations = undefined;
+      }
       this.abort(new Error("RPC was canceled because the RpcPromise was disposed."));
       this.sendRelease();
     }
@@ -729,11 +733,30 @@ class RpcSessionImpl implements Importer, Exporter {
     // Create a proxy WritableStream from the import hook and pump the ReadableStream into it.
     let hook = new RpcImportHook(/*isPromise=*/false, entry);
     let writable = streamImpl.createWritableStreamFromHook(hook);
-    readable.pipeTo(writable).catch(() => {
-      // Errors are handled by the writable stream's error handling -- either the write fails
-      // and the writable side reports it, or the readable side errors and pipeTo aborts the
-      // writable side. Either way, the hook's disposal will handle cleanup.
-    }).finally(() => readableHook.dispose());
+    const reader = readable.getReader();
+    const writer = writable.getWriter();
+    // Some Workers runtimes do not cancel an idle source when pipeTo's destination
+    // errors (even with an AbortSignal). Own the reader so peer failure can cancel it.
+    writer.closed.catch(error => { reader.cancel(error).catch(() => {}); });
+    const pump = async () => {
+      try {
+        for (;;) {
+          const chunk = await reader.read();
+          if (chunk.done) break;
+          await writer.write(chunk.value);
+        }
+        await writer.close();
+      } catch (error) {
+        // Initiate both ends' cleanup without waiting for an arbitrary user cancel callback.
+        void reader.cancel(error).catch(() => {});
+        void writer.abort(error).catch(() => {});
+      } finally {
+        reader.releaseLock();
+        writer.releaseLock();
+        readableHook.dispose();
+      }
+    };
+    void pump();
 
     return importId;
   }
@@ -1064,7 +1087,15 @@ class RpcSessionImpl implements Importer, Exporter {
             // ["readable", importId].
             let { readable, writable } = new TransformStream();
             let hook = streamImpl.createWritableStreamHook(writable);
-            this.exports.push({ hook, refcount: 1, pipeReadable: readable });
+            const exportId = this.exports.length;
+            const entry = { hook, refcount: 1, pipeReadable: readable };
+            this.exports.push(entry);
+            hook.onBroken(error => {
+              // Cancellation must reach an idle source, not wait for its next write.
+              if (this.abortReason || this.exports[exportId] !== entry) return;
+              this.send(["reject", exportId,
+                Devaluator.devaluate(error, undefined, this, undefined, this.encodingLevel)]);
+            });
             continue;
           }
 

src/streams.ts

@@ -117,8 +117,11 @@ class WritableStreamStubHook extends StubHook {
   }
 
   onBroken(callback: (error: any) => void): void {
-    // WritableStream stubs don't really have a "broken" state in the same way.
-    // The caller would notice when write/close/abort fails.
+    const state = this.getState();
+    state.writer.closed.catch(error => {
+      // Releasing the lock also rejects `closed`; that is disposal, not peer failure.
+      if (this.state && !state.closed) callback(error);
+    });
   }
 }
 
@@ -331,6 +334,19 @@ function createWritableStreamFromHook(hook: StubHook): WritableStream {
   };
 
   return new WritableStream({
+    start(controller) {
+      hook.onBroken(error => {
+        if (hookDisposed) return;
+        pendingError = error;
+        controller.error(error);
+        if (windowReject) {
+          windowReject(error);
+          windowResolve = undefined;
+          windowReject = undefined;
+        }
+        disposeHook();
+      });
+    },
     write(chunk, controller) {
       // If we already have an error, fail immediately.
       if (pendingError !== undefined) {