better-webworker
Version:
[English](https://github.com/mchao123/better-webworker/blob/main/README.md) | [中文](https://github.com/mchao123/better-webworker/blob/main/README.zh-CN.md)
320 lines (319 loc) • 10.3 kB
JavaScript
function isTransformedObject(obj) {
return obj && obj._IS_TRANSFORMED_ === true;
}
const transformData = (obj, cache = new Map(), idGen = (function* () {
let id = 0;
while (true) {
yield id++;
}
})()) => {
if (cache.has(obj))
return cache.get(obj);
if (typeof obj !== 'object' && typeof obj !== 'function') {
return obj;
}
if (obj === null)
return obj;
if (isTransformedObject(obj))
return obj;
if (typeof obj === 'function') {
const transformed = {
id: idGen.next().value,
type: 'fn_str',
code: obj.toString(),
_IS_TRANSFORMED_: true
};
cache.set(obj, transformed);
return transformed;
}
const newObj = Array.isArray(obj) ? new Array(obj.length) : {};
cache.set(obj, newObj); // Set cache early to handle circular references
for (const key in obj) {
if (Object.prototype.hasOwnProperty.call(obj, key)) {
// @ts-ignore
newObj[key] = transformData(obj[key], cache, idGen);
}
}
return newObj;
};
const createRequest = (event, name, args, transfer = [], timeout = 5000) => {
const reqid = generateUniqueId(event);
return new Promise((resolve, reject) => {
const timeoutId = setTimeout(() => {
event.promises.delete(reqid);
reject(new Error('Request timed out'));
}, timeout);
event.promises.set(reqid, {
resolve: (value) => {
clearTimeout(timeoutId);
resolve(value);
},
reject: (error) => {
clearTimeout(timeoutId);
reject(error);
}
});
const transferMap = new Map(transfer.map(v => [v, v]));
const transformedArgs = transformData(args, transferMap);
event.thread.postMessage({
__IS_TYPED_WORKER__: true,
isRequest: true,
reqid,
name,
args: transformedArgs,
}, { transfer });
}).finally(() => {
event.promises.delete(reqid);
});
};
const generateUniqueId = (event) => {
let reqid;
do {
reqid = Math.random().toString(36).slice(2);
} while (event.promises.has(reqid));
return reqid;
};
const restoreMessage = (event, root, obj = root.data.data, cache = new Map()) => {
if (typeof obj !== 'object' || obj === null) {
return obj;
}
return new Proxy(obj, {
get(target, prop, receiver) {
const value = Reflect.get(target, prop, receiver);
if (typeof value !== 'object' || value === null) {
return value;
}
if (isTransformedObject(value)) {
switch (value.type) {
case 'fn': {
if (cache.has(value.id)) {
return cache.get(value.id);
}
const fn = function (...args) {
// @ts-ignore
return createRequest(event, value.id, args, fn.transfer, fn.timeout);
};
// @ts-ignore
fn.transfer = [];
fn.timeout = 5000;
cache.set(value.id, fn);
return fn;
}
case 'fn_str': {
if (!cache.has(value.id)) {
cache.set(value.id, new Function('return ' + value.code)());
}
return cache.get(value.id);
}
}
}
return restoreMessage(event, root, value, cache);
}
});
};
const TEMP_FUNCTION_CG = 30 * 1000;
export const TEMP_NAME_PREFIX = 'temp_fn_';
const createRuntime = (thread) => ({
thread,
handlers: new Map(),
promises: new Map(),
cg: new WeakMap(),
});
const messageHandler = (event) => {
const { thread, handlers, promises, cg } = event;
let cgCleanupId;
const cleanupTempHandlers = () => {
const now = Date.now();
for (const [name, fn] of handlers) {
if (!name.startsWith(TEMP_NAME_PREFIX))
continue;
if (!cg.has(fn)) {
cg.set(fn, now);
continue;
}
if (now - cg.get(fn) > TEMP_FUNCTION_CG) {
handlers.delete(name);
cg.delete(fn);
}
}
};
return async (m) => {
var _a;
const data = m.data;
if (!(data === null || data === void 0 ? void 0 : data.__IS_TYPED_WORKER__))
return;
try {
if (data.isRequest) {
const fn = handlers.get(data.name);
if (!fn) {
throw new Error(`Function "${data.name}" not found`);
}
cg.has(fn) && cg.set(fn, Date.now());
const args = restoreMessage(event, m, data.args);
const result = await fn(...args);
thread.postMessage({
__IS_TYPED_WORKER__: true,
isRequest: false,
reqid: data.reqid,
data: transformData(result),
});
}
else {
const promise = promises.get(data.reqid);
if (promise) {
const res = data.data;
const unpacked = (typeof res === 'object' && res !== null)
? restoreMessage(event, m)
: res;
data.isReject ? promise.reject(unpacked) : promise.resolve(unpacked);
}
}
}
catch (error) {
console.error('Worker message handling error:', error);
if (data.isRequest) {
thread.postMessage({
__IS_TYPED_WORKER__: true,
isRequest: false,
reqid: data.reqid,
isReject: true,
data: transformData(error instanceof Error ? error.message : String(error)),
});
}
else {
(_a = promises.get(data.reqid)) === null || _a === void 0 ? void 0 : _a.reject(error);
}
}
cgCleanupId && clearTimeout(cgCleanupId);
cgCleanupId = setTimeout(cleanupTempHandlers, TEMP_FUNCTION_CG);
};
};
/**
* Define message handlers for a Web Worker
*
* @param e - Object containing handler functions to be registered
* @returns A function that returns a WorkerEvent object
*
* @example
* ```ts
* // worker.ts
* import { defineReceive } from 'typed-worker';
*
* export default defineReceive({
* add: (a: number, b: number) => a + b,
* getData: async () => {
* const response = await fetch('https://api.example.com/data');
* return response.json();
* }
* });
* ```
*/
/**
* 定义Worker线程中可用的方法
* @template T - 包含方法签名的对象类型
* @param {T} handlers - 包含Worker方法的对象
* @returns {() => WorkerEvent<T>} - 返回Worker初始化函数
*
* @example
* // worker.ts
* export default defineReceive({
* add(a: number, b: number) {
* return a + b;
* },
* async fetchData(url: string) {
* const response = await fetch(url);
* return response.json();
* }
* });
*/
export const defineReceive = (handlers) => {
const event = createRuntime(self);
Object.entries(handlers).forEach(([name, fn]) => {
event.handlers.set(name, fn);
});
self.onmessage = messageHandler(event);
return null;
};
/**
* Create a typed interface for communicating with a Web Worker
*
* @param worker - Web Worker instance to communicate with
* @returns WorkerEvent object containing typed methods matching the worker's handlers
*
* @example
* ```ts
* // main.ts
* import { useWorker } from 'typed-worker';
*
* const worker = new Worker('worker.js');
*
* interface WorkerAPI {
* add(a: number, b: number): number;
* getData(): Promise<any>;
* }
*
* const { methods } = useWorker<WorkerAPI>(worker);
*
* // Type-safe worker calls
* const sum = await methods.add(1, 2); // Returns: 3
* const data = await methods.getData(); // Returns API data
*
* // Configure timeout for all methods
* methods.timeout = 10000; // 10 seconds
*
* // Configure transferable objects
* const { getData } = methods;
* const buffer = new ArrayBuffer(1024);
* getData.transfer = [buffer];
* ```
*/
export const useWorker = (worker) => {
const event = createRuntime(worker);
const cleanup = () => {
event.handlers.clear();
event.promises.forEach(p => p.reject(new Error('Worker connection terminated')));
event.promises.clear();
};
worker.addEventListener('error', (e) => {
console.error('Worker error:', e);
cleanup();
});
worker.addEventListener('messageerror', (e) => {
console.error('Worker message error:', e);
cleanup();
});
worker.addEventListener('message', messageHandler(event));
// @ts-ignore
return {
worker,
event,
cb: (e, name) => {
let id = name || generateUniqueId(event);
while (!name && event.handlers.has(id)) {
id = TEMP_NAME_PREFIX + Math.random().toString(36).slice(2);
}
event.handlers.set(id, e);
return {
_IS_TRANSFORMED_: true,
type: 'fn',
id, // Use the same id for both handler and transformed object
};
},
methods: new Proxy({
timeout: 5000,
}, {
get(_target, name) {
if (Reflect.has(_target, name))
return Reflect.get(_target, name);
const fn = function (...args) {
// @ts-ignore
return createRequest(event, name, args, fn.transfer, fn.timeout);
};
// @ts-ignore
fn.transfer = [];
fn.timeout = Reflect.get(_target, 'timeout');
return fn;
}
})
};
};