fix(🐛): pump an imported device's async work on the instance that owns it

apps/example/src/Tests.tsx

@@ -9,6 +9,7 @@ import {
 } from "react-native";
 import {
   GPUOffscreenCanvas,
+  importDevice,
   // The flags as the package exports them, not the globals of the same name
   // (see the `exportedFlags` spec in Constants.spec.ts).
   GPUBufferUsage as exportedBufferUsage,
@@ -17,6 +18,12 @@ import {
   GPUShaderStage as exportedShaderStage,
   GPUTextureUsage as exportedTextureUsage,
 } from "react-native-webgpu";
+// Imported after react-native-webgpu on purpose: webgpu installs natively on
+// import and settles on a Dawn instance of its own while Skia's library is not
+// loaded yet. On Android, Skia's device then lives on a second instance, and
+// ImportDevice.spec.ts relies on that split to cover importDevice() across
+// instances.
+import { Skia } from "react-native-skia";
 import { mat4, vec3, mat3 } from "wgpu-matrix";
 
 import { useClient } from "./useClient";
@@ -132,6 +139,10 @@ export const Tests = ({ assets: { di3D, saturn, moon } }: AssetProps) => {
                 },
                 workletDeviceStress,
                 workletSerialization,
+                importSkiaDevice: () =>
+                  Skia.gpu.isAvailable
+                    ? importDevice(Skia.getNativeDevice())
+                    : null,
                 exportedFlags: {
                   GPUBufferUsage: exportedBufferUsage,
                   GPUColorWrite: exportedColorWrite,

packages/webgpu/cpp/rnwgpu/WebGPURuntime.cpp

@@ -9,15 +9,18 @@ namespace rnwgpu {
 namespace {
 
 wgpu::Instance createInstance() {
-  // If a Graphite build of react-native-skia is loaded in this
+  // If a Graphite build of react-native-skia is already loaded in this
   // process, adopt its wgpu::Instance instead of creating our own, so
-  // Graphite's device and every JS-created device live on one instance. That
-  // makes pointer-based handoff (RNWebGPU.importDevice of Skia's device)
-  // sound: the async pump serves the shared instance, so callbacks on
-  // imported devices settle. Discovered at runtime via dlsym so there is no
-  // build-time coupling and native-module init order does not matter. Both
-  // packages must link the same single Dawn copy (they vendor the same
-  // artifact) for the handle to be meaningful.
+  // Graphite's device and every JS-created device live on one instance.
+  // Discovered at runtime via dlsym so there is no build-time coupling. The
+  // adoption is best-effort and decided once, when the instance is first
+  // used: on Android it only succeeds when librnskia.so is loaded before
+  // react-native-webgpu installs (see the RTLD_NOLOAD lookup below), so the
+  // native-module init order decides whether the instance is shared.
+  // RNWebGPU.importDevice does not depend on it: an imported device's async
+  // work is pumped on the instance that owns it (see async::RuntimeContext).
+  // Both packages must link the same single Dawn copy (they vendor the same
+  // artifact) for the handles they exchange to be meaningful.
   using GetInstanceFn = WGPUInstance (*)();
   void *getter = dlsym(RTLD_DEFAULT, "rnskia_getWGPUInstance");
 #ifdef __ANDROID__

packages/webgpu/cpp/rnwgpu/WebGPURuntime.h

@@ -15,8 +15,8 @@ class WebGPURuntime {
 public:
   static WebGPURuntime &shared();
 
-  // The process-wide instance: Skia Graphite's when that is loaded, otherwise
-  // our own. Created on first use. Safe to call from any thread.
+  // The process-wide instance: Skia Graphite's when that is already loaded on
+  // first use, otherwise our own. Safe to call from any thread.
   wgpu::Instance instance();
 
   // Creates a surface on the shared instance for a CAMetalLayer (Apple) or an

packages/webgpu/cpp/rnwgpu/api/RNWebGPU.h

@@ -86,10 +86,12 @@ class RNWebGPU : public NativeObject<RNWebGPU> {
   // Wrap an externally created WGPUDevice (passed as a BigInt pointer, e.g.
   // Skia's Graphite device from Skia.getNativeDevice()) in a GPUDevice.
   // AddRefs the handle, so the original owner keeps its reference. Only sound
-  // because the process links a single Dawn copy and the GPU instance is
-  // shared with the exporter (see the GPU constructor): the calling runtime's
-  // pump serves the instance the device lives on, so its async callbacks
-  // settle normally.
+  // because the process links a single Dawn copy. The device's async
+  // callbacks are pumped on the instance that owns it (RuntimeContext is per
+  // runtime and per instance), which is not necessarily the shared one:
+  // WebGPURuntime adopts Skia's instance only when Skia's getter is reachable
+  // as the instance is first used (on Android, only once librnskia.so is
+  // loaded), and other exporters have their own.
   std::shared_ptr<GPUDevice> importDevice(jsi::Runtime &runtime,
                                           void *pointer) {
     if (pointer == nullptr) {
@@ -99,7 +101,8 @@ class RNWebGPU : public NativeObject<RNWebGPU> {
     auto raw = reinterpret_cast<WGPUDevice>(pointer);
     wgpuDeviceAddRef(raw);
     wgpu::Device device = wgpu::Device::Acquire(raw);
-    auto ctx = async::RuntimeContext::getOrCreate(runtime, _gpu->get());
+    auto ctx = async::RuntimeContext::getOrCreate(
+        runtime, device.GetAdapter().GetInstance());
     return std::make_shared<GPUDevice>(device, ctx, "Imported Device");
   }
 

packages/webgpu/cpp/rnwgpu/async/RuntimeContext.cpp

@@ -2,6 +2,7 @@
 
 #include <memory>
 #include <stdexcept>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -14,7 +15,7 @@ namespace rnwgpu::async {
 
 namespace {
 struct RuntimeData {
-  std::shared_ptr<RuntimeContext> context;
+  std::unordered_map<WGPUInstance, std::shared_ptr<RuntimeContext>> contexts;
 };
 constexpr const char *TAG = "RuntimeContext";
 
@@ -25,12 +26,27 @@ constexpr const char *TAG = "RuntimeContext";
 jsi::Runtime *sMainRuntime = nullptr;
 std::shared_ptr<facebook::react::CallInvoker> sMainInvoker;
 
-// Serializes ProcessEvents() across all runtimes that share a wgpu::Instance.
-// Held only across the ProcessEvents call itself, never while running JS / mailbox
+// ProcessEvents() on one wgpu::Instance is serialized across the runtimes that
+// share it by a mutex per instance, which each context resolves once, on
+// creation. Pumping another library's instance (the one owning a device passed
+// to importDevice) can block on that library's device lock; with one mutex per
+// instance that never stalls the pumps of our own instance. A mutex is held
+// only across the ProcessEvents call itself, never while running JS / mailbox
 // settle-actions, so it cannot deadlock against the per-context mailbox mutex.
-std::mutex &processEventsMutex() {
-  static std::mutex mutex;
-  return mutex;
+std::shared_ptr<std::mutex> processEventsMutexFor(WGPUInstance instance) {
+  static std::mutex registryMutex;
+  static std::unordered_map<WGPUInstance, std::weak_ptr<std::mutex>> registry;
+  std::lock_guard<std::mutex> lock(registryMutex);
+  auto &slot = registry[instance];
+  if (auto existing = slot.lock()) {
+    return existing;
+  }
+  // Contexts retain their instance, so its handle cannot be reused while one
+  // of them is alive: an expired slot belonged to an instance whose runtimes
+  // are all gone, and an instance reusing that address starts afresh.
+  auto created = std::make_shared<std::mutex>();
+  slot = created;
+  return created;
 }
 } // namespace
 
@@ -42,32 +58,32 @@ void RuntimeContext::registerMainRuntime(
 }
 
 RuntimeContext::RuntimeContext(jsi::Runtime &runtime, wgpu::Instance instance)
-    : _runtime(runtime), _instance(std::move(instance)) {
+    : _runtime(runtime), _instance(std::move(instance)),
+      _processEventsMutex(processEventsMutexFor(_instance.Get())) {
   Logger::logToConsole("[%s] Created (runtime=%p)", TAG, &runtime);
 }
 
-std::shared_ptr<RuntimeContext> RuntimeContext::get(jsi::Runtime &runtime) {
-  auto data = runtime.getRuntimeData(runtimeDataUUID());
-  if (!data) {
-    return nullptr;
-  }
-  return std::static_pointer_cast<RuntimeData>(data)->context;
-}
-
 std::shared_ptr<RuntimeContext>
-RuntimeContext::getOrCreate(jsi::Runtime &runtime, wgpu::Instance instance) {
-  if (auto existing = get(runtime)) {
-    return existing;
+RuntimeContext::getOrCreate(jsi::Runtime &runtime,
+                            const wgpu::Instance &instance) {
+  auto data = std::static_pointer_cast<RuntimeData>(
+      runtime.getRuntimeData(runtimeDataUUID()));
+  if (!data) {
+    data = std::make_shared<RuntimeData>();
+    runtime.setRuntimeData(runtimeDataUUID(), data);
   }
-  auto context = std::make_shared<RuntimeContext>(runtime, std::move(instance));
-  // Only the main JS runtime's context carries the CallInvoker; it is used to
-  // deliver spontaneous events (device.lost) without the pump.
-  if (&runtime == sMainRuntime) {
-    context->_callInvoker = sMainInvoker;
+  // Contexts retain their instance, so its handle is a stable key for the
+  // runtime's lifetime, and each instance's callbacks reach the matching pump.
+  // A slot left null by a throwing constructor is simply filled next time.
+  auto &context = data->contexts[instance.Get()];
+  if (!context) {
+    context = std::make_shared<RuntimeContext>(runtime, instance);
+    // Only the main JS runtime's context carries the CallInvoker; it is used to
+    // deliver spontaneous events (device.lost) without the pump.
+    if (&runtime == sMainRuntime) {
+      context->_callInvoker = sMainInvoker;
+    }
   }
-  auto data = std::make_shared<RuntimeData>();
-  data->context = context;
-  runtime.setRuntimeData(runtimeDataUUID(), data);
   return context;
 }
 
@@ -173,7 +189,7 @@ void RuntimeContext::tick() {
   {
     // Serialize ProcessEvents across runtimes sharing this instance. Callbacks
     // fired here only deposit into mailboxes (postSettle), they do not run JS.
-    std::lock_guard<std::mutex> lock(processEventsMutex());
+    std::lock_guard<std::mutex> lock(*_processEventsMutex);
     _instance.ProcessEvents();
   }
   // Settle this runtime's ready promises on this thread, outside the pump lock.

packages/webgpu/cpp/rnwgpu/async/RuntimeContext.h

@@ -22,11 +22,14 @@ class CallInvoker;
 namespace rnwgpu::async {
 
 /**
- * Per-runtime coordinator for asynchronous WebGPU operations.
+ * Per-runtime, per-instance coordinator for asynchronous WebGPU operations.
  *
- * Each JS runtime that uses WebGPU gets its own RuntimeContext, stored in the
- * runtime's runtimeData. Async Dawn operations are registered with
- * CallbackMode::AllowProcessEvents and driven to completion by pumping
+ * Each (JS runtime, Dawn instance) pair gets its own RuntimeContext, stored in
+ * the runtime's runtimeData. A runtime usually has a single one, for the
+ * shared WebGPURuntime instance; a device passed to importDevice() can live on
+ * another instance, which then gets its own context and pump. Async Dawn
+ * operations are registered with CallbackMode::AllowProcessEvents and driven
+ * to completion by pumping
  * `instance.ProcessEvents()` on the runtime's OWN thread via a self-
  * rescheduling tick (scheduled through that runtime's setTimeout). Because
  * ProcessEvents invokes the Dawn callbacks synchronously on the pumping thread,
@@ -56,8 +59,9 @@ namespace rnwgpu::async {
  * deposits a settle-action (a plain C++ closure, no JSI) into the OWNING
  * context's thread-safe mailbox via postSettle(), and each context drains its
  * own mailbox on its own thread during tick(). ProcessEvents() itself is
- * serialized across runtimes by a process-wide mutex, since concurrent
- * ProcessEvents on one instance is not guaranteed reentrant.
+ * serialized across the runtimes sharing an instance by a mutex per instance,
+ * since concurrent ProcessEvents on one instance is not guaranteed reentrant;
+ * pumps of different instances never wait on each other.
  *
  * Threading contract: a RuntimeContext must only be pumped from the runtime it
  * was created for. Request/response async entry points (mapAsync,
@@ -78,9 +82,10 @@ class RuntimeContext : public std::enable_shared_from_this<RuntimeContext> {
 
   RuntimeContext(jsi::Runtime &runtime, wgpu::Instance instance);
 
-  static std::shared_ptr<RuntimeContext> get(jsi::Runtime &runtime);
-  static std::shared_ptr<RuntimeContext> getOrCreate(jsi::Runtime &runtime,
-                                                     wgpu::Instance instance);
+  // The calling runtime's context for `instance`, created on first use and
+  // kept in the runtime's runtimeData for the runtime's lifetime.
+  static std::shared_ptr<RuntimeContext>
+  getOrCreate(jsi::Runtime &runtime, const wgpu::Instance &instance);
 
   // Register the main JS runtime and its CallInvoker. The RuntimeContext
   // created for this runtime gets the invoker (callInvoker() returns it); every
@@ -97,8 +102,8 @@ class RuntimeContext : public std::enable_shared_from_this<RuntimeContext> {
     return _callInvoker;
   }
 
-  // The wgpu::Instance bound to this runtime.
-  wgpu::Instance instance() const { return _instance; }
+  // The wgpu::Instance served by this context on its owning runtime.
+  const wgpu::Instance &instance() const { return _instance; }
 
   // The runtime this context belongs to. Safe for pointer-identity checks from
   // any thread; only touch JSI through it on the runtime's own thread.
@@ -126,6 +131,9 @@ class RuntimeContext : public std::enable_shared_from_this<RuntimeContext> {
 
   jsi::Runtime &_runtime;
   wgpu::Instance _instance;
+  // Serializes ProcessEvents() on _instance across the runtimes that share it;
+  // every context of that instance holds the same mutex.
+  std::shared_ptr<std::mutex> _processEventsMutex;
   // Non-null only for the main JS runtime's context (see registerMainRuntime).
   std::shared_ptr<facebook::react::CallInvoker> _callInvoker;
   std::atomic<std::size_t> _pumpTasks{0};

packages/webgpu/src/__tests__/ImportDevice.spec.ts

@@ -0,0 +1,96 @@
+import { client } from "./setup";
+
+// importDevice() wraps a device another library created, here Skia's Graphite
+// device. Its async operations settle through Dawn's ProcessEvents pump, which
+// has to run on the Dawn instance that owns the device. That is not webgpu's
+// own instance when react-native-skia installs after react-native-webgpu (the
+// example app's order, see apps/example/src/Tests.tsx), so a pump bound to the
+// wrong instance leaves every promise below pending forever. The web and node
+// runners have no device to import and return null, which skips the
+// assertions; so does a native run on a device where Skia has no GPU.
+describe("importDevice", () => {
+  it("settles async operations on a device imported from Skia", async () => {
+    const result = await client.eval(
+      ({ importSkiaDevice, shaders: { triangleVertWGSL, redFragWGSL } }) => {
+        const device = importSkiaDevice ? importSkiaDevice() : null;
+        if (device === null) {
+          return null;
+        }
+        // Each operation races a timer, so a hang names the operation instead
+        // of running into the harness' eval timeout.
+        const within = <T>(promise: Promise<T>) =>
+          Promise.race([
+            promise.then((value) => ({ settled: true, value })),
+            new Promise<{ settled: false; value: null }>((resolve) =>
+              setTimeout(() => resolve({ settled: false, value: null }), 5000),
+            ),
+          ]);
+
+        device.pushErrorScope("validation");
+        // Invalid, maxAnisotropy must be at least 1.
+        device.createSampler({ maxAnisotropy: 0 });
+        const errorScope = within(device.popErrorScope()).then((r) => ({
+          settled: r.settled,
+          hasError: r.settled && r.value !== null,
+        }));
+
+        const pipeline = within(
+          device.createRenderPipelineAsync({
+            layout: "auto",
+            vertex: {
+              module: device.createShaderModule({ code: triangleVertWGSL }),
+            },
+            fragment: {
+              module: device.createShaderModule({ code: redFragWGSL }),
+              targets: [{ format: "rgba8unorm" }],
+            },
+          }),
+        ).then((r) => ({
+          settled: r.settled,
+          isPipeline:
+            r.settled && typeof r.value.getBindGroupLayout === "function",
+        }));
+
+        const src = device.createBuffer({
+          size: 16,
+          usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+        });
+        const dst = device.createBuffer({
+          size: 16,
+          usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+        });
+        device.queue.writeBuffer(src, 0, new Uint32Array([0xc0ffee, 1, 2, 3]));
+        const encoder = device.createCommandEncoder();
+        encoder.copyBufferToBuffer(src, 0, dst, 0, 16);
+        device.queue.submit([encoder.finish()]);
+        const workDone = within(device.queue.onSubmittedWorkDone()).then(
+          (r) => r.settled,
+        );
+        const mapped = within(dst.mapAsync(GPUMapMode.READ)).then((r) => {
+          if (!r.settled) {
+            return { settled: false, first: -1 };
+          }
+          const first = new DataView(dst.getMappedRange()).getUint32(0, true);
+          dst.unmap();
+          return { settled: true, first };
+        });
+
+        return Promise.all([errorScope, pipeline, workDone, mapped]).then(
+          (results) => ({
+            errorScope: results[0],
+            pipeline: results[1],
+            workDoneSettled: results[2],
+            mapped: results[3],
+          }),
+        );
+      },
+    );
+    if (result === null) {
+      return;
+    }
+    expect(result.errorScope).toEqual({ settled: true, hasError: true });
+    expect(result.pipeline).toEqual({ settled: true, isPipeline: true });
+    expect(result.workDoneSettled).toBe(true);
+    expect(result.mapped).toEqual({ settled: true, first: 0xc0ffee });
+  });
+});

packages/webgpu/src/__tests__/setup.ts

@@ -106,6 +106,10 @@ interface GPUTestingContext {
     probeSize: number;
     pixelRatio: number;
   };
+  // Native app harness only: Skia's Graphite device wrapped with
+  // importDevice(), or null when Skia has no GPU. See ImportDevice.spec.ts
+  // and apps/example/src/Tests.tsx.
+  importSkiaDevice?: () => GPUDevice | null;
   mat4: typeof mat4;
   vec3: typeof vec3;
   mat3: typeof mat3;

packages/webgpu/src/importDevice.ts

@@ -9,8 +9,9 @@
  * ```
  *
  * This is only sound when the exporting library links the same single Dawn
- * copy as react-native-webgpu and shares its wgpu::Instance (react-native-skia
- * does both).
+ * copy as react-native-webgpu (react-native-skia does). The device stays on
+ * its owner's Dawn instance and its async operations are pumped there, so the
+ * two libraries do not have to share an instance.
  */
 export const importDevice = (pointer: bigint): GPUDevice => {
   if (typeof RNWebGPU === "undefined") {