Instantiate Emscripten Runtime for python workers separately.

src/pyodide/BUILD.bazel

@@ -211,27 +211,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,7 +264,7 @@ genrule(
         for m in INTERNAL_DATA_MODULES
         if m.endswith(".py")
     ] + [
-        ":pyodide-internal_generated_emscriptenSetup",
+        ":pyodide-internal_generated_emscriptenSetup.js",
         ":pyodide-internal_generated_pyodide.asm.wasm",
         ":pyodide-internal_generated_python_stdlib.zip",
         ":pyodide-internal_generated_pyodide-lock.json",

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');

src/pyodide/internal/python.ts

@@ -18,46 +18,23 @@ import {
   entropyBeforeTopLevel,
   getRandomValues,
 } from 'pyodide-internal:topLevelEntropy/lib';
-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
+ * SetupEmscripten is an internal module defined in setup-emscripten.h the module instantiates
+ * emscripten seperately from this code in another context.
+ * The underlying code for it can be found in pool/emscriptenSetup.ts.
  */
-import pyodideWasmModule from 'pyodide-internal:generated/pyodide.asm.wasm';
+import { default as SetupEmscripten } from 'internal:setup-emscripten';
 
-/**
- * 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';
+import { default as UnsafeEval } from 'internal:unsafe-eval';
+import { simpleRunPython } from 'pyodide-internal:util';
 
 /**
  * After running `instantiateEmscriptenModule` but before calling into any C
  * APIs, we call this function. If `MEMORY` is defined, then we will have passed
  * `noInitialRun: true` and so the C runtime is in an incoherent state until we
  * restore the linear memory from the snapshot.
  */
-async function prepareWasmLinearMemory(Module: Module): Promise<void> {
+function prepareWasmLinearMemory(Module: Module): void {
   // Note: if we are restoring from a snapshot, runtime is not initialized yet.
   mountSitePackages(Module, SITE_PACKAGES.rootInfo);
   entropyMountFiles(Module);
@@ -86,21 +63,22 @@ async function prepareWasmLinearMemory(Module: Module): Promise<void> {
   adjustSysPath(Module);
 }
 
-export async function loadPyodide(
+export function loadPyodide(
   isWorkerd: boolean,
   lockfile: PackageLock,
   indexURL: string
-): Promise<Pyodide> {
-  const Module = await enterJaegerSpan('instantiate_emscripten', () =>
-    instantiateEmscriptenModule(isWorkerd, pythonStdlib, pyodideWasmModule)
+): Pyodide {
+  const Module = enterJaegerSpan('instantiate_emscripten', () =>
+    SetupEmscripten.getModule()
   );
+  Module.API.config.jsglobals = globalThis;
   if (isWorkerd) {
     Module.API.config.indexURL = indexURL;
     Module.API.config.resolveLockFilePromise!(lockfile);
   }
-  setUnsafeEval(UnsafeEval);
-  setGetRandomValues(getRandomValues);
-  await enterJaegerSpan('prepare_wasm_linear_memory', () =>
+  Module.setUnsafeEval(UnsafeEval);
+  Module.setGetRandomValues(getRandomValues);
+  enterJaegerSpan('prepare_wasm_linear_memory', () =>
     prepareWasmLinearMemory(Module)
   );
   maybeSetupSnapshotUpload(Module);
@@ -109,7 +87,7 @@ export async function loadPyodide(
   mountWorkerFiles(Module);
 
   // Finish setting up Pyodide's ffi so we can use the nice Python interface
-  await enterJaegerSpan('finalize_bootstrap', Module.API.finalizeBootstrap);
+  enterJaegerSpan('finalize_bootstrap', Module.API.finalizeBootstrap);
   const pyodide = Module.API.public_api;
   finishSnapshotSetup(pyodide);
   return pyodide;

src/pyodide/python-entrypoint-helper.ts

@@ -31,14 +31,14 @@ function pyimportMainModule(pyodide: Pyodide): PyModule {
   return pyodide.pyimport(mainModuleName);
 }
 
-let pyodidePromise: Promise<Pyodide> | undefined;
-function getPyodide(): Promise<Pyodide> {
+let pyodideInstance: Pyodide | undefined;
+function getPyodide(): Pyodide {
   return enterJaegerSpan('get_pyodide', () => {
-    if (pyodidePromise) {
-      return pyodidePromise;
+    if (pyodideInstance) {
+      return pyodideInstance;
     }
-    pyodidePromise = loadPyodide(IS_WORKERD, LOCKFILE, WORKERD_INDEX_URL);
-    return pyodidePromise;
+    pyodideInstance = loadPyodide(IS_WORKERD, LOCKFILE, WORKERD_INDEX_URL);
+    return pyodideInstance!;
   });
 }
 
@@ -118,7 +118,7 @@ function getMainModule(): Promise<PyModule> {
       return mainModulePromise;
     }
     mainModulePromise = (async function () {
-      const pyodide = await getPyodide();
+      const pyodide = getPyodide();
       await setupPackages(pyodide);
       Limiter.beginStartup();
       try {
@@ -134,7 +134,7 @@ function getMainModule(): Promise<PyModule> {
 }
 
 async function preparePython(): Promise<PyModule> {
-  const pyodide = await getPyodide();
+  const pyodide = getPyodide();
   const mainModule = await getMainModule();
   entropyBeforeRequest(pyodide._module);
   return mainModule;
@@ -177,7 +177,7 @@ try {
     if (IS_WORKERD) {
       // If we're in workerd, we have to do setupPackages in the IoContext, so don't start it yet.
       // TODO: fix this.
-      await getPyodide();
+      getPyodide();
     } else {
       // If we're not in workerd, setupPackages doesn't require IO so we can do it all here.
       await getMainModule();

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;
 }

src/pyodide/types/setup-emscripten.d.ts

@@ -0,0 +1,5 @@
+declare namespace SetupEmscripten {
+  const getModule: () => Module;
+}
+
+export default SetupEmscripten;

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"],

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

@@ -3,6 +3,7 @@
 //     https://opensource.org/licenses/Apache-2.0
 #include "pyodide.h"
 
+#include <workerd/api/pyodide/setup-emscripten.h>
 #include <workerd/util/string-buffer.h>
 #include <workerd/util/strings.h>
 
@@ -484,6 +485,15 @@ void DiskCache::put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data) {
   }
 }
 
+jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
+  js.v8Context()->SetSecurityToken(emscriptenRuntime.contextToken.getHandle(js));
+  return emscriptenRuntime.emscriptenRuntime.getHandle(js);
+}
+
+void SetupEmscripten::visitForGc(jsg::GcVisitor& visitor) {
+  visitor.visit(emscriptenRuntime.emscriptenRuntime);
+}
+
 bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader modules) {
   for (auto module: modules) {
     if (module.isPythonModule()) {

src/workerd/api/pyodide/pyodide.h

@@ -5,6 +5,7 @@
 
 #include "workerd/util/wait-list.h"
 
+#include <workerd/api/pyodide/setup-emscripten.h>
 #include <workerd/io/io-context.h>
 #include <workerd/jsg/jsg.h>
 #include <workerd/jsg/modules-new.h>
@@ -408,6 +409,22 @@ class SimplePythonLimiter: public jsg::Object {
   }
 };
 
+class SetupEmscripten: public jsg::Object {
+public:
+  SetupEmscripten(EmscriptenRuntime emscriptenRuntime)
+      : emscriptenRuntime(kj::mv(emscriptenRuntime)) {};
+
+  jsg::JsValue getModule(jsg::Lock& js);
+
+  JSG_RESOURCE_TYPE(SetupEmscripten) {
+    JSG_METHOD(getModule);
+  }
+
+private:
+  EmscriptenRuntime emscriptenRuntime;
+  void visitForGc(jsg::GcVisitor& visitor);
+};
+
 using Worker = server::config::Worker;
 
 jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(
@@ -419,6 +436,6 @@ bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader module
   api::pyodide::ReadOnlyBuffer, api::pyodide::PyodideMetadataReader,                               \
       api::pyodide::ArtifactBundler, api::pyodide::DiskCache,                                      \
       api::pyodide::DisabledInternalJaeger, api::pyodide::SimplePythonLimiter,                     \
-      api::pyodide::MemorySnapshotResult
+      api::pyodide::MemorySnapshotResult, api::pyodide::SetupEmscripten
 
 }  // namespace workerd::api::pyodide