Instantiate Emscripten Runtime for python workers earlier.

Move ownership of metrics and limitEnforcer to the api type. Currently ownership is shared even though the Isolate class encapsulates the api class. Moving complete ownership to the underlying api class allows the isolate class to be constructed in a different scope to the api class. This is useful for preinitialization of the api class before a request has come in. Add updateConfiguration function to jsg Isolates This can be used to update the given configuration at runtime. Note that while some jsg structs are lazily using the configuration, others can use it at construction and will have the original configuration value.
cloudflare · Nov 7, 2024 · fae3f8e · fae3f8e
1 parent 8bf3af4
commit fae3f8e
Show file tree

Hide file tree

Showing 21 changed files with 346 additions and 89 deletions.
diff --git a/src/pyodide/BUILD.bazel b/src/pyodide/BUILD.bazel
@@ -197,8 +197,6 @@ INTERNAL_MODULES = glob(
     [
         "internal/*.ts",
         "internal/topLevelEntropy/*.ts",
-        # The pool directory is only needed by typescript, it shouldn't be used at runtime.
-        "internal/pool/*.ts",
         "types/*.ts",
         "types/*/*.ts",
     ],
@@ -211,27 +209,27 @@ INTERNAL_DATA_MODULES = glob([
     "internal/*.py",
     "internal/patches/*.py",
     "internal/topLevelEntropy/*.py",
-])
+]) + [
+    "generated/python_stdlib.zip",
+    "generated/pyodide.asm.wasm",
+    "generated/emscriptenSetup.js",
+]
 
 wd_ts_bundle(
     name = "pyodide",
     eslintrc_json = "eslint.config.mjs",
     import_name = "pyodide",
-    internal_data_modules = ["generated/python_stdlib.zip"] + INTERNAL_DATA_MODULES,
+    internal_data_modules = INTERNAL_DATA_MODULES,
     internal_json_modules = [
         "generated/pyodide-lock.json",
         "generated/pyodide-bucket.json",
     ],
-    internal_modules = [
-        "generated/emscriptenSetup.js",
-    ] + INTERNAL_MODULES,
-    internal_wasm_modules = ["generated/pyodide.asm.wasm"],
+    internal_modules = INTERNAL_MODULES,
     js_deps = [
         "generated/emscriptenSetup",
-        "pyodide.asm.js@rule",
         "pyodide.asm.wasm@rule",
-        "pyodide-lock.js@rule",
         "python_stdlib.zip@rule",
+        "pyodide-lock.js@rule",
         "pyodide-bucket.json@rule",
     ],
     lint = False,
@@ -264,9 +262,9 @@ genrule(
         for m in INTERNAL_DATA_MODULES
         if m.endswith(".py")
     ] + [
-        ":pyodide-internal_generated_emscriptenSetup",
-        ":pyodide-internal_generated_pyodide.asm.wasm",
-        ":pyodide-internal_generated_python_stdlib.zip",
+        "pyodide-internal_generated_emscriptenSetup.js",
+        "pyodide-internal_generated_pyodide.asm.wasm",
+        "pyodide-internal_generated_python_stdlib.zip",
         ":pyodide-internal_generated_pyodide-lock.json",
         ":pyodide-internal_generated_pyodide-bucket.json",
     ],
@@ -286,3 +284,66 @@ genrule(
     tools = ["@capnp-cpp//src/capnp:capnp_tool"],
     visibility = ["//visibility:public"],
 )
+
+# copy_file(
+#     name = "emscripten_setup_capnp_file",
+#     src = "emscripten_setup.capnp",
+#     out = "generated/emscripten_setup.capnp",
+# )
+
+# genrule(
+#     name = "emscripten_setup.capnp.bin@rule",
+#     srcs = [
+#         ":generated/emscripten_setup.capnp",
+#         ":generated/emscriptenSetup.js",
+#         ":generated/python_stdlib.zip",
+#         ":generated/pyodide.asm.wasm",
+#     ],
+#     outs = ["emscripten_setup.capnp.bin"],
+#     cmd = " ".join([
+#         "$(execpath @capnp-cpp//src/capnp:capnp_tool)",
+#         "eval",
+#         "$(location generated/emscripten_setup.capnp)",
+#         "emscriptenSetup",
+#         "-o binary",
+#         "> $@",
+#     ]),
+#     tools = ["@capnp-cpp//src/capnp:capnp_tool"],
+#     visibility = ["//visibility:public"],
+# )
+
+# capnp_embed(
+#     name = "emscripten_setup_capnp_file_embed",
+#     src = "generated/emscripten_setup.capnp",
+#     deps = ["emscripten_setup_capnp_file"],
+# )
+
+# capnp_embed(
+#     name = "emscripten_setup_js_file_embed",
+#     src = "generated/emscriptenSetup.js",
+#     deps = ["generated/emscriptenSetup"],
+# )
+
+# capnp_embed(
+#     name = "python_stdlib_zip_file_embed",
+#     src = "generated/python_stdlib.zip",
+#     deps = ["python_stdlib.zip@rule"],
+# )
+
+# capnp_embed(
+#     name = "pyodide_asm_wasm_file_embed",
+#     src = "generated/pyodide.asm.wasm",
+#     deps = ["pyodide.asm.wasm@rule"],
+# )
+
+# cc_capnp_library(
+#     name = "emscripten_setup_capnp",
+#     srcs = ["generated/emscripten_setup.capnp"],
+#     visibility = ["//visibility:public"],
+#     deps = [
+#         ":emscripten_setup_capnp_file_embed",
+#         ":emscripten_setup_js_file_embed",
+#         ":pyodide_asm_wasm_file_embed",
+#         ":python_stdlib_zip_file_embed",
+#     ],
+# )
diff --git a/src/pyodide/emscripten_setup.capnp b/src/pyodide/emscripten_setup.capnp
@@ -0,0 +1,13 @@
+@0xc00ad00cc650fb45;
+
+struct EmscriptenSetup {
+  code @0 :Text;
+  pyodideAsmWasm @1 :Data;
+  pythonStdlibZip @2 :Data;
+}
+
+const emscriptenSetup :EmscriptenSetup = (
+  code = embed "emscriptenSetup.js",
+  pyodideAsmWasm = embed "pyodide.asm.wasm",
+  pythonStdlibZip = embed "python_stdlib.zip",
+);
diff --git a/src/pyodide/internal/pool/emscriptenSetup.ts b/src/pyodide/internal/pool/emscriptenSetup.ts
@@ -13,7 +13,7 @@ import { reportError } from 'pyodide-internal:util';
  */
 import { _createPyodideModule } from 'pyodide-internal:generated/pyodide.asm';
 
-export {
+import {
   setUnsafeEval,
   setGetRandomValues,
 } from 'pyodide-internal:pool/builtin_wrappers';
@@ -56,7 +56,7 @@ function getWaitForDynlibs(resolveReadyPromise: PreRunHook): PreRunHook {
  * This is a simplified version of the `prepareFileSystem` function here:
  * https://github.com/pyodide/pyodide/blob/main/src/js/module.ts
  */
-function getPrepareFileSystem(pythonStdlib: Uint8Array): PreRunHook {
+function getPrepareFileSystem(pythonStdlib: ArrayBuffer): PreRunHook {
   return function prepareFileSystem(Module: Module): void {
     try {
       const pymajor = Module._py_version_major();
@@ -118,7 +118,7 @@ function getInstantiateWasm(
  */
 function getEmscriptenSettings(
   isWorkerd: boolean,
-  pythonStdlib: Uint8Array,
+  pythonStdlib: ArrayBuffer,
   pyodideWasmModule: WebAssembly.Module
 ): EmscriptenSettings {
   const config: PyodideConfig = {
@@ -193,7 +193,7 @@ function* featureDetectionMonkeyPatchesContextManager() {
  */
 export async function instantiateEmscriptenModule(
   isWorkerd: boolean,
-  pythonStdlib: Uint8Array,
+  pythonStdlib: ArrayBuffer,
   wasmModule: WebAssembly.Module
 ): Promise<Module> {
   const emscriptenSettings = getEmscriptenSettings(
@@ -210,6 +210,8 @@ export async function instantiateEmscriptenModule(
 
     // Wait until we've executed all the preRun hooks before proceeding
     const emscriptenModule = await emscriptenSettings.readyPromise;
+    emscriptenModule.setUnsafeEval = setUnsafeEval;
+    emscriptenModule.setGetRandomValues = setGetRandomValues;
     return emscriptenModule;
   } catch (e) {
     console.warn('Error in instantiateEmscriptenModule');

diff --git a/src/pyodide/internal/python.ts b/src/pyodide/internal/python.ts
@@ -18,39 +18,11 @@ import {
   entropyBeforeTopLevel,
   getRandomValues,
 } from 'pyodide-internal:topLevelEntropy/lib';
+import { default as SetupEmscripten } from 'internal:setup-emscripten';
+
 import { default as UnsafeEval } from 'internal:unsafe-eval';
 import { simpleRunPython } from 'pyodide-internal:util';
 
-/**
- * This file is a simplified version of the Pyodide loader:
- * https://github.com/pyodide/pyodide/blob/main/src/js/pyodide.ts
- *
- * In particular, it drops the package lock, which disables
- * `pyodide.loadPackage`. In trade we add memory snapshots here.
- */
-
-/**
- * _createPyodideModule and pyodideWasmModule together are produced by the
- * Emscripten linker
- */
-import pyodideWasmModule from 'pyodide-internal:generated/pyodide.asm.wasm';
-
-/**
- * The Python and Pyodide stdlib zipped together. The zip format is convenient
- * because Python has a "ziploader" that allows one to import directly from a
- * zip file.
- *
- * The ziploader solves bootstrapping problems around unpacking: Python comes
- * with a bunch of C libs to unpack various archive formats, but they need stuff
- * in this zip file to initialize their runtime state.
- */
-import pythonStdlib from 'pyodide-internal:generated/python_stdlib.zip';
-import {
-  instantiateEmscriptenModule,
-  setUnsafeEval,
-  setGetRandomValues,
-} from 'pyodide-internal:generated/emscriptenSetup';
-
 /**
  * After running `instantiateEmscriptenModule` but before calling into any C
  * APIs, we call this function. If `MEMORY` is defined, then we will have passed
@@ -90,14 +62,15 @@ export async function loadPyodide(
   indexURL: string
 ): Promise<Pyodide> {
   const Module = await enterJaegerSpan('instantiate_emscripten', () =>
-    instantiateEmscriptenModule(isWorkerd, pythonStdlib, pyodideWasmModule)
+    SetupEmscripten.getModule()
   );
+  Module.API.config.jsglobals = globalThis;
   if (isWorkerd) {
     Module.API.config.indexURL = indexURL;
     Module.API.config.resolveLockFilePromise!(lockfile);
   }
-  setUnsafeEval(UnsafeEval);
-  setGetRandomValues(getRandomValues);
+  Module.setUnsafeEval(UnsafeEval);
+  Module.setGetRandomValues(getRandomValues);
   await enterJaegerSpan('prepare_wasm_linear_memory', () =>
     prepareWasmLinearMemory(Module)
   );

diff --git a/src/pyodide/types/emscripten.d.ts b/src/pyodide/types/emscripten.d.ts
@@ -68,4 +68,8 @@ interface Module {
   addRunDependency(x: string): void;
   removeRunDependency(x: string): void;
   noInitialRun: boolean;
+  setUnsafeEval(mod: typeof import('internal:unsafe-eval').default): void;
+  setGetRandomValues(
+    func: typeof import('pyodide-internal:topLevelEntropy/lib').getRandomValues
+  ): void;
 }
diff --git a/src/pyodide/types/setup-emscripten.d.ts b/src/pyodide/types/setup-emscripten.d.ts
@@ -0,0 +1,5 @@
+declare namespace SetupEmscripten {
+  const getModule: () => Module;
+}
+
+export default SetupEmscripten;
diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel
@@ -14,6 +14,7 @@ filegroup(
             "html-rewriter.c++",
             "hyperdrive.c++",
             "pyodide/pyodide.c++",
+            "pyodide/setup-emscripten.c++",
             "memory-cache.c++",
             "r2*.c++",
             "rtti.c++",
@@ -37,6 +38,7 @@ filegroup(
             "hyperdrive.h",
             "memory-cache.h",
             "pyodide/pyodide.h",
+            "pyodide/setup-emscripten.h",
             "modules.h",
             "r2*.h",
             "rtti.h",
@@ -126,9 +128,11 @@ wd_cc_library(
     name = "pyodide",
     srcs = [
         "pyodide/pyodide.c++",
+        "pyodide/setup-emscripten.c++",
     ],
     hdrs = [
         "pyodide/pyodide.h",
+        "pyodide/setup-emscripten.h",
         "//src/pyodide:generated/pyodide_extra.capnp.h",
     ],
     implementation_deps = ["//src/workerd/util:string-buffer"],

diff --git a/src/workerd/api/modules.h b/src/workerd/api/modules.h
@@ -49,7 +49,7 @@ void registerBuiltinModules(jsg::modules::ModuleRegistry::Builder& builder, auto
 
   if (featureFlags.getPythonWorkers()) {
     builder.add(pyodide::getExternalPyodideModuleBundle(featureFlags));
-    builder.add(pyodide::getInternalPyodideModuleBundle(featureFlags));
+    builder.add(pyodide::getInternalPyodideModuleBundle<TypeWrapper>(featureFlags));
   }
 
   if (featureFlags.getRttiApi()) {

diff --git a/src/workerd/api/pyodide/pyodide.c++ b/src/workerd/api/pyodide/pyodide.c++
@@ -3,6 +3,8 @@
 //     https://opensource.org/licenses/Apache-2.0
 #include "pyodide.h"
 
+#include <workerd/api/pyodide/setup-emscripten.h>
+#include <workerd/io/worker.h>
 #include <workerd/util/string-buffer.h>
 #include <workerd/util/strings.h>
 
@@ -484,6 +486,24 @@ void DiskCache::put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data) {
   }
 }
 
+jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
+  KJ_IF_SOME(module, emscriptenModule) {
+    return module.getHandle(js);
+  } else {
+    auto& runtime = KJ_ASSERT_NONNULL(workerd::Worker::Api::current().getEmscriptenRuntime());
+    js.v8Context()->SetSecurityToken(runtime.contextToken.getHandle(js));
+    emscriptenModule = runtime.emscriptenRuntime;
+    return KJ_ASSERT_NONNULL(emscriptenModule).getHandle(js);
+  }
+}
+
+void SetupEmscripten::visitForGc(jsg::GcVisitor& visitor) {
+  // const_cast is ok because the GcVisitor doesn't actually change the underlying value of the object.
+  KJ_IF_SOME(module, emscriptenModule) {
+    visitor.visit(const_cast<jsg::JsRef<jsg::JsValue>&>(module));
+  }
+}
+
 bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader modules) {
   for (auto module: modules) {
     if (module.isPythonModule()) {

diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h
@@ -409,6 +409,24 @@ class SimplePythonLimiter: public jsg::Object {
   }
 };
 
+class SetupEmscripten: public jsg::Object {
+public:
+  SetupEmscripten() {};
+  SetupEmscripten(jsg::Lock& js, const jsg::Url&) {}
+
+  jsg::JsValue getModule(jsg::Lock& js);
+
+  JSG_RESOURCE_TYPE(SetupEmscripten) {
+    JSG_METHOD(getModule);
+  }
+
+private:
+  // Reference to the api value of the emscripten module.
+  // Used for visitForGc when no js is currently running.
+  kj::Maybe<const jsg::JsRef<jsg::JsValue>&> emscriptenModule;
+  void visitForGc(jsg::GcVisitor& visitor);
+};
+
 using Worker = server::config::Worker;
 
 jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(
@@ -420,7 +438,7 @@ bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader module
   api::pyodide::PackagesTarReader, api::pyodide::PyodideMetadataReader,                            \
       api::pyodide::ArtifactBundler, api::pyodide::DiskCache,                                      \
       api::pyodide::DisabledInternalJaeger, api::pyodide::SimplePythonLimiter,                     \
-      api::pyodide::MemorySnapshotResult
+      api::pyodide::MemorySnapshotResult, api::pyodide::SetupEmscripten
 
 template <class Registry>
 void registerPyodideModules(Registry& registry, auto featureFlags) {
@@ -431,15 +449,20 @@ void registerPyodideModules(Registry& registry, auto featureFlags) {
   }
   registry.template addBuiltinModule<PackagesTarReader>(
       "pyodide-internal:packages_tar_reader", workerd::jsg::ModuleRegistry::Type::INTERNAL);
+  registry.template addBuiltinModule<SetupEmscripten>(
+      "internal:setup-emscripten", workerd::jsg::ModuleRegistry::Type::INTERNAL);
 }
 
+template <typename TypeWrapper>
 kj::Own<jsg::modules::ModuleBundle> getInternalPyodideModuleBundle(auto featureFlags) {
   jsg::modules::ModuleBundle::BuiltinBuilder builder(
       jsg::modules::ModuleBundle::BuiltinBuilder::Type::BUILTIN_ONLY);
   if (!featureFlags.getPythonExternalBundle() &&
       !util::Autogate::isEnabled(util::AutogateKey::PYTHON_EXTERNAL_BUNDLE)) {
     jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builder, PYODIDE_BUNDLE);
   }
+  static const auto kSpecifier = "internal:setup-emscripten"_url;
+  builder.addObject<SetupEmscripten, TypeWrapper>(kSpecifier);
   return builder.finish();
 }