From 62c9e9e656d34411263365fc1943132bdc72c521 Mon Sep 17 00:00:00 2001
From: Amit Tsur <amit.tsur@gmail.com>
Date: Fri, 22 Dec 2023 11:10:37 +0200
Subject: [PATCH 1/2] Now using one function called handleMessage to handle all
 messages

Instead of creating a new function per call.
---
 src/comlink.ts | 38 ++++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/comlink.ts b/src/comlink.ts
index 8896b71d..ccea421e 100644
--- a/src/comlink.ts
+++ b/src/comlink.ts
@@ -388,10 +388,12 @@ function isMessagePort(endpoint: Endpoint): endpoint is MessagePort {
 }
 
 function closeEndPoint(endpoint: Endpoint) {
+  endpoint.removeEventListener("message", handleMessage);
   if (isMessagePort(endpoint)) endpoint.close();
 }
 
 export function wrap<T>(ep: Endpoint, target?: any): Remote<T> {
+  ep.addEventListener("message", handleMessage);
   return createProxy<T>(ep, [], target) as any;
 }
 
@@ -593,25 +595,33 @@ function fromWireValue(value: WireValue): any {
   }
 }
 
+const messageResolvers: Map<string, (value: WireValue | PromiseLike<WireValue>) => void> = new Map();
+
+function handleMessage(ev: Event) {
+    const { data } = ev as MessageEvent;
+    if (!data || !data.id) {
+      return;
+    }
+    const resolver = messageResolvers.get(data.id);
+    if (resolver) {
+        resolver(data);
+        messageResolvers.delete(data.id);
+    }
+}
+
 function requestResponseMessage(
   ep: Endpoint,
   msg: Message,
   transfers?: Transferable[]
 ): Promise<WireValue> {
-  return new Promise((resolve) => {
-    const id = generateUUID();
-    ep.addEventListener("message", function l(ev: MessageEvent) {
-      if (!ev.data || !ev.data.id || ev.data.id !== id) {
-        return;
-      }
-      ep.removeEventListener("message", l as any);
-      resolve(ev.data);
-    } as any);
-    if (ep.start) {
-      ep.start();
-    }
-    ep.postMessage({ id, ...msg }, transfers);
-  });
+    return new Promise((resolve) => {
+        const id = generateUUID();
+        messageResolvers.set(id, resolve);
+        if (ep.start) {
+          ep.start();
+        }
+        ep.postMessage({ id, ...msg }, transfers);
+    });
 }
 
 function generateUUID(): string {

From dc0881df8061adcb7519b1cfa06adafb7996b336 Mon Sep 17 00:00:00 2001
From: Amit Tsur <amit.tsur@gmail.com>
Date: Fri, 22 Dec 2023 17:37:52 +0200
Subject: [PATCH 2/2] Ensuring resolver removal in case of resolver failure.

---
 src/comlink.ts | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/comlink.ts b/src/comlink.ts
index ccea421e..8dcc8615 100644
--- a/src/comlink.ts
+++ b/src/comlink.ts
@@ -604,8 +604,11 @@ function handleMessage(ev: Event) {
     }
     const resolver = messageResolvers.get(data.id);
     if (resolver) {
+      try {
         resolver(data);
+      } finally {
         messageResolvers.delete(data.id);
+      }
     }
 }