from __future__ import annotations
from dataclasses import dataclass
from multiprocessing import get_context
from queue import Empty
from typing import Any, Callable
@dataclass(frozen=True)
class OracleCallResult:
status: str
value: Any = None
error: str | None = None
def _worker(queue: Any, func: Callable[..., Any], args: tuple[Any, ...]) -> None:
try:
queue.put(OracleCallResult("ok", value=func(*args)))
except BaseException as exc: queue.put(OracleCallResult("error", error=f"{type(exc).__name__}: {exc}"))
def _multiprocessing_context():
try:
return get_context("fork")
except ValueError: return get_context()
def _read_result(queue: Any, exitcode: int | None) -> OracleCallResult:
try:
return queue.get_nowait()
except (Empty, EOFError, OSError, ValueError):
return OracleCallResult("error", error=f"oracle exited with code {exitcode}")
def run_with_timeout(
func: Callable[..., Any],
args: tuple[Any, ...] = (),
timeout_seconds: float | None = None,
) -> OracleCallResult:
if timeout_seconds is None or timeout_seconds <= 0:
try:
return OracleCallResult("ok", value=func(*args))
except BaseException as exc: return OracleCallResult("error", error=f"{type(exc).__name__}: {exc}")
ctx = _multiprocessing_context()
queue = ctx.Queue(maxsize=1)
process = ctx.Process(target=_worker, args=(queue, func, args))
process.start()
process.join(timeout_seconds)
if process.is_alive():
process.terminate()
process.join()
return OracleCallResult("timeout")
return _read_result(queue, process.exitcode)