diff --git a/src/workerd/api/pyodide/pyodide.c++ b/src/workerd/api/pyodide/pyodide.c++ index 9383242d9733..fe78c846b8f9 100644 --- a/src/workerd/api/pyodide/pyodide.c++ +++ b/src/workerd/api/pyodide/pyodide.c++ @@ -27,6 +27,19 @@ void PyodideBundleManager::setPyodideBundleData( kj::mv(version), {.messageReader = kj::mv(messageReader), .bundle = bundle}); } +const kj::Maybe> PyodidePackageManager::getPyodidePackage( + kj::StringPtr id) const { + KJ_IF_SOME(t, packages.lockShared()->find(id)) { + return t.asPtr(); + } + return kj::none; +} + +void PyodidePackageManager::setPyodidePackageData( + kj::String id, kj::Array data) const { + packages.lockExclusive()->insert(kj::mv(id), kj::mv(data)); +} + static int readToTarget( kj::ArrayPtr source, int offset, kj::ArrayPtr buf) { int size = source.size(); @@ -45,6 +58,10 @@ int PackagesTarReader::read(jsg::Lock& js, int offset, kj::Array buf) return readToTarget(PYODIDE_PACKAGES_TAR.get(), offset, buf); } +int SmallPackagesTarReader::read(jsg::Lock& js, int offset, kj::Array buf) { + return readToTarget(source, offset, buf); +} + kj::Array> PyodideMetadataReader::getNames(jsg::Lock& js) { auto builder = kj::heapArrayBuilder>(this->names.size()); for (auto i: kj::zeroTo(builder.capacity())) { diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h index 8ecb23970121..ba98e4946178 100644 --- a/src/workerd/api/pyodide/pyodide.h +++ b/src/workerd/api/pyodide/pyodide.h @@ -29,6 +29,15 @@ class PyodideBundleManager { const kj::MutexGuarded> bundles; }; +class PyodidePackageManager { +public: + void setPyodidePackageData(kj::String id, kj::Array data) const; + const kj::Maybe> getPyodidePackage(kj::StringPtr id) const; + +private: + const kj::MutexGuarded>> packages; +}; + struct PythonConfig { kj::Maybe> packageDiskCacheRoot; kj::Maybe> pyodideDiskCacheRoot; @@ -157,6 +166,21 @@ class PyodideMetadataReader: public jsg::Object { } }; +// Similar to PackagesTarReader, but reads from a dynamic buffer rather than a buffer linked into +// the binary +class SmallPackagesTarReader: public jsg::Object { + kj::ArrayPtr source; + +public: + SmallPackagesTarReader(kj::ArrayPtr src): source(src) {}; + + int read(jsg::Lock& js, int offset, kj::Array buf); + + JSG_RESOURCE_TYPE(SmallPackagesTarReader) { + JSG_METHOD(read); + } +}; + struct MemorySnapshotResult { kj::Array snapshot; kj::Array importedModulesList; @@ -167,15 +191,20 @@ struct MemorySnapshotResult { // CPU architecture-specific artifacts. The logic for loading these is in getArtifacts. class ArtifactBundler: public jsg::Object { public: + typedef kj::HashMap>> PackagePromiseMap; kj::Maybe storedSnapshot; + kj::Own loadedPackages; - ArtifactBundler(kj::Maybe> existingSnapshot) + ArtifactBundler(kj::Own loadedPackages, + kj::Maybe> existingSnapshot = kj::none) : storedSnapshot(kj::none), + loadedPackages(kj::mv(loadedPackages)), existingSnapshot(kj::mv(existingSnapshot)), isValidating(false) {}; ArtifactBundler(bool isValidating = false) : storedSnapshot(kj::none), + loadedPackages(kj::heap()), existingSnapshot(kj::none), isValidating(isValidating) {}; @@ -205,8 +234,9 @@ class ArtifactBundler: public jsg::Object { return isValidating; } - static jsg::Ref makeDisabledBundler() { - return jsg::alloc(); + static jsg::Ref makeDisabledBundler( + kj::Own loadedPackages = kj::heap()) { + return jsg::alloc(kj::mv(loadedPackages)); } void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {