Handle CPU limit exceeded in Python workers

If we call `TerminateExecution()`, it will exit Python execution without
unwinding the stack or cleaning up the runtime state. This leaves the
Python runtime in a permanently messed up state and all further requests
will fail.

This adds a new `cpuLimitNearlyExceededCallback` to the limit enforcer and
hooks it up so that it triggers a SIGINT inside of Python. This can be used
to raise a `CpuLimitExceeded` Python error into the runtime. If this error
is ignored, then we'll hit the hard limit and be terminated. If we ever do
call `TerminateExecution()` on a Python isolate, we should condemn the isolate,
but that is left as a TODO.

To trigger the SIGINT inside of Python, we have to set two addresses:
1. we set `emscripten_signal_clock` to `0` to make the Python eval breaker check
   for a signal on the next tick.
2. we set `_Py_EMSCRIPTEN_SIGNAL_HANDLING` to 1 to make Python check the signal clock.

We also have to set `Module.Py_EmscriptenSignalBuffer` to a buffer with the number
of the signal we wish to trip in it (`SIGINT` aka 2).

When we start a request we set `_Py_EMSCRIPTEN_SIGNAL_HANDLING` to 0 to avoid
ongoing costs of calling out to JavaScript to check the buffer when no signal is set,
and we put a 2 into `Py_EmscriptenSignalBuffer`.

The most annoying aspect of this is that the symbol `emscripten_signal_clock` is not
exported. For Pyodide 0.28.2, I manually located the address of this symbol and hard
coded it. For the next Pyodide, we'll make sure to export it.

Author: hoodmane

src/pyodide/internal/introspection.py

@@ -1,3 +1,4 @@
+import signal
 from inspect import isawaitable, isclass
 from types import FunctionType
 
@@ -86,3 +87,14 @@ async def wrapper_func(relaxed, inst, prop, *args, **kwargs):
         return python_to_rpc(await result)
     else:
         return python_to_rpc(result)
+
+
+class CpuLimitExceeded(BaseException):
+    pass
+
+
+def raise_cpu_limit_exceeded(signum, frame):
+    raise CpuLimitExceeded("Python Worker exceeded CPU time limit")
+
+
+signal.signal(signal.SIGXCPU, raise_cpu_limit_exceeded)

src/pyodide/internal/metadata.ts

@@ -63,3 +63,6 @@ export const LEGACY_GLOBAL_HANDLERS = !NO_GLOBAL_HANDLERS;
 export const LEGACY_VENDOR_PATH = !FORCE_NEW_VENDOR_PATH;
 export const LEGACY_INCLUDE_SDK = !EXTERNAL_SDK;
 export const CHECK_RNG_STATE = !!COMPATIBILITY_FLAGS.python_check_rng_state;
+
+export const setCpuLimitNearlyExceededCallback =
+  MetadataReader.setCpuLimitNearlyExceededCallback.bind(MetadataReader);

src/pyodide/internal/python.ts

@@ -16,7 +16,10 @@ import {
   getRandomValues,
   entropyBeforeRequest,
 } from 'pyodide-internal:topLevelEntropy/lib';
-import { LEGACY_VENDOR_PATH } from 'pyodide-internal:metadata';
+import {
+  LEGACY_VENDOR_PATH,
+  setCpuLimitNearlyExceededCallback,
+} from 'pyodide-internal:metadata';
 import type { PyodideEntrypointHelper } from 'pyodide:python-entrypoint-helper';
 
 /**
@@ -27,7 +30,11 @@ import type { PyodideEntrypointHelper } from 'pyodide:python-entrypoint-helper';
 import { default as SetupEmscripten } from 'internal:setup-emscripten';
 
 import { default as UnsafeEval } from 'internal:unsafe-eval';
-import { PythonWorkersInternalError, reportError } from 'pyodide-internal:util';
+import {
+  PythonWorkersInternalError,
+  reportError,
+  unreachable,
+} from 'pyodide-internal:util';
 import { loadPackages } from 'pyodide-internal:loadPackage';
 import { default as MetadataReader } from 'pyodide-internal:runtime-generated/metadata';
 import { TRANSITIVE_REQUIREMENTS } from 'pyodide-internal:metadata';
@@ -116,20 +123,87 @@ function validatePyodideVersion(pyodide: Pyodide): void {
 }
 
 const origSetTimeout = globalThis.setTimeout.bind(this);
-function setTimeoutTopLevelPatch(
-  handler: () => void,
-  timeout: number | undefined
-): number {
-  // Redirect top level setTimeout(cb, 0) to queueMicrotask().
-  // If we don't know how to handle it, call normal setTimeout() to force failure.
-  if (typeof handler === 'string') {
-    return origSetTimeout(handler, timeout);
+
+function makeSetTimeout(Module: Module): typeof setTimeout {
+  return function setTimeoutTopLevelPatch(
+    handler: () => void,
+    timeout: number | undefined
+  ): number {
+    // Redirect top level setTimeout(cb, 0) to queueMicrotask().
+    // If we don't know how to handle it, call normal setTimeout() to force failure.
+    if (typeof handler === 'string') {
+      return origSetTimeout(handler, timeout);
+    }
+    function wrappedHandler() {
+      // In case an Exceeded CPU occurred just as Python was exiting, there may be one waiting that
+      // will interrupt the wrong task. Clear signals before entering the task.
+      // This is covered by cpu-limit-exceeded.ew-test "async_trip" test.
+      clearSignals(Module);
+      handler();
+    }
+    if (timeout) {
+      return origSetTimeout(wrappedHandler, timeout);
+    }
+    queueMicrotask(wrappedHandler);
+    return 0;
+  } as typeof setTimeout;
+}
+
+function getSignalClockAddr(Module: Module): number {
+  if (Module.API.version !== '0.28.2') {
+    throw new PythonWorkersInternalError(
+      'getSignalClockAddr only supported in 0.28.2'
+    );
   }
-  if (timeout) {
-    return origSetTimeout(handler, timeout);
+  // This is the address here:
+  // https://github.com/python/cpython/blob/main/Python/emscripten_signal.c#L42
+  //
+  // Since the symbol isn't exported, we can't access it directly. Instead, we used wasm-objdump and
+  // searched for the call site to _Py_CheckEmscriptenSignals_Helper(), then read the offset out of
+  // the assembly code.
+  //
+  // TODO: Export this symbol in the next Pyodide release so we can stop using the magic number.
+  const emscripten_signal_clock_offset = 3171536;
+  return Module.___memory_base.value + emscripten_signal_clock_offset;
+}
+
+function setupRuntimeSignalHandling(Module: Module): void {
+  Module.Py_EmscriptenSignalBuffer = new Uint8Array(1);
+  const version = Module.API.version;
+  if (version === '0.26.0a2') {
+    return;
+  }
+  if (version === '0.28.2') {
+    // The callback sets signal_clock to 0 and signal_handling to 1. It has to be in C++ because we
+    // don't hold the isolate lock when we call it. JS code would be:
+    //
+    // function callback() { Module.HEAP8[getSignalClockAddr(Module)] = 0;
+    //    Module.HEAP8[Module._Py_EMSCRIPTEN_SIGNAL_HANDLING] = 1;
+    // }
+    setCpuLimitNearlyExceededCallback(
+      Module.HEAP8,
+      getSignalClockAddr(Module),
+      Module._Py_EMSCRIPTEN_SIGNAL_HANDLING
+    );
+    return;
+  }
+  unreachable(version);
+}
+
+const SIGXCPU = 24;
+
+export function clearSignals(Module: Module): void {
+  if (Module.API.version === '0.28.2') {
+    // In case the previous request was aborted, make sure that:
+    // 1. a sigint is waiting in the signal buffer
+    // 2. signal handling is off
+    //
+    // We will turn signal handling on as part of triggering the interrupt, having it on otherwise
+    // just wastes cycles.
+    Module.Py_EmscriptenSignalBuffer[0] = SIGXCPU;
+    Module.HEAPU32[getSignalClockAddr(Module)] = 1;
+    Module.HEAPU32[Module._Py_EMSCRIPTEN_SIGNAL_HANDLING / 4] = 0;
   }
-  queueMicrotask(handler);
-  return 0;
 }
 
 export function loadPyodide(
@@ -150,7 +224,7 @@ export function loadPyodide(
     Module.setUnsafeEval(UnsafeEval);
     Module.setGetRandomValues(getRandomValues);
     Module.setSetTimeout(
-      setTimeoutTopLevelPatch as typeof setTimeout,
+      makeSetTimeout(Module),
       clearTimeout,
       setInterval,
       clearInterval
@@ -193,6 +267,7 @@ export function loadPyodide(
       }
     );
     setupPythonSearchPath(pyodide);
+    setupRuntimeSignalHandling(Module);
     return pyodide;
   } catch (e) {
     // In edgeworker test suite, without this we get the file name and line number of the exception

src/pyodide/internal/util.ts

@@ -78,3 +78,7 @@ export function invalidateCaches(Module: Module): void {
     `from importlib import invalidate_caches; invalidate_caches(); del invalidate_caches`
   );
 }
+
+export function unreachable(msg: never): never {
+  throw new PythonWorkersInternalError(`Unreachable: ${msg}`);
+}

src/pyodide/python-entrypoint-helper.ts

@@ -2,7 +2,11 @@
 // This file is a BUILTIN module that provides the actual implementation for the
 // python-entrypoint.js USER module.
 
-import { beforeRequest, loadPyodide } from 'pyodide-internal:python';
+import {
+  beforeRequest,
+  loadPyodide,
+  clearSignals,
+} from 'pyodide-internal:python';
 import { enterJaegerSpan } from 'pyodide-internal:jaeger';
 import { patchLoadPackage } from 'pyodide-internal:setupPackages';
 import {
@@ -292,6 +296,8 @@ async function doPyCallHelper(
   pyfunc: PyCallable,
   args: any[]
 ): Promise<any> {
+  const pyodide = await getPyodide();
+  clearSignals(pyodide._module);
   try {
     if (pyfunc.callWithOptions) {
       return await pyfunc.callWithOptions(

src/pyodide/types/emscripten.d.ts

@@ -38,6 +38,7 @@ interface API {
   };
   serializeHiwireState(serializer: (obj: any) => any): SnapshotConfig;
   pyVersionTuple: [number, number, number];
+  scheduleCallback: (callback: () => void, timeout: number) => void;
 }
 
 interface LDSO {
@@ -74,9 +75,7 @@ interface EmscriptenSettings {
   ) => WebAssembly.Exports;
   reportUndefinedSymbolsNoOp: () => void;
   noInitialRun?: boolean;
-  API: {
-    config: API['config'];
-  };
+  API: Pick<API, 'config'>;
   readyPromise: Promise<Module>;
   rejectReadyPromise: (e: any) => void;
 }
@@ -132,4 +131,7 @@ interface Module {
   getEmptyTableSlot(): number;
   freeTableIndexes: number[];
   LD_LIBRARY_PATH: string;
+  Py_EmscriptenSignalBuffer: Uint8Array;
+  _Py_EMSCRIPTEN_SIGNAL_HANDLING: number;
+  ___memory_base: WebAssembly.Global<'i32'>;
 }

src/pyodide/types/runtime-generated/metadata.d.ts

@@ -30,6 +30,11 @@ declare namespace MetadataReader {
   const read: (index: number, position: number, buffer: Uint8Array) => number;
   const getTransitiveRequirements: () => Set<string>;
   const getCompatibilityFlags: () => CompatibilityFlags;
+  const setCpuLimitNearlyExceededCallback: (
+    buf: Uint8Array,
+    sig_clock: number,
+    sig_flag: number
+  ) => void;
   const constructor: {
     getBaselineSnapshotImports(): string[];
   };

src/workerd/api/pyodide/pyodide.c++

@@ -91,6 +91,19 @@ kj::Array<kj::StringPtr> PyodideMetadataReader::getNames(
   return builder.releaseAsArray();
 }
 
+void PyodideMetadataReader::setCpuLimitNearlyExceededCallback(
+    jsg::Lock& js, kj::Array<kj::byte> wasm_memory, int sig_clock, int sig_flag) {
+  // This callback has to be implemented in C++ because we don't hold the isolate lock when we call
+  // it. It also has to be signal safe since we call it from the cpu time limiter.
+  Worker::Isolate::from(js).setCpuLimitNearlyExceededCallback(
+      [wasm_memory = kj::mv(wasm_memory), sig_clock, sig_flag]() mutable {
+    // Set signal handling clock to fire on the next check.
+    wasm_memory[sig_clock] = 0;
+    // Set signal handling to on
+    wasm_memory[sig_flag] = 1;
+  });
+}
+
 kj::Array<kj::String> PythonModuleInfo::getPythonFileContents() {
   auto builder = kj::Vector<kj::String>(names.size());
   for (auto i: kj::zeroTo(names.size())) {

src/workerd/api/pyodide/pyodide.h

@@ -238,6 +238,12 @@ class PyodideMetadataReader: public jsg::Object {
 
   static kj::Array<kj::StringPtr> getBaselineSnapshotImports();
 
+  // We call this during Python setup with the wasm memory and the addresses of the signal clock and
+  // the flag to indicate whether signal handling is on or off. It sets up the isolate
+  // CpuLimitNearlyExceeded callback to trigger a signal in Python.
+  void setCpuLimitNearlyExceededCallback(
+      jsg::Lock& js, kj::Array<kj::byte> wasm_memory, int sig_clock, int sig_flag);
+
   // Similar to Cloudflare::::getCompatibilityFlags in global-scope.c++, but the key difference is
   // that it returns experimental flags even if `experimental` is not enabled. This avoids a gotcha
   // where an experimental compat flag is enabled in our C++ code, but not in our JS code.
@@ -266,6 +272,7 @@ class PyodideMetadataReader: public jsg::Object {
     JSG_METHOD(getTransitiveRequirements);
     JSG_METHOD(getCompatibilityFlags);
     JSG_STATIC_METHOD(getBaselineSnapshotImports);
+    JSG_METHOD(setCpuLimitNearlyExceededCallback);
   }
 
   void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {

src/workerd/io/io-context.c++

@@ -192,6 +192,9 @@ IoContext::IoContext(ThreadContext& thread,
 
     return promise;
   };
+  KJ_IF_SOME(cb, this->worker->getIsolate().getCpuLimitNearlyExceededCallback()) {
+    limitEnforcer->setCpuLimitNearlyExceededCallback(kj::mv(cb));
+  }
 
   // Arrange to abort when limits expire.
   abortWhen(makeLimitsPromise());