From 9e0b76c7edc7bc063dd9331c112cc1b8ab12a475 Mon Sep 17 00:00:00 2001
From: Garrett Gu <garrett@cloudflare.com>
Date: Tue, 10 Sep 2024 12:12:56 +0000
Subject: [PATCH] Address Hood feedback

---
 src/pyodide/internal/loadPackage.ts     | 22 ++++++++++------------
 src/pyodide/internal/setupPackages.ts   |  7 +++++++
 src/pyodide/python-entrypoint-helper.ts |  1 -
 src/workerd/api/pyodide/pyodide.h       |  4 +++-
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/pyodide/internal/loadPackage.ts b/src/pyodide/internal/loadPackage.ts
index c7c5c70c314e..d78cb43aa6a7 100644
--- a/src/pyodide/internal/loadPackage.ts
+++ b/src/pyodide/internal/loadPackage.ts
@@ -40,16 +40,14 @@ async function decompressArrayBuffer(
 
 // loadBundleFromR2 loads the package from the internet (through fetch) and uses the DiskCache as
 // a backing store. This is only used in local dev.
-async function loadBundleFromR2(
-  requirement: string
-): Promise<[string, Reader]> {
+async function loadBundleFromR2(requirement: string): Promise<Reader> {
   // first check if the disk cache has what we want
   const filename = LOCKFILE['packages'][requirement]['file_name'];
   const cached = DiskCache.get(filename);
   if (cached) {
     const decompressed = await decompressArrayBuffer(cached);
     const reader = new ArrayBufferReader(decompressed);
-    return [requirement, reader];
+    return reader;
   }
 
   // we didn't find it in the disk cache, continue with original fetch
@@ -61,22 +59,22 @@ async function loadBundleFromR2(
 
   DiskCache.put(filename, compressed);
   const reader = new ArrayBufferReader(decompressed);
-  return [requirement, reader];
+  return reader;
 }
 
 async function loadBundleFromArtifactBundler(
   requirement: string
-): Promise<[string, Reader]> {
+): Promise<Reader> {
   const packagesVersion = PACKAGES_VERSION;
   const filename = LOCKFILE['packages'][requirement]['file_name'];
   const fullPath = 'python-package-bucket/' + packagesVersion + '/' + filename;
   const reader = ArtifactBundler.getPackage(fullPath);
-  if (!reader)
+  if (!reader) {
     throw new Error(
       'Failed to get package ' + fullPath + ' from ArtifactBundler'
     );
-  return Promise.resolve([requirement, reader]);
-  // ^ this is okay to do during startup since it resolves immediately
+  }
+  return reader;
 }
 
 /**
@@ -102,13 +100,13 @@ class ArrayBufferReader {
 async function loadPackagesImpl(
   Module: Module,
   requirements: Set<string>,
-  loadBundle: (req: string) => Promise<[string, Reader]>
+  loadBundle: (req: string) => Promise<Reader>
 ) {
-  let loadPromises = [];
+  let loadPromises: Promise<[string, Reader]>[] = [];
   let loading = [];
   for (const req of requirements) {
     if (SITE_PACKAGES.loadedRequirements.has(req)) continue;
-    loadPromises.push(loadBundle(req));
+    loadPromises.push(loadBundle(req).then((r) => [req, r]));
     loading.push(req);
   }
 
diff --git a/src/pyodide/internal/setupPackages.ts b/src/pyodide/internal/setupPackages.ts
index 783124a32eb0..fe386f714a42 100644
--- a/src/pyodide/internal/setupPackages.ts
+++ b/src/pyodide/internal/setupPackages.ts
@@ -130,6 +130,13 @@ export function buildSitePackages(requirements: Set<string>): SitePackagesDir {
   const [bigTarInfo, bigTarSoFiles] = parseTarInfo();
 
   let requirementsInBigBundle = new Set([...STDLIB_PACKAGES]);
+
+  // Currently, we include all packages within the big bundle in Edgeworker.
+  // During this transitionary period, we add the option (via autogate)
+  // to load packages from GCS (in which case they are accessible through the ArtifactBundler)
+  // or to simply use the packages within the big bundle. The latter is not ideal
+  // since we're locked to a specific packages version, so we will want to move away
+  // from it eventually.
   if (!LOAD_WHEELS_FROM_R2 && !LOAD_WHEELS_FROM_ARTIFACT_BUNDLER) {
     requirements.forEach((r) => requirementsInBigBundle.add(r));
   }
diff --git a/src/pyodide/python-entrypoint-helper.ts b/src/pyodide/python-entrypoint-helper.ts
index b718631c6954..b48346329c45 100644
--- a/src/pyodide/python-entrypoint-helper.ts
+++ b/src/pyodide/python-entrypoint-helper.ts
@@ -188,7 +188,6 @@ try {
       if (typeof mainModule[pyHandlerName] === 'function') {
         handlers[handlerName] = makeHandler(pyHandlerName);
       }
-      handlers[handlerName] = makeHandler(pyHandlerName);
     }
   }
   /**
diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h
index 7db85b229371..e93e34c4f036 100644
--- a/src/workerd/api/pyodide/pyodide.h
+++ b/src/workerd/api/pyodide/pyodide.h
@@ -204,7 +204,9 @@ struct MemorySnapshotResult {
 class ArtifactBundler: public jsg::Object {
 public:
   kj::Maybe<const PyodidePackageManager&> packageManager;
-  // ^ lifetime should be static since there is normally one worker set for the whole process. see worker-set.h
+  // ^ lifetime should be contained by lifetime of ArtifactBundler since there is normally one worker set for the whole process. see worker-set.h
+  // In other words:
+  // WorkerSet lifetime = PackageManager lifetime and Worker lifetime = ArtifactBundler lifetime and WorkerSet owns and will outlive Worker, so PackageManager outlives ArtifactBundler
   kj::Maybe<MemorySnapshotResult> storedSnapshot;
 
   ArtifactBundler(kj::Maybe<const PyodidePackageManager&> packageManager,