Proxy WebSocket communication through the main thread (#576)

* Add Worker proxy support for webR channels

* Add `async: false` option for webR Worker proxy

* Update NEWS.md

* Fix proxy worker under Node

* Add tests for proxy worker

* Don't add main thread proxies in PostMessage mode

Author: georgestagg

NEWS.md

@@ -6,6 +6,8 @@
 
 * Added CSS media query for system dark mode (#572).
 
+* Proxy Web Worker communication through the main thread when using the `SharedArrayBuffer` communication channel (#576).
+
 # webR 0.5.5
 
 ## New features

src/.eslintrc.js

@@ -24,6 +24,7 @@ module.exports = {
     "esbuild.js",
     "jest.config.js",
     "tests/webr.config.js",
+    "tests/scripts/proxy-worker.worker.js",
     "tests/packages.config.js"
   ],
   plugins: ["@typescript-eslint", "jest", "jsdoc", "react"],

src/tests/scripts/proxy-worker.worker.js

@@ -0,0 +1,5 @@
+const { parentPort } = require('node:worker_threads');
+
+parentPort.onmessage = () => {
+  parentPort.postMessage("pong");
+};

src/tests/webR/chan/websocket.test.ts renamed to src/tests/webR/chan/proxy-websocket.test.ts

@@ -0,0 +1,38 @@
+import { WebR } from '../../../webR/webr-main';
+import { ChannelType } from '../../../webR/chan/channel-common';
+
+const webR = new WebR({
+  channelType: ChannelType.SharedArrayBuffer,
+  baseUrl: '../dist/',
+  RArgs: ['--quiet'],
+});
+
+describe('Communicate over Worker proxy', () => {
+  test('Initialises successfully', async () => {
+    await expect(webR.init()).resolves.not.toThrow();
+  });
+
+  test('Use Worker from webR worker thread', async () => {
+    await webR.init();
+    const res = (await webR.evalRString(`
+      webr::eval_js('
+        var worker = new Worker("./tests/scripts/proxy-worker.worker.js");
+        worker.onmessage = (ev) => globalThis.gotReply = ev.data;
+        worker.postMessage("ping");
+        globalThis.gotReply = null;
+      ')
+
+      while (is.na(webr::eval_js("globalThis.gotReply"))) {
+        Sys.sleep(1)
+      }
+
+      webr::eval_js("worker.terminate()")
+      webr::eval_js("globalThis.gotReply")
+    `));
+    expect(res).toEqual("pong");
+  });
+});
+
+afterAll(() => {
+  return webR.close();
+});

src/tests/webR/chan/proxy-worker.test.ts

@@ -142,6 +142,10 @@ export class PostMessageChannelWorker {
   #parked = new Map<string, ResolveFn<Message>>();
   #dispatch: (msg: Message) => void = () => 0;
   #promptDepth = 0;
+  
+  // Main thread proxies only work with SharedBufferChannelWorker for now
+  WebSocketProxy = IN_NODE ? undefined : WebSocket;
+  WorkerProxy = IN_NODE ? undefined : Worker;
 
   constructor() {
     this.#ep = (IN_NODE ? require('worker_threads').parentPort : globalThis) as Endpoint;

src/webR/chan/channel-postmessage.ts

@@ -1,5 +1,5 @@
 import { promiseHandles, newCrossOriginWorker, isCrossOrigin } from '../utils';
-import { EventMessage, Message, Response, SyncRequest, WebSocketCloseMessage, WebSocketMessage, WebSocketOpenMessage } from './message';
+import { EventMessage, Message, PostMessageWorkerMessage, Response, SyncRequest, WebSocketCloseMessage, WebSocketMessage, WebSocketOpenMessage, WorkerErrorMessage, WorkerMessage, WorkerMessageErrorMessage } from './message';
 import { Endpoint } from './task-common';
 import { syncResponse } from './task-main';
 import { ChannelMain, ChannelWorker } from './channel';
@@ -145,6 +145,25 @@ export class SharedBufferChannelMain extends ChannelMain {
             await syncResponse(worker, reqData, { type: 'eval-response', data });
             break;
           }
+          case 'post-message-worker': {
+            const message = payload.data as PostMessageWorkerMessage['data'];
+            message.handles = promiseHandles();
+            this.systemQueue.put({ type: 'postMessageWorker', data: message });
+
+            if (message.async) {
+              await syncResponse(worker, reqData, { type: 'post-message-response' });
+            } else {
+              message.handles.promise.then(
+                (value) => {
+                  void syncResponse(worker, reqData, { type: 'post-message-response', data: { result: value } });
+                },
+                (error) => {
+                  void syncResponse(worker, reqData, { type: 'post-message-response', data: { error: String(error) } });
+                }
+              );
+            }
+            break;
+          }
           default:
             throw new WebRChannelError(`Unsupported request type '${payload.type}'.`);
         }
@@ -162,11 +181,14 @@ export class SharedBufferChannelMain extends ChannelMain {
 
 import { setEventBuffer, setEventsHandler, SyncTask } from './task-worker';
 import { Module } from '../emscripten';
-import { WebSocketProxy, WebSocketProxyFactory } from './websocket';
+import { WebSocketProxy, WebSocketProxyFactory } from './proxy-websocket';
+import { WorkerProxy, WorkerProxyFactory } from './proxy-worker';
 
 export class SharedBufferChannelWorker implements ChannelWorker {
   WebSocketProxy: typeof WebSocket;
-  proxies: Map<string, WebSocketProxy>;
+  WorkerProxy: typeof Worker;
+  ws: Map<string, WebSocketProxy>;
+  workers: Map<string, WorkerProxy>;
   #ep: Endpoint;
   #dispatch: (msg: Message) => void = () => 0;
   #eventBuffer = new Int32Array(new SharedArrayBuffer(4));
@@ -178,23 +200,26 @@ export class SharedBufferChannelWorker implements ChannelWorker {
     setEventBuffer(this.#eventBuffer.buffer);
     setEventsHandler(() => this.handleEvents());
 
+    // Event functionality to be handled via proxy to main thread
     this.WebSocketProxy = WebSocketProxyFactory.proxy(this);
-    this.proxies = new Map();
+    this.ws = new Map();
+    this.WorkerProxy = WorkerProxyFactory.proxy(this);
+    this.workers = new Map();
   }
 
   resolve() {
     this.write({ type: 'resolve', data: this.#eventBuffer.buffer });
   }
 
-  write(msg: Message, transfer?: [Transferable]) {
+  write(msg: Message, transfer?: Transferable[]) {
     this.#ep.postMessage(msg, transfer);
   }
 
-  writeSystem(msg: Message, transfer?: [Transferable]) {
+  writeSystem(msg: Message, transfer?: Transferable[]) {
     this.#ep.postMessage({ type: 'system', data: msg }, transfer);
   }
 
-  syncRequest(msg: Message, transfer?: [Transferable]): Message {
+  syncRequest(msg: Message, transfer?: Transferable[]): Message {
     const task = new SyncTask(this.#ep, msg, transfer);
     return task.syncify() as Message;
   }
@@ -240,22 +265,37 @@ export class SharedBufferChannelWorker implements ChannelWorker {
             break;
           case 'websocket-open': {
             const message = response.data.msg as WebSocketOpenMessage;
-            this.proxies.get(message.data.uuid)?._accept();
+            this.ws.get(message.data.uuid)?._accept();
             break;
           }
           case 'websocket-message': {
             const message = response.data.msg as WebSocketMessage;
-            this.proxies.get(message.data.uuid)?._recieve(message.data.data);
+            this.ws.get(message.data.uuid)?._recieve(message.data.data);
             break;
           }
           case 'websocket-close': {
             const message = response.data.msg as WebSocketCloseMessage;
-            this.proxies.get(message.data.uuid)?._close(message.data.code, message.data.reason);
+            this.ws.get(message.data.uuid)?._close(message.data.code, message.data.reason);
             break;
           }
           case 'websocket-error': {
             const message = response.data.msg as WebSocketMessage;
-            this.proxies.get(message.data.uuid)?._error();
+            this.ws.get(message.data.uuid)?._error();
+            break;
+          }
+          case 'worker-message': {
+            const message = response.data.msg as WorkerMessage;
+            this.workers.get(message.data.uuid)?._message(message.data.data);
+            break;
+          }
+          case 'worker-messageerror': {
+            const message = response.data.msg as WorkerMessageErrorMessage;
+            this.workers.get(message.data.uuid)?._messageerror(message.data.data);
+            break;
+          }
+          case 'worker-error': {
+            const message = response.data.msg as WorkerErrorMessage;
+            this.workers.get(message.data.uuid)?._error();
             break;
           }
           default:

src/webR/chan/channel-shared.ts

@@ -101,6 +101,8 @@ export abstract class ChannelMain {
 }
 
 export interface ChannelWorker {
+  WebSocketProxy: typeof WebSocket | undefined;
+  WorkerProxy: typeof Worker | undefined;
   resolve(): void;
   write(msg: Message, transfer?: [Transferable]): void;
   writeSystem(msg: Message, transfer?: [Transferable]): void;
@@ -112,7 +114,6 @@ export interface ChannelWorker {
   inputOrDispatch: () => number;
   setDispatchHandler: (dispatch: (msg: Message) => void) => void;
   resolveRequest: (msg: Message) => void;
-  WebSocketProxy?: typeof WebSocket;
 }
 
 /**

src/webR/chan/channel.ts

@@ -2,6 +2,7 @@
  * WebR communication channel messaging and request types.
  * @module Message
  */
+import { PromiseHandles } from '../utils';
 import { generateUUID, transfer, UUID } from './task-common';
 
 /** A webR communication channel message. */
@@ -153,6 +154,74 @@ export interface WebSocketCloseMessage {
   };
 }
 
+/** A webR communication channel `proxyWorker` message.
+ * @internal
+ */
+export interface ProxyWorkerMessage {
+  type: 'proxyWorker';
+  data: {
+    uuid: string;
+    url: string;
+    options?: WorkerOptions;
+  };
+}
+
+/** A webR communication channel `postMessageWorker` message.
+ * @internal
+ */
+export interface PostMessageWorkerMessage {
+  type: 'postMessageWorker';
+  data: {
+    uuid: string;
+    data: unknown;
+    async: boolean;
+    transfer?: Transferable[];
+    handles?: PromiseHandles<unknown>;
+  };
+}
+
+/** A webR communication channel `terminateWorker` message.
+ * @internal
+ */
+export interface TerminateWorkerMessage {
+  type: 'terminateWorker';
+  data: {
+    uuid: string;
+  };
+}
+
+/** A webR communication channel `worker-message` message.
+ * @internal
+ */
+export interface WorkerMessage {
+  type: 'worker-message';
+  data: {
+    uuid: string;
+    data: any;
+  };
+}
+
+/** A webR communication channel `worker-messageerror` message.
+ * @internal
+ */
+export interface WorkerMessageErrorMessage {
+  type: 'worker-messageerror';
+  data: {
+    uuid: string;
+    data: any;
+  };
+}
+
+/** A webR communication channel `worker-error` message.
+ * @internal
+ */
+export interface WorkerErrorMessage {
+  type: 'worker-error';
+  data: {
+    uuid: string;
+  };
+}
+
 /** A webR communication channel sync-request.
  * @internal
  */

src/webR/chan/message.ts

@@ -97,7 +97,7 @@ export class WebSocketProxyFactory {
           type: 'proxyWebSocket',
           data: { uuid: this.uuid, url: this.url, protocol: this.protocol }
         });
-        chan.proxies.set(this.uuid, this);
+        chan.ws.set(this.uuid, this);
       }
 
       send(data: string | ArrayBufferLike | Blob | ArrayBufferView): void {
@@ -129,7 +129,7 @@ export class WebSocketProxyFactory {
         const ev = new CloseEvent('close', { code, reason });
         this.dispatchEvent(ev);
         this.onclose?.(ev);
-        chan.proxies.delete(this.uuid);
+        chan.ws.delete(this.uuid);
       }
 
       _error(): void {