Add PackagePromiseMap to ArtifactBundler

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

@@ -27,6 +27,19 @@ void PyodideBundleManager::setPyodideBundleData(
       kj::mv(version), {.messageReader = kj::mv(messageReader), .bundle = bundle});
 }
 
+const kj::Maybe<kj::ArrayPtr<const unsigned char>> PyodidePackageManager::getPyodidePackage(
+    kj::StringPtr id) const {
+  KJ_IF_SOME(t, packages.lockShared()->find(id)) {
+    return t.asPtr();
+  }
+  return kj::none;
+}
+
+void PyodidePackageManager::setPyodidePackageData(
+    kj::String id, kj::Array<unsigned char> data) const {
+  packages.lockExclusive()->insert(kj::mv(id), kj::mv(data));
+}
+
 static int readToTarget(
     kj::ArrayPtr<const kj::byte> source, int offset, kj::ArrayPtr<kj::byte> buf) {
   int size = source.size();
@@ -45,6 +58,10 @@ int PackagesTarReader::read(jsg::Lock& js, int offset, kj::Array<kj::byte> buf)
   return readToTarget(PYODIDE_PACKAGES_TAR.get(), offset, buf);
 }
 
+int SmallPackagesTarReader::read(jsg::Lock& js, int offset, kj::Array<kj::byte> buf) {
+  return readToTarget(source, offset, buf);
+}
+
 kj::Array<jsg::JsRef<jsg::JsString>> PyodideMetadataReader::getNames(jsg::Lock& js) {
   auto builder = kj::heapArrayBuilder<jsg::JsRef<jsg::JsString>>(this->names.size());
   for (auto i: kj::zeroTo(builder.capacity())) {

src/workerd/api/pyodide/pyodide.h

@@ -29,6 +29,15 @@ class PyodideBundleManager {
   const kj::MutexGuarded<kj::HashMap<kj::String, MessageBundlePair>> bundles;
 };
 
+class PyodidePackageManager {
+public:
+  void setPyodidePackageData(kj::String id, kj::Array<unsigned char> data) const;
+  const kj::Maybe<kj::ArrayPtr<const unsigned char>> getPyodidePackage(kj::StringPtr id) const;
+
+private:
+  const kj::MutexGuarded<kj::HashMap<kj::String, kj::Array<unsigned char>>> packages;
+};
+
 struct PythonConfig {
   kj::Maybe<kj::Own<const kj::Directory>> packageDiskCacheRoot;
   kj::Maybe<kj::Own<const kj::Directory>> pyodideDiskCacheRoot;
@@ -157,6 +166,21 @@ class PyodideMetadataReader: public jsg::Object {
   }
 };
 
+// Similar to PackagesTarReader, but reads from a dynamic buffer rather than a buffer linked into
+// the binary
+class SmallPackagesTarReader: public jsg::Object {
+  kj::ArrayPtr<const kj::byte> source;
+
+public:
+  SmallPackagesTarReader(kj::ArrayPtr<const kj::byte> src): source(src) {};
+
+  int read(jsg::Lock& js, int offset, kj::Array<kj::byte> buf);
+
+  JSG_RESOURCE_TYPE(SmallPackagesTarReader) {
+    JSG_METHOD(read);
+  }
+};
+
 struct MemorySnapshotResult {
   kj::Array<kj::byte> snapshot;
   kj::Array<kj::String> importedModulesList;
@@ -167,15 +191,20 @@ struct MemorySnapshotResult {
 // CPU architecture-specific artifacts. The logic for loading these is in getArtifacts.
 class ArtifactBundler: public jsg::Object {
 public:
+  typedef kj::HashMap<kj::String, kj::Promise<kj::Own<SmallPackagesTarReader>>> PackagePromiseMap;
   kj::Maybe<MemorySnapshotResult> storedSnapshot;
+  kj::Own<PackagePromiseMap> loadedPackages;
 
-  ArtifactBundler(kj::Maybe<kj::Array<kj::byte>> existingSnapshot)
+  ArtifactBundler(kj::Own<PackagePromiseMap> loadedPackages,
+      kj::Maybe<kj::Array<kj::byte>> existingSnapshot = kj::none)
       : storedSnapshot(kj::none),
+        loadedPackages(kj::mv(loadedPackages)),
         existingSnapshot(kj::mv(existingSnapshot)),
         isValidating(false) {};
 
   ArtifactBundler(bool isValidating = false)
       : storedSnapshot(kj::none),
+        loadedPackages(kj::heap<PackagePromiseMap>()),
         existingSnapshot(kj::none),
         isValidating(isValidating) {};
 
@@ -205,8 +234,9 @@ class ArtifactBundler: public jsg::Object {
     return isValidating;
   }
 
-  static jsg::Ref<ArtifactBundler> makeDisabledBundler() {
-    return jsg::alloc<ArtifactBundler>();
+  static jsg::Ref<ArtifactBundler> makeDisabledBundler(
+      kj::Own<PackagePromiseMap> loadedPackages = kj::heap<PackagePromiseMap>()) {
+    return jsg::alloc<ArtifactBundler>(kj::mv(loadedPackages));
   }
 
   void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {