@xylabs/threads 4.7.7 → 4.7.9

package/src/master/invocation-proxy.ts

@@ -0,0 +1,151 @@
+/* eslint-disable import-x/no-internal-modules */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/*
+ * This source file contains the code for proxying calls in the master thread to calls in the workers
+ * by `.postMessage()`-ing.
+ *
+ * Keep in mind that this code can make or break the program's performance! Need to optimize more…
+ */
+
+import DebugLogger from 'debug'
+import { multicast, Observable } from 'observable-fns'
+
+import { deserialize, serialize } from '../common.ts'
+import { ObservablePromise } from '../observable-promise.ts'
+import { isTransferDescriptor } from '../transferable.ts'
+import type {
+  ModuleMethods, ModuleProxy, ProxyableFunction, Worker as WorkerType,
+} from '../types/master.ts'
+import type {
+  MasterJobCancelMessage,
+  MasterJobRunMessage,
+  WorkerJobErrorMessage,
+  WorkerJobResultMessage,
+  WorkerJobStartMessage,
+} from '../types/messages.ts'
+import {
+  MasterMessageType,
+  WorkerMessageType,
+} from '../types/messages.ts'
+
+const debugMessages = DebugLogger('threads:master:messages')
+
+let nextJobUID = 1
+
+const dedupe = <T>(array: T[]): T[] => [...new Set(array)]
+
+const isJobErrorMessage = (data: any): data is WorkerJobErrorMessage => data && data.type === WorkerMessageType.error
+const isJobResultMessage = (data: any): data is WorkerJobResultMessage => data && data.type === WorkerMessageType.result
+const isJobStartMessage = (data: any): data is WorkerJobStartMessage => data && data.type === WorkerMessageType.running
+
+function createObservableForJob<ResultType>(worker: WorkerType, jobUID: number): Observable<ResultType> {
+  return new Observable((observer) => {
+    let asyncType: 'observable' | 'promise' | undefined
+
+    const messageHandler = ((event: MessageEvent) => {
+      debugMessages('Message from worker:', event.data)
+      if (!event.data || event.data.uid !== jobUID) return
+
+      if (isJobStartMessage(event.data)) {
+        asyncType = event.data.resultType
+      } else if (isJobResultMessage(event.data)) {
+        if (asyncType === 'promise') {
+          if (event.data.payload !== undefined) {
+            observer.next(deserialize(event.data.payload))
+          }
+          observer.complete()
+          worker.removeEventListener('message', messageHandler)
+        } else {
+          if (event.data.payload) {
+            observer.next(deserialize(event.data.payload))
+          }
+          if (event.data.complete) {
+            observer.complete()
+            worker.removeEventListener('message', messageHandler)
+          }
+        }
+      } else if (isJobErrorMessage(event.data)) {
+        const error = deserialize(event.data.error as any)
+        if (asyncType === 'promise' || !asyncType) {
+          observer.error(error)
+        } else {
+          observer.error(error)
+        }
+        worker.removeEventListener('message', messageHandler)
+      }
+    }) as EventListener
+
+    worker.addEventListener('message', messageHandler)
+
+    return () => {
+      if (asyncType === 'observable' || !asyncType) {
+        const cancelMessage: MasterJobCancelMessage = {
+          type: MasterMessageType.cancel,
+          uid: jobUID,
+        }
+        worker.postMessage(cancelMessage)
+      }
+      worker.removeEventListener('message', messageHandler)
+    }
+  })
+}
+
+function prepareArguments(rawArgs: any[]): { args: any[]; transferables: Transferable[] } {
+  if (rawArgs.length === 0) {
+    // Exit early if possible
+    return {
+      args: [],
+      transferables: [],
+    }
+  }
+
+  const args: any[] = []
+  const transferables: Transferable[] = []
+
+  for (const arg of rawArgs) {
+    if (isTransferDescriptor(arg)) {
+      args.push(serialize(arg.send))
+      transferables.push(...arg.transferables)
+    } else {
+      args.push(serialize(arg))
+    }
+  }
+
+  return {
+    args,
+    transferables: transferables.length === 0 ? transferables : dedupe(transferables),
+  }
+}
+
+export function createProxyFunction<Args extends any[], ReturnType>(worker: WorkerType, method?: string) {
+  return ((...rawArgs: Args) => {
+    const uid = nextJobUID++
+    const { args, transferables } = prepareArguments(rawArgs)
+    const runMessage: MasterJobRunMessage = {
+      args,
+      method,
+      type: MasterMessageType.run,
+      uid,
+    }
+
+    debugMessages('Sending command to run function to worker:', runMessage)
+
+    try {
+      worker.postMessage(runMessage, transferables)
+    } catch (error) {
+      return ObservablePromise.from(Promise.reject(error))
+    }
+
+    return ObservablePromise.from(multicast(createObservableForJob<ReturnType>(worker, uid)))
+  }) as any as ProxyableFunction<Args, ReturnType>
+}
+
+export function createProxyModule<Methods extends ModuleMethods>(worker: WorkerType, methodNames: string[]): ModuleProxy<Methods> {
+  const proxy: any = {}
+
+  for (const methodName of methodNames) {
+    proxy[methodName] = createProxyFunction(worker, methodName)
+  }
+
+  return proxy
+}

package/src/master/pool-browser.ts

@@ -0,0 +1,400 @@
+/* eslint-disable import-x/export */
+/* eslint-disable unicorn/no-thenable */
+
+/* eslint-disable @typescript-eslint/member-ordering */
+/* eslint-disable unicorn/no-array-reduce */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/* eslint-disable @typescript-eslint/no-namespace */
+
+/// <reference lib="esnext" />
+
+import DebugLogger from 'debug'
+import {
+  multicast, Observable, Subject,
+} from 'observable-fns'
+
+import { allSettled } from '../ponyfills.ts'
+import { defaultPoolSize } from './implementation.browser.ts'
+import type {
+  PoolEvent, QueuedTask, TaskRunFunction, WorkerDescriptor,
+} from './pool-types.ts'
+import { PoolEventType } from './pool-types.ts'
+import { Thread } from './thread.ts'
+
+export declare namespace Pool {
+  type Event<ThreadType extends Thread = any> = PoolEvent<ThreadType>
+  type EventType = PoolEventType
+}
+
+let nextPoolID = 1
+
+function createArray(size: number): number[] {
+  const array: number[] = []
+  for (let index = 0; index < size; index++) {
+    array.push(index)
+  }
+  return array
+}
+
+function delay(ms: number) {
+  return new Promise(resolve => setTimeout(resolve, ms))
+}
+
+function flatMap<In, Out>(array: In[], mapper: (element: In) => Out[]): Out[] {
+  return array.reduce<Out[]>((flattened, element) => [...flattened, ...mapper(element)], [])
+}
+
+function slugify(text: string) {
+  return text.replaceAll(/\W/g, ' ').trim().replaceAll(/\s+/g, '-')
+}
+
+function spawnWorkers<ThreadType extends Thread>(spawnWorker: () => Promise<ThreadType>, count: number): Array<WorkerDescriptor<ThreadType>> {
+  return createArray(count).map(
+    (): WorkerDescriptor<ThreadType> => ({
+      init: spawnWorker(),
+      runningTasks: [],
+    }),
+  )
+}
+
+/**
+ * Thread pool managing a set of worker threads.
+ * Use it to queue tasks that are run on those threads with limited
+ * concurrency.
+ */
+export interface Pool<ThreadType extends Thread> {
+  /**
+   * Returns a promise that resolves once the task queue is emptied.
+   * Promise will be rejected if any task fails.
+   *
+   * @param allowResolvingImmediately Set to `true` to resolve immediately if task queue is currently empty.
+   */
+  completed(allowResolvingImmediately?: boolean): Promise<any>
+
+  /**
+   * Returns a promise that resolves once the task queue is emptied.
+   * Failing tasks will not cause the promise to be rejected.
+   *
+   * @param allowResolvingImmediately Set to `true` to resolve immediately if task queue is currently empty.
+   */
+  settled(allowResolvingImmediately?: boolean): Promise<Error[]>
+
+  /**
+   * Returns an observable that yields pool events.
+   */
+  events(): Observable<PoolEvent<ThreadType>>
+
+  /**
+   * Queue a task and return a promise that resolves once the task has been dequeued,
+   * started and finished.
+   *
+   * @param task An async function that takes a thread instance and invokes it.
+   */
+  queue<Return>(task: TaskRunFunction<ThreadType, Return>): QueuedTask<ThreadType, Return>
+
+  /**
+   * Terminate all pool threads.
+   *
+   * @param force Set to `true` to kill the thread even if it cannot be stopped gracefully.
+   */
+  terminate(force?: boolean): Promise<void>
+}
+
+interface PoolOptions {
+  /** Maximum no. of tasks to run on one worker thread at a time. Defaults to one. */
+  concurrency?: number
+
+  /** Maximum no. of jobs to be queued for execution before throwing an error. */
+  maxQueuedJobs?: number
+
+  /** Gives that pool a name to be used for debug logging, letting you distinguish between log output of different pools. */
+  name?: string
+
+  /** No. of worker threads to spawn and to be managed by the pool. */
+  size?: number
+}
+
+class WorkerPool<ThreadType extends Thread> implements Pool<ThreadType> {
+  static EventType = PoolEventType
+
+  private readonly debug: DebugLogger.Debugger
+  private readonly eventObservable: Observable<PoolEvent<ThreadType>>
+  private readonly options: PoolOptions
+  private readonly workers: Array<WorkerDescriptor<ThreadType>>
+
+  private readonly eventSubject = new Subject<PoolEvent<ThreadType>>()
+  private initErrors: Error[] = []
+  private isClosing = false
+  private nextTaskID = 1
+  private taskQueue: Array<QueuedTask<ThreadType, any>> = []
+
+  constructor(spawnWorker: () => Promise<ThreadType>, optionsOrSize?: number | PoolOptions) {
+    const options: PoolOptions = typeof optionsOrSize === 'number' ? { size: optionsOrSize } : optionsOrSize || {}
+
+    const { size = defaultPoolSize } = options
+
+    this.debug = DebugLogger(`threads:pool:${slugify(options.name || String(nextPoolID++))}`)
+    this.options = options
+    this.workers = spawnWorkers(spawnWorker, size)
+
+    this.eventObservable = multicast(Observable.from(this.eventSubject))
+
+    Promise.all(this.workers.map(worker => worker.init)).then(
+      () =>
+        this.eventSubject.next({
+          size: this.workers.length,
+          type: PoolEventType.initialized,
+        }),
+      (error) => {
+        this.debug('Error while initializing pool worker:', error)
+        this.eventSubject.error(error)
+        this.initErrors.push(error)
+      },
+    )
+  }
+
+  private findIdlingWorker(): WorkerDescriptor<ThreadType> | undefined {
+    const { concurrency = 1 } = this.options
+    return this.workers.find(worker => worker.runningTasks.length < concurrency)
+  }
+
+  private async runPoolTask(worker: WorkerDescriptor<ThreadType>, task: QueuedTask<ThreadType, any>) {
+    const workerID = this.workers.indexOf(worker) + 1
+
+    this.debug(`Running task #${task.id} on worker #${workerID}...`)
+    this.eventSubject.next({
+      taskID: task.id,
+      type: PoolEventType.taskStart,
+      workerID,
+    })
+
+    try {
+      const returnValue = await task.run(await worker.init)
+
+      this.debug(`Task #${task.id} completed successfully`)
+      this.eventSubject.next({
+        returnValue,
+        taskID: task.id,
+        type: PoolEventType.taskCompleted,
+        workerID,
+      })
+    } catch (ex) {
+      const error = ex as Error
+      this.debug(`Task #${task.id} failed`)
+      this.eventSubject.next({
+        error,
+        taskID: task.id,
+        type: PoolEventType.taskFailed,
+        workerID,
+      })
+    }
+  }
+
+  private run(worker: WorkerDescriptor<ThreadType>, task: QueuedTask<ThreadType, any>) {
+    const runPromise = (async () => {
+      const removeTaskFromWorkersRunningTasks = () => {
+        worker.runningTasks = worker.runningTasks.filter(someRunPromise => someRunPromise !== runPromise)
+      }
+
+      // Defer task execution by one tick to give handlers time to subscribe
+      await delay(0)
+
+      try {
+        await this.runPoolTask(worker, task)
+      } finally {
+        removeTaskFromWorkersRunningTasks()
+
+        if (!this.isClosing) {
+          this.scheduleWork()
+        }
+      }
+    })()
+
+    worker.runningTasks.push(runPromise)
+  }
+
+  private scheduleWork() {
+    this.debug('Attempt de-queueing a task in order to run it...')
+
+    const availableWorker = this.findIdlingWorker()
+    if (!availableWorker) return
+
+    const nextTask = this.taskQueue.shift()
+    if (!nextTask) {
+      this.debug('Task queue is empty')
+      this.eventSubject.next({ type: PoolEventType.taskQueueDrained })
+      return
+    }
+
+    this.run(availableWorker, nextTask)
+  }
+
+  private taskCompletion(taskID: number) {
+    return new Promise<any>((resolve, reject) => {
+      const eventSubscription = this.events().subscribe((event) => {
+        if (event.type === PoolEventType.taskCompleted && event.taskID === taskID) {
+          eventSubscription.unsubscribe()
+          resolve(event.returnValue)
+        } else if (event.type === PoolEventType.taskFailed && event.taskID === taskID) {
+          eventSubscription.unsubscribe()
+          reject(event.error)
+        } else if (event.type === PoolEventType.terminated) {
+          eventSubscription.unsubscribe()
+          reject(new Error('Pool has been terminated before task was run.'))
+        }
+      })
+    })
+  }
+
+  async settled(allowResolvingImmediately: boolean = false): Promise<Error[]> {
+    const getCurrentlyRunningTasks = () => flatMap(this.workers, worker => worker.runningTasks)
+
+    const taskFailures: Error[] = []
+
+    const failureSubscription = this.eventObservable.subscribe((event) => {
+      if (event.type === PoolEventType.taskFailed) {
+        taskFailures.push(event.error)
+      }
+    })
+
+    if (this.initErrors.length > 0) {
+      throw this.initErrors[0]
+    }
+    if (allowResolvingImmediately && this.taskQueue.length === 0) {
+      await allSettled(getCurrentlyRunningTasks())
+      return taskFailures
+    }
+
+    await new Promise<void>((resolve, reject) => {
+      const subscription = this.eventObservable.subscribe({
+        error: reject,
+        next(event) {
+          if (event.type === PoolEventType.taskQueueDrained) {
+            subscription.unsubscribe()
+            resolve(void 0)
+          }
+        }, // make a pool-wide error reject the completed() result promise
+      })
+    })
+
+    await allSettled(getCurrentlyRunningTasks())
+    failureSubscription.unsubscribe()
+
+    return taskFailures
+  }
+
+  async completed(allowResolvingImmediately: boolean = false) {
+    const settlementPromise = this.settled(allowResolvingImmediately)
+
+    const earlyExitPromise = new Promise<Error[]>((resolve, reject) => {
+      const subscription = this.eventObservable.subscribe({
+        error: reject,
+        next(event) {
+          if (event.type === PoolEventType.taskQueueDrained) {
+            subscription.unsubscribe()
+            resolve(settlementPromise)
+          } else if (event.type === PoolEventType.taskFailed) {
+            subscription.unsubscribe()
+            reject(event.error)
+          }
+        }, // make a pool-wide error reject the completed() result promise
+      })
+    })
+
+    const errors = await Promise.race([settlementPromise, earlyExitPromise])
+
+    if (errors.length > 0) {
+      throw errors[0]
+    }
+  }
+
+  events() {
+    return this.eventObservable
+  }
+
+  queue(taskFunction: TaskRunFunction<ThreadType, any>) {
+    const { maxQueuedJobs = Number.POSITIVE_INFINITY } = this.options
+
+    if (this.isClosing) {
+      throw new Error('Cannot schedule pool tasks after terminate() has been called.')
+    }
+    if (this.initErrors.length > 0) {
+      throw this.initErrors[0]
+    }
+
+    const taskID = this.nextTaskID++
+    const taskCompletion = this.taskCompletion(taskID)
+
+    taskCompletion.catch((error) => {
+      // Prevent unhandled rejections here as we assume the user will use
+      // `pool.completed()`, `pool.settled()` or `task.catch()` to handle errors
+      this.debug(`Task #${taskID} errored:`, error)
+    })
+
+    const task: QueuedTask<ThreadType, any> = {
+      cancel: () => {
+        if (!this.taskQueue.includes(task)) return
+        this.taskQueue = this.taskQueue.filter(someTask => someTask !== task)
+        this.eventSubject.next({
+          taskID: task.id,
+          type: PoolEventType.taskCanceled,
+        })
+      },
+      id: taskID,
+      run: taskFunction,
+      then: taskCompletion.then.bind(taskCompletion),
+    }
+
+    if (this.taskQueue.length >= maxQueuedJobs) {
+      throw new Error(
+        'Maximum number of pool tasks queued. Refusing to queue another one.\n'
+        + 'This usually happens for one of two reasons: We are either at peak '
+        + "workload right now or some tasks just won't finish, thus blocking the pool.",
+      )
+    }
+
+    this.debug(`Queueing task #${task.id}...`)
+    this.taskQueue.push(task)
+
+    this.eventSubject.next({
+      taskID: task.id,
+      type: PoolEventType.taskQueued,
+    })
+
+    this.scheduleWork()
+    return task
+  }
+
+  async terminate(force?: boolean) {
+    this.isClosing = true
+    if (!force) {
+      await this.completed(true)
+    }
+    this.eventSubject.next({
+      remainingQueue: [...this.taskQueue],
+      type: PoolEventType.terminated,
+    })
+    this.eventSubject.complete()
+    await Promise.all(this.workers.map(async worker => Thread.terminate(await worker.init)))
+  }
+}
+
+/**
+ * Thread pool constructor. Creates a new pool and spawns its worker threads.
+ */
+function PoolConstructor<ThreadType extends Thread>(spawnWorker: () => Promise<ThreadType>, optionsOrSize?: number | PoolOptions) {
+  // The function exists only so we don't need to use `new` to create a pool (we still can, though).
+  // If the Pool is a class or not is an implementation detail that should not concern the user.
+  return new WorkerPool(spawnWorker, optionsOrSize)
+}
+
+;(PoolConstructor as any).EventType = PoolEventType
+
+/**
+ * Thread pool constructor. Creates a new pool and spawns its worker threads.
+ */
+export const Pool = PoolConstructor as typeof PoolConstructor & { EventType: typeof PoolEventType }
+
+export type { PoolEvent, QueuedTask } from './pool-types.ts'
+export { PoolEventType } from './pool-types.ts'
+export { Thread } from './thread.ts'