How to replicate data from the main thread to a web worker using ImmerJS

Introduction

Web Workers are an established browser technology for running Javascript tasks in a background thread. They're the gold standard for executing long-running, CPU-intensive tasks in the browser. At Flux.ai, we successfully harnessed Web Workers, paired with ImmerJS patches, to minimize data transfer and deliver an ultra-fast user experience. This post will take you through our journey of using Web Workers and ImmerJS for data replication in our web-based EDA tool.

The Problem

Flux.ai, an innovative web-based EDA tool, needs to compute the layout of thousands of electronic components simultaneously for its unique PCB layouting system. This process must adhere to user-defined rules. Our initial prototype revealed that layouting could take several seconds, leading us to explore the capabilities of Web Workers to parallelize this process and unblock the UI.

At bottom, the web worker API is extremely simple. A single method, postMessage, sends data to a web woker, and the same postMessage method is used to send data back to the main thread. We use a popular abstraction layer on top of postMessage, Comlink, developed several years ago by Google, that makes it possible to call one of your functions in a web worker as if it existed in the main thread. Newer, better or similar abstractions may exist. We did learn in using Comlink that it can easily blow up your JavaScript bundle size.

The trouble with using a web worker in a pure RPC style is that you most likely have a lot of data to pass through postMessage which is as slow as JSON.stringify, as a rule of thumb. This was definitely true in our case. We calculated that it would take 100ms at our desired level of scale just to transfer the layouting data each way, eating into the benefit of a parallel web worker.

Exploring SharedArrayBuffer for Data Transfer

A potential solution to the data transfer problem could be using SharedArrayBuffer, recommended for use with web workers. However, SharedArrayBuffer "represents a generic raw binary data buffer" meaning that a) it is of fixed size and b) it does not accept JS objects, strings, or other typical application data. Our investigations led us to conclude that the performance benefits were offset by the encoding and decoding costs in SharedArrayBuffer. One hope for the future is a Stage 3 ECMAScript proposal for growable ArrayBuffers.

The Solution

We decided instead to populate our web worker with all the data on initial load of a Flux document (while the user is already waiting) and update it with changes as they happened. An added benefit of this approach is that the functions designed to run inside the web worker can also be run in the main thread with the flip of a global variable. You might want to do this for Jest tests, for example, which do not support web workers by default.

We got our changes in document data from ImmerJS, something we were already using as part of Redux Toolkit. Immer is an extremely popular library that enables copy-on-write for built-in data types via a Proxy. A lesser-known feature of Immer is Patches. The function produceWithPatches will return a sequence of patches that represent the changes to the original input.

We made a function that wraps produceWithPatches and assigns the patches back into the document for use downstream.

//
// in helpers.ts
//
export function withPatches(
  document: IDocumentState,
  mutationFn: Parameters[1]
): IDocumentState {
  const [newDocument, forward] = produceWithPatches(document, mutationFn);
  if (forward.length === 0) {
    return newDocument;
  } else {
    return produce(newDocument, (draftDoc) => {
      draftDoc.latestPatchSeq = forward;
    });
  }
}

//
// in reducers.ts
//
const documentReducer = (
    state: IDocumentState | null = documentInitialState,
    action: IDocumentReduxAction,
): IDocumentState | null => {
    if (!state) {
        // note, we don't create patches on the first load of the document
        if (action.type === Actions.loadDocumentActions.successType) {
            return action.response
        }
        return state;
    } else {
        return withPatches(
            state,
            (draftState) => {
                if (isAnyOf(Actions.setSubjectProperties)(action)) {
                // ... do mutations
            }
        )
    }
}

With the patches in hand, we could then complete our data flow from main thread to web worker and back again. The main thread calls the worker functions from middleware after every global state change. In Flux, we use redux-observable middleware.

More Code Samples

In the code, the relevant functions look like this, assuming you are using Comlink.

//
// in LayoutEngineInMain.ts
//
import Comlink from "comlink-loader!./LayoutEngineInWorker";
import { Patch } from "immer";

const comlink = new Comlink();

export async function setInitialDocumentState(
  documentState: DocumentState
): void {
  comlink.setInitialDocumentState(documentState);
}

export function applyDocumentPatches(patches: Patch[]): Patch[] {
  const layoutPatches = comlink.applyDocumentPatches(patches);
  // apply these patches to the global state in middleware
  return layoutPatches;
}

//
// in LayoutEngineInWorker.ts
//
import { Patch, applyPatches } from "immer";
import { LayoutEngine, DocumentState } from "./LayoutEngine";

let documentState: DocumentState | undefined;

export function setInitialDocumentState(state: DocumentState): void {
  documentState = state;
}

export function applyDocumentPatches(patches: Patch[]): Patch[] {
  if (documentState === undefined) {
    throw new Error("First call setInitialDocumentState");
  }
  documentState = applyPatches(documentState, patches);
  return new LayoutEngine(documentState).recomputeLayout();
}

Results: Speedy Data Replication with Web Workers and ImmerJS

The result of our use of Web Workers and ImmerJS patches was a significant reduction in workload on every document change and the ability for users to continue interacting with the application during a large re-layout - a priceless benefit in our web-based EDA tool.

Extra Credit: Boosting Speed with ImmerJS

For extra speed in our web worker, we forked the Immer applyPatches function. The original version was too slow for our needs. So, we adapted applyPatches to skip the draft step and mutate the target object in-place, resulting in a 10X speedup.

In conclusion, Web Workers and ImmerJS have proven to be powerful tools for efficient data replication in Javascript, particularly in the context of our web-based EDA tool, Flux.ai. They offer a potent combination for handling complex, CPU-intensive tasks, and improving user experience through faster data transfer and processing.

‍