@lightningjs/threadx
Version:
A web browser-based JavaScript library that helps manage the communcation of data between one or more web worker threads.
696 lines (653 loc) • 21.6 kB
text/typescript
/*
* Copyright 2023 Comcast Cable Communications Management, LLC
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* SPDX-License-Identifier: Apache-2.0
*/
import { SharedObject } from './SharedObject.js';
import { stringifyTypeId } from './buffer-struct-utils.js';
import { assertTruthy, resolvedGlobal } from './utils.js';
interface ThreadXOptions {
/**
* The ID of the worker. Must be unique across all workers.
*
* Should be an integer value between 1 and 899.
*
* @internalRemarks
* The reason for the 899 limit is the way we generate unique IDs for
* BufferStructs. See `BufferStruct.ts` for more details.
*/
workerId: number;
workerName: string;
sharedObjectFactory?: (buffer: SharedArrayBuffer) => SharedObject | null;
// TOOD: Ultimately replace this with a more generic event handler system
onObjectShared?: (sharedObject: SharedObject) => void;
onBeforeObjectForgotten?: (sharedObject: SharedObject) => void;
onMessage?: (message: any) => Promise<any>;
}
declare global {
class DedicatedWorkerGlobalScope {
DedicatedWorkerGlobalScope: typeof DedicatedWorkerGlobalScope;
postMessage(message: any, transfer?: Transferable[]): void;
addEventListener<K extends keyof WindowEventHandlersEventMap>(
type: K,
listener: (
this: WindowEventHandlers,
ev: WindowEventHandlersEventMap[K],
) => any,
options?: boolean | AddEventListenerOptions,
): void;
addEventListener(
type: string,
listener: EventListenerOrEventListenerObject | null,
options?: boolean | AddEventListenerOptions,
): void;
}
interface WindowOrWorkerGlobalScope {
THREADX?: ThreadX;
DedicatedWorkerGlobalScope?: typeof DedicatedWorkerGlobalScope;
}
}
/**
* Created to define a common interface for both Worker parents (`self`) and
* Worker instances
*/
interface WorkerCommon {
postMessage(message: any, transfer?: Transferable[]): void;
addEventListener<K extends keyof WindowEventHandlersEventMap>(
type: K,
listener: (
this: WindowEventHandlers,
ev: WindowEventHandlersEventMap[K],
) => any,
options?: boolean | AddEventListenerOptions,
): void;
addEventListener(
type: string,
listener: EventListenerOrEventListenerObject | null,
options?: boolean | AddEventListenerOptions,
): void;
terminate?(): void;
}
interface ForgetOptions {
/**
* If true, no warning will be logged if the object is not found.
*/
silent?: boolean;
}
interface ThreadXMessage {
threadXMessageType: string;
}
interface ReadyMessage extends ThreadXMessage {
threadXMessageType: 'ready';
}
interface ShareObjectsMessage extends ThreadXMessage {
threadXMessageType: 'shareObjects';
buffers: SharedArrayBuffer[];
}
interface ForgetObjectsMessage extends ThreadXMessage {
threadXMessageType: 'forgetObjects';
objectIds: number[];
}
interface SharedObjectEmitMessage extends ThreadXMessage {
threadXMessageType: 'sharedObjectEmit';
sharedObjectId: number;
eventName: string;
data: Record<string, unknown>;
}
interface ResponseMessage extends ThreadXMessage {
threadXMessageType: 'response';
asyncMsgId: number;
error?: true;
data: Record<string, unknown>;
}
interface CloseMessage extends ThreadXMessage {
threadXMessageType: 'close';
}
function isMessage(
messageType: 'ready',
message: unknown,
): message is ReadyMessage;
function isMessage(
messageType: 'shareObjects',
message: unknown,
): message is ShareObjectsMessage;
function isMessage(
messageType: 'forgetObjects',
message: unknown,
): message is ForgetObjectsMessage;
function isMessage(
messageType: 'sharedObjectEmit',
message: unknown,
): message is SharedObjectEmitMessage;
function isMessage(
messageType: 'response',
message: unknown,
): message is ResponseMessage;
function isMessage(
messageType: 'close',
message: unknown,
): message is CloseMessage;
function isMessage(
messageType: string,
message: unknown,
): message is MessageEvent {
return (
typeof message === 'object' &&
message !== null &&
'threadXMessageType' in message &&
message.threadXMessageType === messageType
);
}
function isWebWorker(selfObj: any): selfObj is DedicatedWorkerGlobalScope {
// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
return typeof selfObj.DedicatedWorkerGlobalScope === 'function';
}
export class ThreadX {
static init(options: ThreadXOptions) {
if (resolvedGlobal.THREADX) {
throw new Error('ThreadX.init(): ThreadX already initialized.');
}
const threadX = new ThreadX(options);
resolvedGlobal.THREADX = threadX;
return threadX;
}
static destroy() {
if (!resolvedGlobal.THREADX) {
console.warn('ThreadX.destroy(): ThreadX is not initialized.');
return;
}
delete resolvedGlobal.THREADX;
return;
}
/**
* Get the Worker ID of the current worker
*
* @remarks
* This is only valid after ThreadX.init() has been called.
*/
static get workerId(): number {
if (!resolvedGlobal.THREADX) {
throw new Error('ThreadX not initialized');
}
return resolvedGlobal.THREADX.workerId;
}
/**
* Get the Worker Name of the current thread
*
* @remarks
* This is only valid after ThreadX.init() has been called.
*/
static get workerName(): string {
if (!resolvedGlobal.THREADX) {
throw new Error('ThreadX not initialized');
}
return resolvedGlobal.THREADX.workerName;
}
static get instance(): ThreadX {
if (!resolvedGlobal.THREADX) {
throw new Error('ThreadX not initialized');
}
return resolvedGlobal.THREADX;
}
readonly workerId: number;
readonly workerName: string;
readonly sharedObjectFactory?: (
buffer: SharedArrayBuffer,
) => SharedObject | null;
private readonly onSharedObjectCreated?: (sharedObject: SharedObject) => void;
private readonly onBeforeObjectForgotten?: (
sharedObject: SharedObject,
) => void;
/**
* User-defined message handler
*/
private readonly onUserMessage?: (message: any) => Promise<void>;
readonly sharedObjects = new Map<number, SharedObject>();
/**
* WeakMap of SharedObjects to additional metadata
*/
private sharedObjectData = new WeakMap<
SharedObject,
{
workerName: string;
/**
* Whether the SharedObject has been confirmed to be shared with the other worker
*/
shareConfirmed: boolean;
/**
* Queue of messages to emit on the SharedObject once it is confirmed to be shared
*/
emitQueue: Array<
[eventName: string, data: Record<string, unknown>]
> | null;
}
>();
readonly workers = new Map<string, WorkerCommon>();
private workerReadyPromises = new Map<
string,
{
promise: Promise<void> | null;
resolve: () => void;
}
>();
private pendingAsyncMsgs = new Map<
number,
{
resolve: (data: any) => void;
reject: (data: any) => void;
}
>();
private nextAsyncMsgId = 0;
private nextUniqueId = 0;
/**
* Suppress emitting events from SharedObjects
*
* @remarks
* This is used to prevent infinite loops when emitting events from a SharedObject
* that is shared with another worker.
*
* We set this to true when we receive a SharedObjectEmitMessage from another worker
* and set it back to false after we have emitted the event on the SharedObject.
*/
private suppressSharedObjectEmit = false;
private constructor(options: ThreadXOptions) {
this.workerId = options.workerId;
this.workerName = options.workerName;
this.nextUniqueId = options.workerId * 10000000000000 + 1;
this.sharedObjectFactory = options.sharedObjectFactory;
this.onSharedObjectCreated = options.onObjectShared;
this.onBeforeObjectForgotten = options.onBeforeObjectForgotten;
this.onUserMessage = options.onMessage;
const mySelf: unknown = resolvedGlobal;
if (isWebWorker(mySelf)) {
this.registerWorker('parent', mySelf);
this.sendMessage('parent', {
threadXMessageType: 'ready',
} satisfies ReadyMessage);
}
}
registerWorker(workerName: string, worker: WorkerCommon) {
this.workers.set(workerName, worker);
// Set up a promise that will resolve when the worker sends the
// 'ready' message
let readyResolve!: () => void;
let readyPromise: Promise<void>;
if (workerName === 'parent') {
// parent worker is always ready
readyPromise = Promise.resolve();
readyResolve = () => {
// do nothing
};
} else {
readyPromise = new Promise<void>((resolve) => {
readyResolve = resolve;
});
}
this.workerReadyPromises.set(workerName, {
promise: readyPromise,
resolve: readyResolve,
});
this.listenForWorkerMessages(workerName, worker);
}
closeWorker(workerName: string) {
if (!this.workers.has(workerName)) {
throw new Error(`Worker ${workerName} not registered.`);
}
this.closeWorkerAsync(workerName).catch(console.error);
}
async closeWorkerAsync(
workerName: string,
timeout = 5000,
): Promise<'graceful' | 'forced'> {
const worker = this.workers.get(workerName);
if (!worker) {
throw new Error(`Worker ${workerName} not registered.`);
}
const result = await Promise.race([
new Promise((resolve) => {
setTimeout(() => {
resolve(false);
}, timeout);
}),
this.sendMessageAsync(workerName, {
threadXMessageType: 'close',
} satisfies CloseMessage) as Promise<boolean>,
]);
this.workers.delete(workerName);
this.workerReadyPromises.delete(workerName);
if (!result) {
console.warn(
`threadX.closeWorkerAsync(): Worker "${workerName}" did not respond to "close" message within ${timeout}ms. Forcing termination.`,
);
worker.terminate?.();
return 'forced';
}
return 'graceful';
}
private listenForWorkerMessages(workerName: string, worker: WorkerCommon) {
worker.addEventListener('message', (event) => {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
const { data } = event as { data: Record<string, unknown> };
// Process only if message is a ThreadX message
const asyncMsgId = data.__asyncMsgId as number | undefined;
this.onMessage(workerName, data)
.then((response) => {
if (asyncMsgId !== undefined) {
worker.postMessage({
threadXMessageType: 'response',
asyncMsgId: asyncMsgId,
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
data: response,
} satisfies ResponseMessage);
}
})
.catch((error) => {
if (asyncMsgId !== undefined) {
worker.postMessage({
threadXMessageType: 'response',
asyncMsgId: asyncMsgId,
error: true,
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
data: error,
} satisfies ResponseMessage);
}
});
});
}
/**
* Share a SharedObject with a worker
*
* @param workerName Worker to share with
* @param sharedObject
*/
async shareObjects(workerName: string, sharedObjects: SharedObject[]) {
for (const sharedObject of sharedObjects) {
if (this.sharedObjects.get(sharedObject.id)) {
// Currently we only support sharing objects with only a single worker
// TODO: Support sharing objects with multiple workers?
// - Do we really need to do this?
console.warn(
`ThreadX.shareObject(): SharedObject ${
sharedObject.id
} (TypeID: ${stringifyTypeId(
sharedObject.typeId,
)}) is already shared.`,
);
} else {
this.sharedObjects.set(sharedObject.id, sharedObject);
this.sharedObjectData.set(sharedObject, {
workerName: workerName,
shareConfirmed: false,
emitQueue: null,
});
}
}
await this.sendMessageAsync(workerName, {
threadXMessageType: 'shareObjects',
buffers: sharedObjects.map((so) => {
return SharedObject.extractBuffer(so);
}),
} satisfies ShareObjectsMessage);
for (const sharedObject of sharedObjects) {
const soData = this.sharedObjectData.get(sharedObject);
if (soData) {
soData.shareConfirmed = true;
const { emitQueue } = soData;
if (emitQueue) {
for (const event of emitQueue) {
this.__sharedObjectEmit(sharedObject, event[0], event[1]);
}
soData.emitQueue = null;
}
}
}
// TODO: Handle case where worker fails to create shared object on its end
// - We could issue you an error event back to the sharer
}
/**
* Tell ThreadX to forget about SharedObjects
*
* @remarks
* This causes ThreadX on the current worker and the worker that the object
* is shared with to forget about the object. It is up to the worker code to
* actually make sure that no other references to the SharedObjects exist so
* that they can be garbage collected.
*
* A worker can implement the onObjectForgotten() callback to be notified
* when a SharedObject is forgotten.
*
* @param sharedObject
* @param options Options
*/
async forgetObjects(
sharedObjects: SharedObject[],
options: ForgetOptions = {},
) {
/**
* Map of worker name to array of SharedObjects
*
* @remarks
* We group the shared objects by worker so that we can send a single message
* to forget all of the objects shared with each worker.
*/
const objectsByWorker = new Map<string, SharedObject[]>();
for (const sharedObject of sharedObjects) {
if (!this.sharedObjects.has(sharedObject.id)) {
// Currently we only support sharing objects with only a single worker
if (!options.silent) {
console.warn(
`ThreadX.forgetObject(): SharedObject ${
sharedObject.id
} (TypeID: ${stringifyTypeId(sharedObject.typeId)}) is not shared.`,
);
}
} else {
const soData = this.sharedObjectData.get(sharedObject);
assertTruthy(soData);
let objectsInWorker = objectsByWorker.get(soData.workerName);
if (!objectsInWorker) {
objectsInWorker = [];
objectsByWorker.set(soData.workerName, objectsInWorker);
}
objectsInWorker.push(sharedObject);
this.sharedObjects.delete(sharedObject.id);
this.sharedObjectData.delete(sharedObject);
}
}
const promises: Promise<void>[] = [];
for (const [workerName, objectsInWorker] of objectsByWorker) {
promises.push(
// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
this.sendMessageAsync(workerName, {
threadXMessageType: 'forgetObjects',
objectIds: objectsInWorker.map((so) => so.id),
} satisfies ForgetObjectsMessage),
);
}
await Promise.all(promises);
}
sendMessage(
workerName: string,
message: Record<string, unknown>,
transfer?: Transferable[] | undefined,
): void {
const worker = this.workers.get(workerName);
if (!worker) {
throw new Error(
`ThreadX.sendMessage(): Worker '${workerName}' not registered.`,
);
}
this.sendMessageAsync(workerName, message, transfer, {
skipResponseWait: true,
}).catch(console.error);
}
async sendMessageAsync(
workerName: string,
message: Record<string, unknown>,
transfer?: Transferable[] | undefined,
options: { skipResponseWait?: boolean } = {},
): Promise<any> {
const worker = this.workers.get(workerName);
if (!worker) {
throw new Error(
`ThreadX.execMessage(): Worker '${workerName}' not registered.`,
);
}
// Wait for the worker to be ready (if it isn't already)
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
await this.workerReadyPromises.get(workerName)!.promise;
if (options.skipResponseWait) {
worker.postMessage(message, transfer);
return;
}
const asyncMsgId = this.nextAsyncMsgId++;
const promise = new Promise((resolve, reject) => {
this.pendingAsyncMsgs.set(asyncMsgId, {
resolve,
reject,
});
});
message.__asyncMsgId = asyncMsgId;
worker.postMessage(message, transfer);
return promise;
}
private async onMessage(srcWorkerName: string, message: any): Promise<any> {
if (isMessage('shareObjects', message)) {
message.buffers.forEach((buffer: SharedArrayBuffer) => {
const sharedObject = this.sharedObjectFactory?.(buffer);
if (!sharedObject) {
throw new Error(
'ThreadX.onMesasge(): Failed to create shared object.',
);
}
this.sharedObjects.set(sharedObject.id, sharedObject);
this.sharedObjectData.set(sharedObject, {
workerName: srcWorkerName,
shareConfirmed: true,
emitQueue: null,
});
this.onSharedObjectCreated?.(sharedObject);
});
} else if (isMessage('forgetObjects', message)) {
message.objectIds.forEach((id: number) => {
const sharedObject = this.sharedObjects.get(id);
if (!sharedObject) {
// If we can't find the SharedObject then it wasn't shared with this
// worker. Just ignore the message.
return;
}
this.onBeforeObjectForgotten?.(sharedObject);
this.sharedObjects.delete(id);
sharedObject.destroy();
});
} else if (isMessage('sharedObjectEmit', message)) {
const sharedObject = this.sharedObjects.get(message.sharedObjectId);
if (!sharedObject) {
// If we can't find the SharedObject then it wasn't shared with this
// worker. Just ignore the message.
return;
}
// Prevent emitting the event back to the worker that sent it.
this.suppressSharedObjectEmit = true;
sharedObject.emit(message.eventName, message.data);
this.suppressSharedObjectEmit = false;
} else if (isMessage('response', message)) {
const response = this.pendingAsyncMsgs.get(message.asyncMsgId);
if (!response) {
throw new Error(
`ThreadX.onMessage(): Received response for unknown request (ID: ${message.asyncMsgId})`,
);
}
this.pendingAsyncMsgs.delete(message.asyncMsgId);
if (message.error) {
response.reject(message.data);
} else {
response.resolve(message.data);
}
} else if (isMessage('close', message)) {
resolvedGlobal.close();
return true;
} else if (isMessage('ready', message)) {
// Resolve the worker ready promise
this.workerReadyPromises.get(srcWorkerName)?.resolve();
return true;
} else if (this.onUserMessage) {
return await this.onUserMessage(message);
}
}
getSharedObjectById(id: number): SharedObject | null {
return this.sharedObjects.get(id) || null;
}
/**
* Generates an ID that is unique across all ThreadX workers.
*
* @remarks
* The ID is based on the `workerId` set in the `ThreadXOptions` and an
* incrementing counter. For the ID to actually be unique the `workerId` must
* also be unique.
*
* @returns A unique ID
*/
generateUniqueId(): number {
return this.nextUniqueId++;
}
/**
* Emit an event from a SharedObject to all other workers
*
* @internalRemarks
* For internal ThreadX use only.
*
* Since we aren't sure what workers are sharing a SharedObject we need to
* emit the event to all workers. (TODO: Possible optimization?)
*
* @param sharedObject
* @param eventName
* @param data
* @returns
*/
__sharedObjectEmit(
sharedObject: SharedObject,
eventName: string,
data: Record<string, unknown>,
) {
// If we are currently emitting an event from a SharedObject that originated
// from another worker then we don't want to emit the event again.
if (this.suppressSharedObjectEmit) {
return;
}
const soData = this.sharedObjectData.get(sharedObject);
if (!soData) {
// Object isn't shared with any workers yet. Not even in process to do so.
// Just ignore the emit.
return;
}
if (!soData.shareConfirmed) {
// Object is in the process of being shared with other workers. Queue the
// emit until the share is confirmed.
if (!soData.emitQueue) {
soData.emitQueue = [];
}
soData.emitQueue.push([eventName, data]);
return;
}
const worker = this.workers.get(soData.workerName);
assertTruthy(worker, 'Worker not found');
worker.postMessage({
threadXMessageType: 'sharedObjectEmit',
sharedObjectId: sharedObject.id,
eventName,
data,
});
}
}