From 8fcd67250c25439ec3fd5190c7bfd302229ca5aa Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Tue, 23 Jul 2024 15:27:12 +0200
Subject: [PATCH] Load Pyodide runtime from external capnproto file

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
---
 src/workerd/api/modules.h           |  4 +++-
 src/workerd/api/pyodide/pyodide.c++ | 10 ++++++++++
 src/workerd/api/pyodide/pyodide.h   | 24 ++++++++++++++++++------
 src/workerd/util/autogate.c++       |  2 ++
 src/workerd/util/autogate.h         |  1 +
 5 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/src/workerd/api/modules.h b/src/workerd/api/modules.h
index 419aae1fbb8..666e51263f8 100644
--- 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++
index a491ec87b2b..ab6fcbde523 100644
--- 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
index 6e7d9d9f3f8..19d692fa5da 100644
--- 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++
index 6430dd067cb..e21b5d916f6 100644
--- 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
index 8a97eb1fec3..0a7598c0d0b 100644
--- 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.
 };