Skip to content

Commit

Permalink
Instantiate Emscripten Runtime for python workers earlier.
Browse files Browse the repository at this point in the history
  • Loading branch information
danlapid committed Dec 7, 2024
1 parent 52168e3 commit 9272eae
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 104 deletions.
2 changes: 1 addition & 1 deletion src/workerd/api/pyodide/pyodide.c++
Original file line number Diff line number Diff line change
Expand Up @@ -490,7 +490,7 @@ jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
}

void SetupEmscripten::visitForGc(jsg::GcVisitor& visitor) {
visitor.visit(emscriptenRuntime.emscriptenRuntime);
visitor.visit(const_cast<EmscriptenRuntime&>(emscriptenRuntime).emscriptenRuntime);
}

bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader modules) {
Expand Down
6 changes: 3 additions & 3 deletions src/workerd/api/pyodide/pyodide.h
Original file line number Diff line number Diff line change
Expand Up @@ -411,8 +411,8 @@ class SimplePythonLimiter: public jsg::Object {

class SetupEmscripten: public jsg::Object {
public:
SetupEmscripten(EmscriptenRuntime emscriptenRuntime)
: emscriptenRuntime(kj::mv(emscriptenRuntime)) {};
SetupEmscripten(const EmscriptenRuntime& emscriptenRuntime)
: emscriptenRuntime(emscriptenRuntime) {};

jsg::JsValue getModule(jsg::Lock& js);

Expand All @@ -421,7 +421,7 @@ class SetupEmscripten: public jsg::Object {
}

private:
EmscriptenRuntime emscriptenRuntime;
const EmscriptenRuntime& emscriptenRuntime;
void visitForGc(jsg::GcVisitor& visitor);
};

Expand Down
206 changes: 106 additions & 100 deletions src/workerd/server/workerd-api.c++
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,92 @@ static const PythonConfig defaultConfig{
};
} // namespace

namespace {
kj::Path getPyodideBundleFileName(kj::StringPtr version) {
return kj::Path(kj::str("pyodide_", version, ".capnp.bin"));
}

kj::Maybe<kj::Own<const kj::ReadableFile>> getPyodideBundleFile(
const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir, kj::StringPtr version) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto file = dir->tryOpenFile(filename);

return file;
}

return kj::none;
}

void writePyodideBundleFileToDisk(const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir,
kj::StringPtr version,
kj::ArrayPtr<byte> bytes) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto replacer = dir->replaceFile(filename, kj::WriteMode::CREATE | kj::WriteMode::MODIFY);

replacer->get().writeAll(bytes);
replacer->commit();
}
}

kj::Maybe<jsg::Bundle::Reader> fetchPyodideBundle(
const api::pyodide::PythonConfig& pyConfig, kj::StringPtr version) {
if (pyConfig.pyodideBundleManager.getPyodideBundle(version) != kj::none) {
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

auto maybePyodideBundleFile = getPyodideBundleFile(pyConfig.pyodideDiskCacheRoot, version);
KJ_IF_SOME(pyodideBundleFile, maybePyodideBundleFile) {
auto body = pyodideBundleFile->readAllBytes();
pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

if (version == "dev") {
// the "dev" version is special and indicates we're using the tip-of-tree version built for testing
// so we shouldn't fetch it from the internet, only check for its existence in the disk cache
return kj::none;
}

{
KJ_LOG(INFO, "Loading Pyodide package from internet...");
kj::Thread([&]() {
kj::AsyncIoContext io = kj::setupAsyncIo();
kj::HttpHeaderTable table;

kj::TlsContext::Options options;
options.useSystemTrustStore = true;

kj::Own<kj::TlsContext> tls = kj::heap<kj::TlsContext>(kj::mv(options));
auto& network = io.provider->getNetwork();
auto tlsNetwork = tls->wrapNetwork(network);
auto& timer = io.provider->getTimer();

auto client = kj::newHttpClient(timer, table, network, *tlsNetwork);

kj::HttpHeaders headers(table);

kj::String url =
kj::str("https://pyodide.runtime-playground.workers.dev/pyodide-capnp-bin/pyodide_",
version, ".capnp.bin");

auto req = client->request(kj::HttpMethod::GET, url.asPtr(), headers);

auto res = req.response.wait(io.waitScope);
auto body = res.body->readAllBytes().wait(io.waitScope);

writePyodideBundleFileToDisk(pyConfig.pyodideDiskCacheRoot, version, body);

pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
});
}

KJ_LOG(INFO, "Loaded Pyodide package from internet");
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}
} // namespace

struct WorkerdApi::Impl final {
kj::Own<CompatibilityFlags::Reader> features;
kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> maybeOwnedModuleRegistry;
Expand All @@ -137,6 +223,7 @@ struct WorkerdApi::Impl final {
JsgWorkerdIsolate jsgIsolate;
api::MemoryCacheProvider& memoryCacheProvider;
const PythonConfig& pythonConfig;
kj::Maybe<api::pyodide::EmscriptenRuntime> maybeEmscriptenRuntime;

class Configuration {
public:
Expand Down Expand Up @@ -175,8 +262,20 @@ struct WorkerdApi::Impl final {
limitEnforcer->getCreateParams()),
memoryCacheProvider(memoryCacheProvider),
pythonConfig(pythonConfig) {
jsgIsolate.runInLockScope(
[&](JsgWorkerdIsolate::Lock& lock) { limitEnforcer->customizeIsolate(lock.v8Isolate); });
jsgIsolate.runInLockScope([&](JsgWorkerdIsolate::Lock& lock) {
limitEnforcer->customizeIsolate(lock.v8Isolate);
if (features->getPythonWorkers()) {
auto pythonRelease = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(*features));
auto version = getPythonBundleName(pythonRelease);
auto bundle = KJ_ASSERT_NONNULL(
fetchPyodideBundle(pythonConfig, version), "Failed to get Pyodide bundle");
auto context = lock.newContext<api::ServiceWorkerGlobalScope>(lock.v8Isolate);
v8::Context::Scope scope(context.getHandle(lock));
// Init emscripten synchronously, the python script will import setup-emscripten and
// call setEmscriptenModele
maybeEmscriptenRuntime = api::pyodide::EmscriptenRuntime::initialize(lock, true, bundle);
}
});
}

static v8::Local<v8::String> compileTextGlobal(
Expand Down Expand Up @@ -437,92 +536,6 @@ kj::Maybe<jsg::ModuleRegistry::ModuleInfo> WorkerdApi::tryCompileModule(jsg::Loc
KJ_UNREACHABLE;
}

namespace {
kj::Path getPyodideBundleFileName(kj::StringPtr version) {
return kj::Path(kj::str("pyodide_", version, ".capnp.bin"));
}

kj::Maybe<kj::Own<const kj::ReadableFile>> getPyodideBundleFile(
const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir, kj::StringPtr version) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto file = dir->tryOpenFile(filename);

return file;
}

return kj::none;
}

void writePyodideBundleFileToDisk(const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir,
kj::StringPtr version,
kj::ArrayPtr<byte> bytes) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto replacer = dir->replaceFile(filename, kj::WriteMode::CREATE | kj::WriteMode::MODIFY);

replacer->get().writeAll(bytes);
replacer->commit();
}
}

kj::Maybe<jsg::Bundle::Reader> fetchPyodideBundle(
const api::pyodide::PythonConfig& pyConfig, kj::StringPtr version) {
if (pyConfig.pyodideBundleManager.getPyodideBundle(version) != kj::none) {
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

auto maybePyodideBundleFile = getPyodideBundleFile(pyConfig.pyodideDiskCacheRoot, version);
KJ_IF_SOME(pyodideBundleFile, maybePyodideBundleFile) {
auto body = pyodideBundleFile->readAllBytes();
pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

if (version == "dev") {
// the "dev" version is special and indicates we're using the tip-of-tree version built for testing
// so we shouldn't fetch it from the internet, only check for its existence in the disk cache
return kj::none;
}

{
KJ_LOG(INFO, "Loading Pyodide package from internet...");
kj::Thread([&]() {
kj::AsyncIoContext io = kj::setupAsyncIo();
kj::HttpHeaderTable table;

kj::TlsContext::Options options;
options.useSystemTrustStore = true;

kj::Own<kj::TlsContext> tls = kj::heap<kj::TlsContext>(kj::mv(options));
auto& network = io.provider->getNetwork();
auto tlsNetwork = tls->wrapNetwork(network);
auto& timer = io.provider->getTimer();

auto client = kj::newHttpClient(timer, table, network, *tlsNetwork);

kj::HttpHeaders headers(table);

kj::String url =
kj::str("https://pyodide.runtime-playground.workers.dev/pyodide-capnp-bin/pyodide_",
version, ".capnp.bin");

auto req = client->request(kj::HttpMethod::GET, url.asPtr(), headers);

auto res = req.response.wait(io.waitScope);
auto body = res.body->readAllBytes().wait(io.waitScope);

writePyodideBundleFileToDisk(pyConfig.pyodideDiskCacheRoot, version, body);

pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
});
}

KJ_LOG(INFO, "Loaded Pyodide package from internet");
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}
} // namespace

void WorkerdApi::compileModules(jsg::Lock& lockParam,
config::Worker::Reader conf,
Worker::ValidationErrorReporter& errorReporter,
Expand All @@ -537,22 +550,15 @@ void WorkerdApi::compileModules(jsg::Lock& lockParam,
if (hasPythonModules(confModules)) {
KJ_REQUIRE(featureFlags.getPythonWorkers(),
"The python_workers compatibility flag is required to use Python.");
// Inject SetupEmscripten module
modules->addBuiltinModule("internal:setup-emscripten",
jsg::alloc<SetupEmscripten>(KJ_ASSERT_NONNULL(impl->maybeEmscriptenRuntime)),
workerd::jsg::ModuleRegistry::Type::INTERNAL);

auto pythonRelease = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
auto version = getPythonBundleName(pythonRelease);
auto bundle = KJ_ASSERT_NONNULL(
fetchPyodideBundle(impl->pythonConfig, version), "Failed to get Pyodide bundle");
// Inject SetupEmscripten module
{
auto& lock = kj::downcast<JsgWorkerdIsolate::Lock>(lockParam);
auto context = lock.newContext<api::ServiceWorkerGlobalScope>({}, lock.v8Isolate);
v8::Context::Scope scope(context.getHandle(lock));
// Init emscripten synchronously, the python script will import setup-emscripten and
// call setEmscriptenModele
auto emscriptenRuntime = api::pyodide::EmscriptenRuntime::initialize(lock, true, bundle);
modules->addBuiltinModule("internal:setup-emscripten",
jsg::alloc<SetupEmscripten>(kj::mv(emscriptenRuntime)),
workerd::jsg::ModuleRegistry::Type::INTERNAL);
}

// Inject Pyodide bundle
modules->addBuiltinBundle(bundle, kj::none);
Expand Down

0 comments on commit 9272eae

Please sign in to comment.