Snapshots: allow snapshotting some user code

Makefile.envs

@@ -164,52 +164,53 @@ export MAIN_MODULE_LDFLAGS= $(LDFLAGS_BASE) \
 	-lsdl.js \
 	-sGL_WORKAROUND_SAFARI_GETCONTEXT_BUG=0
 
-EXPORTS= _main\
-		\
-		,_free \
-		\
-		,_hiwire_new \
-		,_hiwire_intern \
-		,_hiwire_num_refs \
-		,_hiwire_get \
-		,_hiwire_incref \
-		,_hiwire_decref \
-		,_hiwire_pop \
-		,__hiwire_set \
-		,__hiwire_immortal_add \
-		,_jslib_init \
-		,_init_pyodide_proxy \
-		\
-		,_PyBuffer_Release \
-		,_Py_DecRef \
-		,_PyDict_New \
-		,_PyDict_SetItem \
-		,__PyErr_CheckSignals \
-		,_PyErr_CheckSignals \
-		,_PyErr_Clear \
-		,_PyErr_Occurred \
-		,_PyErr_Print \
-		,_PyErr_SetRaisedException \
-		,_PyErr_SetString \
-		,_PyEval_SaveThread \
-		,_PyEval_RestoreThread \
-		,_PyFloat_FromDouble \
-		,_PyGILState_Check \
-		,_Py_IncRef \
-		,_PyList_New \
-		,_PyList_SetItem \
-		,__PyLong_FromByteArray \
-		,_PyLong_FromDouble \
-		,_PyMem_Free \
-		,_PyObject_GetAIter \
-		,_PyObject_GetIter \
-		,_PyObject_Size \
-		,_PyRun_SimpleString \
-		,_PySet_Add \
-		,_PySet_New \
-		,__PyTraceback_Add \
-		,_PyUnicode_Data \
-		,_PyUnicode_New \
+EXPORTS=_main \
+   ,_free \
+   \
+   ,_hiwire_new \
+   ,_hiwire_intern \
+   ,_hiwire_num_refs \
+   ,_hiwire_get \
+   ,_hiwire_incref \
+   ,_hiwire_decref \
+   ,_hiwire_pop \
+   ,__hiwire_get \
+   ,__hiwire_set \
+   ,__hiwire_immortal_get \
+   ,__hiwire_immortal_add \
+   ,_jslib_init \
+   ,_init_pyodide_proxy \
+   \
+   ,_PyBuffer_Release \
+   ,_Py_DecRef \
+   ,_PyDict_New \
+   ,_PyDict_SetItem \
+   ,__PyErr_CheckSignals \
+   ,_PyErr_CheckSignals \
+   ,_PyErr_Clear \
+   ,_PyErr_Occurred \
+   ,_PyErr_Print \
+   ,_PyErr_SetRaisedException \
+   ,_PyErr_SetString \
+   ,_PyEval_SaveThread \
+   ,_PyEval_RestoreThread \
+   ,_PyFloat_FromDouble \
+   ,_PyGILState_Check \
+   ,_Py_IncRef \
+   ,_PyList_New \
+   ,_PyList_SetItem \
+   ,__PyLong_FromByteArray \
+   ,_PyLong_FromDouble \
+   ,_PyMem_Free \
+   ,_PyObject_GetAIter \
+   ,_PyObject_GetIter \
+   ,_PyObject_Size \
+   ,_PyRun_SimpleString \
+   ,_PySet_Add \
+   ,_PySet_New \
+   ,__PyTraceback_Add \
+   ,_PyUnicode_Data \
+   ,_PyUnicode_New \
 
 
 ifeq ($(DISABLE_DYLINK), 1)

src/core/_pyodide_core.c

@@ -65,6 +65,11 @@ EM_JS(void, set_pyodide_module, (JsVal mod), {
 int
 init_pyodide_proxy()
 {
+  EM_ASM({
+    // sourmash needs open64 to mean the same thing as open.
+    // Emscripten 3.1.44 seems to have removed it??
+    wasmImports["open64"] = wasmImports["open"];
+  });
   bool success = false;
   // Enable JavaScript access to the _pyodide module.
   PyObject* _pyodide = PyImport_ImportModule("_pyodide");
@@ -83,12 +88,6 @@ EM_JS_DEPS(pyodide_core_deps, "stackAlloc,stackRestore,stackSave");
 PyObject*
 PyInit__pyodide_core(void)
 {
-  EM_ASM({
-    // sourmash needs open64 to mean the same thing as open.
-    // Emscripten 3.1.44 seems to have removed it??
-    wasmImports["open64"] = wasmImports["open"];
-  });
-
   bool success = false;
   PyObject* _pyodide = NULL;
   PyObject* core_module = NULL;

src/js/api.ts

@@ -10,6 +10,12 @@ import { scheduleCallback } from "./scheduler";
 import { TypedArray } from "./types";
 import { IN_NODE, detectEnvironment } from "./environments";
 import "./literal-map.js";
+import {
+  makeGlobalsProxy,
+  SnapshotConfig,
+  syncUpSnapshotLoad1,
+  syncUpSnapshotLoad2,
+} from "./snapshot";
 
 // Exported for micropip
 API.loadBinaryFile = loadBinaryFile;
@@ -645,6 +651,15 @@ export class PyodideAPI {
     API.debug_ffi = debug;
     return orig;
   }
+
+  static makeMemorySnapshot(): Uint8Array {
+    if (!API.config._makeSnapshot) {
+      throw new Error(
+        "Can only use pyodide.makeMemorySnapshot if the _makeSnapshot option is passed to loadPyodide",
+      );
+    }
+    return API.makeSnapshot();
+  }
 }
 
 /** @hidden */
@@ -699,7 +714,7 @@ API.bootstrapFinalizedPromise = new Promise<void>(
   (r) => (bootstrapFinalized = r),
 );
 
-function jsFinderHook(o: object) {
+export function jsFinderHook(o: object) {
   if ("__all__" in o) {
     return;
   }
@@ -713,63 +728,17 @@ function jsFinderHook(o: object) {
   });
 }
 
-/**
- * Set up some of the JavaScript state that is normally set up by C initialization code. TODO:
- * adjust C code to simplify.
- *
- * This is divided up into two parts: syncUpSnapshotLoad1 has to happen at the beginning of
- * finalizeBootstrap before the public API is setup, syncUpSnapshotLoad2 happens near the end.
- *
- * This code is quite sensitive to the details of our setup, so it might break if we move stuff
- * around far away in the code base. Ideally over time we can structure the code to make it less
- * brittle.
- */
-function syncUpSnapshotLoad1() {
-  // hiwire init puts a null at the beginning of both the mortal and immortal tables.
-  Module.__hiwire_set(0, null);
-  Module.__hiwire_immortal_add(null);
-  // Usually importing _pyodide_core would trigger jslib_init but we need to manually call it.
-  Module._jslib_init();
-  // Puts deduplication map into the immortal table.
-  // TODO: Add support for snapshots to hiwire and move this to a hiwire_snapshot_init function.
-  Module.__hiwire_immortal_add(new Map());
-  // An interned JS string.
-  // TODO: Better system for handling interned strings.
-  Module.__hiwire_immortal_add(
-    "This borrowed proxy was automatically destroyed at the end of a function call. Try using create_proxy or create_once_callable.",
-  );
-  // Set API._pyodide to a proxy of the _pyodide module.
-  // Normally called by import _pyodide.
-  Module._init_pyodide_proxy();
-}
-
-/**
- * Fill in the JsRef table.
- */
-function syncUpSnapshotLoad2() {
-  [
-    null,
-    jsFinderHook,
-    API.config.jsglobals,
-    API.public_api,
-    Module.API,
-    scheduleCallback,
-    Module.API,
-    {},
-    null,
-    null,
-  ].forEach((v, idx) => Module.__hiwire_set(idx, v));
-}
-
 /**
  * This function is called after the emscripten module is finished initializing,
  * so eval_code is newly available.
  * It finishes the bootstrap so that once it is complete, it is possible to use
  * the core `pyodide` apis. (But package loading is not ready quite yet.)
  * @private
  */
-API.finalizeBootstrap = function (fromSnapshot?: boolean): PyodideInterface {
-  if (fromSnapshot) {
+API.finalizeBootstrap = function (
+  snapshotConfig?: SnapshotConfig,
+): PyodideInterface {
+  if (snapshotConfig) {
     syncUpSnapshotLoad1();
   }
   let [err, captured_stderr] = API.rawRun("import _pyodide_core");
@@ -802,11 +771,15 @@ API.finalizeBootstrap = function (fromSnapshot?: boolean): PyodideInterface {
   // Set up key Javascript modules.
   let importhook = API._pyodide._importhook;
   let pyodide = makePublicAPI();
-  if (fromSnapshot) {
-    syncUpSnapshotLoad2();
+  if (API.config._makeSnapshot) {
+    API.config.jsglobals = makeGlobalsProxy(API.config.jsglobals);
+  }
+  const jsglobals = API.config.jsglobals;
+  if (snapshotConfig) {
+    syncUpSnapshotLoad2(jsglobals, snapshotConfig);
   } else {
     importhook.register_js_finder.callKwargs({ hook: jsFinderHook });
-    importhook.register_js_module("js", API.config.jsglobals);
+    importhook.register_js_module("js", jsglobals);
     importhook.register_js_module("pyodide_js", pyodide);
   }
 

src/js/pyodide.ts

@@ -16,6 +16,7 @@ import type { PyodideInterface } from "./api.js";
 import type { TypedArray, Module } from "./types";
 import type { EmscriptenSettings } from "./emscripten-settings";
 import type { PackageData } from "./load-package";
+import { SnapshotConfig } from "./snapshot";
 export type { PyodideInterface, TypedArray };
 
 export { version, type PackageData };
@@ -42,6 +43,7 @@ export type ConfigType = {
   _node_mounts: string[];
   env: { [key: string]: string };
   packages: string[];
+  _makeSnapshot: boolean;
 };
 
 /**
@@ -213,11 +215,13 @@ export async function loadPyodide(
     await loadScript(scriptSrc);
   }
 
-  let snapshot;
+  let snapshot: Uint8Array | undefined = undefined;
   if (options._loadSnapshot) {
-    snapshot = await options._loadSnapshot;
-    if (snapshot?.constructor?.name === "ArrayBuffer") {
-      snapshot = new Uint8Array(snapshot);
+    const snp = await options._loadSnapshot;
+    if (ArrayBuffer.isView(snp)) {
+      snapshot = snp;
+    } else {
+      snapshot = new Uint8Array(snp);
     }
     emscriptenSettings.noInitialRun = true;
     // @ts-ignore
@@ -248,17 +252,12 @@ If you updated the Pyodide version, make sure you also updated the 'indexURL' pa
     throw new Error("Didn't expect to load any more file_packager files!");
   };
 
+  let snapshotConfig: SnapshotConfig | undefined = undefined;
   if (snapshot) {
-    // @ts-ignore
-    Module.HEAP8.set(snapshot);
+    snapshotConfig = API.restoreSnapshot(snapshot);
   }
   // runPython works starting after the call to finalizeBootstrap.
-  const pyodide = API.finalizeBootstrap(!!snapshot);
-
-  if (options._makeSnapshot) {
-    // @ts-ignore
-    pyodide._snapshot = Module.HEAP8.slice();
-  }
+  const pyodide = API.finalizeBootstrap(snapshotConfig);
   API.sys.path.insert(0, API.config.env.HOME);
 
   if (!pyodide.version.includes("dev")) {