Unsubscribe from worker observables when master unsubscribes (#263)

* Proxy `.unsubscribe()` calls to worker observables

* Fix tiny-worker message handling bug

Author: andywer

src/master/invocation-proxy.ts

@@ -17,6 +17,7 @@ import {
   Worker as WorkerType
 } from "../types/master"
 import {
+  MasterJobCancelMessage,
   MasterJobRunMessage,
   MasterMessageType,
   WorkerJobErrorMessage,
@@ -71,8 +72,19 @@ function createObservableForJob<ResultType>(worker: WorkerType, jobUID: number):
         worker.removeEventListener("message", messageHandler)
       }
     }) as EventListener
+
     worker.addEventListener("message", messageHandler)
-    return () => worker.removeEventListener("message", messageHandler)
+
+    return () => {
+      if (asyncType === "observable" || !asyncType) {
+        const cancelMessage: MasterJobCancelMessage = {
+          type: MasterMessageType.cancel,
+          uid: jobUID
+        }
+        worker.postMessage(cancelMessage)
+      }
+      worker.removeEventListener("message", messageHandler)
+    }
   })
 }
 

src/types/messages.ts

@@ -9,17 +9,23 @@ export interface SerializedError {
 // Messages sent by master:
 
 export enum MasterMessageType {
+  cancel = "cancel",
   run = "run"
 }
 
+export type MasterJobCancelMessage = {
+  type: MasterMessageType.cancel,
+  uid: number
+}
+
 export type MasterJobRunMessage = {
   type: MasterMessageType.run,
   uid: number,
   method?: string,
   args: any[]
 }
 
-export type MasterSentMessage = MasterJobRunMessage
+export type MasterSentMessage = MasterJobCancelMessage | MasterJobRunMessage
 
 ////////////////////////////
 // Messages sent by worker:

src/worker/implementation.tiny-worker.ts

@@ -24,14 +24,22 @@ const postMessageToMaster: AbstractedWorkerAPI["postMessageToMaster"] = function
   self.postMessage(data)
 }
 
+let muxingHandlerSetUp = false
+const messageHandlers = new Set<(data: any) => void>()
+
 const subscribeToMasterMessages: AbstractedWorkerAPI["subscribeToMasterMessages"] = function subscribeToMasterMessages(onMessage) {
-  const messageHandler = (messageEvent: MessageEvent) => {
-    onMessage(messageEvent.data)
-  }
-  const unsubscribe = () => {
-    self.removeEventListener("message", messageHandler as EventListener)
+  if (!muxingHandlerSetUp) {
+    // We have one multiplexing message handler as tiny-worker's
+    // addEventListener() only allows you to set a single message handler
+    self.addEventListener("message", ((event: MessageEvent) => {
+      messageHandlers.forEach(handler => handler(event.data))
+    }) as EventListener)
+    muxingHandlerSetUp = true
   }
-  self.addEventListener("message", messageHandler as EventListener)
+
+  messageHandlers.add(onMessage)
+
+  const unsubscribe = () => messageHandlers.delete(onMessage)
   return unsubscribe
 }
 

src/worker/index.ts

@@ -1,8 +1,9 @@
 import isSomeObservable from "is-observable"
-import { Observable } from "observable-fns"
+import { Observable, Subscription } from "observable-fns"
 import { deserialize, serialize } from "../common"
 import { isTransferDescriptor, TransferDescriptor } from "../transferable"
 import {
+  MasterJobCancelMessage,
   MasterJobRunMessage,
   MasterMessageType,
   SerializedError,
@@ -24,6 +25,9 @@ export const isWorkerRuntime = Implementation.isWorkerRuntime
 
 let exposeCalled = false
 
+const activeSubscriptions = new Map<number, Subscription<any>>()
+
+const isMasterJobCancelMessage = (thing: any): thing is MasterJobCancelMessage => thing && thing.type === MasterMessageType.cancel
 const isMasterJobRunMessage = (thing: any): thing is MasterJobRunMessage => thing && thing.type === MasterMessageType.run
 
 /**
@@ -124,11 +128,18 @@ async function runFunction(jobUID: number, fn: WorkerFunction, args: any[]) {
   postJobStartMessage(jobUID, resultType)
 
   if (isObservable(syncResult)) {
-    syncResult.subscribe(
+    const subscription = syncResult.subscribe(
       value => postJobResultMessage(jobUID, false, serialize(value)),
-      error => postJobErrorMessage(jobUID, serialize(error) as any),
-      () => postJobResultMessage(jobUID, true)
+      error => {
+        postJobErrorMessage(jobUID, serialize(error) as any)
+        activeSubscriptions.delete(jobUID)
+      },
+      () => {
+        postJobResultMessage(jobUID, true)
+        activeSubscriptions.delete(jobUID)
+      }
     )
+    activeSubscriptions.set(jobUID, subscription)
   } else {
     try {
       const result = await syncResult
@@ -174,6 +185,18 @@ export function expose(exposed: WorkerFunction | WorkerModule<any>) {
   } else {
     throw Error(`Invalid argument passed to expose(). Expected a function or an object, got: ${exposed}`)
   }
+
+  Implementation.subscribeToMasterMessages(messageData => {
+    if (isMasterJobCancelMessage(messageData)) {
+      const jobUID = messageData.uid
+      const subscription = activeSubscriptions.get(jobUID)
+
+      if (subscription) {
+        subscription.unsubscribe()
+        activeSubscriptions.delete(jobUID)
+      }
+    }
+  })
 }
 
 if (typeof self !== "undefined" && typeof self.addEventListener === "function" && Implementation.isWorkerRuntime()) {