chore(🗿): improve async WebGPU model (#3897)

Author: wcandillon

packages/skia/android/CMakeLists.txt

@@ -120,9 +120,8 @@ if(SK_GRAPHITE)
         "${PROJECT_SOURCE_DIR}/../cpp/jsi2/Promise.cpp"
 
         # WebGPU async system
-        "${PROJECT_SOURCE_DIR}/../cpp/rnwgpu/async/AsyncRunner.cpp"
         "${PROJECT_SOURCE_DIR}/../cpp/rnwgpu/async/AsyncTaskHandle.cpp"
-        "${PROJECT_SOURCE_DIR}/../cpp/rnwgpu/async/JSIMicrotaskDispatcher.cpp"
+        "${PROJECT_SOURCE_DIR}/../cpp/rnwgpu/async/RuntimeContext.cpp"
 
         # WebGPU API
         "${PROJECT_SOURCE_DIR}/../cpp/rnwgpu/api/GPU.cpp"

packages/skia/cpp/api/JsiSkApi.h

@@ -9,7 +9,7 @@
 #ifdef SK_GRAPHITE
 #include "rnskia/RNDawnContext.h"
 #include "rnwgpu/api/GPUDevice.h"
-#include "rnwgpu/async/AsyncRunner.h"
+#include "rnwgpu/async/RuntimeContext.h"
 #endif
 
 #include "JsiNativeBuffer.h"
@@ -164,12 +164,12 @@ class JsiSkApi : public JsiSkHostObject {
         "getDevice", JSI_HOST_FUNCTION_LAMBDA {
 #ifdef SK_GRAPHITE
           auto &dawnContext = DawnContext::getInstance();
-          auto asyncRunner = rnwgpu::async::AsyncRunner::get(runtime);
-          if (!asyncRunner) {
-            throw jsi::JSError(runtime, "AsyncRunner not initialized");
-          }
+          // Per-runtime context: async ops on this device resolve on the calling
+          // runtime's own thread (via its ProcessEvents pump).
+          auto context = rnwgpu::async::RuntimeContext::getOrCreate(
+              runtime, dawnContext.getWGPUInstance());
           auto device = std::make_shared<rnwgpu::GPUDevice>(
-              dawnContext.getWGPUDevice(), asyncRunner, "Skia Device");
+              dawnContext.getWGPUDevice(), context, "Skia Device");
           return rnwgpu::GPUDevice::create(runtime, device);
 #else
       throw jsi::JSError(runtime,

packages/skia/cpp/jsi2/NativeObject.h

@@ -432,6 +432,29 @@ class NativeObject : public jsi::NativeState,
     prototype.setProperty(runtime, name, func);
   }
 
+  /**
+   * Install a method whose native implementation needs the calling jsi::Runtime
+   * as its first parameter. Used by entry points that must act per-runtime
+   * (e.g. GPU::requestAdapter, which creates a per-runtime RuntimeContext).
+   */
+  template <typename ReturnType, typename... Args>
+  static void
+  installMethodWithRuntime(jsi::Runtime &runtime, jsi::Object &prototype,
+                           const char *name,
+                           ReturnType (Derived::*method)(jsi::Runtime &,
+                                                         Args...)) {
+    auto func = jsi::Function::createFromHostFunction(
+        runtime, jsi::PropNameID::forUtf8(runtime, name), sizeof...(Args),
+        [method](jsi::Runtime &rt, const jsi::Value &thisVal,
+                 const jsi::Value *args, size_t count) -> jsi::Value {
+          auto native = Derived::fromValue(rt, thisVal);
+          return callMethodWithRuntime(native.get(), method, rt, args,
+                                       std::index_sequence_for<Args...>{},
+                                       count);
+        });
+    prototype.setProperty(runtime, name, func);
+  }
+
   /**
    * Install a getter on the prototype.
    */
@@ -567,6 +590,22 @@ class NativeObject : public jsi::NativeState,
   }
 
 private:
+  // Helper to call a method that takes the calling jsi::Runtime as its first
+  // parameter, with JSI argument conversion for the rest and JSI conversion of
+  // the result.
+  template <typename ReturnType, typename... Args, size_t... Is>
+  static jsi::Value
+  callMethodWithRuntime(Derived *obj,
+                        ReturnType (Derived::*method)(jsi::Runtime &, Args...),
+                        jsi::Runtime &runtime, const jsi::Value *args,
+                        std::index_sequence<Is...>, size_t count) {
+    ReturnType result = (obj->*method)(
+        runtime, rnwgpu::JSIConverter<std::decay_t<Args>>::fromJSI(
+                     runtime, args[Is], Is >= count)...);
+    return rnwgpu::JSIConverter<std::decay_t<ReturnType>>::toJSI(
+        runtime, std::move(result));
+  }
+
   // Helper to call a method with JSI argument conversion
   template <typename ReturnType, typename... Args, size_t... Is>
   static jsi::Value callMethod(Derived *obj,

packages/skia/cpp/rnskia/RNSkManager.cpp

@@ -23,6 +23,8 @@
 #include "rnwgpu/api/descriptors/GPUMapMode.h"
 #include "rnwgpu/api/descriptors/GPUShaderStage.h"
 #include "rnwgpu/api/descriptors/GPUTextureUsage.h"
+#include "rnwgpu/api/WebGPUConstants.h"
+#include "rnwgpu/async/RuntimeContext.h"
 #include "jsi2/Promise.h"
 
 #include "include/core/SkData.h"
@@ -82,6 +84,12 @@ void RNSkManager::installBindings() {
       jsi::Object::createFromHostObject(*_jsRuntime, _viewApi));
 
 #ifdef SK_GRAPHITE
+  // Register the main runtime + its CallInvoker so spontaneous events
+  // (device.lost / uncapturederror) on main-runtime devices can be delivered to
+  // the JS thread without the ProcessEvents pump. Worklet-runtime devices have
+  // no invoker (best-effort; see the RuntimeContext "Threading model" doc).
+  rnwgpu::async::RuntimeContext::registerMainRuntime(_jsRuntime, _jsCallInvoker);
+
   // Install WebGPU constructors
   rnwgpu::GPU::installConstructor(*_jsRuntime);
   rnwgpu::GPUUncapturedErrorEvent::installConstructor(*_jsRuntime);
@@ -104,18 +112,26 @@ void RNSkManager::installBindings() {
                                      std::move(navigator));
   }
 
-  // Install WebGPU constant objects as plain JS objects
-  _jsRuntime->global().setProperty(*_jsRuntime, "GPUBufferUsage",
-                                   rnwgpu::GPUBufferUsage::create(*_jsRuntime));
-  _jsRuntime->global().setProperty(*_jsRuntime, "GPUColorWrite",
-                                   rnwgpu::GPUColorWrite::create(*_jsRuntime));
-  _jsRuntime->global().setProperty(*_jsRuntime, "GPUMapMode",
-                                   rnwgpu::GPUMapMode::create(*_jsRuntime));
-  _jsRuntime->global().setProperty(*_jsRuntime, "GPUShaderStage",
-                                   rnwgpu::GPUShaderStage::create(*_jsRuntime));
+  // Install WebGPU constant objects as plain JS objects on the main runtime.
+  rnwgpu::installWebGPUConstants(*_jsRuntime);
+
+  // Install a global `installWebGPU()` host function so worklet runtimes can get
+  // the same constants. A host function captured into a worklet is serialized as
+  // a SerializableHostFunction and re-created on the worklet runtime, so the body
+  // runs there (its `rt` is the worklet runtime) and installs the constants on
+  // that runtime. The constants come from the native wgpu::*Usage enums, so the
+  // values stay a single source of truth across every runtime. Calling it on a
+  // runtime that already has the globals is a safe, idempotent no-op.
   _jsRuntime->global().setProperty(
-      *_jsRuntime, "GPUTextureUsage",
-      rnwgpu::GPUTextureUsage::create(*_jsRuntime));
+      *_jsRuntime, "installWebGPU",
+      jsi::Function::createFromHostFunction(
+          *_jsRuntime, jsi::PropNameID::forAscii(*_jsRuntime, "installWebGPU"),
+          0,
+          [](jsi::Runtime &rt, const jsi::Value & /*thisVal*/,
+             const jsi::Value * /*args*/, size_t /*count*/) -> jsi::Value {
+            rnwgpu::installWebGPUConstants(rt);
+            return jsi::Value::undefined();
+          }));
 
   // Install RNWebGPU global object for WebGPU Canvas support
   auto rnWebGPU = std::make_shared<rnwgpu::RNWebGPU>(gpu, nullptr);

packages/skia/cpp/rnwgpu/SurfaceRegistry.h

@@ -7,6 +7,12 @@
 
 #include "webgpu/webgpu_cpp.h"
 
+#ifdef __APPLE__
+namespace dawn::native::metal {
+void WaitForCommandsToBeScheduled(WGPUDevice device);
+} // namespace dawn::native::metal
+#endif
+
 namespace rnwgpu {
 
 struct NativeInfo {
@@ -112,13 +118,39 @@ class SurfaceInfo {
     height = newHeight;
   }
 
-  void present() {
+  // Present the current surface texture. Called synchronously from the thread
+  // that did getCurrentTexture / submit (via GPUCanvasContext::present), so it
+  // preserves Dawn surface thread-affinity. No-op when offscreen / unconfigured
+  // (no surface).
+  void presentFrame() {
+#ifdef __APPLE__
+    // Ensure command buffers are scheduled before presenting. Read the device
+    // under a shared lock, then wait without holding it (the wait can block).
+    // The device may be reconfigured between the two locks; that is safe because
+    // present() is called on the rendering thread right after submit(), the wait
+    // just flushes that thread's already-submitted work, and the Present() below
+    // re-checks `surface` under the unique lock before touching it.
+    wgpu::Device device;
+    {
+      std::shared_lock<std::shared_mutex> lock(_mutex);
+      device = config.device;
+    }
+    if (device) {
+      dawn::native::metal::WaitForCommandsToBeScheduled(device.Get());
+    }
+#endif
     std::unique_lock<std::shared_mutex> lock(_mutex);
     if (surface) {
       surface.Present();
     }
   }
 
+  // True when an on-screen wgpu::Surface is attached (vs offscreen texture).
+  bool hasSurface() {
+    std::shared_lock<std::shared_mutex> lock(_mutex);
+    return surface != nullptr;
+  }
+
   wgpu::Texture getCurrentTexture() {
     std::shared_lock<std::shared_mutex> lock(_mutex);
     if (surface) {

packages/skia/cpp/rnwgpu/api/GPU.cpp

@@ -9,17 +9,15 @@
 
 #include "Convertors.h"
 #include "jsi2/JSIConverter.h"
-#include "rnwgpu/async/JSIMicrotaskDispatcher.h"
+#include "rnwgpu/async/RuntimeContext.h"
 
 namespace rnwgpu {
 
-GPU::GPU(jsi::Runtime &runtime, wgpu::Instance instance)
-    : NativeObject(CLASS_NAME), _instance(instance) {
-  auto dispatcher = std::make_shared<async::JSIMicrotaskDispatcher>(runtime);
-  _async = async::AsyncRunner::getOrCreate(runtime, _instance, dispatcher);
-}
+GPU::GPU(jsi::Runtime & /*runtime*/, wgpu::Instance instance)
+    : NativeObject(CLASS_NAME), _instance(instance) {}
 
 async::AsyncTaskHandle GPU::requestAdapter(
+    jsi::Runtime &runtime,
     std::optional<std::shared_ptr<GPURequestAdapterOptions>> options) {
   wgpu::RequestAdapterOptions aOptions;
   Convertor conv;
@@ -32,21 +30,26 @@ async::AsyncTaskHandle GPU::requestAdapter(
   constexpr auto kDefaultBackendType = wgpu::BackendType::Vulkan;
 #endif
   aOptions.backendType = kDefaultBackendType;
-  return _async->postTask(
-      [this, aOptions](const async::AsyncTaskHandle::ResolveFunction &resolve,
-                       const async::AsyncTaskHandle::RejectFunction &reject) {
+
+  // Per-runtime context: async ops requested on this runtime resolve on this
+  // runtime's own thread (via its ProcessEvents pump).
+  auto context = async::RuntimeContext::getOrCreate(runtime, _instance);
+  return context->postTask(
+      [this, aOptions,
+       context](const async::AsyncTaskHandle::ResolveFunction &resolve,
+                const async::AsyncTaskHandle::RejectFunction &reject) {
         _instance.RequestAdapter(
             &aOptions, wgpu::CallbackMode::AllowProcessEvents,
-            [asyncRunner = _async, resolve,
+            [context, resolve,
              reject](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter,
                      wgpu::StringView message) {
               if (message.length) {
                 fprintf(stderr, "%s", message.data);
               }
 
               if (status == wgpu::RequestAdapterStatus::Success && adapter) {
-                auto adapterHost = std::make_shared<GPUAdapter>(
-                    std::move(adapter), asyncRunner);
+                auto adapterHost =
+                    std::make_shared<GPUAdapter>(std::move(adapter), context);
                 auto result =
                     std::variant<std::nullptr_t, std::shared_ptr<GPUAdapter>>(
                         adapterHost);

packages/skia/cpp/rnwgpu/api/GPU.h

@@ -9,8 +9,8 @@
 
 #include "jsi2/NativeObject.h"
 
-#include "rnwgpu/async/AsyncRunner.h"
 #include "rnwgpu/async/AsyncTaskHandle.h"
+#include "rnwgpu/async/RuntimeContext.h"
 
 #include "webgpu/webgpu_cpp.h"
 
@@ -32,15 +32,19 @@ class GPU : public NativeObject<GPU> {
 public:
   std::string getBrand() { return CLASS_NAME; }
 
+  // requestAdapter needs the calling runtime so each runtime gets its own
+  // RuntimeContext (and ProcessEvents pump on its own thread).
   async::AsyncTaskHandle requestAdapter(
+      jsi::Runtime &runtime,
       std::optional<std::shared_ptr<GPURequestAdapterOptions>> options);
   wgpu::TextureFormat getPreferredCanvasFormat();
 
   std::unordered_set<std::string> getWgslLanguageFeatures();
 
   static void definePrototype(jsi::Runtime &runtime, jsi::Object &prototype) {
     installGetter(runtime, prototype, "__brand", &GPU::getBrand);
-    installMethod(runtime, prototype, "requestAdapter", &GPU::requestAdapter);
+    installMethodWithRuntime(runtime, prototype, "requestAdapter",
+                             &GPU::requestAdapter);
     installMethod(runtime, prototype, "getPreferredCanvasFormat",
                   &GPU::getPreferredCanvasFormat);
     installGetter(runtime, prototype, "wgslLanguageFeatures",
@@ -51,7 +55,6 @@ class GPU : public NativeObject<GPU> {
 
 private:
   wgpu::Instance _instance;
-  std::shared_ptr<async::AsyncRunner> _async;
 };
 
 } // namespace rnwgpu

packages/skia/cpp/rnwgpu/api/GPUAdapter.cpp

@@ -138,7 +138,7 @@ async::AsyncTaskHandle GPUAdapter::requestDevice(
         }
         _instance.RequestDevice(
             &deviceDesc, wgpu::CallbackMode::AllowProcessEvents,
-            [asyncRunner = _async, resolve, reject, label, creationRuntime,
+            [context = _async, resolve, reject, label, creationRuntime,
              deviceLostBinding](wgpu::RequestDeviceStatus status,
                                 wgpu::Device device,
                                 wgpu::StringView message) {
@@ -190,7 +190,7 @@ async::AsyncTaskHandle GPUAdapter::requestDevice(
                   creationRuntime);
 
               auto deviceHost = std::make_shared<GPUDevice>(std::move(device),
-                                                            asyncRunner, label);
+                                                            context, label);
               *deviceLostBinding = deviceHost;
               resolve([deviceHost = std::move(deviceHost)](
                           jsi::Runtime &runtime) mutable {

packages/skia/cpp/rnwgpu/api/GPUAdapter.h

@@ -8,8 +8,8 @@
 
 #include "jsi2/NativeObject.h"
 
-#include "rnwgpu/async/AsyncRunner.h"
 #include "rnwgpu/async/AsyncTaskHandle.h"
+#include "rnwgpu/async/RuntimeContext.h"
 
 #include "webgpu/webgpu_cpp.h"
 
@@ -27,7 +27,7 @@ class GPUAdapter : public NativeObject<GPUAdapter> {
   static constexpr const char *CLASS_NAME = "GPUAdapter";
 
   explicit GPUAdapter(wgpu::Adapter instance,
-                      std::shared_ptr<async::AsyncRunner> async)
+                      std::shared_ptr<async::RuntimeContext> async)
       : NativeObject(CLASS_NAME), _instance(instance), _async(async) {}
 
 public:
@@ -53,7 +53,7 @@ class GPUAdapter : public NativeObject<GPUAdapter> {
 
 private:
   wgpu::Adapter _instance;
-  std::shared_ptr<async::AsyncRunner> _async;
+  std::shared_ptr<async::RuntimeContext> _async;
 };
 
 } // namespace rnwgpu

packages/skia/cpp/rnwgpu/api/GPUBuffer.h

@@ -9,8 +9,8 @@
 
 #include "jsi2/NativeObject.h"
 
-#include "rnwgpu/async/AsyncRunner.h"
 #include "rnwgpu/async/AsyncTaskHandle.h"
+#include "rnwgpu/async/RuntimeContext.h"
 
 #include "webgpu/webgpu_cpp.h"
 
@@ -25,7 +25,7 @@ class GPUBuffer : public NativeObject<GPUBuffer> {
   static constexpr const char *CLASS_NAME = "GPUBuffer";
 
   explicit GPUBuffer(wgpu::Buffer instance,
-                     std::shared_ptr<async::AsyncRunner> async,
+                     std::shared_ptr<async::RuntimeContext> async,
                      std::string label)
       : NativeObject(CLASS_NAME), _instance(instance), _async(async),
         _label(label) {}
@@ -71,7 +71,7 @@ class GPUBuffer : public NativeObject<GPUBuffer> {
 
 private:
   wgpu::Buffer _instance;
-  std::shared_ptr<async::AsyncRunner> _async;
+  std::shared_ptr<async::RuntimeContext> _async;
   std::string _label;
   struct Mapping {
     uint64_t start;