Load Pyodide runtime from external capnproto file (#2430)

These changes are needed to load the Pyodide runtime from an external capnproto file. In workerd, we'll keep loading Pyodide in the same way, but we moved this to WorkerdAPI. In our internal codebase, we now download the bundle and call `setPyodideBundleData` when we load the first Python worker. Then when setting up the module registry, we add the bundle from `pyodideBundleGlobal`. TODO in followups: 1. Make workerd load the bundle in the same way 2. Add support for multiple versions of the bundle 3. Testing?
cloudflare · Jul 26, 2024 · 1bf1c83 · 1bf1c83
1 parent 32d67fd
commit 1bf1c83
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 7 deletions.
diff --git a/src/workerd/api/modules.h b/src/workerd/api/modules.h
@@ -20,7 +20,9 @@ namespace workerd::api {
 template <class Registry>
 void registerModules(Registry& registry, auto featureFlags) {
   node::registerNodeJsCompatModules(registry, featureFlags);
-  pyodide::registerPyodideModules(registry, featureFlags);
+  if (featureFlags.getPythonWorkers()) {
+    pyodide::registerPyodideModules(registry, featureFlags);
+  }
   registerUnsafeModules(registry, featureFlags);
   if (featureFlags.getRttiApi()) {
     registerRTTIModule(registry);

diff --git a/src/workerd/api/pyodide/pyodide.c++ b/src/workerd/api/pyodide/pyodide.c++
@@ -7,6 +7,16 @@
 
 namespace workerd::api::pyodide {
 
+kj::Maybe<kj::Array<unsigned char>> pyodideBundleDataGlobal = kj::none;
+kj::Maybe<kj::Own<capnp::FlatArrayMessageReader>> pyodideBundleReaderGlobal = kj::none;
+kj::Maybe<jsg::Bundle::Reader> pyodideBundleGlobal = kj::none;
+
+void setPyodideBundleData(kj::Array<unsigned char> data) {
+  pyodideBundleReaderGlobal = kj::heap<capnp::FlatArrayMessageReader>(kj::arrayPtr(
+        reinterpret_cast<const capnp::word*>(data.begin()), data.size() / sizeof(capnp::word))).attach(kj::mv(data));
+  pyodideBundleGlobal = KJ_REQUIRE_NONNULL(pyodideBundleReaderGlobal)->getRoot<jsg::Bundle>();
+}
+
 static int readToTarget(kj::ArrayPtr<const kj::byte> source, int offset, kj::ArrayPtr<kj::byte> buf) {
   int size = source.size();
   if (offset >= size || offset < 0) {

diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h
@@ -11,9 +11,17 @@
 #include <workerd/jsg/url.h>
 #include <workerd/server/workerd.capnp.h>
 #include <workerd/io/io-context.h>
+#include <workerd/util/autogate.h>
+#include "capnp/serialize.h"
 
 namespace workerd::api::pyodide {
 
+// singleton that owns bundle
+extern kj::Maybe<jsg::Bundle::Reader> pyodideBundleGlobal;
+
+void setPyodideBundleData(kj::Array<unsigned char> data);
+
+
 struct PythonConfig {
   kj::Maybe<kj::Own<const kj::Directory>> diskCacheRoot;
   bool createSnapshot;
@@ -344,25 +352,29 @@ bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader module
   api::pyodide::SimplePythonLimiter
 
 template <class Registry> void registerPyodideModules(Registry& registry, auto featureFlags) {
-  if (featureFlags.getPythonWorkers()) {
-    // We add `pyodide:` packages here including python-entrypoint-helper.js.
+  // We add `pyodide:` packages here including python-entrypoint-helper.js.
+  if (!util::Autogate::isEnabled(util::AutogateKey::PYODIDE_LOAD_EXTERNAL)) {
     registry.addBuiltinBundle(PYODIDE_BUNDLE, kj::none);
-    registry.template addBuiltinModule<PackagesTarReader>(
-        "pyodide-internal:packages_tar_reader", workerd::jsg::ModuleRegistry::Type::INTERNAL);
   }
+  registry.template addBuiltinModule<PackagesTarReader>(
+      "pyodide-internal:packages_tar_reader", workerd::jsg::ModuleRegistry::Type::INTERNAL);
 }
 
 kj::Own<jsg::modules::ModuleBundle> getInternalPyodideModuleBundle(auto featureFlags) {
   jsg::modules::ModuleBundle::BuiltinBuilder builder(
       jsg::modules::ModuleBundle::BuiltinBuilder::Type::BUILTIN_ONLY);
-  jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builder, PYODIDE_BUNDLE);
+  if (!util::Autogate::isEnabled(util::AutogateKey::PYODIDE_LOAD_EXTERNAL)) {
+    jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builder, PYODIDE_BUNDLE);
+  }
   return builder.finish();
 }
 
 kj::Own<jsg::modules::ModuleBundle> getExternalPyodideModuleBundle(auto featureFlags) {
   jsg::modules::ModuleBundle::BuiltinBuilder builder(
       jsg::modules::ModuleBundle::BuiltinBuilder::Type::BUILTIN);
-  jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builder, PYODIDE_BUNDLE);
+  if (!util::Autogate::isEnabled(util::AutogateKey::PYODIDE_LOAD_EXTERNAL)) {
+    jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builder, PYODIDE_BUNDLE);
+  }
   return builder.finish();
 }
 

diff --git a/src/workerd/util/autogate.c++ b/src/workerd/util/autogate.c++
@@ -15,6 +15,8 @@ kj::StringPtr KJ_STRINGIFY(AutogateKey key) {
   switch (key) {
     case AutogateKey::TEST_WORKERD:
       return "test-workerd"_kj;
+    case AutogateKey::PYODIDE_LOAD_EXTERNAL:
+      return "pyodide-load-external"_kj;
     case AutogateKey::NumOfKeys:
       KJ_FAIL_ASSERT("NumOfKeys should not be used in getName");
   }

diff --git a/src/workerd/util/autogate.h b/src/workerd/util/autogate.h
@@ -13,6 +13,7 @@ namespace workerd::util {
 // Workerd-specific list of autogate keys (can also be used in internal repo).
 enum class AutogateKey {
   TEST_WORKERD,
+  PYODIDE_LOAD_EXTERNAL,
   NumOfKeys // Reserved for iteration.
 };