Skip to content
rachel/workerd
File

Blob: src/pyodide/internal/snapshot.ts

typescript967 lines
1// Copyright (c) 2026 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5import { enterJaegerSpan } from 'pyodide-internal:jaeger';
6import { default as ArtifactBundler } from 'pyodide-internal:artifacts';
7import { default as UnsafeEval } from 'internal:unsafe-eval';
8import { default as DiskCache } from 'pyodide-internal:disk_cache';
9import { type FilePath, VIRTUALIZED_DIR } from 'pyodide-internal:setupPackages';
10import { default as EmbeddedPackagesTarReader } from 'pyodide-internal:packages_tar_reader';
11import {
12  SHOULD_SNAPSHOT_TO_DISK,
13  IS_CREATING_BASELINE_SNAPSHOT,
14  MEMORY_SNAPSHOT_READER,
15  REQUIREMENTS,
16  IS_CREATING_SNAPSHOT,
17  IS_EW_VALIDATING,
18  IS_DYNAMIC_WORKER,
19  IS_DEDICATED_SNAPSHOT_ENABLED,
20  COMPATIBILITY_FLAGS,
21  type CompatibilityFlags,
22  IS_SECOND_VALIDATION_PHASE,
23} from 'pyodide-internal:metadata';
24import {
25  invalidateCaches,
26  PythonWorkersInternalError,
27  PythonUserError,
28  simpleRunPython,
29  unreachable,
30} from 'pyodide-internal:util';
31import { default as MetadataReader } from 'pyodide-internal:runtime-generated/metadata';
32import type { PyodideEntrypointHelper } from 'pyodide:python-entrypoint-helper';
33import { entropyAfterSnapshot } from 'pyodide-internal:topLevelEntropy/lib';
34import {
35  deserializeJsModule,
36  maybeSerializeJsModule,
37  type SerializedJsModule,
38} from 'pyodide-internal:serializeJsModule';
39import { PyodideVersion } from 'pyodide-internal:const';
40 
41// A handle is the pointer into the linear memory returned by dlopen. Multiple dlopens will return
42// multiple pointers.
43type DsoHandles = {
44  [name: string]: { handles: number[] };
45};
46 
47// This is the info about Dsos that we record in getMemoryPatched. Namely the load order and where
48// their metadata is allocated. It would be natural to also have dsoHandles in here, but we don't
49// need a global variable to calculate dsoHandles since it is calculable from the information that
50// Emscripten stores in Module.LDSO (see recordDsoHandles).
51type DsoLoadInfo = {
52  readonly loadOrder: string[];
53  readonly soMemoryBases: { [name: string]: number };
54  readonly soTableBases?: { [name: string]: number };
55};
56 
57// This is the old wire format, where "settings" is mixed with the DsoHandles information and
58// DsoLoadInfo is not present because we used to preload all dynamic libraries in a standard order.
59// The "version" field is never present in the old wire format, but at runtime accessing it will
60// give undefined. We need to include it here so that we can use the version field as the
61// descriminator for `OldSnapshotMeta | SnapshotMeta`.
62type OldSnapshotMeta = DsoHandles & {
63  readonly settings?: { readonly baselineSnapshot?: boolean };
64  readonly version?: undefined;
65};
66 
67type LoadedSnapshotSettings = {
68  readonly snapshotType: ArtifactBundler.SnapshotType;
69  readonly compatFlags: CompatibilityFlags;
70};
71 
72type SnapshotSettings = {
73  readonly baselineSnapshot?: boolean;
74} & Partial<LoadedSnapshotSettings>;
75 
76// The new wire format, with additional information about the hiwire state, the order that dsos were
77// loaded in, and their memory bases. We also moved settings out of the dsoHandles.
78type SnapshotMeta = {
79  // We just store importedModulesList to help with testing and introspection
80  readonly importedModulesList: ReadonlyArray<string> | undefined;
81  readonly hiwire: SnapshotConfig | undefined;
82  readonly dsoHandles: DsoHandles;
83  readonly settings: SnapshotSettings;
84  readonly version: 1;
85  readonly jsModuleNames?: ReadonlyArray<string>;
86} & DsoLoadInfo;
87 
88// MEMORY_SNAPSHOT_READER has type SnapshotReader | undefined
89type SnapshotReader = {
90  readMemorySnapshot: (offset: number, buf: Uint32Array | Uint8Array) => void;
91  getMemorySnapshotSize: () => number;
92  disposeMemorySnapshot: () => void;
93};
94 
95// Extra info that isn't stored in the snapshot but is calculated at runtime: snapshotSize and
96// snapshotOffset would be awkward to store in the snapshot. snapshotReader is equal to
97// MEMORY_SNAPSHOT_READER, but typescript knows it is defined.
98type LoadedSnapshotExtras = {
99  snapshotSize: number;
100  snapshotOffset: number;
101  snapshotReader: SnapshotReader;
102};
103 
104type LoadedSnapshotMeta = SnapshotMeta &
105  LoadedSnapshotExtras & { readonly settings: LoadedSnapshotSettings };
106 
107/**
108 * Constants
109 */
110// "\x00snp"
111const SNAPSHOT_MAGIC = 0x706e7300;
112const CREATE_SNAPSHOT_VERSION = 2;
113const HEADER_SIZE = 4 * 4;
114 
115/**
116 * Global variables for the memory snapshot.
117 */
118const LOADED_SNAPSHOT_META: LoadedSnapshotMeta | undefined = decodeSnapshot(
119  MEMORY_SNAPSHOT_READER
120);
121let JS_MODULES: Record<string, any>;
122 
123export async function fillSnapshotJsModules(
124  doAnImport: (mod: string) => Promise<any>
125): Promise<void> {
126  JS_MODULES = await importJsModulesFromSnapshot(
127    doAnImport,
128    LOADED_SNAPSHOT_META?.jsModuleNames
129  );
130}
131const CREATED_SNAPSHOT_META: Required<DsoLoadInfo> = {
132  soMemoryBases: {},
133  soTableBases: {},
134  loadOrder: [],
135};
136if (LOADED_SNAPSHOT_META) {
137  // Make sure we include the soMemoryBases and loadOrder from the baseline snapshot when we are
138  // generating stacked snapshots.
139  Object.assign(
140    CREATED_SNAPSHOT_META.soMemoryBases,
141    LOADED_SNAPSHOT_META.soMemoryBases
142  );
143  Object.assign(
144    CREATED_SNAPSHOT_META.soTableBases,
145    LOADED_SNAPSHOT_META.soTableBases
146  );
147  CREATED_SNAPSHOT_META.loadOrder.push(...LOADED_SNAPSHOT_META.loadOrder);
148}
149export const LOADED_SNAPSHOT_TYPE = LOADED_SNAPSHOT_META?.settings.snapshotType;
150 
151/**
152 * Preload a dynamic library.
153 *
154 * Emscripten would usually figure out all of these details for us
155 * automatically. These defaults work for shared libs that are configured as
156 * standard Python extensions. This naive approach will not work for libraries
157 * like scipy, shapely, geos...
158 * TODO(someday) fix this.
159 */
160function loadDynlib(
161  Module: Module,
162  path: string,
163  wasmModuleData: Uint8Array
164): void {
165  const wasmModule = UnsafeEval.newWasmModule(wasmModuleData);
166  const dso = Module.newDSO(path, undefined, 'loading');
167  // even though these are used via dlopen, we are allocating them in an arena
168  // outside the heap and the memory cannot be reclaimed. So I don't think it
169  // would help us to allow them to be dealloc'd.
170  dso.refcount = Infinity;
171  // Hopefully they are used with dlopen
172  dso.global = false;
173  const options = {};
174  // Passing this empty object as dylibLocalScope fixes symbol lookup in dependent shared libraries
175  // that are not loaded globally, thus fixing one of our problems with the upstream shift away from
176  // RLTD_GLOBAL. Emscripten should probably be updated so that if dylibLocalScope is undefined it
177  // will give the dynamic library a new empty loading scope.
178  const dylibLocalScope = {};
179  dso.exports = Module.loadWebAssemblyModule(
180    wasmModule,
181    options,
182    path,
183    dylibLocalScope
184  );
185  // "handles" are dlopen handles. There will be one entry in the `handles` list
186  // for each dlopen handle that has not been dlclosed. We need to keep track of
187  // these across
188  const { handles } = LOADED_SNAPSHOT_META?.dsoHandles[path] || { handles: [] };
189  for (const handle of handles) {
190    Module.LDSO.loadedLibsByHandle[handle] = dso;
191  }
192  Module.LDSO.loadedLibsByName[path.split('/').at(-1)!] = dso;
193}
194 
195/**
196 * This function is used to ensure the order in which we load SO_FILES stays the same. It is only
197 * used for 0.26.0a2, later we look at SNAPSHOT_META.loadOrder to decide what order to load libs.
198 *
199 * The sort always puts _lzma.so and _ssl.so first, because these SO_FILES are loaded in the
200 * baseline snapshot, and if we want to generate a package snapshot while a baseline snapshot is
201 * loaded we need them to be first. The rest of the files are sorted alphabetically.
202 *
203 * The `filePaths` list is of the form [["folder", "file.so"], ["file.so"]], so each element in it
204 * is effectively a file path.
205 */
206function sortSoFiles(filePaths: FilePath[]): FilePath[] {
207  let result = [];
208  let hasLzma = false;
209  let hasSsl = false;
210  const lzmaFile = '_lzma.so';
211  const sslFile = '_ssl.so';
212  for (const path of filePaths) {
213    if (path.length == 1 && path[0] == lzmaFile) {
214      hasLzma = true;
215    } else if (path.length == 1 && path[0] == sslFile) {
216      hasSsl = true;
217    } else {
218      result.push(path);
219    }
220  }
221 
222  // JS might handle sorting lists of lists fine, but I'd rather be explicit here and make it compare
223  // strings.
224  result = result
225    .map((x) => x.join('/'))
226    .sort()
227    .map((x) => x.split('/'));
228  if (hasSsl) {
229    result.unshift([sslFile]);
230  }
231  if (hasLzma) {
232    result.unshift([lzmaFile]);
233  }
234 
235  return result;
236}
237 
238/**
239 * Used to ensure that the memoryBase of the dynamic library is stable when restoring snapshots.
240 * If the dynamic library was loaded in a memory snapshot,
241 */
242function getMemoryPatched(
243  Module: Module,
244  libPath: string,
245  size: number
246): number {
247  if (Module.API.version === PyodideVersion.V0_26_0a2) {
248    return Module.getMemory(size);
249  }
250  // Sometimes the module is loaded once by path and once by name, in either order. I'm not really
251  // sure why. But let's check if we snapshoted a load of the library by name.
252  const libName = libPath.split('/').at(-1)!;
253  // 1. Is it loaded in the snapshot? Replay the memory base.
254  {
255    const { soMemoryBases, soTableBases } = LOADED_SNAPSHOT_META ?? {};
256    // If we loaded this library before taking the snapshot, we already allocated the memory and the
257    // allocator remembers because its state is in the linear memory. We just have to look it up.
258    const tableBase = Module.wasmTable.length;
259    const expectedTableBase =
260      soTableBases?.[libPath] ?? soTableBases?.[libName];
261    if (expectedTableBase && tableBase !== expectedTableBase) {
262      // If this happens, we will segfault if we ever try to use this dynamic library.
263      // Save ourselves some debugging pain by crashing early.
264      throw new PythonWorkersInternalError(
265        `Error loading ${libName}: Expected table base ${expectedTableBase} but got table base ${tableBase}`
266      );
267    }
268    const memoryBase = soMemoryBases?.[libPath] ?? soMemoryBases?.[libName];
269    if (memoryBase) {
270      return memoryBase;
271    }
272  }
273  // 2. It's not loaded in the snapshot. Record
274  {
275    const { loadOrder, soMemoryBases, soTableBases } = CREATED_SNAPSHOT_META;
276    // Okay, we didn't load this before so we need to allocate new memory for it. Also record what we
277    // did in case someone makes a snapshot from this run.
278    loadOrder.push(libPath);
279    const memoryBase = Module.getMemory(size);
280    // Track both by full path and by name. That gives us a chance to resolve conflicts in name by the
281    // full path.
282    soMemoryBases[libPath] = memoryBase;
283    soMemoryBases[libName] = memoryBase;
284    soTableBases[libPath] = Module.wasmTable.length;
285    soTableBases[libName] = Module.wasmTable.length;
286    return memoryBase;
287  }
288}
289 
290function loadDynlibFromTarFs(
291  Module: Module,
292  base: string,
293  node: TarFSInfo | undefined,
294  soFile: string[]
295): void {
296  for (const part of soFile) {
297    node = node?.children?.get(part);
298  }
299  if (!node?.contentsOffset) {
300    node = VIRTUALIZED_DIR.getDynlibRoot();
301    for (const part of soFile) {
302      node = node?.children?.get(part);
303    }
304  }
305  if (!node?.contentsOffset) {
306    throw Error(`fs node could not be found for ${soFile.join('/')}`);
307  }
308  const { contentsOffset, size } = node;
309  if (contentsOffset === undefined) {
310    throw Error(`contentsOffset not defined for ${soFile.join('/')}`);
311  }
312  const wasmModuleData = new Uint8Array(size);
313  (node.reader ?? EmbeddedPackagesTarReader).read(
314    contentsOffset,
315    wasmModuleData
316  );
317  const path = base + soFile.join('/');
318  loadDynlib(Module, path, wasmModuleData);
319}
320 
321function loadDynlibFromVendor(
322  Module: Module,
323  soFile: string[],
324  userBundleNames: string[]
325): void {
326  const path = soFile.slice(3).join('/');
327  const index = userBundleNames.indexOf(path);
328  if (index == -1) {
329    throw new PythonWorkersInternalError(
330      `Could not find ${path} in user bundle, which is required by the snapshot.`
331    );
332  }
333  // The MetadataReader holds the user bundle's contents.
334  const buffer = new Uint8Array(MetadataReader.getSizes()[index]!);
335  MetadataReader.read(index, 0, buffer);
336  loadDynlib(Module, path, buffer);
337}
338 
339/**
340 * Preloading for the legacy 0.26 version. This loads all dynamic libraries
341 * visible in the site-packages directory. They are loaded before the runtime is
342 * initialized outside of the heap, using the same mechanism for DT_NEEDED libs
343 * (i.e., the libs that are loaded before the program starts because you passed
344 * them as linker args).
345 */
346function preloadDynamicLibs026(Module: Module): void {
347  const sitePackages = Module.FS.sessionSitePackages + '/';
348  const sitePackagesRoot = VIRTUALIZED_DIR.getSitePackagesRoot();
349  const loadedBaselineSnapshot =
350    LOADED_SNAPSHOT_META?.settings?.baselineSnapshot;
351  let SO_FILES_TO_LOAD: string[][];
352  if (IS_CREATING_BASELINE_SNAPSHOT || loadedBaselineSnapshot) {
353    SO_FILES_TO_LOAD = [['_lzma.so'], ['_ssl.so']];
354  } else {
355    SO_FILES_TO_LOAD = sortSoFiles(VIRTUALIZED_DIR.getSoFilesToLoad());
356  }
357  for (const soFile of SO_FILES_TO_LOAD) {
358    loadDynlibFromTarFs(Module, sitePackages, sitePackagesRoot, soFile);
359  }
360}
361 
362/**
363 * If we're restoring from a snapshot, we need to preload dynamic libraries so that any function
364 * pointers that point into the dylib symbols work correctly.
365 * Load the dynamic libraries in loadOrder. Mostly logic dealing with paths.
366 */
367function preloadDynamicLibsMain(Module: Module, loadOrder: string[]): void {
368  const sitePackages = Module.FS.sessionSitePackages + '/';
369  const sitePackagesRoot = VIRTUALIZED_DIR.getSitePackagesRoot();
370  const dynlibRoot = VIRTUALIZED_DIR.getDynlibRoot();
371  const dynlibPath = '/usr/lib/';
372  const userBundleNames = MetadataReader.getNames();
373  for (let path of loadOrder) {
374    let root = sitePackagesRoot;
375    let base = '';
376    if (path.startsWith(sitePackages)) {
377      path = path.slice(sitePackages.length);
378      base = sitePackages;
379    } else if (path.startsWith(dynlibPath)) {
380      path = path.slice(dynlibPath.length);
381      root = dynlibRoot;
382      base = dynlibPath;
383    }
384 
385    const pathSplit = Module.PATH.normalizeArray(path.split('/'), true);
386    if (pathSplit[0] == '') {
387      // This is a file path beginning with `/`, like /session/metadata/vendor/pkg/lib.so. So we
388      // are loading the vendored package's dynlibs here.
389      //
390      // TODO(EW-9508): support .so's in user bundle outside vendor dir.
391      loadDynlibFromVendor(Module, pathSplit, userBundleNames);
392    } else {
393      // This is a file path relative to the site-packages directory, like pkg/lib.so. So we are
394      // loading the built-in package's dynlibs here.
395      loadDynlibFromTarFs(Module, base, root, pathSplit);
396    }
397  }
398}
399 
400function preloadDynamicLibs(Module: Module): void {
401  if (Module.API.version === PyodideVersion.V0_26_0a2) {
402    // In 0.26.0a2 we need to preload dynamic libraries even if we aren't restoring a snapshot.
403    preloadDynamicLibs026(Module);
404    return;
405  }
406  //
407  const loadOrder = LOADED_SNAPSHOT_META?.loadOrder;
408  if (!loadOrder) {
409    // In newer versions we only need to do the preloading if there is a snapshot to restore.
410    return;
411  }
412  // In Pyodide 0.28 we switched from using top level EM_JS to initialize the CountArgs function
413  // pointer to using an initializer to work around a regression in Emscripten 4.0.3 and 4.0.4. We
414  // could drop this patch because we are now on Emscripten 4.0.9.
415  // https://github.com/pyodide/pyodide/blob/main/cpython/patches/0008-Fix-Emscripten-call-trampoline-compatibility-with-Em.patch
416  //
417  // Unfortunately, this initializer allocates a function table slot and is called before dynamic
418  // loading when taking the snapshot but after when restoring the snapshot. Thus, when restoring a
419  // snapshot, before loading dynamic libraries we reserve a function pointer for the
420  // CountArgsPointer and after loading dynamic libraries, we put it in the free list so it will be
421  // used at the right moment.
422  const PyEMCountArgsPtr = Module.getEmptyTableSlot();
423  preloadDynamicLibsMain(Module, loadOrder);
424  Module.freeTableIndexes.push(PyEMCountArgsPtr);
425}
426 
427/**
428 * This records which dynamic libraries have open handles (handed out by dlopen,
429 * not yet dlclosed). We'll need to track this information so that we don't
430 * crash if we dlsym the handle after restoring from the snapshot
431 */
432function recordDsoHandles(Module: Module): DsoHandles {
433  const dylinkInfo: DsoHandles = {};
434  for (const [h, { name }] of Object.entries(Module.LDSO.loadedLibsByHandle)) {
435    const handle = Number(h);
436    if (handle === 0) {
437      continue;
438    }
439    dylinkInfo[name] ??= { handles: [] };
440    dylinkInfo[name].handles.push(handle);
441  }
442  return dylinkInfo;
443}
444 
445/**
446 * Python modules do a lot of work the first time they are imported. The memory
447 * snapshot will save more time the more of this work is included. However, we
448 * can't snapshot the JS runtime state so we have no ffi. Thus some imports from
449 * user code will fail.
450 *
451 * If we are doing a baseline snapshot, just import everything from
452 * baselineSnapshotImports. These will all succeed.
453 *
454 * If doing a more dedicated "package" snap shot, also try to import each
455 * user import that is importing non-vendored modules.
456 *
457 * All of this is being done in the __main__ global scope, so be careful not to
458 * pollute it with extra included-by-default names (user code is executed in its
459 * own separate module scope though so it's not _that_ important).
460 *
461 * This function returns a list of modules that have been imported.
462 */
463function memorySnapshotDoImports(Module: Module): string[] {
464  const baselineSnapshotImports =
465    MetadataReader.constructor.getBaselineSnapshotImports();
466  const toImport = baselineSnapshotImports.join(',');
467  const toDelete = Array.from(
468    new Set(baselineSnapshotImports.map((x) => x.split('.', 1)[0]))
469  ).join(',');
470 
471  simpleRunPython(Module, `import ${toImport}`);
472  simpleRunPython(Module, 'sysconfig.get_config_vars()');
473  // Delete to avoid polluting globals
474  simpleRunPython(Module, `del ${toDelete}`);
475  if (IS_CREATING_BASELINE_SNAPSHOT) {
476    // We've done all the imports for the baseline snapshot.
477    return [];
478  }
479  if (REQUIREMENTS.length == 0) {
480    // Don't attempt to scan for package imports if the Worker has specified no package
481    // requirements, as this means their code isn't going to be importing any modules that we need
482    // to include in a snapshot.
483    return [];
484  }
485 
486  // The `importedModules` list will contain all modules that have been imported, including local
487  // modules, the usual `js` and other stdlib modules. We want to filter out local imports, so we
488  // grab them and put them into a set for fast filtering.
489  const importedModules: string[] = MetadataReader.getPackageSnapshotImports(
490    Module.API.version
491  );
492  const deduplicatedModules = [...new Set(importedModules)];
493 
494  // Import the modules list so they are included in the snapshot.
495  if (deduplicatedModules.length > 0) {
496    simpleRunPython(Module, 'import ' + deduplicatedModules.join(','));
497  }
498 
499  return deduplicatedModules;
500}
501 
502function describeValue(val: any): string {
503  try {
504    const out = [];
505 
506    const type = typeof val;
507    const isObject = type === 'object';
508    out.push(`Value: ${val}`);
509    out.push(`Type: ${type}`);
510    try {
511      const constructorName = val?.constructor?.name; // eslint-disable-line
512      if (constructorName) {
513        out.push(`Constructor name: ${constructorName}`);
514      }
515    } catch {
516      // Ignore errors when getting keys
517    }
518 
519    if (val && isObject) {
520      try {
521        out.push(
522          `Keys: ${Object.keys(val as Record<string, unknown>)
523            .slice(0, 10)
524            .join(', ')}`
525        );
526      } catch {
527        // Ignore errors when getting keys
528      }
529    }
530 
531    try {
532      if (val && isObject && typeof (val as Error).stack === 'string') {
533        out.push(`Stack:\n${(val as Error).stack}`);
534      }
535    } catch {
536      // Ignore errors when getting stack
537    }
538 
539    try {
540      out.push(`toStringTag: ${Object.prototype.toString.call(val)}`);
541    } catch {
542      // Ignore errors when getting object type
543    }
544 
545    try {
546      const contents = JSON.stringify(val);
547      if (contents?.length > 100) {
548        out.push(`Contents: ${contents.slice(0, 100)}...`);
549      } else if (contents) {
550        out.push(`Contents: ${contents}`);
551      }
552    } catch {
553      // Ignore JSON stringify errors
554    }
555 
556    return out.join('\n');
557  } catch (err) {
558    return `Error describing value: ${err}`;
559  }
560}
561 
562/**
563 * When we create a dedicated memory snapshot, we capture all the globals in the user's top-level
564 * scope. If those globals refer to JS objects, then we may fail to serialise them. This function
565 * creates a user error to inform the user of this issue.
566 *
567 * It's important that we give the user as much information about this as possible, so they can
568 * understand what they need to change in order to resolve the problem on their end.
569 */
570function createUnserializableObjectError(obj: any): PythonUserError {
571  // TODO: Create docs for this and link them here.
572  const error = `Can't serialize top-level variable.
573Please review any global variables you or your imported modules create at the top-level of your code, and consider deleting them.
574
575Description of the value:
576${describeValue(obj)}
577`;
578  return new PythonUserError(error);
579}
580 
581async function importJsModulesFromSnapshot(
582  doAnImport: (mod: string) => Promise<any>,
583  jsModuleNames: ReadonlyArray<string> | undefined
584): Promise<Record<string, any>> {
585  if (jsModuleNames === undefined) {
586    return {};
587  }
588  async function doImport(x: string): Promise<any> {
589    if (x === 'global this') {
590      return globalThis;
591    }
592    return await doAnImport(x);
593  }
594  return Object.fromEntries(
595    await Promise.all(
596      jsModuleNames.map(
597        async (x): Promise<[string, any]> => [x, await doImport(x)]
598      )
599    )
600  );
601}
602 
603type CustomSerialized =
604  | { pyodide_entrypoint_helper: true }
605  | { cloudflare_compat_flags: true }
606  | SerializedJsModule;
607/**
608 * Global objects that need a custom serializer
609 */
610export type CustomSerializedObjects = {
611  pyodide_entrypoint_helper: PyodideEntrypointHelper;
612  cloudflare_compat_flags: CompatibilityFlags;
613};
614 
615function getHiwireSerializer(
616  globalObj: CustomSerializedObjects,
617  modules: Set<string>
618): (obj: any) => CustomSerialized {
619  return function serializer(obj: any): CustomSerialized {
620    if (obj === globalObj.pyodide_entrypoint_helper) {
621      return { pyodide_entrypoint_helper: true };
622    } else if (obj === globalObj.cloudflare_compat_flags) {
623      return { cloudflare_compat_flags: true };
624    }
625    const serializedModule = maybeSerializeJsModule(obj, modules);
626    if (serializedModule) {
627      return serializedModule;
628    }
629    throw createUnserializableObjectError(obj);
630  };
631}
632 
633function getHiwireDeserializer(
634  globalObj: CustomSerializedObjects
635): (obj: CustomSerialized) => any {
636  return function deserializer(obj) {
637    if ('pyodide_entrypoint_helper' in obj) {
638      return globalObj.pyodide_entrypoint_helper;
639    } else if ('cloudflare_compat_flags' in obj) {
640      return globalObj.cloudflare_compat_flags;
641    }
642    if ('jsModule' in obj) {
643      return deserializeJsModule(obj, JS_MODULES);
644    }
645    unreachable(obj, `Can't deserialize ${obj}`);
646  };
647}
648 
649/**
650 * Create memory snapshot by importing SNAPSHOT_IMPORTS to ensure these packages
651 * are initialized in the linear memory snapshot and then saving a copy of the
652 * linear memory into MEMORY.
653 */
654function makeLinearMemorySnapshot(
655  Module: Module,
656  importedModulesList: string[],
657  customSerializedObjects: CustomSerializedObjects,
658  snapshotType: ArtifactBundler.SnapshotType
659): Uint8Array {
660  const dsoHandles = recordDsoHandles(Module);
661  let hiwire: SnapshotConfig | undefined;
662  const jsModuleNames: Set<string> = new Set();
663  if (Module.API.version !== PyodideVersion.V0_26_0a2) {
664    hiwire = Module.API.serializeHiwireState(
665      getHiwireSerializer(customSerializedObjects, jsModuleNames)
666    );
667  }
668  const settings: SnapshotSettings = {
669    baselineSnapshot: IS_CREATING_BASELINE_SNAPSHOT,
670    snapshotType,
671    compatFlags: COMPATIBILITY_FLAGS,
672  };
673  return encodeSnapshot(Module.HEAP8, {
674    version: 1,
675    dsoHandles,
676    hiwire,
677    importedModulesList,
678    jsModuleNames: Array.from(jsModuleNames),
679    settings,
680    ...CREATED_SNAPSHOT_META,
681  });
682}
683 
684/**
685 * Encode heap and dsoJSON into the memory snapshot artifact that we'll upload
686 */
687function encodeSnapshot(heap: Uint8Array, meta: SnapshotMeta): Uint8Array {
688  const json = JSON.stringify(meta);
689  let snapshotOffset = HEADER_SIZE + 2 * json.length;
690  // align to 8 bytes
691  snapshotOffset = Math.ceil(snapshotOffset / 8) * 8;
692  const toUpload = new Uint8Array(snapshotOffset + heap.length);
693  const encoder = new TextEncoder();
694  const { written: jsonByteLength } = encoder.encodeInto(
695    json,
696    toUpload.subarray(HEADER_SIZE)
697  );
698  const header = new Uint32Array(toUpload.buffer);
699  header[0] = SNAPSHOT_MAGIC;
700  header[1] = CREATE_SNAPSHOT_VERSION;
701  header[2] = snapshotOffset;
702  header[3] = jsonByteLength;
703  toUpload.subarray(snapshotOffset).set(heap);
704  return toUpload;
705}
706 
707/**
708 * Decode heap and dsoJSON from the memory snapshot reader
709 */
710function decodeSnapshot(
711  reader: SnapshotReader | undefined
712): LoadedSnapshotMeta | undefined {
713  if (!reader) {
714    return undefined;
715  }
716  if (reader.getMemorySnapshotSize() === 0) {
717    throw new PythonWorkersInternalError(
718      `SnapshotReader returned memory snapshot size of 0`
719    );
720  }
721  const header = new Uint32Array(4);
722  reader.readMemorySnapshot(0, header);
723  if (header[0] !== SNAPSHOT_MAGIC) {
724    throw new PythonWorkersInternalError(
725      `Invalid magic number ${header[0]}, expected ${SNAPSHOT_MAGIC}`
726    );
727  }
728  // buf[1] is SNAPSHOT_VERSION (unused currently)
729  const snapshotOffset = header[2]!;
730  const jsonByteLength = header[3]!;
731 
732  const snapshotSize = reader.getMemorySnapshotSize() - snapshotOffset;
733  const jsonBuf = new Uint8Array(jsonByteLength);
734  const offset = header.byteLength; // the json starts after the header
735  reader.readMemorySnapshot(offset, jsonBuf);
736  const json = new TextDecoder().decode(jsonBuf);
737  const meta = JSON.parse(json) as OldSnapshotMeta | SnapshotMeta;
738  const extras: LoadedSnapshotExtras = {
739    snapshotSize,
740    snapshotOffset,
741    snapshotReader: reader,
742  };
743  if (!meta?.version) {
744    return {
745      version: 1,
746      importedModulesList: undefined,
747      dsoHandles: meta,
748      hiwire: undefined,
749      loadOrder: [],
750      soMemoryBases: {},
751      settings: {
752        snapshotType: meta.settings?.baselineSnapshot ? 'baseline' : 'package',
753        compatFlags: {},
754        ...meta.settings,
755      },
756      jsModuleNames: [],
757      ...extras,
758    };
759  }
760  return {
761    ...meta,
762    ...extras,
763    settings: {
764      ...meta.settings,
765      snapshotType:
766        meta.settings.snapshotType ??
767        (meta.settings.baselineSnapshot ? 'baseline' : 'package'),
768      compatFlags: meta.settings.compatFlags ?? {},
769    },
770  };
771}
772 
773export function isRestoringSnapshot(): boolean {
774  return !!LOADED_SNAPSHOT_META;
775}
776 
777function checkSnapshotType(snapshotType: string): void {
778  if (SHOULD_SNAPSHOT_TO_DISK) {
779    return;
780  }
781  // Dynamic workers don't yet support dedicated snapshots, so they will receive a baseline
782  // snapshot from GCS even when the dedicated snapshot compat flag is enabled. Skip the
783  // snapshot type validation in this case.
784  if (IS_DYNAMIC_WORKER) {
785    return;
786  }
787  if (
788    !IS_EW_VALIDATING &&
789    snapshotType === 'dedicated' &&
790    !IS_DEDICATED_SNAPSHOT_ENABLED
791  ) {
792    throw new PythonWorkersInternalError(
793      'Received dedicated snapshot but compat flag for dedicated snapshots is not enabled'
794    );
795  }
796 
797  if (
798    !IS_EW_VALIDATING &&
799    snapshotType !== 'dedicated' &&
800    IS_DEDICATED_SNAPSHOT_ENABLED
801  ) {
802    throw new PythonWorkersInternalError(
803      'Received non-dedicated snapshot but compat flag for dedicated snapshots is enabled'
804    );
805  }
806 
807  // If we have a snapshot in the bundle and the dedicated snapshot flag is enabled, then we
808  // should verify that the snapshot in the bundle is a dedicated snapshot. If it is not
809  // we should fail with an error.
810  if (snapshotType !== 'dedicated' && IS_SECOND_VALIDATION_PHASE) {
811    throw new PythonWorkersInternalError(
812      'The second validation phase should receive a dedicated snapshot, got ' +
813        snapshotType
814    );
815  }
816}
817 
818export function maybeRestoreSnapshot(Module: Module): void {
819  Module.noInitialRun = isRestoringSnapshot();
820  Module.getMemoryPatched = getMemoryPatched;
821  // Make sure memory is large enough
822  Module.growMemory(LOADED_SNAPSHOT_META?.snapshotSize ?? 0);
823  enterJaegerSpan('preload_dynamic_libs', () => {
824    preloadDynamicLibs(Module);
825  });
826  // TODO: Remove the jaeger span here
827  enterJaegerSpan('remove_run_dependency', () => {
828    Module.removeRunDependency('dynlibs');
829  });
830  if (!LOADED_SNAPSHOT_META) {
831    return;
832  }
833  const { snapshotSize, snapshotOffset, snapshotReader, settings } =
834    LOADED_SNAPSHOT_META;
835  checkSnapshotType(settings.snapshotType);
836 
837  Module.growMemory(snapshotSize);
838  snapshotReader.readMemorySnapshot(snapshotOffset, Module.HEAP8);
839  snapshotReader.disposeMemorySnapshot();
840  // Invalidate caches if we have a snapshot because the contents of site-packages
841  // may have changed.
842  invalidateCaches(Module);
843}
844 
845function collectSnapshot(
846  Module: Module,
847  importedModulesList: string[],
848  customSerializedObjects: CustomSerializedObjects,
849  snapshotType: ArtifactBundler.SnapshotType
850): void {
851  if (!IS_EW_VALIDATING && !SHOULD_SNAPSHOT_TO_DISK) {
852    throw new PythonWorkersInternalError(
853      "Attempted to collect snapshot outside of context where it's supported."
854    );
855  }
856  const snapshot = makeLinearMemorySnapshot(
857    Module,
858    importedModulesList,
859    customSerializedObjects,
860    snapshotType
861  );
862  entropyAfterSnapshot(Module);
863  if (IS_EW_VALIDATING) {
864    ArtifactBundler.storeMemorySnapshot({
865      snapshot,
866      importedModulesList,
867      snapshotType,
868    });
869  } else if (SHOULD_SNAPSHOT_TO_DISK) {
870    DiskCache.putSnapshot('snapshot.bin', snapshot);
871  } else {
872    throw new PythonWorkersInternalError('Unreachable');
873  }
874}
875 
876/**
877 * Collects a dedicated snapshot. This is only called after the top-level of the worker has been
878 * run.
879 */
880export function maybeCollectDedicatedSnapshot(
881  Module: Module,
882  customSerializedObjects: CustomSerializedObjects | null
883): void {
884  if (!IS_CREATING_SNAPSHOT) {
885    return;
886  }
887 
888  if (!IS_DEDICATED_SNAPSHOT_ENABLED) {
889    return;
890  }
891 
892  if (Module.API.version == PyodideVersion.V0_26_0a2) {
893    // 0.26.0a2 does not support serialisation of the hiwire state, so it cannot support dedicated
894    // snapshots.
895    throw new PythonWorkersInternalError(
896      'Dedicated snapshot is not supported for Python runtime version 0.26.0a2'
897    );
898  }
899 
900  if (!customSerializedObjects) {
901    throw new PythonWorkersInternalError(
902      'customSerializedObjects is required for dedicated snapshot'
903    );
904  }
905  collectSnapshot(Module, [], customSerializedObjects, 'dedicated');
906}
907 
908/**
909 * Collects either a baseline or package snapshot. This is called prior to running the top-level
910 * of the worker and crucially before the worker files are mounted.
911 *
912 * Dedicated snapshots are collected in `maybeCollectDedicatedSnapshot`.
913 */
914export function maybeCollectSnapshot(
915  Module: Module,
916  customSerializedObjects: CustomSerializedObjects
917): void {
918  // In order to surface any problems that occur in `memorySnapshotDoImports` to
919  // users in local development, always call it even if we aren't actually
920  const importedModulesList = memorySnapshotDoImports(Module);
921  if (!IS_CREATING_SNAPSHOT) {
922    return;
923  }
924 
925  if (IS_DEDICATED_SNAPSHOT_ENABLED) {
926    // We are not interested in collecting a baseline/package snapshot here if this feature flag
927    // is enabled.
928    return;
929  }
930 
931  collectSnapshot(
932    Module,
933    importedModulesList,
934    customSerializedObjects,
935    IS_CREATING_BASELINE_SNAPSHOT ? 'baseline' : 'package'
936  );
937}
938 
939export function finalizeBootstrap(
940  Module: Module,
941  customSerializedObjects: CustomSerializedObjects
942): void {
943  Module.API.config._makeSnapshot =
944    IS_CREATING_SNAPSHOT &&
945    Module.API.version !== PyodideVersion.V0_26_0a2 &&
946    Module.API.version !== PyodideVersion.V0_28_2;
947  enterJaegerSpan('finalize_bootstrap', () => {
948    Module.API.finalizeBootstrap(
949      LOADED_SNAPSHOT_META?.hiwire,
950      getHiwireDeserializer(customSerializedObjects)
951    );
952  });
953  // finalizeBootstrap overrides LD_LIBRARY_PATH. Restore it.
954  simpleRunPython(
955    Module,
956    `import os; os.environ["LD_LIBRARY_PATH"] += ":/session/metadata/python_modules/lib/"; del os`
957  );
958  if (IS_CREATING_SNAPSHOT) {
959    return;
960  }
961  Module.API.public_api.registerJsModule('_cf_internal_snapshot_info', {
962    loadedSnapshot: !!LOADED_SNAPSHOT_META,
963    loadedBaselineSnapshot: LOADED_SNAPSHOT_META?.settings.baselineSnapshot,
964    importedModulesList: LOADED_SNAPSHOT_META?.importedModulesList,
965  });
966}