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