diff --git a/src/pyodide/python-entrypoint-helper.ts b/src/pyodide/python-entrypoint-helper.ts index e06b202576f..f29807198b2 100644 --- a/src/pyodide/python-entrypoint-helper.ts +++ b/src/pyodide/python-entrypoint-helper.ts @@ -35,11 +35,34 @@ import { } from 'pyodide-internal:util'; import { PyodideVersion } from 'pyodide-internal:const'; import { default as introspectionSource } from 'pyodide-internal:introspection.py'; -export { createImportProxy } from 'pyodide-internal:serializeJsModule'; +import { createImportProxy } from 'pyodide-internal:serializeJsModule'; +export { createImportProxy }; type PyFuture = Promise & { copy(): PyFuture; destroy(): void }; const waitUntilPatched = new WeakSet(); +function getPatchedWaitUntil(ctx: { + waitUntil: (p: Promise | PyFuture) => void; +}): (p: Promise | PyFuture) => void { + const origWaitUntil: (p: Promise) => void = ctx.waitUntil.bind(ctx); + function waitUntil(p: Promise | PyFuture): void { + origWaitUntil( + (async function (): Promise { + if ('copy' in p) { + p = p.copy(); + } + try { + await p; + } finally { + if ('destroy' in p) { + p.destroy(); + } + } + })() + ); + } + return waitUntil; +} function patchWaitUntil(ctx: { waitUntil: (p: Promise | PyFuture) => void; @@ -54,21 +77,7 @@ function patchWaitUntil(ctx: { if (waitUntilPatched.has(ctx)) { return; } - const origWaitUntil: (p: Promise) => void = ctx.waitUntil.bind(ctx); - function waitUntil(p: Promise | PyFuture): void { - origWaitUntil( - (async function (): Promise { - if ('copy' in p) { - p = p.copy(); - } - await p; - if ('destroy' in p) { - p.destroy(); - } - })() - ); - } - ctx.waitUntil = waitUntil; + ctx.waitUntil = getPatchedWaitUntil(ctx); waitUntilPatched.add(ctx); } @@ -93,13 +102,35 @@ function get_pyodide_entrypoint_helper(): PyodideEntrypointHelper { return _pyodide_entrypoint_helper; } +async function getCloudflareWorkersModule( + doAnImport: (mod: string) => Promise +): Promise<{ env: any }> { + const cloudflareWorkersModule = await doAnImport('cloudflare:workers'); + const waitUntil = createImportProxy( + 'cloudflare:workers', + getPatchedWaitUntil(cloudflareWorkersModule), + ['waitUntil'] + ); + return new Proxy(cloudflareWorkersModule, { + get(_target: any, prop: string | symbol): any { + if (prop === 'waitUntil') { + return waitUntil; + } + // @ts-expect-error untyped Reflect.get + // eslint-disable-next-line prefer-rest-params + return Reflect.get(...arguments); + }, + }); +} + export async function setDoAnImport( doAnImport: (mod: string) => Promise, workerEntrypoint: any ): Promise { + const cloudflareWorkersModule = await getCloudflareWorkersModule(doAnImport); _pyodide_entrypoint_helper = { doAnImport, - cloudflareWorkersModule: await doAnImport('cloudflare:workers'), + cloudflareWorkersModule, cloudflareSocketsModule: await doAnImport('cloudflare:sockets'), workerEntrypoint, patchWaitUntil, @@ -112,7 +143,12 @@ export async function setDoAnImport( throw new PythonWorkersInternalError(message); }, }; - await fillSnapshotJsModules(doAnImport); + await fillSnapshotJsModules(async (mod: string): Promise => { + if (mod === 'cloudflare:workers') { + return cloudflareWorkersModule; + } + return await doAnImport(mod); + }); } function handleSrcImport(pyodide: Pyodide, e: any): never { diff --git a/src/workerd/server/tests/python/python-rpc/worker.py b/src/workerd/server/tests/python/python-rpc/worker.py index f7c3b7d3d61..37f45ae3505 100644 --- a/src/workerd/server/tests/python/python-rpc/worker.py +++ b/src/workerd/server/tests/python/python-rpc/worker.py @@ -10,14 +10,16 @@ from unittest import TestCase import js -from workers import Blob, Request, Response, WorkerEntrypoint, handler +from workers import Blob, Request, Response, WorkerEntrypoint, handler, waitUntil -from pyodide.ffi import JsException, JsProxy, to_js +from pyodide.ffi import JsException, JsProxy, create_proxy, to_js assertRaises = TestCase().assertRaises assertRaisesRegex = TestCase().assertRaisesRegex testFuture = Future() +globalWaitUntilFuture = Future() +explicitProxyWaitUntilFuture = Future() class PythonRpcTester(WorkerEntrypoint): @@ -66,6 +68,20 @@ async def sleep_then_set_result(self): async def test_wait_until_coroutine_lifetime(self): self.ctx.waitUntil(self.sleep_then_set_result()) + async def explicit_proxy_wait_until(self): + async def set_result(): + await sleep(0.1) + explicitProxyWaitUntilFuture.set_result(100) + + self.ctx.waitUntil(create_proxy(set_result())) + + async def sleep_then_set_global_wait_until_result(self): + await sleep(0.1) + globalWaitUntilFuture.set_result(100) + + async def test_global_wait_until_coroutine_lifetime(self): + waitUntil(self.sleep_then_set_global_wait_until_result()) + class CustomType: def __init__(self, x): @@ -280,3 +296,15 @@ def my_func(): assert not testFuture.done() await sleep(0.2) assert testFuture.result() == 100 + + # Check that waitUntil() accepts a PyProxy explicitly created by the user. + await env.PythonRpc.explicit_proxy_wait_until() + assert not explicitProxyWaitUntilFuture.done() + await sleep(0.2) + assert explicitProxyWaitUntilFuture.result() == 100 + + # Check that the module-level waitUntil() also keeps its coroutine alive. + await env.PythonRpc.test_global_wait_until_coroutine_lifetime() + assert not globalWaitUntilFuture.done() + await sleep(0.2) + assert globalWaitUntilFuture.result() == 100