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 cpu_limit_exceeded(signum, frame):
+    raise CpuLimitExceeded("Worker exceeded CPU time limit")
+
+
+signal.signal(signal.SIGINT, 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';
 
 /**
@@ -132,6 +135,44 @@ function setTimeoutTopLevelPatch(
   return 0;
 }
 
+function getSignalClockAddr(Module: Module) {
+  if (Module.API.version !== '0.28.2') {
+    throw new PythonWorkersInternalError('Should not happen');
+  }
+  // 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) {
+  Module.Py_EmscriptenSignalBuffer = new Uint8Array([2]);
+  if (Module.API.version === '0.28.2') {
+    setCpuLimitNearlyExceededCallback(
+      Module.HEAP8,
+      getSignalClockAddr(Module),
+      Module._Py_EMSCRIPTEN_SIGNAL_HANDLING
+    );
+  }
+}
+
+export function setupRequestSignalHandling(Module: Module) {
+  // 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] = 2;
+  Module.HEAPU32[Module._Py_EMSCRIPTEN_SIGNAL_HANDLING / 4] = 0;
+}
+
 export function loadPyodide(
   isWorkerd: boolean,
   lockfile: PackageLock,
@@ -193,6 +234,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/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();
+  await 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,17 @@ 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) {
+  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());

src/workerd/io/limit-enforcer.h

@@ -156,6 +156,7 @@ class LimitEnforcer {
   // Returns a promise that will reject if and when a limit is exceeded that prevents further
   // JavaScript execution, such as the CPU or memory limit.
   virtual kj::Promise<void> onLimitsExceeded() = 0;
+  virtual void setCpuLimitNearlyExceededCallback(kj::Function<void(void)>) = 0;
 
   // Throws an exception if a limit has already been exceeded which prevents further JavaScript
   // execution, such as the CPU or memory limit.