Skip to content
rachel/workerd
File

Blob: src/workerd/api/pyodide/pyodide.h

cpp554 lines
1// Copyright (c) 2017-2022 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#pragma once
5 
6#include <workerd/io/compatibility-date.capnp.h>
7#include <workerd/jsg/jsg.h>
8#include <workerd/jsg/modules-new.h>
9#include <workerd/util/strong-bool.h>
10 
11#include <pyodide/generated/pyodide_extra.capnp.h>
12#include <pyodide/pyodide_static.capnp.h>
13 
14#include <capnp/serialize.h>
15#include <kj/array.h>
16#include <kj/common.h>
17#include <kj/compat/http.h>
18#include <kj/filesystem.h>
19#include <kj/function.h>
20#include <kj/string.h>
21#include <kj/table.h>
22#include <kj/timer.h>
23 
24namespace workerd::api::pyodide {
25 
26WD_STRONG_BOOL(CreateBaselineSnapshot);
27WD_STRONG_BOOL(IsTracing);
28WD_STRONG_BOOL(IsValidating);
29WD_STRONG_BOOL(IsWorkerd);
30WD_STRONG_BOOL(SnapshotToDisk);
31 
32const auto PYTHON_PACKAGES_URL = "https://pyodide-capnp-bin.edgeworker.net/";
33class PyodideBundleManager {
34 public:
35  void setPyodideBundleData(kj::String version, kj::Array<unsigned char> data) const;
36  const kj::Maybe<jsg::Bundle::Reader> getPyodideBundle(kj::StringPtr version) const;
37 
38 private:
39  struct MessageBundlePair {
40    kj::Own<capnp::FlatArrayMessageReader> messageReader;
41    jsg::Bundle::Reader bundle;
42  };
43  const kj::MutexGuarded<kj::HashMap<kj::String, MessageBundlePair>> bundles;
44};
45 
46class PyodidePackageManager {
47 public:
48  void setPyodidePackageData(kj::String id, kj::Array<unsigned char> data) const;
49  const kj::Maybe<const kj::Array<unsigned char>&> getPyodidePackage(kj::StringPtr id) const;
50 
51 private:
52  const kj::MutexGuarded<kj::HashMap<kj::String, kj::Array<unsigned char>>> packages;
53};
54 
55struct PythonConfig {
56  kj::Maybe<kj::Own<const kj::Directory>> packageDiskCacheRoot;
57  kj::Maybe<kj::Own<const kj::Directory>> pyodideDiskCacheRoot;
58  kj::Maybe<kj::Own<const kj::Directory>> snapshotDirectory;
59  const PyodideBundleManager pyodideBundleManager;
60  const PyodidePackageManager pyodidePackageManager;
61  bool createSnapshot;
62  bool createBaselineSnapshot;
63  kj::Maybe<kj::String> loadSnapshotFromDisk;
64};
65 
66// A function to read a segment of the tar file into a buffer
67// Set up this way to avoid copying files that aren't accessed.
68class ReadOnlyBuffer: public jsg::Object {
69  kj::ArrayPtr<const kj::byte> source;
70 
71 public:
72  ReadOnlyBuffer(kj::ArrayPtr<const kj::byte> src): source(src) {};
73 
74  int read(jsg::Lock& js, int offset, kj::Array<kj::byte> buf);
75 
76  JSG_RESOURCE_TYPE(ReadOnlyBuffer) {
77    JSG_METHOD(read);
78  }
79};
80 
81class PythonModuleInfo {
82 public:
83  PythonModuleInfo(kj::Array<kj::String> names, kj::Array<kj::Array<kj::byte>> contents)
84      : names(kj::mv(names)),
85        contents(kj::mv(contents)) {
86    KJ_REQUIRE(this->names.size() == this->contents.size());
87  }
88  kj::Array<kj::String> names;
89  kj::Array<kj::Array<kj::byte>> contents;
90 
91  PythonModuleInfo clone() const {
92    auto clonedContents =
93        KJ_MAP(content, this->contents) { return kj::heapArray<kj::byte>(content); };
94    auto clonedNames = KJ_MAP(name, this->names) { return kj::str(name); };
95    return PythonModuleInfo(kj::mv(clonedNames), kj::mv(clonedContents));
96  }
97 
98  // Return the list of names to import into a package snapshot.
99  kj::Array<kj::String> getPackageSnapshotImports(kj::StringPtr version);
100  // Takes in a list of Python files (their contents). Parses these files to find the import
101  // statements, then returns a list of modules imported via those statements.
102  //
103  // For example:
104  // import a, b, c
105  // from z import x
106  // import t.y.u
107  // from . import k
108  //
109  // -> ["a", "b", "c", "z", "t.y.u"]
110  //
111  // Package relative imports are ignored.
112  static kj::Array<kj::String> parsePythonScriptImports(kj::Array<kj::String> files);
113  kj::HashSet<kj::String> getWorkerModuleSet();
114  kj::Array<kj::String> getPythonFileContents();
115  static kj::Array<kj::String> filterPythonScriptImports(kj::HashSet<kj::String> workerModules,
116      kj::ArrayPtr<kj::String> imports,
117      kj::StringPtr version);
118};
119 
120// A class wrapping the information stored in a WorkerBundle, in particular the Python source files
121// and metadata about the worker.
122//
123// This is done this way to avoid copying files as much as possible. We set up a Metadata File
124// System which reads the contents as they are needed.
125class PyodideMetadataReader: public jsg::Object {
126 public:
127  //
128  struct State {
129    kj::String mainModule;
130    PythonModuleInfo moduleInfo;
131    kj::Array<kj::String> requirements;
132    kj::String pyodideVersion;
133    kj::String packagesVersion;
134    kj::String packagesLock;
135    bool isWorkerdFlag;
136    bool isTracingFlag;
137    bool snapshotToDisk;
138    bool createBaselineSnapshot;
139    kj::Maybe<kj::Array<kj::byte>> memorySnapshot;
140 
141    State(kj::String mainModule,
142        kj::Array<kj::String> names,
143        kj::Array<kj::Array<kj::byte>> contents,
144        kj::Array<kj::String> requirements,
145        kj::String pyodideVersion,
146        kj::String packagesVersion,
147        kj::String packagesLock,
148        IsWorkerd isWorkerd,
149        IsTracing isTracing,
150        SnapshotToDisk snapshotToDisk,
151        CreateBaselineSnapshot createBaselineSnapshot,
152        kj::Maybe<kj::Array<kj::byte>> memorySnapshot)
153        : mainModule(kj::mv(mainModule)),
154          moduleInfo(kj::mv(names), kj::mv(contents)),
155          requirements(kj::mv(requirements)),
156          pyodideVersion(kj::mv(pyodideVersion)),
157          packagesVersion(kj::mv(packagesVersion)),
158          packagesLock(kj::mv(packagesLock)),
159          isWorkerdFlag(isWorkerd),
160          isTracingFlag(isTracing),
161          snapshotToDisk(snapshotToDisk),
162          createBaselineSnapshot(createBaselineSnapshot),
163          memorySnapshot(kj::mv(memorySnapshot)) {
164      verifyNoMainModuleInVendor();
165    }
166 
167    State(const State& other);
168 
169    void verifyNoMainModuleInVendor();
170 
171    kj::Own<State> clone();
172  };
173 
174  PyodideMetadataReader(kj::Own<State> state): state(kj::mv(state)) {}
175 
176  bool isWorkerd() {
177    return state->isWorkerdFlag;
178  }
179 
180  bool isTracing() {
181    return state->isTracingFlag;
182  }
183 
184  bool shouldSnapshotToDisk() {
185    return state->snapshotToDisk;
186  }
187 
188  bool isCreatingBaselineSnapshot() {
189    return state->createBaselineSnapshot;
190  }
191 
192  // Returns whether the python-abort-isolate-on-fatal-error autogate is enabled. When true, the
193  // Python on_fatal handler should call abortIsolate() to terminate the isolate after reporting.
194  bool shouldAbortIsolateOnFatalError();
195 
196  kj::StringPtr getMainModule() {
197    return state->mainModule;
198  }
199 
200  // Returns the filenames of the files inside of the WorkerBundle that end with the specified
201  // file extension.
202  // TODO: Remove this.
203  kj::Array<kj::StringPtr> getNames(jsg::Lock& js, jsg::Optional<kj::String> maybeExtFilter);
204  kj::Array<int> getSizes(jsg::Lock& js);
205 
206  // Return the list of names to import into a package snapshot.
207  kj::Array<kj::String> getPackageSnapshotImports(kj::String version);
208 
209  kj::Array<jsg::JsRef<jsg::JsString>> getRequirements(jsg::Lock& js);
210 
211  int read(jsg::Lock& js, int index, int offset, kj::Array<kj::byte> buf);
212 
213  bool hasMemorySnapshot() {
214    return state->memorySnapshot != kj::none;
215  }
216  int getMemorySnapshotSize() {
217    if (state->memorySnapshot == kj::none) {
218      return 0;
219    }
220    return KJ_REQUIRE_NONNULL(state->memorySnapshot).size();
221  }
222 
223  void disposeMemorySnapshot() {
224    state->memorySnapshot = kj::none;
225  }
226  int readMemorySnapshot(int offset, kj::Array<kj::byte> buf);
227 
228  kj::StringPtr getPyodideVersion() {
229    return state->pyodideVersion;
230  }
231 
232  kj::StringPtr getPackagesVersion() {
233    return state->packagesVersion;
234  }
235 
236  kj::StringPtr getPackagesLock() {
237    return state->packagesLock;
238  }
239 
240  kj::HashSet<kj::String> getTransitiveRequirements();
241 
242  static kj::Array<kj::StringPtr> getBaselineSnapshotImports();
243 
244  // We call this during Python setup with the wasm memory and the addresses of the signal clock and
245  // the flag to indicate whether signal handling is on or off. It sets up the isolate
246  // CpuLimitNearlyExceeded callback to trigger a signal in Python.
247  void setCpuLimitNearlyExceededCallback(
248      jsg::Lock& js, kj::Array<kj::byte> wasm_memory, int sig_clock, int sig_flag);
249 
250  // Similar to Cloudflare::::getCompatibilityFlags in global-scope.c++, but the key difference is
251  // that it returns experimental flags even if `experimental` is not enabled. This avoids a gotcha
252  // where an experimental compat flag is enabled in our C++ code, but not in our JS code.
253  //
254  // This is only for use by our Python runtime.
255  jsg::JsObject getCompatibilityFlags(jsg::Lock& js);
256 
257  JSG_RESOURCE_TYPE(PyodideMetadataReader) {
258    JSG_METHOD(isWorkerd);
259    JSG_METHOD(isTracing);
260    JSG_METHOD(getMainModule);
261    JSG_METHOD(getRequirements);
262    JSG_METHOD(getNames);
263    JSG_METHOD(getSizes);
264    JSG_METHOD(getPackageSnapshotImports);
265    JSG_METHOD(read);
266    JSG_METHOD(hasMemorySnapshot);
267    JSG_METHOD(getMemorySnapshotSize);
268    JSG_METHOD(readMemorySnapshot);
269    JSG_METHOD(disposeMemorySnapshot);
270    JSG_METHOD(shouldSnapshotToDisk);
271    JSG_METHOD(getPyodideVersion);
272    JSG_METHOD(getPackagesVersion);
273    JSG_METHOD(getPackagesLock);
274    JSG_METHOD(isCreatingBaselineSnapshot);
275    JSG_METHOD(shouldAbortIsolateOnFatalError);
276    JSG_METHOD(getTransitiveRequirements);
277    JSG_METHOD(getCompatibilityFlags);
278    JSG_STATIC_METHOD(getBaselineSnapshotImports);
279    JSG_METHOD(setCpuLimitNearlyExceededCallback);
280  }
281 
282  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
283    tracker.trackField("mainModule", state->mainModule);
284    for (const auto& name: state->moduleInfo.names) {
285      tracker.trackField("name", name);
286    }
287    for (const auto& content: state->moduleInfo.contents) {
288      tracker.trackField("content", content);
289    }
290    for (const auto& requirement: state->requirements) {
291      tracker.trackField("requirement", requirement);
292    }
293  }
294 
295 private:
296  kj::Own<State> state;
297};
298 
299struct MemorySnapshotResult {
300  kj::Array<kj::byte> snapshot;
301  kj::Array<kj::String> importedModulesList;
302  kj::String snapshotType;
303  JSG_STRUCT(snapshot, importedModulesList, snapshotType);
304};
305 
306// This used to be declared nested as ArtifactBundler::State, but then there was a need to
307// forward-declare it, so here we are.
308struct ArtifactBundler_State {
309  kj::Maybe<const PyodidePackageManager&> packageManager;
310  // ^ lifetime should be contained by lifetime of ArtifactBundler since there is normally one worker set for the whole process. see worker-set.h
311  // In other words:
312  // WorkerSet lifetime = PackageManager lifetime and Worker lifetime = ArtifactBundler lifetime and WorkerSet owns and will outlive Worker, so PackageManager outlives ArtifactBundler
313 
314  // The storedSnapshot is only used while isValidating is true.
315  kj::Maybe<MemorySnapshotResult> storedSnapshot;
316 
317  // A memory snapshot of the state of the Python interpreter after initialization. Used to speed
318  // up cold starts.
319  kj::Maybe<kj::Array<const kj::byte>> existingSnapshot;
320 
321  // Set only when the validator is running. This is used to determine if it is appropriate
322  // to store a memory snapshot.
323  bool isValidating;
324 
325  // Set when the worker is a dynamically-loaded worker. Dynamic workers don't support dedicated
326  // snapshots yet, so the Python runtime uses this to skip snapshot type validation.
327  bool isDynamicWorkerFlag;
328 
329  ArtifactBundler_State(kj::Maybe<const PyodidePackageManager&> packageManager,
330      kj::Maybe<kj::Array<const kj::byte>> existingSnapshot,
331      bool isValidating = false,
332      bool isDynamicWorker = false)
333      : packageManager(packageManager),
334        storedSnapshot(kj::none),
335        existingSnapshot(kj::mv(existingSnapshot)),
336        isValidating(isValidating),
337        isDynamicWorkerFlag(isDynamicWorker) {};
338 
339  kj::Own<ArtifactBundler_State> clone() {
340    return kj::heap<ArtifactBundler_State>(packageManager,
341        existingSnapshot.map(
342            [](kj::Array<const kj::byte>& data) { return kj::heapArray<const kj::byte>(data); }),
343        isValidating, isDynamicWorkerFlag);
344  }
345};
346 
347// A loaded bundle of artifacts for a particular script id. It can also contain V8 version and
348// CPU architecture-specific artifacts. The logic for loading these is in getArtifacts.
349class ArtifactBundler: public jsg::Object {
350 public:
351  using State = ArtifactBundler_State;
352 
353  ArtifactBundler(kj::Own<State> inner): inner(kj::mv(inner)) {};
354 
355  void storeMemorySnapshot(jsg::Lock& js, MemorySnapshotResult snapshot) {
356    KJ_REQUIRE(inner->isValidating);
357    inner->storedSnapshot = kj::mv(snapshot);
358  }
359 
360  bool hasMemorySnapshot() {
361    return inner->existingSnapshot != kj::none;
362  }
363 
364  int getMemorySnapshotSize() {
365    if (inner->existingSnapshot == kj::none) {
366      return 0;
367    }
368    return KJ_REQUIRE_NONNULL(inner->existingSnapshot).size();
369  }
370 
371  int readMemorySnapshot(int offset, kj::Array<kj::byte> buf);
372  void disposeMemorySnapshot() {
373    inner->existingSnapshot = kj::none;
374  }
375 
376  // Determines whether this ArtifactBundler was created inside the validator.
377  bool isEwValidating() {
378    return inner->isValidating;
379  }
380 
381  // Determines whether this ArtifactBundler belongs to a dynamically-loaded worker.
382  bool isDynamicWorker() {
383    return inner->isDynamicWorkerFlag;
384  }
385 
386  static kj::Own<State> makeDisabledBundler() {
387    return kj::heap<State>(kj::none, kj::none);
388  }
389 
390  // Creates an ArtifactBundler that only grants access to packages, and not a memory snapshot.
391  static kj::Own<State> makePackagesOnlyBundler(kj::Maybe<const PyodidePackageManager&> manager) {
392    return kj::heap<State>(manager, kj::none);
393  }
394 
395  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
396    KJ_IF_SOME(snap, inner->existingSnapshot) {
397      tracker.trackFieldWithSize("snapshot", snap.size());
398    }
399  }
400 
401  bool isEnabled() {
402    return false;  // TODO(later): Remove this function once we regenerate the bundle.
403  }
404 
405  kj::Maybe<jsg::Ref<ReadOnlyBuffer>> getPackage(jsg::Lock& js, kj::String path) {
406    KJ_IF_SOME(pacman, inner->packageManager) {
407      KJ_IF_SOME(ptr, pacman.getPyodidePackage(path)) {
408        return js.alloc<ReadOnlyBuffer>(ptr);
409      }
410    }
411 
412    return kj::none;
413  }
414 
415  JSG_RESOURCE_TYPE(ArtifactBundler) {
416    JSG_METHOD(hasMemorySnapshot);
417    JSG_METHOD(getMemorySnapshotSize);
418    JSG_METHOD(readMemorySnapshot);
419    JSG_METHOD(disposeMemorySnapshot);
420    JSG_METHOD(isEwValidating);
421    JSG_METHOD(isDynamicWorker);
422    JSG_METHOD(storeMemorySnapshot);
423    JSG_METHOD(isEnabled);
424    JSG_METHOD(getPackage);
425  }
426 
427 private:
428  kj::Own<State> inner;
429};
430 
431class DisabledInternalJaeger: public jsg::Object {
432 public:
433  static jsg::Ref<DisabledInternalJaeger> create(jsg::Lock& js) {
434    return js.alloc<DisabledInternalJaeger>();
435  }
436  JSG_RESOURCE_TYPE(DisabledInternalJaeger) {}
437};
438 
439// This cache is used by Pyodide to store wheels fetched over the internet across workerd restarts in local dev only
440class DiskCache: public jsg::Object {
441 private:
442  static const kj::Maybe<kj::Own<const kj::Directory>> NULL_CACHE_ROOT;  // always set to kj::none
443 
444  const kj::Maybe<kj::Own<const kj::Directory>>& cacheRoot;
445  const kj::Maybe<kj::Own<const kj::Directory>>& snapshotRoot;
446 
447 public:
448  DiskCache(): cacheRoot(NULL_CACHE_ROOT), snapshotRoot(NULL_CACHE_ROOT) {};  // Disabled disk cache
449  DiskCache(const kj::Maybe<kj::Own<const kj::Directory>>& cacheRoot,
450      const kj::Maybe<kj::Own<const kj::Directory>>& snapshotRoot)
451      : cacheRoot(cacheRoot),
452        snapshotRoot(snapshotRoot) {};
453 
454  jsg::Optional<kj::Array<kj::byte>> get(jsg::Lock& js, kj::String key);
455  void put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data);
456  void putSnapshot(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data);
457 
458  JSG_RESOURCE_TYPE(DiskCache) {
459    JSG_METHOD(get);
460    JSG_METHOD(put);
461    JSG_METHOD(putSnapshot);
462  }
463};
464 
465// Reports worker fatal errors to the request observer for Runtime Analytics.
466// This is exposed to the Python runtime as a module so that the on_fatal callback
467// can report fatal errors.
468class WorkerFatalReporter: public jsg::Object {
469 public:
470  WorkerFatalReporter() {}
471 
472  void reportFatal(jsg::Lock& js, kj::String error);
473  void reportPythonWorkersInternalError(jsg::Lock& js);
474 
475  JSG_RESOURCE_TYPE(WorkerFatalReporter) {
476    JSG_METHOD(reportFatal);
477    JSG_METHOD(reportPythonWorkersInternalError);
478  }
479};
480 
481// A limiter which will throw if the startup is found to exceed limits. The script will still be
482// able to run for longer than the limit, but an error will be thrown as soon as the startup
483// finishes. This way we can enforce a Python-specific startup limit.
484//
485// TODO(later): stop execution as soon limit is reached, instead of doing so after the fact.
486class SimplePythonLimiter: public jsg::Object {
487 private:
488  int startupLimitMs;
489  kj::Maybe<kj::Function<kj::TimePoint()>> getTimeCb;
490 
491  kj::Maybe<kj::TimePoint> startTime;
492 
493 public:
494  SimplePythonLimiter(): startupLimitMs(0), getTimeCb(kj::none) {}
495 
496  SimplePythonLimiter(int startupLimitMs, kj::Function<kj::TimePoint()> getTimeCb)
497      : startupLimitMs(startupLimitMs),
498        getTimeCb(kj::mv(getTimeCb)) {}
499 
500  static jsg::Ref<SimplePythonLimiter> makeDisabled(jsg::Lock& js) {
501    return js.alloc<SimplePythonLimiter>();
502  }
503 
504  void beginStartup() {
505    KJ_IF_SOME(cb, getTimeCb) {
506      JSG_REQUIRE(startTime == kj::none, TypeError, "Cannot call `beginStartup` multiple times.");
507      startTime = cb();
508    }
509  }
510 
511  void finishStartup(kj::Maybe<kj::String> snapshotType) {
512    KJ_IF_SOME(cb, getTimeCb) {
513      JSG_REQUIRE(startTime != kj::none, TypeError, "Need to call `beginStartup` first.");
514      auto endTime = cb();
515      kj::Duration diff = endTime - KJ_ASSERT_NONNULL(startTime);
516      auto diffMs = diff / kj::MILLISECONDS;
517 
518      JSG_REQUIRE(diffMs <= startupLimitMs, TypeError, "Python Worker startup exceeded CPU limit ",
519          diffMs, "<=", startupLimitMs, " with snapshot ", snapshotType.orDefault(kj::str("none")));
520    }
521  }
522 
523  JSG_RESOURCE_TYPE(SimplePythonLimiter) {
524    JSG_METHOD(beginStartup);
525    JSG_METHOD(finishStartup);
526  }
527};
528 
529kj::Maybe<kj::String> getPyodideLock(PythonSnapshotRelease::Reader pythonSnapshotRelease);
530 
531// Returns a list of filenames we need to fetch according to the pyodide-lock.json file
532// in addition to the requirements argument, we also must include all "stdlib" packages
533// as well as any transitive dependencies needed
534kj::Array<kj::String> getPythonPackageFiles(kj::StringPtr lockFileContents,
535    kj::ArrayPtr<kj::String> requirements,
536    kj::StringPtr packagesVersion);
537 
538// Constructs the path to a Python package in the package repository
539kj::String getPyodidePackagePath(kj::StringPtr packagesVersion, kj::StringPtr filename);
540 
541#define EW_PYODIDE_ISOLATE_TYPES                                                                   \
542  api::pyodide::ReadOnlyBuffer, api::pyodide::PyodideMetadataReader,                               \
543      api::pyodide::ArtifactBundler, api::pyodide::DiskCache,                                      \
544      api::pyodide::DisabledInternalJaeger, api::pyodide::SimplePythonLimiter,                     \
545      api::pyodide::WorkerFatalReporter, api::pyodide::MemorySnapshotResult
546 
547}  // namespace workerd::api::pyodide
548 
549namespace workerd {
550kj::Maybe<PythonSnapshotRelease::Reader> getPythonSnapshotRelease(
551    CompatibilityFlags::Reader featureFlags);
552kj::String getPythonBundleName(PythonSnapshotRelease::Reader pyodideRelease);
553}  // namespace workerd