Implement PackagePromiseMap in a way that avoids storing kj::Promises…

… on the v8 isolate heap. Note that we can't just wrap the promise in an IoOwn because we don't yet have an IoContext when the promise is created

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

@@ -211,4 +211,43 @@ bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader module
   return false;
 }
 
+void PackagePromiseMap::insert(
+    kj::String path, kj::Promise<kj::Own<SmallPackagesTarReader>> promise) {
+  KJ_ASSERT(waitlists.find(kj::str(path)) == kj::none);
+  waitlists.insert(kj::str(path), CrossThreadWaitList());
+
+  promise
+      .then([this, p = kj::mv(path)](auto a) {
+    auto path = kj::str(p);
+    auto maybeWaitlist = waitlists.find(p);
+    KJ_IF_SOME(waitlist, maybeWaitlist) {
+      fetchedPackages.insert(kj::mv(path), kj::mv(a));
+      waitlist.fulfill();
+    } else {
+      JSG_FAIL_REQUIRE(Error, "Failed to get waitlist for package", path);
+    }
+  }).detach([](kj::Exception&& exception) {
+    JSG_FAIL_REQUIRE(Error, "Failed to get package", exception);
+  });
+}
+
+kj::Promise<kj::Own<SmallPackagesTarReader>> PackagePromiseMap::getPromise(kj::StringPtr path) {
+  auto maybeWaitlist = waitlists.find(path);
+
+  KJ_IF_SOME(waitlist, maybeWaitlist) {
+    co_await waitlist.addWaiter();
+    auto maybeEntry = fetchedPackages.findEntry(path);
+    KJ_IF_SOME(entryRef, maybeEntry) {
+      auto entry = fetchedPackages.release(entryRef);
+
+      co_return kj::mv(entry.value);
+    } else {
+      JSG_FAIL_REQUIRE(Error, "Failed to get package when trying to get promise", path);
+    }
+
+  } else {
+    JSG_FAIL_REQUIRE(Error, "Failed to get waitlist for package when trying to get promise", path);
+  }
+}
+
 }  // namespace workerd::api::pyodide

src/workerd/api/pyodide/pyodide.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "workerd/util/wait-list.h"
 #include <kj/array.h>
 #include <kj/debug.h>
 #include <kj/common.h>
@@ -187,12 +188,26 @@ struct MemorySnapshotResult {
   JSG_STRUCT(snapshot, importedModulesList);
 };
 
+// Before a Pyodide isolate starts up, all its packages begin loading. This struct provides an
+// interface for the Pyodide startup code to request promises for each package that was loaded.
+struct PackagePromiseMap {
+  // This implementation is complex because we're not allowed to store kj::Promises on an object
+  // stored on the V8 heap. Ideally, we'd like to simply store a map from package paths to Promises.
+  // Instead, we store a CrossThreadWaitList and get a new Promise from the waitlist as needed.
+  kj::HashMap<kj::String, kj::Own<SmallPackagesTarReader>> fetchedPackages;
+  kj::HashMap<kj::String, CrossThreadWaitList> waitlists;
+
+  void insert(kj::String path, kj::Promise<kj::Own<SmallPackagesTarReader>> promise);
+
+  // Note it is explicitly not allowed to call getPromise more than once on the same package
+  // This is because allowing this causes ownership of the SmallPackagesTarReader to be muddied.
+  kj::Promise<kj::Own<SmallPackagesTarReader>> getPromise(kj::StringPtr path);
+};
+
 // A loaded bundle of artifacts for a particular script id. It can also contain V8 version and
 // CPU architecture-specific artifacts. The logic for loading these is in getArtifacts.
 class ArtifactBundler: public jsg::Object {
 public:
-  typedef kj::HashMap<kj::String, IoOwn<kj::Promise<kj::Own<SmallPackagesTarReader>>>>
-      PackagePromiseMap;
   kj::Maybe<MemorySnapshotResult> storedSnapshot;
   kj::Own<PackagePromiseMap> loadedPackages;
 
@@ -251,6 +266,10 @@ class ArtifactBundler: public jsg::Object {
     return false;  // TODO(later): Remove this function once we regenerate the bundle.
   }
 
+  kj::Promise<kj::Own<SmallPackagesTarReader>> getPackage(kj::StringPtr path) {
+    return loadedPackages->getPromise(path);
+  }
+
   JSG_RESOURCE_TYPE(ArtifactBundler) {
     JSG_METHOD(hasMemorySnapshot);
     JSG_METHOD(getMemorySnapshotSize);