Enable JSPI (#3259)

This enables JSPI in Python workers. It will not be usable until we
also update to the newer Pyodide runtime.

src/workerd/api/pyodide/pyodide.c++

@@ -485,6 +485,7 @@ void DiskCache::put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data) {
 }
 
 jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
+  js.installJspi();
   js.v8Context()->SetSecurityToken(emscriptenRuntime.contextToken.getHandle(js));
   return emscriptenRuntime.emscriptenRuntime.getHandle(js);
 }

src/workerd/api/pyodide/setup-emscripten.c++

@@ -72,6 +72,7 @@ EmscriptenRuntime EmscriptenRuntime::initialize(
   kj::Maybe<capnp::Data::Reader> emsciptenSetupJsReader;
   kj::Maybe<capnp::Data::Reader> pythonStdlibZipReader;
   kj::Maybe<capnp::Data::Reader> pyodideAsmWasmReader;
+  js.installJspi();
   for (auto module: bundle.getModules()) {
     if (module.getName().endsWith("emscriptenSetup.js")) {
       emsciptenSetupJsReader = module.getData();

src/workerd/jsg/jsg.c++

@@ -187,6 +187,12 @@ void Lock::setAllowEval(bool allow) {
   IsolateBase::from(v8Isolate).setAllowEval({}, allow);
 }
 
+void Lock::installJspi() {
+  IsolateBase::from(v8Isolate).setJspiEnabled({}, true);
+  v8Isolate->InstallConditionalFeatures(v8Context());
+  IsolateBase::from(v8Isolate).setJspiEnabled({}, false);
+}
+
 void Lock::setCaptureThrowsAsRejections(bool capture) {
   IsolateBase::from(v8Isolate).setCaptureThrowsAsRejections({}, capture);
 }

src/workerd/jsg/jsg.h

@@ -2528,6 +2528,9 @@ class Lock {
   // Use to enable/disable dynamic code evaluation (via eval(), new Function(), or WebAssembly).
   void setAllowEval(bool allow);
 
+  // Install JSPI on the current context. Currently used only for Python workers.
+  void installJspi();
+
   void setCaptureThrowsAsRejections(bool capture);
   void setCommonJsExportDefault(bool exportDefault);
 

src/workerd/jsg/setup.c++

@@ -343,6 +343,7 @@ IsolateBase::IsolateBase(const V8System& system,
 
     ptr->SetModifyCodeGenerationFromStringsCallback(&modifyCodeGenCallback);
     ptr->SetAllowWasmCodeGenerationCallback(&allowWasmCallback);
+    ptr->SetWasmJSPIEnabledCallback(&jspiEnabledCallback);
 
     // We don't support SharedArrayBuffer so Atomics.wait() doesn't make sense, and might allow DoS
     // attacks.
@@ -461,6 +462,12 @@ bool IsolateBase::allowWasmCallback(v8::Local<v8::Context> context, v8::Local<v8
   return self->evalAllowed;
 }
 
+bool IsolateBase::jspiEnabledCallback(v8::Local<v8::Context> context) {
+  IsolateBase* self =
+      static_cast<IsolateBase*>(context->GetIsolate()->GetData(SET_DATA_ISOLATE_BASE));
+  return self->jspiEnabled;
+}
+
 void IsolateBase::jitCodeEvent(const v8::JitCodeEvent* event) noexcept {
   // We register this callback with V8 in order to build a mapping of code addresses to source
   // code locations, which we use when reporting stack traces during crashes.

src/workerd/jsg/setup.h

@@ -122,6 +122,9 @@ class IsolateBase {
   inline void setAllowEval(kj::Badge<Lock>, bool allow) {
     evalAllowed = allow;
   }
+  inline void setJspiEnabled(kj::Badge<Lock>, bool enabled) {
+    jspiEnabled = enabled;
+  }
   inline void setCaptureThrowsAsRejections(kj::Badge<Lock>, bool capture) {
     captureThrowsAsRejections = capture;
   }
@@ -236,6 +239,7 @@ class IsolateBase {
   v8::Isolate* ptr;
   kj::Maybe<kj::String> uuid;
   bool evalAllowed = false;
+  bool jspiEnabled = false;
 
   // The Web Platform API specifications require that any API that returns a JavaScript Promise
   // should never throw errors synchronously. Rather, they are supposed to capture any synchronous
@@ -320,6 +324,7 @@ class IsolateBase {
   static v8::ModifyCodeGenerationFromStringsResult modifyCodeGenCallback(
       v8::Local<v8::Context> context, v8::Local<v8::Value> source, bool isCodeLike);
   static bool allowWasmCallback(v8::Local<v8::Context> context, v8::Local<v8::String> source);
+  static bool jspiEnabledCallback(v8::Local<v8::Context> context);
 
   static void jitCodeEvent(const v8::JitCodeEvent* event) noexcept;
 

src/workerd/server/tests/python/hello/worker.py

@@ -1,4 +1,8 @@
 def test():
+    from js.WebAssembly import Suspending
+
+    Suspending  # noqa: B018
+
     # This just tests that nothing raises when we run this. It isn't great though
     # because we don't test whether we printed anything.
     # TODO: update this to test that something happened