Merge pull request #678 from foxglove/achim/fix-messageport-close

Fix message ports not being closed when proxy is relased

Author: benjamind

src/comlink.ts

@@ -231,6 +231,10 @@ type PendingListenersMap = Map<
   string,
   (value: WireValue | PromiseLike<WireValue>) => void
 >;
+type EndpointWithPendingListeners = {
+  endpoint: Endpoint;
+  pendingListeners: PendingListenersMap;
+};
 
 /**
  * Internal transfer handler to handle thrown exceptions.
@@ -415,7 +419,7 @@ export function wrap<T>(ep: Endpoint, target?: any): Remote<T> {
     }
   });
 
-  return createProxy<T>(ep, pendingListeners, [], target) as any;
+  return createProxy<T>({ endpoint: ep, pendingListeners }, [], target) as any;
 }
 
 function throwIfProxyReleased(isReleased: boolean) {
@@ -424,11 +428,11 @@ function throwIfProxyReleased(isReleased: boolean) {
   }
 }
 
-function releaseEndpoint(ep: Endpoint) {
-  return requestResponseMessage(ep, new Map(), {
+function releaseEndpoint(epWithPendingListeners: EndpointWithPendingListeners) {
+  return requestResponseMessage(epWithPendingListeners, {
     type: MessageType.RELEASE,
   }).then(() => {
-    closeEndPoint(ep);
+    closeEndPoint(epWithPendingListeners.endpoint);
   });
 }
 
@@ -441,24 +445,31 @@ interface FinalizationRegistry<T> {
   ): void;
   unregister(unregisterToken: object): void;
 }
-declare var FinalizationRegistry: FinalizationRegistry<Endpoint>;
+declare var FinalizationRegistry: FinalizationRegistry<EndpointWithPendingListeners>;
 
-const proxyCounter = new WeakMap<Endpoint, number>();
+const proxyCounter = new WeakMap<EndpointWithPendingListeners, number>();
 const proxyFinalizers =
   "FinalizationRegistry" in globalThis &&
-  new FinalizationRegistry((ep: Endpoint) => {
-    const newCount = (proxyCounter.get(ep) || 0) - 1;
-    proxyCounter.set(ep, newCount);
-    if (newCount === 0) {
-      releaseEndpoint(ep);
+  new FinalizationRegistry(
+    (epWithPendingListeners: EndpointWithPendingListeners) => {
+      const newCount = (proxyCounter.get(epWithPendingListeners) || 0) - 1;
+      proxyCounter.set(epWithPendingListeners, newCount);
+      if (newCount === 0) {
+        releaseEndpoint(epWithPendingListeners).finally(() => {
+          epWithPendingListeners.pendingListeners.clear();
+        });
+      }
     }
-  });
+  );
 
-function registerProxy(proxy: object, ep: Endpoint) {
-  const newCount = (proxyCounter.get(ep) || 0) + 1;
-  proxyCounter.set(ep, newCount);
+function registerProxy(
+  proxy: object,
+  epWithPendingListeners: EndpointWithPendingListeners
+) {
+  const newCount = (proxyCounter.get(epWithPendingListeners) || 0) + 1;
+  proxyCounter.set(epWithPendingListeners, newCount);
   if (proxyFinalizers) {
-    proxyFinalizers.register(proxy, ep, proxy);
+    proxyFinalizers.register(proxy, epWithPendingListeners, proxy);
   }
 }
 
@@ -469,8 +480,7 @@ function unregisterProxy(proxy: object) {
 }
 
 function createProxy<T>(
-  ep: Endpoint,
-  pendingListeners: PendingListenersMap,
+  epWithPendingListeners: EndpointWithPendingListeners,
   path: (string | number | symbol)[] = [],
   target: object = function () {}
 ): Remote<T> {
@@ -481,31 +491,31 @@ function createProxy<T>(
       if (prop === releaseProxy) {
         return () => {
           unregisterProxy(proxy);
-          releaseEndpoint(ep);
-          pendingListeners.clear();
+          releaseEndpoint(epWithPendingListeners).finally(() => {
+            epWithPendingListeners.pendingListeners.clear();
+          });
           isProxyReleased = true;
         };
       }
       if (prop === "then") {
         if (path.length === 0) {
           return { then: () => proxy };
         }
-        const r = requestResponseMessage(ep, pendingListeners, {
+        const r = requestResponseMessage(epWithPendingListeners, {
           type: MessageType.GET,
           path: path.map((p) => p.toString()),
         }).then(fromWireValue);
         return r.then.bind(r);
       }
-      return createProxy(ep, pendingListeners, [...path, prop]);
+      return createProxy(epWithPendingListeners, [...path, prop]);
     },
     set(_target, prop, rawValue) {
       throwIfProxyReleased(isProxyReleased);
       // FIXME: ES6 Proxy Handler `set` methods are supposed to return a
       // boolean. To show good will, we return true asynchronously ¯\_(ツ)_/¯
       const [value, transferables] = toWireValue(rawValue);
       return requestResponseMessage(
-        ep,
-        pendingListeners,
+        epWithPendingListeners,
         {
           type: MessageType.SET,
           path: [...path, prop].map((p) => p.toString()),
@@ -518,18 +528,17 @@ function createProxy<T>(
       throwIfProxyReleased(isProxyReleased);
       const last = path[path.length - 1];
       if ((last as any) === createEndpoint) {
-        return requestResponseMessage(ep, pendingListeners, {
+        return requestResponseMessage(epWithPendingListeners, {
           type: MessageType.ENDPOINT,
         }).then(fromWireValue);
       }
       // We just pretend that `bind()` didn’t happen.
       if (last === "bind") {
-        return createProxy(ep, pendingListeners, path.slice(0, -1));
+        return createProxy(epWithPendingListeners, path.slice(0, -1));
       }
       const [argumentList, transferables] = processArguments(rawArgumentList);
       return requestResponseMessage(
-        ep,
-        pendingListeners,
+        epWithPendingListeners,
         {
           type: MessageType.APPLY,
           path: path.map((p) => p.toString()),
@@ -542,8 +551,7 @@ function createProxy<T>(
       throwIfProxyReleased(isProxyReleased);
       const [argumentList, transferables] = processArguments(rawArgumentList);
       return requestResponseMessage(
-        ep,
-        pendingListeners,
+        epWithPendingListeners,
         {
           type: MessageType.CONSTRUCT,
           path: path.map((p) => p.toString()),
@@ -553,7 +561,7 @@ function createProxy<T>(
       ).then(fromWireValue);
     },
   });
-  registerProxy(proxy, ep);
+  registerProxy(proxy, epWithPendingListeners);
   return proxy as any;
 }
 
@@ -622,11 +630,12 @@ function fromWireValue(value: WireValue): any {
 }
 
 function requestResponseMessage(
-  ep: Endpoint,
-  pendingListeners: PendingListenersMap,
+  epWithPendingListeners: EndpointWithPendingListeners,
   msg: Message,
   transfers?: Transferable[]
 ): Promise<WireValue> {
+  const ep = epWithPendingListeners.endpoint;
+  const pendingListeners = epWithPendingListeners.pendingListeners;
   return new Promise((resolve) => {
     const id = generateUUID();
     pendingListeners.set(id, resolve);

tests/node/main.mjs

@@ -24,5 +24,17 @@ describe("node", () => {
       const otherProxy = Comlink.wrap(otherEp);
       expect(await otherProxy(20, 1)).to.equal(21);
     });
+
+    it("releaseProxy closes MessagePort created by createEndpoint", async function () {
+      const proxy = Comlink.wrap(nodeEndpoint(this.worker));
+      const otherEp = await proxy[Comlink.createEndpoint]();
+      const otherProxy = Comlink.wrap(otherEp);
+      expect(await otherProxy(20, 1)).to.equal(21);
+
+      await new Promise((resolve) => {
+        otherEp.close = resolve; // Resolve the promise when the MessagePort is closed.
+        otherProxy[Comlink.releaseProxy](); // Release the proxy, which should close the MessagePort.
+      });
+    });
   });
 });

tests/worker.comlink.test.js

@@ -33,4 +33,16 @@ describe("Comlink across workers", function () {
     const otherProxy = Comlink.wrap(otherEp);
     expect(await otherProxy(20, 1)).to.equal(21);
   });
+
+  it("releaseProxy closes MessagePort created by createEndpoint", async function () {
+    const proxy = Comlink.wrap(this.worker);
+    const otherEp = await proxy[Comlink.createEndpoint]();
+    const otherProxy = Comlink.wrap(otherEp);
+    expect(await otherProxy(20, 1)).to.equal(21);
+
+    await new Promise((resolve) => {
+      otherEp.close = resolve; // Resolve the promise when the MessagePort is closed.
+      otherProxy[Comlink.releaseProxy](); // Release the proxy, which should close the MessagePort.
+    });
+  });
 });