webgpu: ensure api restrictions

The webgpu API can only be used in the context of a durable object and when the "webgpu" compatibility flag is set.
cloudflare · Aug 18, 2023 · da2bad2 · da2bad2
1 parent baf7c0b
commit da2bad2
Show file tree

Hide file tree

Showing 12 changed files with 130 additions and 21 deletions.
diff --git a/src/workerd/api/global-scope.h b/src/workerd/api/global-scope.h
@@ -44,7 +44,9 @@ class Navigator: public jsg::Object {
 public:
   kj::StringPtr getUserAgent() { return "Cloudflare-Workers"_kj; }
 #ifdef WORKERD_EXPERIMENTAL_ENABLE_WEBGPU
-  jsg::Ref<api::gpu::GPU> getGPU() { return jsg::alloc<api::gpu::GPU>(); }
+  jsg::Ref<api::gpu::GPU> getGPU(CompatibilityFlags::Reader flags) {
+    return jsg::alloc<api::gpu::GPU>(flags);
+  }
 #endif
 
   JSG_RESOURCE_TYPE(Navigator) {

diff --git a/src/workerd/api/gpu/gpu.c++ b/src/workerd/api/gpu/gpu.c++
@@ -3,6 +3,7 @@
 //     https://opensource.org/licenses/Apache-2.0
 
 #include "gpu.h"
+#include "workerd/jsg/exception.h"
 #include <dawn/dawn_proc.h>
 
 namespace workerd::api::gpu {
@@ -21,7 +22,17 @@ void initialize() {
   dawnProcSetProcs(&dawn::native::GetProcs());
 }
 
-GPU::GPU() {
+GPU::GPU(CompatibilityFlags::Reader flags) {
+  // is this a durable object?
+  KJ_IF_MAYBE (actor, IoContext::current().getActor()) {
+    JSG_REQUIRE(actor->getPersistent() != nullptr, TypeError,
+                "webgpu api is only available in Durable Objects (no storage)");
+  } else {
+    JSG_FAIL_REQUIRE(TypeError, "webgpu api is only available in Durable Objects (no actor)");
+  };
+
+  JSG_REQUIRE(flags.getWebgpu(), TypeError, "webgpu needs the webgpu compatibility flag set");
+
   instance_.DiscoverDefaultAdapters();
 }
 

diff --git a/src/workerd/api/gpu/gpu.h b/src/workerd/api/gpu/gpu.h
@@ -39,7 +39,7 @@ struct GPURequestAdapterOptions {
 
 class GPU : public jsg::Object {
 public:
-  explicit GPU();
+  explicit GPU(CompatibilityFlags::Reader flags);
   JSG_RESOURCE_TYPE(GPU) {
     JSG_METHOD(requestAdapter);
   }

diff --git a/src/workerd/api/gpu/webgpu-buffer-test.gpu-wd-test b/src/workerd/api/gpu/webgpu-buffer-test.gpu-wd-test
@@ -7,8 +7,15 @@ const unitTests :Workerd.Config = (
         modules = [
           (name = "worker", esModule = embed "webgpu-buffer-test.js")
         ],
+        durableObjectNamespaces = [
+          (className = "DurableObjectExample", uniqueKey = "210bd0cbd803ef7883a1ee9d86cce06e"),
+        ],
+        durableObjectStorage = (inMemory = void),
+        bindings = [
+          (name = "ns", durableObjectNamespace = "DurableObjectExample"),
+        ],
         compatibilityDate = "2023-01-15",
-        compatibilityFlags = ["experimental", "nodejs_compat"],
+        compatibilityFlags = ["experimental", "nodejs_compat", "webgpu"],
       )
     ),
   ],

diff --git a/src/workerd/api/gpu/webgpu-buffer-test.js b/src/workerd/api/gpu/webgpu-buffer-test.js
@@ -1,10 +1,14 @@
-import { deepEqual, ok } from "node:assert";
+import { deepEqual, ok, equal } from "node:assert";
 
 // run manually for now
 // bazel run --//src/workerd/io:enable_experimental_webgpu //src/workerd/server:workerd -- test `realpath ./src/workerd/api/gpu/webgpu-buffer-test.gpu-wd-test` --verbose --experimental
 
-export const read_sync_stack = {
-  async test(ctrl, env, ctx) {
+export class DurableObjectExample {
+  constructor(state) {
+    this.state = state;
+  }
+
+  async fetch() {
     ok(navigator.gpu);
     const adapter = await navigator.gpu.requestAdapter();
     ok(adapter);
@@ -56,6 +60,18 @@ export const read_sync_stack = {
     const copyArrayBuffer = gpuReadBuffer.getMappedRange();
     ok(copyArrayBuffer);
 
-    deepEqual(new Uint8Array(copyArrayBuffer), new Uint8Array([ 0, 1, 2, 3 ]));
+    deepEqual(new Uint8Array(copyArrayBuffer), new Uint8Array([0, 1, 2, 3]));
+
+    return new Response("OK");
+  }
+}
+
+export const buffer_mapping = {
+  async test(ctrl, env, ctx) {
+    let id = env.ns.idFromName("A");
+    let obj = env.ns.get(id);
+    let res = await obj.fetch("http://foo/test");
+    let text = await res.text();
+    equal(text, "OK");
   },
 };
diff --git a/src/workerd/api/gpu/webgpu-compute-test.gpu-wd-test b/src/workerd/api/gpu/webgpu-compute-test.gpu-wd-test
@@ -7,8 +7,15 @@ const unitTests :Workerd.Config = (
         modules = [
           (name = "worker", esModule = embed "webgpu-compute-test.js")
         ],
+        durableObjectNamespaces = [
+          (className = "DurableObjectExample", uniqueKey = "210bd0cbd803ef7883a1ee9d86cce06e"),
+        ],
+        durableObjectStorage = (inMemory = void),
+        bindings = [
+          (name = "ns", durableObjectNamespace = "DurableObjectExample"),
+        ],
         compatibilityDate = "2023-01-15",
-        compatibilityFlags = ["experimental", "nodejs_compat"],
+        compatibilityFlags = ["experimental", "nodejs_compat", "webgpu"],
       )
     ),
   ],

diff --git a/src/workerd/api/gpu/webgpu-compute-test.js b/src/workerd/api/gpu/webgpu-compute-test.js
@@ -1,10 +1,14 @@
-import { deepEqual, ok } from "node:assert";
+import { deepEqual, ok, equal } from "node:assert";
 
 // run manually for now
 // bazel run --//src/workerd/io:enable_experimental_webgpu //src/workerd/server:workerd -- test `realpath ./src/workerd/api/gpu/webgpu-compute-test.gpu-wd-test` --verbose --experimental
 
-export const read_sync_stack = {
-  async test(ctrl, env, ctx) {
+export class DurableObjectExample {
+  constructor(state) {
+    this.state = state;
+  }
+
+  async fetch() {
     ok(navigator.gpu);
     if (!("gpu" in navigator)) {
       console.log(
@@ -271,5 +275,17 @@ export const read_sync_stack = {
       new Float32Array(arrayBuffer),
       new Float32Array([2, 2, 50, 60, 114, 140])
     );
+
+    return new Response("OK");
+  }
+}
+
+export const compute_shader = {
+  async test(ctrl, env, ctx) {
+    let id = env.ns.idFromName("A");
+    let obj = env.ns.get(id);
+    let res = await obj.fetch("http://foo/test");
+    let text = await res.text();
+    equal(text, "OK");
   },
 };
diff --git a/src/workerd/api/gpu/webgpu-errors-test.gpu-wd-test b/src/workerd/api/gpu/webgpu-errors-test.gpu-wd-test
@@ -7,8 +7,15 @@ const unitTests :Workerd.Config = (
         modules = [
           (name = "worker", esModule = embed "webgpu-errors-test.js")
         ],
+        durableObjectNamespaces = [
+          (className = "DurableObjectExample", uniqueKey = "210bd0cbd803ef7883a1ee9d86cce06e"),
+        ],
+        durableObjectStorage = (inMemory = void),
+        bindings = [
+          (name = "ns", durableObjectNamespace = "DurableObjectExample"),
+        ],
         compatibilityDate = "2023-01-15",
-        compatibilityFlags = ["experimental", "nodejs_compat"],
+        compatibilityFlags = ["experimental", "nodejs_compat", "webgpu"],
       )
     ),
   ],

diff --git a/src/workerd/api/gpu/webgpu-errors-test.js b/src/workerd/api/gpu/webgpu-errors-test.js
@@ -1,10 +1,14 @@
-import { ok } from "node:assert";
+import { ok, equal } from "node:assert";
 
 // run manually for now
 // bazel run --//src/workerd/io:enable_experimental_webgpu //src/workerd/server:workerd -- test `realpath ./src/workerd/api/gpu/webgpu-errors-test.gpu-wd-test` --verbose --experimental
 
-export const read_sync_stack = {
-  async test(ctrl, env, ctx) {
+export class DurableObjectExample {
+  constructor(state) {
+    this.state = state;
+  }
+
+  async fetch() {
     ok(navigator.gpu);
 
     const adapter = await navigator.gpu.requestAdapter();
@@ -79,5 +83,17 @@ export const read_sync_stack = {
 
     // ensure callback with error was indeed called
     ok(callbackCalled);
+
+    return new Response("OK");
+  }
+}
+
+export const error_handling = {
+  async test(ctrl, env, ctx) {
+    let id = env.ns.idFromName("A");
+    let obj = env.ns.get(id);
+    let res = await obj.fetch("http://foo/test");
+    let text = await res.text();
+    equal(text, "OK");
   },
 };
diff --git a/src/workerd/api/gpu/webgpu-write-test.gpu-wd-test b/src/workerd/api/gpu/webgpu-write-test.gpu-wd-test
@@ -7,8 +7,15 @@ const unitTests :Workerd.Config = (
         modules = [
           (name = "worker", esModule = embed "webgpu-write-test.js")
         ],
+        durableObjectNamespaces = [
+          (className = "DurableObjectExample", uniqueKey = "210bd0cbd803ef7883a1ee9d86cce06e"),
+        ],
+        durableObjectStorage = (inMemory = void),
+        bindings = [
+          (name = "ns", durableObjectNamespace = "DurableObjectExample"),
+        ],
         compatibilityDate = "2023-01-15",
-        compatibilityFlags = ["experimental", "nodejs_compat"],
+        compatibilityFlags = ["experimental", "nodejs_compat", "webgpu"],
       )
     ),
   ],

diff --git a/src/workerd/api/gpu/webgpu-write-test.js b/src/workerd/api/gpu/webgpu-write-test.js
@@ -1,10 +1,14 @@
-import { ok, deepEqual } from "node:assert";
+import { ok, deepEqual, equal } from "node:assert";
 
 // run manually for now
 // bazel run --//src/workerd/io:enable_experimental_webgpu //src/workerd/server:workerd -- test `realpath ./src/workerd/api/gpu/webgpu-write-test.gpu-wd-test` --verbose --experimental
 
-export const read_sync_stack = {
-  async test(ctrl, env, ctx) {
+export class DurableObjectExample {
+  constructor(state) {
+    this.state = state;
+  }
+
+  async fetch() {
     ok(navigator.gpu);
     const adapter = await navigator.gpu.requestAdapter();
     ok(adapter);
@@ -23,6 +27,18 @@ export const read_sync_stack = {
 
     // Write bytes to buffer.
     new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);
-    deepEqual(new Uint8Array(arrayBuffer), new Uint8Array([ 0, 1, 2, 3 ]));
+    deepEqual(new Uint8Array(arrayBuffer), new Uint8Array([0, 1, 2, 3]));
+
+    return new Response("OK");
+  }
+}
+
+export const buffer_write = {
+  async test(ctrl, env, ctx) {
+    let id = env.ns.idFromName("A");
+    let obj = env.ns.get(id);
+    let res = await obj.fetch("http://foo/test");
+    let text = await res.text();
+    equal(text, "OK");
   },
 };
diff --git a/src/workerd/io/compatibility-date.capnp b/src/workerd/io/compatibility-date.capnp
@@ -340,4 +340,8 @@ struct CompatibilityFlags @0x8f8c1b68151b6cef {
       $compatEnableFlag("rtti_api")
       $experimental;
   # Enables the `workerd:rtti` module for querying runtime-type-information from JavaScript.
+
+  webgpu @35 :Bool
+      $compatEnableFlag("webgpu")
+      $experimental;
 }