Add WebWorkerTaskExecutor

[kateinoigakukun]

WebWorkerTaskExecutor is an implementation of TaskExecutor protocol, which is introduced by SE-0417 since Swift 6.0. This task executor runs tasks on Worker threads, which is useful for offloading computationally expensive tasks from the main thread.

The WebWorkerTaskExecutor is designed to work with Web Workers API and Node.js's worker_threads module.

This depends on swiftlang/swift#75008

Example

Try it out: https://swiftwasm-threading-example.vercel.app/

Source: https://github.com/kateinoigakukun/swiftwasm-threading-example/

Usage

import JavaScriptEventLoop

JavaScriptEventLoop.installGlobalExecutor()
WebWorkerTaskExecutor.installGlobalExecutor()

func render() async {
  let executor = WebWorkerTaskExecutor(numberOfThreads: 8)
  defer { executor.terminate() }

  await withTaskGroup(of: Void.self) { group in
      let yStride = scene.height / concurrency
      for i in 0..<concurrency {
          let yRange = i * yStride..<(i + 1) * yStride
          let work = Work(scene: scene, imageView: imageView, yRange: yRange)
          group.addTask(executorPreference: executor) { work.run() }
      }
      if scene.height % concurrency != 0 {
          let work = Work(scene: scene, imageView: imageView, yRange: (concurrency * yStride)..<scene.height)
          group.addTask(executorPreference: executor) { work.run() }
      }
  }
}

Also JavaScript-side needs some tweaks:

// --- main.js
class ThreadRegistry {
  workers = new Map();
  nextTid = 1;

  spawnThread(worker, module, memory, startArg) {
    const tid = this.nextTid++;
    this.workers.set(tid, worker);
    worker.postMessage({ module, memory, tid, startArg });
    return tid;
  }

  listenMainJobFromWorkerThread(tid, listener) {
    const worker = this.workers.get(tid);
    worker.onmessage = (event) => {
      listener(event.data);
    };
  }

  wakeUpWorkerThread(tid, data) {
    const worker = this.workers.get(tid);
    worker.postMessage(data);
  }
}

const threads = new ThreadRegistry();
const swift = new SwiftRuntime({
  threadChannel: {
    wakeUpWorkerThread: threads.wakeUpWorkerThread.bind(threads),
    listenMainJobFromWorkerThread: threads.listenMainJobFromWorkerThread.bind(threads)
  }
});

const module = await WebAssembly.compile("./main.wasm");
const memory = new WebAssembly.Memory({ initial: 256, maximum: 16384, shared: true })

const importObject = {
  wasi_snapshot_preview1: wasi.wasiImport,
  javascript_kit: swift.wasmImports,
  env: { memory },
  wasi: {
    "thread-spawn": (startArg) => {
      const worker = new Worker("./worker.js", { type: "module" });
      return threads.spawnThread(worker, module, memory, startArg);
    }
  }
}

// Instantiate the WebAssembly file
const instance = await WebAssembly.instantiate(module, importObject);

swift.setInstance(instance);
// Start the WebAssembly WASI instance!
wasi.start(instance, swift);


// --- worker.js 

self.onmessage = (event) => {
  self.onmessage = null;

  const swift = new SwiftRuntime({
    threadChannel: {
      wakeUpMainThread: (unownedJob) => {
        // Send the job to the main thread
        postMessage(unownedJob);
      },
      listenWakeEventFromMainThread: (listener) => {
        self.onmessage = (event) => listener(event.data);
      }
    }
  }


  const { module, memory, tid, startArg } = event.data;

  const importObject = {
    wasi_snapshot_preview1: wasi.wasiImport,
    javascript_kit: swift.wasmImports,
    env: { memory },
    wasi: {
      "thread-spawn": (startArg) => {
         throw new Error("Cannot spawn a new thread from a worker thread")
      }
    }
  }

  const instance = await WebAssembly.instantiate(module, importObject);
  swift.setInstance(instance);
  wasi.setInstance(instance);
  swift.startThread(tid, startArg);
}

Dataflow overview

Known limitations

@MainActor does not hop back to main thread

Due to an issue in Cooperative global executor, @MainActor and MainActor.run don't switch execution thread when WebWorkerTaskExecutor is preferred.

JSObject instance cannot cross the thread boundary

Due to the underlying Web Worker limitation, JavaScript objects (JSObject) cannot be shared with or transferred to another thread. You need to convert it into Swift-native objects to represent it in shared memory space.

let canvas = JSObject.global.document.getElementById("my-canvas")
Task(executorPreference: executor) {
  let context = canvas.getContext("2d") // Invalid!
}

TODO

• Job scheduling should take into account Task priority
• Better scheduling algorithm
  • e.g. Schedule jobs to other workers when the current worker thread is busy (the current scheduler enqueues jobs to the current worker thread)
• Rewrite JSClosure not to use global storage
• Add a way to transfer or clone objects across workers
• Worker-backed SerialiExecutor implementation?
  • It would be useful to bind an actor to a dedicated worker. In such actor context, we can safely access JS objects

[github-actions]

Time Change: +399ms (4%)

Total Time: 9,570ms

Test name	Duration	Change
Serialization/JavaScript function call through Wasm import	24ms	+5ms (18%)	⚠️
Serialization/JavaScript function call through Wasm import with int	17ms	+3ms (18%)	⚠️
Serialization/JavaScript function call from Swift	110ms	+15ms (13%)	⚠️
Serialization/Swift Int to JavaScript with assignment	343ms	+28ms (8%)	🔍
Serialization/JavaScript Number to Swift Int	320ms	+20ms (6%)	🔍

View Unchanged

Test name	Duration	Change
Serialization/Swift Int to JavaScript with call	976ms	+45ms (4%)
Serialization/Swift String to JavaScript with assignment	390ms	+14ms (3%)
Serialization/Swift String to JavaScript with call	997ms	+8ms (0%)
Serialization/JavaScript String to Swift String	3,653ms	+137ms (3%)
Object heap/Increment and decrement RC	2,723ms	+121ms (4%)

View Baselines

Test name	Duration
Serialization/Call JavaScript function directly	4ms
Serialization/Assign JavaScript number directly	2ms
Serialization/Call with JavaScript number directly	3ms
Serialization/Write JavaScript string directly	2ms
Serialization/Call with JavaScript string directly	5ms

[kateinoigakukun]

CC: @ephemer might be interested in :)

[ephemer]

Hi @kateinoigakukun thanks, just saw this in the JSKit code and randomly saw your ping here after investigating more closely. I don't get notifications from GitHub, sorry!

This is interesting, and definitely makes sense to me as an approach for Swift concurrency. Our use case is to put an ML workload onto a background thread. That workload is managed by tfjs, and we can't do much more with splitting up the work into chunks like in the example here. I think we'd have to use "real" threading for our use case to benefit from this direction of work. I think generally that would be the best outcome for Swift Wasm too (using "real" Wasm threads). I know that not all runtime environments support it though (including ours – we need to support browsers back to ~Safari 12)