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.SIGINT, 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';
@@ -132,6 +139,60 @@ function setTimeoutTopLevelPatch(
   return 0;
 }
 
+function getSignalClockAddr(Module: Module): number {
+  if (Module.API.version !== '0.28.2') {
+    throw new PythonWorkersInternalError(
+      'getSignalClockAddr only supported in 0.28.2'
+    );
+  }
+  // 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 SIGINT = 2;
+
+export function setupRequestSignalHandling(Module: Module): void {
+  // 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] = SIGINT;
+  Module.HEAPU32[Module._Py_EMSCRIPTEN_SIGNAL_HANDLING / 4] = 0;
+}
+
 export function loadPyodide(
   isWorkerd: boolean,
   lockfile: PackageLock,
@@ -193,6 +254,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,
+  setupRequestSignalHandling,
+} 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();
+  setupRequestSignalHandling(pyodide._module);
   try {
     if (pyfunc.callWithOptions) {
       return await pyfunc.callWithOptions(

src/pyodide/types/emscripten.d.ts

@@ -132,4 +132,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> buf, 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.
+  Worker::Isolate::from(js).setCpuLimitNearlyExceededCallback(
+      [buf = kj::mv(buf), sig_clock, sig_flag]() mutable {
+    // Set signal handling clock to fire on the next check.
+    buf[sig_clock] = 0;
+    // Set signal handling to on
+    buf[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,9 @@ class PyodideMetadataReader: public jsg::Object {
 
   static kj::Array<kj::StringPtr> getBaselineSnapshotImports();
 
+  void setCpuLimitNearlyExceededCallback(
+      jsg::Lock& js, kj::Array<kj::byte> buf, int addr, int addr2);
+
   // 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 +269,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++

@@ -193,6 +193,8 @@ IoContext::IoContext(ThreadContext& thread,
 
     return promise;
   };
+  limitEnforcer->setCpuLimitNearlyExceededCallback(
+      [this]() { this->worker->getIsolate().cpuLimitNearlyExceeded(); });
 
   // Arrange to abort when limits expire.
   abortWhen(makeLimitsPromise());