| b69ab31 | | | 1 | /** |
| b69ab31 | | | 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
| b69ab31 | | | 3 | * |
| b69ab31 | | | 4 | * This source code is licensed under the MIT license found in the |
| b69ab31 | | | 5 | * LICENSE file in the root directory of this source tree. |
| b69ab31 | | | 6 | */ |
| b69ab31 | | | 7 | |
| b69ab31 | | | 8 | import type {CancellationToken} from 'shared/CancellationToken'; |
| b69ab31 | | | 9 | |
| b69ab31 | | | 10 | type WorkerModuleType<Request, Response> = { |
| b69ab31 | | | 11 | handleMessage: (callback: (msg: Response) => void, msg: MessageEvent<Request>) => void; |
| b69ab31 | | | 12 | }; |
| b69ab31 | | | 13 | /** |
| b69ab31 | | | 14 | * Some environments, like jest tests, don't support WebWorkers. |
| b69ab31 | | | 15 | * Let's just import the equivalent file dynamically and call the functions |
| b69ab31 | | | 16 | * instead of doing message passing, so we can validate syntax highlighting in tests. |
| b69ab31 | | | 17 | */ |
| b69ab31 | | | 18 | export class SynchronousWorker<Request, Response> { |
| b69ab31 | | | 19 | private importedModulePromise: Promise<WorkerModuleType<Request, Response>> | undefined; |
| b69ab31 | | | 20 | constructor( |
| b69ab31 | | | 21 | private getImportedModulePromise: () => Promise<WorkerModuleType<Request, Response>>, |
| b69ab31 | | | 22 | ) {} |
| b69ab31 | | | 23 | |
| b69ab31 | | | 24 | public getImportedModule(): Promise<WorkerModuleType<Request, Response>> { |
| b69ab31 | | | 25 | if (this.importedModulePromise) { |
| b69ab31 | | | 26 | return this.importedModulePromise; |
| b69ab31 | | | 27 | } |
| b69ab31 | | | 28 | this.importedModulePromise = this.getImportedModulePromise(); |
| b69ab31 | | | 29 | return this.importedModulePromise; |
| b69ab31 | | | 30 | } |
| b69ab31 | | | 31 | |
| b69ab31 | | | 32 | public onmessage = (_e: MessageEvent) => null; |
| b69ab31 | | | 33 | public postMessage(msg: Request) { |
| b69ab31 | | | 34 | this.getImportedModule().then(module => { |
| b69ab31 | | | 35 | module.handleMessage( |
| b69ab31 | | | 36 | (msg: Response) => { |
| b69ab31 | | | 37 | this.onmessage({data: msg} as MessageEvent<Response>); |
| b69ab31 | | | 38 | }, |
| b69ab31 | | | 39 | {data: msg} as MessageEvent<Request>, |
| b69ab31 | | | 40 | ); |
| b69ab31 | | | 41 | }); |
| b69ab31 | | | 42 | } |
| b69ab31 | | | 43 | |
| b69ab31 | | | 44 | public dispose(): void { |
| b69ab31 | | | 45 | return undefined; |
| b69ab31 | | | 46 | } |
| b69ab31 | | | 47 | } |
| b69ab31 | | | 48 | |
| b69ab31 | | | 49 | export class WorkerApi<Request extends {type: string}, Response extends {type: string}> { |
| b69ab31 | | | 50 | private id = 0; |
| b69ab31 | | | 51 | private requests = new Map<number, (response: Response) => void>(); |
| b69ab31 | | | 52 | private listeners = new Map<Response['type'], (msg: Response) => void>(); |
| b69ab31 | | | 53 | |
| b69ab31 | | | 54 | constructor(public worker: Worker) { |
| b69ab31 | | | 55 | type ResponseWithId = Response & {id: number}; |
| b69ab31 | | | 56 | this.worker.onmessage = e => { |
| b69ab31 | | | 57 | const msg = e.data as Response; |
| b69ab31 | | | 58 | const id = (msg as ResponseWithId).id; |
| b69ab31 | | | 59 | const callback = this.requests.get(id); |
| b69ab31 | | | 60 | if (callback) { |
| b69ab31 | | | 61 | callback(msg); |
| b69ab31 | | | 62 | this.requests.delete(id); |
| b69ab31 | | | 63 | } |
| b69ab31 | | | 64 | |
| b69ab31 | | | 65 | const listener = this.listeners.get(msg.type); |
| b69ab31 | | | 66 | if (listener) { |
| b69ab31 | | | 67 | listener(msg); |
| b69ab31 | | | 68 | } |
| b69ab31 | | | 69 | }; |
| b69ab31 | | | 70 | } |
| b69ab31 | | | 71 | |
| b69ab31 | | | 72 | /** Send a message, then wait for a reply */ |
| b69ab31 | | | 73 | request<T extends Request['type']>( |
| b69ab31 | | | 74 | msg: Request & {type: T}, |
| b69ab31 | | | 75 | cancellationToken?: CancellationToken, |
| b69ab31 | | | 76 | ): Promise<Response & {type: T}> { |
| b69ab31 | | | 77 | return new Promise<Response & {type: T}>(resolve => { |
| b69ab31 | | | 78 | const id = this.id++; |
| b69ab31 | | | 79 | this.worker.postMessage({...msg, id}); |
| b69ab31 | | | 80 | |
| b69ab31 | | | 81 | const disposeOnCancel = cancellationToken?.onCancel(() => { |
| b69ab31 | | | 82 | this.worker.postMessage({type: 'cancel', idToCancel: id}); |
| b69ab31 | | | 83 | }); |
| b69ab31 | | | 84 | this.requests.set(id, result => { |
| b69ab31 | | | 85 | (resolve as (response: Response) => void)(result); |
| b69ab31 | | | 86 | disposeOnCancel?.(); |
| b69ab31 | | | 87 | }); |
| b69ab31 | | | 88 | }); |
| b69ab31 | | | 89 | } |
| b69ab31 | | | 90 | |
| b69ab31 | | | 91 | /** listen for messages from the server of a given type */ |
| b69ab31 | | | 92 | listen<T extends Response['type']>( |
| b69ab31 | | | 93 | type: T, |
| b69ab31 | | | 94 | listener: (msg: Response & {type: T}) => void, |
| b69ab31 | | | 95 | ): () => void { |
| b69ab31 | | | 96 | this.listeners.set(type, listener as (msg: Response) => void); |
| b69ab31 | | | 97 | return () => this.listeners.delete(type); |
| b69ab31 | | | 98 | } |
| b69ab31 | | | 99 | |
| b69ab31 | | | 100 | dispose() { |
| b69ab31 | | | 101 | this.worker.terminate(); |
| b69ab31 | | | 102 | } |
| b69ab31 | | | 103 | } |