Fix disposed RPC callback registrations while preserving notification order

.changeset/rpc-callback-lifecycle.md

@@ -0,0 +1,5 @@
+---
+"capnweb": patch
+---
+
+Release disposed RPC brokenness callback registrations while preserving notification order and reentrant abort behavior.

__tests__/on-broken.node.mjs

@@ -0,0 +1,105 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import {RpcSession, RpcTarget} from "../dist/index.js";
+const settle = async () => { for (let i=0; i<4; i++) await new Promise(setImmediate); };
+function transportPair() {
+  const boxes = Array.from({length:2}, () => ({queue:[], waiter:null, closed:null}));
+  return boxes.map((inbox,i) => ({
+    send(message) { const out=boxes[1-i]; if(out.waiter) { out.waiter.resolve(message); out.waiter=null; } else out.queue.push(message); },
+    receive() { if(inbox.queue.length) return Promise.resolve(inbox.queue.shift()); if(inbox.closed) return Promise.reject(inbox.closed); inbox.waiter=Promise.withResolvers(); return inbox.waiter.promise; },
+    failReceive(reason) { inbox.closed=reason; inbox.waiter?.reject(reason); inbox.waiter=null; },
+    abort(reason) { this.failReceive(reason); this.onAbort?.(reason); },
+  }));
+}
+class Thing extends RpcTarget { ping() { return 'ok'; } }
+class Main extends RpcTarget { make() { return new Thing(); } delayed() { return this.gate.promise; } fail() { throw Error('remote failure'); } }
+function fixture() {
+  const [clientTransport, serverTransport] = transportPair();
+  const target = new Main(); target.gate=Promise.withResolvers();
+  const server = new RpcSession(serverTransport,target);
+  const client = new RpcSession(clientTransport);
+  const main=client.getRemoteMain();
+  return {client, server, main, target, clientTransport,
+    async breakClient() { clientTransport.failReceive(Error('connection lost')); await settle(); },
+    async close() { clientTransport.failReceive(Error('test done')); serverTransport.failReceive(Error('test done')); await settle(); },
+  };
+}
+test('final reference disposal unregisters callbacks and remains idempotent', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  let fired=0; const thing=await f.main.make(); thing.onRpcBroken(()=>fired++);
+  const second=thing.dup(); thing[Symbol.dispose](); second[Symbol.dispose](); second[Symbol.dispose]();
+  await settle(); assert.deepEqual(f.client.getStats(),{imports:1,exports:1});
+  await f.breakClient(); assert.equal(fired,0);
+ } finally { await f.close(); }
+});
+test('same-session resolution preserves global callback registration order', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  const seen=[]; const pending=f.main.delayed(); pending.onRpcBroken(()=>seen.push('first'));
+  const other=await f.main.make(); other.onRpcBroken(()=>seen.push('middle'));
+  pending.onRpcBroken(()=>seen.push('last'));
+  f.target.gate.resolve(new Thing()); const resolved=await pending;
+  assert.equal(await resolved.ping(),'ok'); await f.breakClient();
+  assert.deepEqual(seen,['first','middle','last']);
+ } finally { await f.close(); }
+});
+test('session abort notifies pending entries in order despite reentrant disposal and consumer throw', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  const seen=[]; const a=await f.main.make(), b=await f.main.make(), c=await f.main.make();
+  a.onRpcBroken(()=>{seen.push('a'); b[Symbol.dispose](); c[Symbol.dispose]();});
+  b.onRpcBroken(()=>{seen.push('b'); throw Error('consumer failure');});
+  c.onRpcBroken(()=>seen.push('c'));
+  await f.breakClient(); assert.deepEqual(seen,['a','b','c']);
+ } finally { await f.close(); }
+});
+
+test('same-session transferred callbacks unregister after all resolved references dispose', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  let fired=0; const pending=f.main.delayed(); pending.onRpcBroken(()=>fired++);
+  f.target.gate.resolve(new Thing()); const resolved=await pending; const second=resolved.dup();
+  resolved[Symbol.dispose](); second[Symbol.dispose](); await settle();
+  assert.deepEqual(f.client.getStats(),{imports:1,exports:1});
+  await f.breakClient(); assert.equal(fired,0);
+ } finally { await f.close(); }
+});
+test('resolve callback consumer throw does not strand the RPC result or abort other capabilities', {timeout:2000}, async()=>{
+ const f=fixture(); let timer; try {
+  const other=await f.main.make(), seen=[];
+  const pending=f.main.fail(); pending.onRpcBroken(()=>{seen.push('first'); throw Error('consumer failure');});
+  pending.onRpcBroken(()=>seen.push('last'));
+  const result=await Promise.race([
+   Promise.resolve(pending).then(()=>({status:'fulfilled'}),e=>({status:'rejected',message:e.message})),
+   new Promise(resolve=>{timer=setTimeout(()=>resolve({status:'timeout'}),150);}),
+  ]);
+  clearTimeout(timer); assert.deepEqual(result,{status:'rejected',message:'remote failure'});
+  assert.deepEqual(seen,['first','last']); assert.equal(await other.ping(),'ok');
+ } finally { clearTimeout(timer); await f.close(); }
+});
+test('resolve callback reentrant session abort notifies later callbacks exactly once in original order', {timeout:2000}, async()=>{
+ const f=fixture(); let timer; try {
+  const other=await f.main.make(), seen=[];
+  const pending=f.main.fail(); pending.onRpcBroken(()=>{seen.push('first'); f.main[Symbol.dispose]();});
+  other.onRpcBroken(()=>seen.push('middle')); pending.onRpcBroken(()=>seen.push('last'));
+  const result=await Promise.race([
+   Promise.resolve(pending).then(()=> 'fulfilled',()=> 'rejected'),
+   new Promise(resolve=>{timer=setTimeout(()=>resolve('timeout'),150);}),
+  ]);
+  clearTimeout(timer); assert.equal(result,'rejected');
+  assert.deepEqual(seen,['first','middle','last']);
+ } finally { clearTimeout(timer); await f.close(); }
+});
+test('transport abort reentrant disposal cannot remove pending notifications', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  const seen=[]; const a=await f.main.make(), b=await f.main.make();
+  a.onRpcBroken(()=>seen.push('a')); b.onRpcBroken(()=>seen.push('b'));
+  f.clientTransport.onAbort=()=>b[Symbol.dispose]();
+  await f.breakClient(); assert.deepEqual(seen,['a','b']);
+ } finally { await f.close(); }
+});
+test('registering a broken callback during session abort reports the error immediately', {timeout:2000}, async()=>{
+ const f=fixture(); try {
+  const seen=[]; const a=await f.main.make(), b=await f.main.make();
+  a.onRpcBroken(()=>{seen.push('a'); b.onRpcBroken(()=>seen.push('new'));});
+  b.onRpcBroken(()=>seen.push('b'));
+  await f.breakClient(); assert.deepEqual(seen,['a','new','b']);
+ } finally { await f.close(); }
+});

src/rpc.ts

@@ -182,6 +182,8 @@ type ExportTableEntry = {
   pipeReadable?: ReadableStream
 };
 
+type OnBrokenRegistration = {callback: (error: any) => void, index: number};
+
 // Entry on the imports table.
 class ImportTableEntry {
   constructor(public session: RpcSessionImpl, public importId: number, pulling: boolean) {
@@ -196,9 +198,9 @@ class ImportTableEntry {
   private activePull?: PromiseWithResolvers<void>;
   public resolution?: StubHook;
 
-  // List of integer indexes into session.onBrokenCallbacks which are callbacks registered on
+  // List of registrations in session.onBrokenCallbacks which are callbacks registered on
   // this import. Initialized on first use (so `undefined` is the same as an empty list).
-  private onBrokenRegistrations?: number[];
+  private onBrokenRegistrations?: OnBrokenRegistration[];
 
   resolve(resolution: StubHook) {
     // TODO: Need embargo handling here? PayloadStubHook needs to be wrapped in a
@@ -218,26 +220,25 @@ class ImportTableEntry {
     this.resolution = resolution;
     this.sendRelease();
 
-    if (this.onBrokenRegistrations) {
-      // Delete all our callback registrations from this session and re-register them on the
-      // target stub.
-      for (let i of this.onBrokenRegistrations) {
-        let callback = this.session.onBrokenCallbacks[i];
+    let registrations = this.onBrokenRegistrations;
+    this.onBrokenRegistrations = undefined;
+    if (registrations) {
+      for (let registration of registrations) {
+        let {index, callback} = registration;
+        // A previous callback may have aborted the session and already notified the rest.
+        if (this.session.onBrokenCallbacks[index] !== registration) continue;
         let endIndex = this.session.onBrokenCallbacks.length;
+        // Unregister before calling application code, which may throw or abort the session.
+        delete this.session.onBrokenCallbacks[index];
         resolution.onBroken(callback);
-        if (this.session.onBrokenCallbacks[endIndex] === callback) {
-          // Oh, calling onBroken() just registered the callback back on this connection again.
-          // But when the connection dies, we want all the callbacks to be called in the order in
-          // which they were registered. So we don't want this one pushed to the back of the line
-          // here. So, let's remove the newly-added registration and keep the original.
-          // TODO: This is quite hacky, think about whether this is really the right answer.
+        let moved = this.session.onBrokenCallbacks[endIndex];
+        if (moved?.callback === callback) {
+          // Keep global registration order and update the new owner's cleanup index.
+          moved.index = index;
+          this.session.onBrokenCallbacks[index] = moved;
           delete this.session.onBrokenCallbacks[endIndex];
-        } else {
-          // The callback is now registered elsewhere, so delete it from our session.
-          delete this.session.onBrokenCallbacks[i];
         }
       }
-      this.onBrokenRegistrations = undefined;
     }
 
     if (this.activePull) {
@@ -273,8 +274,9 @@ class ImportTableEntry {
         this.activePull = undefined;
       }
 
-      // The RpcSession itself will have called all our callbacks so we don't need to track the
-      // registrations anymore.
+      for (let registration of this.onBrokenRegistrations ?? []) {
+        delete this.session.onBrokenCallbacks[registration.index];
+      }
       this.onBrokenRegistrations = undefined;
     }
   }
@@ -283,11 +285,11 @@ class ImportTableEntry {
     if (this.resolution) {
       this.resolution.onBroken(callback);
     } else {
-      let index = this.session.onBrokenCallbacks.length;
-      this.session.onBrokenCallbacks.push(callback);
-
-      if (!this.onBrokenRegistrations) this.onBrokenRegistrations = [];
-      this.onBrokenRegistrations.push(index);
+      let registration = this.session.registerOnBroken(callback);
+      if (registration) {
+        if (!this.onBrokenRegistrations) this.onBrokenRegistrations = [];
+        this.onBrokenRegistrations.push(registration);
+      }
     }
   }
 
@@ -486,7 +488,7 @@ class RpcSessionImpl implements Importer, Exporter {
 
   // Sparse array of onBrokenCallback registrations. Items are strictly appended to the end but
   // may be deleted from the middle (hence leaving the array sparse).
-  onBrokenCallbacks: ((error: any) => void)[] = [];
+  onBrokenCallbacks: OnBrokenRegistration[] = [];
 
   // Encoding level from the transport (defaults to "string")
   private encodingLevel: EncodingLevel;
@@ -526,6 +528,16 @@ class RpcSessionImpl implements Importer, Exporter {
     this.readLoop().catch(err => this.abort(err));
   }
 
+  registerOnBroken(callback: (error: any) => void): OnBrokenRegistration | undefined {
+    if (this.abortReason !== undefined) {
+      new ErrorStubHook(this.abortReason).onBroken(callback);
+      return;
+    }
+    let registration = {callback, index: this.onBrokenCallbacks.length};
+    this.onBrokenCallbacks.push(registration);
+    return registration;
+  }
+
   // Should only be called once immediately after construction.
   getMainImport(): RpcImportHook {
     return new RpcMainHook(this.imports[0]);
@@ -952,6 +964,9 @@ class RpcSessionImpl implements Importer, Exporter {
     }
 
     this.abortReason = error;
+    // Detach before application code can re-enter disposal and delete pending callbacks.
+    let onBrokenCallbacks = this.onBrokenCallbacks;
+    this.onBrokenCallbacks = [];
     if (this.onBatchDone) {
       this.onBatchDone.reject(error);
     }
@@ -968,9 +983,9 @@ class RpcSessionImpl implements Importer, Exporter {
 
     // WATCH OUT: these are sparse arrays. `for/let/of` will iterate only positive indexes
     // including deleted indexes -- bad. We need to use `for/let/in` instead.
-    for (let i in this.onBrokenCallbacks) {
+    for (let i in onBrokenCallbacks) {
       try {
-        this.onBrokenCallbacks[i](error);
+        onBrokenCallbacks[i].callback(error);
       } catch (err) {
         // Treat as unhandled rejection.
         Promise.resolve(err);