Python snapshots: Add autogate to use snapshots from disk for develop…

…ment

src/pyodide/internal/python.js

@@ -75,7 +75,9 @@ export function getMemoryToUpload() {
   if (!MEMORY_TO_UPLOAD) {
     throw new TypeError("Expected MEMORY_TO_UPLOAD to be set");
   }
-  return MEMORY_TO_UPLOAD;
+  const tmp = MEMORY_TO_UPLOAD;
+  MEMORY_TO_UPLOAD = undefined;
+  return tmp;
 }
 
 /**
@@ -193,7 +195,11 @@ const PRELOADED_SO_FILES = [];
 function preloadDynamicLibs(Module) {
   let SO_FILES_TO_LOAD = SITE_PACKAGES_SO_FILES;
   if (IS_CREATING_BASELINE_SNAPSHOT || DSO_METADATA?.settings?.baselineSnapshot) {
-    SO_FILES_TO_LOAD = [[ '_lzma.so' ], [ '_ssl.so' ]];
+    // Ideally this should be just
+    // [[ '_lzma.so' ], [ '_ssl.so' ]]
+    // but we put a few more because we messed up the memory snapshot...
+    SO_FILES_TO_LOAD = [["_hashlib.so"],["_lzma.so"],["_sqlite3.so"],["_ssl.so"]];
+    // SO_FILES_TO_LOAD = [[ '_lzma.so' ], [ '_ssl.so' ]];
   }
   try {
     const sitePackages = getSitePackagesPath(Module);
@@ -417,7 +423,7 @@ function makeLinearMemorySnapshot(Module) {
   memorySnapshotDoImports(Module);
   const dsoJSON = recordDsoHandles(Module);
   if (IS_CREATING_BASELINE_SNAPSHOT) {
-    checkLoadedSoFiles(dsoJSON);
+    // checkLoadedSoFiles(dsoJSON);
   }
   return encodeSnapshot(Module.HEAP8, dsoJSON);
 }

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

@@ -3,6 +3,10 @@
 #include "kj/common.h"
 #include "kj/debug.h"
 
+#if !_WIN32
+#include <fcntl.h>
+#endif
+
 namespace workerd::api::pyodide {
 
 static int readToTarget(kj::ArrayPtr<const kj::byte> source, int offset, kj::ArrayPtr<kj::byte> buf) {
@@ -120,4 +124,126 @@ jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(Worker::Reader conf) {
                                            false /* isTracing */, false /* createBaselineSnapshot */, kj::none /* memorySnapshot */);
 }
 
+using namespace util;
+
+bool
+shouldStoreSnapshotToDisk() {
+  if (!Autogate::isEnabled(AutogateKey::LOCAL_DEV_PYTHON_SNAPSHOT)) {
+    return false;
+  }
+#if _WIN32
+  return false;
+#else
+  int h = open("/tmp/snapshot.bin", O_RDONLY);
+  if (h >= 0) {
+    close(h);
+    return false;
+  }
+  if (errno == ENOENT) {
+    errno = 0;
+    return true;
+  }
+  KJ_LOG(WARNING, "Local dev Python snapshots enabled, but got unexpected error opening /tmp/snapshot.bin");
+  perror("open");
+  errno = 0;
+  return false;
+#endif
+}
+
+bool
+shouldLoadSnapshotFromDisk() {
+  if (!Autogate::isEnabled(AutogateKey::LOCAL_DEV_PYTHON_SNAPSHOT)) {
+    return false;
+  }
+#if _WIN32
+  return false;
+#else
+  int h = open("/tmp/snapshot.bin", O_RDONLY);
+  if (h >= 0) {
+    close(h);
+    return true;
+  }
+  if (errno != ENOENT) {
+    KJ_LOG(WARNING, "Error opening snapshot for reading:");
+    perror("open");
+  }
+  errno = 0;
+  return false;
+#endif
+}
+
+void
+maybeStoreSnapshotToDisk(kj::ArrayPtr<kj::byte> array) {
+  if (!Autogate::isEnabled(AutogateKey::LOCAL_DEV_PYTHON_SNAPSHOT)) {
+    return;
+  }
+#if _WIN32
+  return;
+#else
+  KJ_LOG(WARNING, "Local Dev Python Snapshot enabled");
+  KJ_LOG(WARNING, "Storing snapshot to /tmp/snapshot.bin");
+  int h = open("/tmp/snapshot.bin", O_CREAT | O_WRONLY);
+  if (h == -1) {
+    KJ_LOG(WARNING, "Failed to open /tmp/snapshot.bin for writing");
+    perror("open");
+    errno = 0;
+    return;
+  }
+  int res = write(h, array.begin(), array.size());
+  if (res == -1) {
+    KJ_LOG(WARNING, "Error writing snapshot");
+    perror("write");
+    errno = 0;
+    close(h);
+    return;
+  }
+  if (res < array.size()) {
+    KJ_LOG(WARNING, "Only wrote", res, "bytes expected to write", array.size());
+  }
+  KJ_LOG(WARNING, "Stored snapshot of size", res, "to /tmp/snapshot.bin");
+#endif
+}
+
+kj::Maybe<kj::Array<byte>>
+maybeLoadSnapshotFromDisk() {
+  if (!Autogate::isEnabled(AutogateKey::LOCAL_DEV_PYTHON_SNAPSHOT)) {
+    return kj::none;
+  }
+#if _WIN32
+  return kj::none;
+#else
+  KJ_LOG(WARNING, "Local Dev Python Snapshot enabled");
+  KJ_LOG(WARNING, "Checking for snapshot in /tmp/snapshot.bin");
+  int h = open("/tmp/snapshot.bin", O_RDONLY);
+  if (h == -1) {
+    if (errno == ENOENT) {
+      KJ_LOG(WARNING, "No snapshot found");
+    } else {
+      KJ_LOG(WARNING, "Error opening file for reading:");
+      perror("open");
+    }
+    errno = 0;
+    return kj::none;
+  }
+
+  // Find snapshot length
+  int len = lseek(h, 0L, SEEK_END);
+  lseek(h, 0L, SEEK_SET);
+  // allocate result
+  auto result = kj::heapArray<kj::byte>(len);
+  // read snapshot into result
+  int res = read(h, result.begin(), len);
+  if (res == -1) {
+    KJ_LOG(WARNING, "Error reading snapshot");
+    perror("read");
+    close(h);
+    errno = 0;
+    return kj::none;
+  }
+  KJ_LOG(WARNING, "Read snapshot of length", len);
+  close(h);
+  return result;
+#endif
+}
+
 } // namespace workerd::api::pyodide

src/workerd/api/pyodide/pyodide.h

@@ -12,6 +12,18 @@
 
 namespace workerd::api::pyodide {
 
+bool
+shouldLoadSnapshotFromDisk();
+
+bool
+shouldStoreSnapshotToDisk();
+
+kj::Maybe<kj::Array<byte>>
+maybeLoadSnapshotFromDisk();
+
+void
+maybeStoreSnapshotToDisk(kj::ArrayPtr<kj::byte> array);
+
 // A function to read a segment of the tar file into a buffer
 // Set up this way to avoid copying files that aren't accessed.
 class PackagesTarReader : public jsg::Object {
@@ -55,6 +67,12 @@ class PyodideMetadataReader : public jsg::Object {
   }
 
   bool isTracing() {
+    if (this->isTracingFlag) {
+      return true;
+    }
+    // If we are loading the snapshot in local dev mode from disk, force
+    // isTracing to true in case we want to run the v8 profiler.
+    this->isTracingFlag = shouldLoadSnapshotFromDisk();
     return this->isTracingFlag;
   }
 
@@ -75,8 +93,13 @@ class PyodideMetadataReader : public jsg::Object {
   int read(jsg::Lock& js, int index, int offset, kj::Array<kj::byte> buf);
 
   bool hasMemorySnapshot() {
+    if (memorySnapshot != kj::none) {
+      return true;
+    }
+    memorySnapshot = maybeLoadSnapshotFromDisk();
     return memorySnapshot != kj::none;
   }
+
   int getMemorySnapshotSize() {
     if (memorySnapshot == kj::none) {
       return 0;
@@ -87,6 +110,7 @@ class PyodideMetadataReader : public jsg::Object {
   void disposeMemorySnapshot() {
     memorySnapshot = kj::none;
   }
+
   int readMemorySnapshot(int offset, kj::Array<kj::byte> buf);
 
   JSG_RESOURCE_TYPE(PyodideMetadataReader) {
@@ -170,6 +194,7 @@ class ArtifactBundler : public jsg::Object {
 
   void storeMemorySnapshot(jsg::Lock& js, kj::Array<kj::byte> snapshot) {
     KJ_REQUIRE(isValidating);
+    maybeStoreSnapshotToDisk(snapshot);
     storedSnapshot = kj::mv(snapshot);
   }
 
@@ -196,6 +221,12 @@ class ArtifactBundler : public jsg::Object {
 
   // Determines whether this ArtifactBundler was created inside the validator.
   bool isEwValidating() {
+    if (isValidating) {
+      return true;
+    }
+    // If we are loading from disk, force validating mode on to get JS code to
+    // give us the memory snapshot
+    isValidating = shouldStoreSnapshotToDisk();
     return isValidating;
   }
 

src/workerd/util/autogate.c++

@@ -13,6 +13,8 @@ kj::StringPtr KJ_STRINGIFY(AutogateKey key) {
   switch (key) {
     case AutogateKey::TEST_WORKERD:
       return "test-workerd"_kj;
+    case AutogateKey::LOCAL_DEV_PYTHON_SNAPSHOT:
+      return "local-dev-python-snapshot"_kj;
     case AutogateKey::NumOfKeys:
       KJ_FAIL_ASSERT("NumOfKeys should not be used in getName");
   }

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,
+  LOCAL_DEV_PYTHON_SNAPSHOT,
   NumOfKeys // Reserved for iteration.
 };