From 4c3e4504683a708f75805b81139f4be0a16c84d6 Mon Sep 17 00:00:00 2001
From: pelikhan <jhalleux@microsoft.com>
Date: Mon, 13 Feb 2023 11:03:04 -0800
Subject: [PATCH] fix: add token to unregister await handler

---
 src/jdom/eventsource.ts | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/jdom/eventsource.ts b/src/jdom/eventsource.ts
index 74e5b22aa..393ef32c4 100644
--- a/src/jdom/eventsource.ts
+++ b/src/jdom/eventsource.ts
@@ -63,6 +63,25 @@ export function dependencyId(nodes: IEventSource | IEventSource[]) {
 
 let nextNodeId = 0
 
+/**
+ * Collects unsubscribe handlers and unmounts them
+ */
+export class JDCancellationToken {
+    private subscriptions: (() => void)[] = []
+
+    constructor() {}
+
+    mount(...items: (() => void)[]) {
+        this.subscriptions.push(...items)
+    }
+
+    unmount() {
+        const us = this.subscriptions
+        this.subscriptions = []
+        us.forEach(unsub => unsub())
+    }
+}
+
 /**
  * Base class for evented nodes in Jacdac
  * @category JDOM
@@ -141,12 +160,16 @@ export class JDEventSource implements IEventSource {
      * @param eventName
      * @param timeout
      */
-    async awaitOnce(eventName: string | string[]): Promise<void> {
+    async awaitOnce(
+        eventName: string | string[],
+        token?: JDCancellationToken
+    ): Promise<void> {
         if (!eventName) return
 
         const p = new Promise<void>(resolve => {
             const handler = () => resolve?.()
             this.once(eventName, handler)
+            token?.mount(() => this.off(eventName, handler))
         })
         return p
     }