diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f3a91f..9f0d9e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,49 @@ # Changelog +## Unreleased + +### Added + +- **Interrupt running Python** - `py:interrupt/1` and `py_context:interrupt/1` + raise `KeyboardInterrupt` in the thread executing a context; the in-flight call + returns `{error, interrupted}`. `py_context:call/eval/exec` now interrupt + automatically when their timeout expires, so `{error, timeout}` stops the Python + code instead of only abandoning the reply while the thread kept burning CPU. + Works in both `worker` and `owngil` modes, and is callable while the context + process is blocked in a NIF. CPython delivers async exceptions at bytecode + boundaries, so code blocked in a C call (`time.sleep`, a numpy kernel, a socket + read) is interrupted once that call returns. See `docs/interrupts.md`. +- **Per-context memory caps** - `py_context:new(#{mode => owngil, memory_limit => Bytes})` + caps memory allocated by one context; exceeding it raises `MemoryError` there and + leaves other contexts untouched. Requires `{enable_memory_limits, true}` in the + application environment, since the allocator is hooked before Python starts. + Accounting covers obmalloc traffic only: allocations over 512 bytes (large + binaries, numpy buffers) are not counted, granularity is one 1 MB arena, and + `worker` mode returns `{error, memory_limit_requires_owngil}` because those + contexts share the main interpreter. `py_nif:context_memory_usage/1` reports + usage. See `docs/memory.md`. +- **Async generator streaming** - `py:stream_start/3,4` accepts async generators + again, driving them on a private event loop and emitting the same + `{py_stream, Ref, ...}` events. The docs claimed this since 3.0.0 without the + code behind it. `py:stream/4` with kwargs and `py:stream_eval/1,2` remain + sync-only, which is now stated explicitly. + +### Changed + +- **Callback results cross as external term format** - results returned from an + Erlang callback into Python are encoded with `term_to_binary` and decoded by the + same `term_to_py` converter used for call arguments, replacing the Python-repr + string that was parsed with `ast.literal_eval`. This fixes binaries containing + backslashes, quotes, newlines or tabs (which produced an unparseable literal and + were silently handed to Python as the raw repr text), `[]` arriving as `''`, + float precision loss, and the base64 round-trip for pids and references, which + now cross as native `Pid` and `Ref` objects. + + Breaking: an Erlang string returned from a callback (`"abc"`, a list of + integers) now reaches Python as `[97, 98, 99]` rather than `'abc'`, the same + conversion call arguments have always used. Return a binary (`<<"abc">>`) for a + Python `str`. + ## 3.1.1 (2026-05-31) ### Changed diff --git a/c_src/py_callback.c b/c_src/py_callback.c index 4aa87df..69d822b 100644 --- a/c_src/py_callback.c +++ b/c_src/py_callback.c @@ -946,6 +946,30 @@ static int copy_callback_results_to_nested(suspended_context_state_t *nested, return 0; } +/** + * Decode an external-term-format binary into a Python object. + * + * Returns a NEW reference on success, NULL on failure (decode error or a + * conversion error, in which case a Python exception may be set). + */ +static PyObject *etf_binary_to_py(const char *data, size_t len) { + ErlNifEnv *tmp_env = enif_alloc_env(); + if (tmp_env == NULL) { + return NULL; + } + + ERL_NIF_TERM term; + if (enif_binary_to_term(tmp_env, (unsigned char *)data, len, &term, + ERL_NIF_BIN2TERM_SAFE) == 0) { + enif_free_env(tmp_env); + return NULL; + } + + PyObject *result = term_to_py(tmp_env, term); + enif_free_env(tmp_env); + return result; +} + /** * Helper to convert __etf__:base64 strings to Python objects. * Used for encoding pids and references in callback responses. @@ -1003,28 +1027,8 @@ static PyObject *decode_etf_string(const char *str, Py_ssize_t len) { return NULL; } - /* Create a temporary NIF environment to decode the term */ - ErlNifEnv *tmp_env = enif_alloc_env(); - if (tmp_env == NULL) { - Py_DECREF(decoded); - return NULL; - } - - /* Decode the ETF binary to an Erlang term */ - ERL_NIF_TERM term; - if (enif_binary_to_term(tmp_env, (unsigned char *)bin_data, bin_len, &term, ERL_NIF_BIN2TERM_SAFE) == 0) { - /* Decoding failed */ - enif_free_env(tmp_env); - Py_DECREF(decoded); - return NULL; - } - + PyObject *result = etf_binary_to_py(bin_data, (size_t)bin_len); Py_DECREF(decoded); - - /* Convert the term to a Python object */ - PyObject *result = term_to_py(tmp_env, term); - enif_free_env(tmp_env); - return result; } @@ -1150,7 +1154,12 @@ static PyObject *convert_etf_strings(PyObject *obj) { /** * Helper to parse callback response data into a Python object. - * Response format: status_byte (0=ok, 1=error) + python_repr_string + * + * Response format: status_byte + payload + * 0 = ok, payload is a Python source literal (legacy; kept so responses + * produced by an older Erlang side still decode) + * 1 = error, payload is an error message + * 2 = ok, payload is an external-term-format binary (current encoding) */ static PyObject *parse_callback_response(unsigned char *response_data, size_t response_len) { if (response_len < 1) { @@ -1161,7 +1170,7 @@ static PyObject *parse_callback_response(unsigned char *response_data, size_t re uint8_t status = response_data[0]; if (response_len < 2) { - if (status == 0) { + if (status == 0 || status == 2) { Py_RETURN_NONE; } else { PyErr_SetString(PyExc_RuntimeError, "Erlang callback failed"); @@ -1173,7 +1182,18 @@ static PyObject *parse_callback_response(unsigned char *response_data, size_t re size_t result_len = response_len - 1; PyObject *result = NULL; - if (status == 0) { + if (status == 2) { + /* Current encoding: external term format, decoded by the same + * converter that handles call arguments. */ + result = etf_binary_to_py(result_str, result_len); + if (result == NULL) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_RuntimeError, + "Failed to decode callback result"); + } + return NULL; + } + } else if (status == 0) { /* Try to evaluate the result string as Python literal. * Import ast.literal_eval fresh to support subinterpreters * (the cached g_ast_literal_eval may be from a different interpreter). */ diff --git a/c_src/py_convert.c b/c_src/py_convert.c index cbd9bab..0d27b0e 100644 --- a/c_src/py_convert.c +++ b/c_src/py_convert.c @@ -825,6 +825,16 @@ static ERL_NIF_TERM make_py_error(ErlNifEnv *env) { enif_make_tuple2(env, ATOM_STOP_ITERATION, ATOM_NONE)); } + /* KeyboardInterrupt means py_nif:context_interrupt/1 injected an async + * exception, so report it as a distinct reason rather than a Python error. */ + if (PyErr_GivenExceptionMatches(type, PyExc_KeyboardInterrupt)) { + PyErr_Clear(); + Py_XDECREF(type); + Py_XDECREF(value); + Py_XDECREF(traceback); + return enif_make_tuple2(env, ATOM_ERROR, enif_make_atom(env, "interrupted")); + } + /* Get exception message. PyUnicode_AsUTF8 can return NULL (e.g. a str with * lone surrogates); never pass NULL to enif_make_string. */ PyObject *str = PyObject_Str(value); diff --git a/c_src/py_mem_limit.c b/c_src/py_mem_limit.c new file mode 100644 index 0000000..5e2d7dc --- /dev/null +++ b/c_src/py_mem_limit.c @@ -0,0 +1,280 @@ +/* Copyright 2026 Benoit Chesneau + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file py_mem_limit.c + * @brief Optional per-interpreter memory caps via obmalloc arena accounting. + * + * Why arenas: PyObjectArenaAllocator's free hook receives the block size, + * so allocations and frees can be attributed exactly without adding a size + * header to every object. PyMem_SetAllocator's free hook does not, which is + * why object-level accounting would require growing every allocation. + * + * Enforcement raises MemoryError asynchronously in the thread that crossed the + * cap, the same mechanism py_nif:context_interrupt/1 uses. Returning NULL from + * the arena allocator does NOT work: obmalloc treats a failed arena as a reason + * to fall back to PyMem_RawMalloc (Objects/obmalloc.c, _PyObject_Malloc), so + * the allocation would silently succeed off-arena and the cap would stop + * counting instead of stopping the code. + * + * Because an async exception lands at the next bytecode boundary, usage can + * overshoot the cap slightly before the code stops. The cap re-arms once usage + * drops back below it. + * + * Scope and limits (documented in doc/features.md): + * - Attribution is per *interpreter*. All worker-mode contexts share the main + * interpreter, so a per-context cap is only meaningful in owngil mode. + * - Allocations larger than obmalloc's small-object threshold (512 bytes) go + * straight to malloc and are NOT counted: large bytes objects, numpy + * buffers, and anything using its own allocator. + * - Granularity is one arena (1 MB on current CPython). + * + * The hooks are only installed when the application asks for them, so the + * default build path is untouched. + */ + +#ifdef HAVE_OWNGIL + +/** @brief Maximum number of interpreters tracked at once */ +#define PY_MEM_LIMIT_SLOTS 64 + +typedef struct { + /** @brief Interpreter this slot accounts for (NULL = free slot) */ + PyInterpreterState *interp; + + /** @brief Arena bytes currently allocated to this interpreter */ + size_t used; + + /** @brief Cap in bytes (0 = unlimited, accounting only) */ + size_t limit; + + /** @brief True once the cap was hit and MemoryError was injected */ + bool tripped; +} py_mem_slot_t; + +static py_mem_slot_t g_mem_slots[PY_MEM_LIMIT_SLOTS]; +static pthread_mutex_t g_mem_slots_mutex = PTHREAD_MUTEX_INITIALIZER; +static PyObjectArenaAllocator g_base_arena; +static bool g_mem_limits_enabled = false; + +/** + * Find the slot for @p interp, optionally creating it. + * + * Caller must hold g_mem_slots_mutex. Never acquires the GIL, so the + * allocator (which runs holding a GIL) cannot deadlock against registration. + */ +static py_mem_slot_t *mem_slot_find(PyInterpreterState *interp, bool create) { + py_mem_slot_t *free_slot = NULL; + + for (int i = 0; i < PY_MEM_LIMIT_SLOTS; i++) { + if (g_mem_slots[i].interp == interp) { + return &g_mem_slots[i]; + } + if (free_slot == NULL && g_mem_slots[i].interp == NULL) { + free_slot = &g_mem_slots[i]; + } + } + + if (!create || free_slot == NULL) { + return NULL; + } + + free_slot->interp = interp; + free_slot->used = 0; + free_slot->limit = 0; + free_slot->tripped = false; + return free_slot; +} + +/** @brief Current interpreter, or NULL when no thread state is attached */ +static PyInterpreterState *mem_current_interp(void) { + PyThreadState *tstate = PyThreadState_GetUnchecked(); + return (tstate != NULL) ? PyThreadState_GetInterpreter(tstate) : NULL; +} + +static void *py_mem_arena_alloc(void *ctx, size_t size) { + PyInterpreterState *interp = mem_current_interp(); + py_mem_slot_t *slot = NULL; + bool inject = false; + + if (interp != NULL) { + pthread_mutex_lock(&g_mem_slots_mutex); + slot = mem_slot_find(interp, true); + if (slot != NULL) { + slot->used += size; + if (slot->limit != 0 && slot->used > slot->limit && !slot->tripped) { + slot->tripped = true; + inject = true; + } + } + pthread_mutex_unlock(&g_mem_slots_mutex); + } + + if (inject) { + /* Raise MemoryError at the next bytecode boundary in this thread. We + * hold this interpreter's GIL (obmalloc runs under it), and + * SetAsyncExc neither allocates nor re-enters the allocator. + * Done outside g_mem_slots_mutex to keep that lock leaf-level. */ + PyThreadState_SetAsyncExc(PyThread_get_thread_ident(), PyExc_MemoryError); + } + + void *ptr = g_base_arena.alloc(ctx, size); + + if (ptr == NULL && slot != NULL) { + /* Roll back the reservation on a genuine allocation failure */ + pthread_mutex_lock(&g_mem_slots_mutex); + slot = mem_slot_find(interp, false); + if (slot != NULL) { + slot->used = (slot->used > size) ? slot->used - size : 0; + } + pthread_mutex_unlock(&g_mem_slots_mutex); + } + + return ptr; +} + +static void py_mem_arena_free(void *ctx, void *ptr, size_t size) { + PyInterpreterState *interp = mem_current_interp(); + + if (interp != NULL) { + pthread_mutex_lock(&g_mem_slots_mutex); + py_mem_slot_t *slot = mem_slot_find(interp, false); + if (slot != NULL) { + slot->used = (slot->used > size) ? slot->used - size : 0; + if (slot->tripped && slot->limit != 0 && slot->used < slot->limit) { + /* Back under the cap, re-arm enforcement */ + slot->tripped = false; + } + } + pthread_mutex_unlock(&g_mem_slots_mutex); + } + + g_base_arena.free(ctx, ptr, size); +} + +/** + * @brief Install the accounting arena allocator + * + * Must be called before Py_Initialize, since arenas are allocated during + * interpreter startup and the base allocator is captured here. + */ +static void py_mem_limit_install(void) { + if (g_mem_limits_enabled) { + return; + } + + PyObject_GetArenaAllocator(&g_base_arena); + if (g_base_arena.alloc == NULL || g_base_arena.free == NULL) { + return; + } + + PyObjectArenaAllocator wrapper = g_base_arena; + wrapper.alloc = py_mem_arena_alloc; + wrapper.free = py_mem_arena_free; + PyObject_SetArenaAllocator(&wrapper); + + g_mem_limits_enabled = true; +} + +static bool py_mem_limit_enabled(void) { + return g_mem_limits_enabled; +} + +/** + * @brief Set the cap for an interpreter + * @param limit Bytes, or 0 to remove the cap (accounting continues) + * @return 0 on success, -1 if no slot is available + */ +static int py_mem_limit_set(PyInterpreterState *interp, size_t limit) { + if (!g_mem_limits_enabled || interp == NULL) { + return -1; + } + + pthread_mutex_lock(&g_mem_slots_mutex); + py_mem_slot_t *slot = mem_slot_find(interp, true); + if (slot != NULL) { + slot->limit = limit; + slot->tripped = false; + } + pthread_mutex_unlock(&g_mem_slots_mutex); + + return (slot != NULL) ? 0 : -1; +} + +/** + * @brief Read current usage and cap for an interpreter + * @return 0 on success, -1 if the interpreter is not tracked + */ +static int py_mem_limit_get(PyInterpreterState *interp, size_t *used, size_t *limit) { + if (!g_mem_limits_enabled || interp == NULL) { + return -1; + } + + pthread_mutex_lock(&g_mem_slots_mutex); + py_mem_slot_t *slot = mem_slot_find(interp, false); + if (slot != NULL) { + *used = slot->used; + *limit = slot->limit; + } + pthread_mutex_unlock(&g_mem_slots_mutex); + + return (slot != NULL) ? 0 : -1; +} + +/** + * @brief Release the slot for a destroyed interpreter + * + * Must be called after Py_EndInterpreter, so the arenas freed during teardown + * are still accounted. Without this a recycled PyInterpreterState pointer + * would inherit the dead interpreter's usage. + */ +static void py_mem_limit_forget(PyInterpreterState *interp) { + if (!g_mem_limits_enabled || interp == NULL) { + return; + } + + pthread_mutex_lock(&g_mem_slots_mutex); + py_mem_slot_t *slot = mem_slot_find(interp, false); + if (slot != NULL) { + slot->interp = NULL; + slot->used = 0; + slot->limit = 0; + slot->tripped = false; + } + pthread_mutex_unlock(&g_mem_slots_mutex); +} + +#else /* !HAVE_OWNGIL */ + +/* Memory limits require OWN_GIL subinterpreters (Python 3.14+): without them + * every context shares the main interpreter, so a per-context cap has no + * meaning. These stubs keep the call sites free of #ifdef. */ + +static void py_mem_limit_install(void) {} +static bool py_mem_limit_enabled(void) { return false; } + +#ifdef HAVE_SUBINTERPRETERS +static int py_mem_limit_set(PyInterpreterState *interp, size_t limit) { + (void)interp; (void)limit; + return -1; +} +static int py_mem_limit_get(PyInterpreterState *interp, size_t *used, size_t *limit) { + (void)interp; (void)used; (void)limit; + return -1; +} +static void py_mem_limit_forget(PyInterpreterState *interp) { (void)interp; } +#endif + +#endif /* HAVE_OWNGIL */ diff --git a/c_src/py_nif.c b/c_src/py_nif.c index 0b33bca..f7443ba 100644 --- a/c_src/py_nif.c +++ b/c_src/py_nif.c @@ -261,6 +261,7 @@ static int is_inline_schedule_marker(PyObject *obj); * ============================================================================ */ #include "py_util.c" +#include "py_mem_limit.c" #include "py_convert.c" #include "py_exec.c" #include "py_logging.c" @@ -351,6 +352,13 @@ static void context_destructor(ErlNifEnv *env, void *obj) { /* Close callback pipes if open */ close_pipe_pair(ctx->callback_pipe); + /* Refcount is zero here, so no interrupt can be in flight. Contexts that + * leaked an unresponsive thread keep a reference and never reach this. */ + if (ctx->interrupt_mutex_init) { + pthread_mutex_destroy(&ctx->interrupt_mutex); + ctx->interrupt_mutex_init = false; + } + /* Skip if already destroyed by nif_context_destroy */ if (ctx->destroyed) { return; @@ -1011,6 +1019,19 @@ static ERL_NIF_TERM nif_py_init(ErlNifEnv *env, int argc, const ERL_NIF_TERM arg } #endif + /* Per-context memory caps hook the obmalloc arena allocator, which has to + * be installed before Python allocates anything. Opt-in: the default path + * leaves the allocator untouched. */ + if (argc > 0 && enif_is_map(env, argv[0])) { + ERL_NIF_TERM mem_value; + if (enif_get_map_value(env, argv[0], + enif_make_atom(env, "enable_memory_limits"), + &mem_value) && + enif_is_identical(mem_value, ATOM_TRUE)) { + py_mem_limit_install(); + } + } + /* Initialize Python with thread support. * If Python is already initialized (e.g., after app restart without * calling Py_Finalize), skip initialization to avoid corruption. */ @@ -3170,10 +3191,14 @@ static void *worker_context_thread_main(void *arg) { ctx->response_ok = false; ctx->response_term = 0; - /* Acquire GIL and process the request */ + /* Acquire GIL and process the request. + * exec_enter before / exec_leave after the GIL (see the locking + * invariant on py_context::interrupt_mutex). */ + py_context_exec_enter(ctx); gstate = PyGILState_Ensure(); owngil_execute_request(ctx); /* Reuse execute functions */ PyGILState_Release(gstate); + py_context_exec_leave(ctx); /* Copy response to request struct */ req->result_env = enif_alloc_env(); @@ -3708,10 +3733,14 @@ static void *owngil_context_thread_main(void *arg) { ctx->response_ok = false; ctx->response_term = 0; - /* Acquire our GIL and process the request */ + /* Acquire our GIL and process the request. + * exec_enter before / exec_leave after the GIL (see the locking + * invariant on py_context::interrupt_mutex). */ + py_context_exec_enter(ctx); PyEval_RestoreThread(ctx->own_gil_tstate); owngil_execute_request(ctx); PyEval_SaveThread(); + py_context_exec_leave(ctx); /* Copy response to request struct */ req->result_env = enif_alloc_env(); @@ -3751,10 +3780,15 @@ static void *owngil_context_thread_main(void *arg) { ctx->module_cache = NULL; /* End interpreter - this releases our GIL and cleans up */ + PyInterpreterState *ended_interp = ctx->own_gil_interp; Py_EndInterpreter(ctx->own_gil_tstate); ctx->own_gil_tstate = NULL; ctx->own_gil_interp = NULL; + /* Release the memory accounting slot only after teardown, so the arenas + * freed by Py_EndInterpreter are still attributed to this interpreter. */ + py_mem_limit_forget(ended_interp); + /* Don't call PyGILState_Release(gstate) here! * After Py_NewInterpreterFromConfig switched us to the OWN_GIL interpreter, * the original gstate is no longer valid. Py_EndInterpreter handles cleanup. */ @@ -4718,6 +4752,16 @@ static ERL_NIF_TERM nif_context_create(ErlNifEnv *env, int argc, const ERL_NIF_T ctx->module_cache = NULL; ctx->uses_worker_thread = false; + /* Interrupt support */ + ctx->interrupt_mutex_init = (pthread_mutex_init(&ctx->interrupt_mutex, NULL) == 0); + if (!ctx->interrupt_mutex_init) { + enif_release_resource(ctx); + return make_error(env, "mutex_init_failed"); + } + atomic_store(&ctx->exec_in_flight, false); + atomic_store(&ctx->interrupt_pending, false); + ctx->exec_thread_id = 0; + /* Create callback pipe for blocking callback responses */ if (pipe(ctx->callback_pipe) < 0) { enif_release_resource(ctx); @@ -4731,6 +4775,8 @@ static ERL_NIF_TERM nif_context_create(ErlNifEnv *env, int argc, const ERL_NIF_T #ifdef HAVE_SUBINTERPRETERS ctx->uses_own_gil = false; + ctx->own_gil_tstate = NULL; + ctx->own_gil_interp = NULL; if (use_owngil) { /* OWN_GIL mode: create dedicated pthread with OWN_GIL subinterpreter */ @@ -4764,6 +4810,155 @@ static ERL_NIF_TERM nif_context_create(ErlNifEnv *env, int argc, const ERL_NIF_T return enif_make_tuple3(env, ATOM_OK, ref, enif_make_uint(env, ctx->interp_id)); } +/** + * @brief Set a memory cap for a context + * + * nif_context_set_memory_limit(ContextRef, Bytes) -> ok | {error, Reason} + * + * Bytes = 0 removes the cap. Requires owngil mode: accounting is per + * interpreter, and every worker-mode context shares the main interpreter. + * Requires the runtime to have been started with enable_memory_limits. + */ +static ERL_NIF_TERM nif_context_set_memory_limit(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]) { + (void)argc; + py_context_t *ctx; + ErlNifUInt64 limit; + + if (!enif_get_resource(env, argv[0], PY_CONTEXT_RESOURCE_TYPE, (void **)&ctx)) { + return make_error(env, "invalid_context"); + } + if (!enif_get_uint64(env, argv[1], &limit)) { + return make_error(env, "invalid_limit"); + } + if (atomic_load(&ctx->destroyed)) { + return make_error(env, "context_destroyed"); + } + if (!py_mem_limit_enabled()) { + return make_error(env, "memory_limits_disabled"); + } + +#ifdef HAVE_SUBINTERPRETERS + if (!ctx->uses_own_gil || ctx->own_gil_interp == NULL) { + return make_error(env, "memory_limit_requires_owngil"); + } + if (py_mem_limit_set(ctx->own_gil_interp, (size_t)limit) != 0) { + return make_error(env, "memory_limit_unavailable"); + } + return ATOM_OK; +#else + return make_error(env, "memory_limit_requires_owngil"); +#endif +} + +/** + * @brief Report accounted memory usage for a context + * + * nif_context_memory_usage(ContextRef) -> {ok, Used, Limit} | {error, Reason} + * + * Used counts obmalloc arena bytes for this context's interpreter. It does + * not include allocations that bypass obmalloc (over 512 bytes). + */ +static ERL_NIF_TERM nif_context_memory_usage(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]) { + (void)argc; + py_context_t *ctx; + + if (!enif_get_resource(env, argv[0], PY_CONTEXT_RESOURCE_TYPE, (void **)&ctx)) { + return make_error(env, "invalid_context"); + } + if (atomic_load(&ctx->destroyed)) { + return make_error(env, "context_destroyed"); + } + if (!py_mem_limit_enabled()) { + return make_error(env, "memory_limits_disabled"); + } + +#ifdef HAVE_SUBINTERPRETERS + if (!ctx->uses_own_gil || ctx->own_gil_interp == NULL) { + return make_error(env, "memory_limit_requires_owngil"); + } + + size_t used = 0, limit = 0; + if (py_mem_limit_get(ctx->own_gil_interp, &used, &limit) != 0) { + return make_error(env, "not_tracked"); + } + return enif_make_tuple3(env, ATOM_OK, + enif_make_uint64(env, (ErlNifUInt64)used), + enif_make_uint64(env, (ErlNifUInt64)limit)); +#else + return make_error(env, "memory_limit_requires_owngil"); +#endif +} + +/** + * @brief Interrupt Python code currently running in a context + * + * nif_context_interrupt(ContextRef) -> ok | not_running + * + * Raises KeyboardInterrupt asynchronously in whichever thread is executing + * this context. KeyboardInterrupt is a BaseException, so ordinary + * `except Exception:` handlers in user code do not swallow it, and it is a + * static builtin valid in every subinterpreter. + * + * CPython delivers an async exception at the next bytecode boundary, so a + * thread blocked inside a C call (time.sleep, a numpy kernel, a socket read) + * is not interrupted until that call returns. + * + * Dirty IO-bound: this blocks on the GIL, which the running thread only + * yields at switch-interval boundaries. + */ +static ERL_NIF_TERM nif_context_interrupt(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { + (void)argc; + py_context_t *ctx; + + if (!enif_get_resource(env, argv[0], PY_CONTEXT_RESOURCE_TYPE, (void **)&ctx)) { + return make_error(env, "invalid_context"); + } + + if (!runtime_is_running() || atomic_load(&ctx->destroyed)) { + return enif_make_atom(env, "not_running"); + } + + /* Held across the GIL acquisition below (see the locking invariant on + * py_context::interrupt_mutex). This thread must not already hold the GIL. */ + pthread_mutex_lock(&ctx->interrupt_mutex); + + if (!atomic_load(&ctx->exec_in_flight) || atomic_load(&ctx->destroyed)) { + pthread_mutex_unlock(&ctx->interrupt_mutex); + return enif_make_atom(env, "not_running"); + } + + unsigned long tid = ctx->exec_thread_id; + bool injected = false; + +#ifdef HAVE_SUBINTERPRETERS + if (ctx->uses_own_gil && ctx->own_gil_interp != NULL) { + /* Attach to this context's subinterpreter. The thread state is created + * on THIS thread so the 3.12+ tstate/thread binding assertions hold. */ + PyThreadState *tstate = PyThreadState_New(ctx->own_gil_interp); + if (tstate != NULL) { + PyEval_RestoreThread(tstate); + injected = (PyThreadState_SetAsyncExc(tid, PyExc_KeyboardInterrupt) > 0); + PyThreadState_Clear(tstate); + PyThreadState_DeleteCurrent(); /* detaches and drops the OWN_GIL */ + } + } else +#endif + { + PyGILState_STATE gstate = PyGILState_Ensure(); + injected = (PyThreadState_SetAsyncExc(tid, PyExc_KeyboardInterrupt) > 0); + PyGILState_Release(gstate); + } + + if (injected) { + atomic_store(&ctx->interrupt_pending, true); + } + pthread_mutex_unlock(&ctx->interrupt_mutex); + + return injected ? ATOM_OK : enif_make_atom(env, "not_running"); +} + /** * @brief Destroy a Python context * @@ -7987,6 +8182,9 @@ static ErlNifFunc nif_funcs[] = { /* Process-per-context API (no mutex) */ {"context_create", 1, nif_context_create, 0}, {"context_destroy", 1, nif_context_destroy, 0}, + {"context_interrupt", 1, nif_context_interrupt, ERL_NIF_DIRTY_JOB_IO_BOUND}, + {"context_set_memory_limit", 2, nif_context_set_memory_limit, 0}, + {"context_memory_usage", 1, nif_context_memory_usage, 0}, {"context_call", 5, nif_context_call, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"context_eval", 3, nif_context_eval, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"context_exec", 2, nif_context_exec, ERL_NIF_DIRTY_JOB_CPU_BOUND}, diff --git a/c_src/py_nif.h b/c_src/py_nif.h index 8977b8f..002418a 100644 --- a/c_src/py_nif.h +++ b/c_src/py_nif.h @@ -1025,6 +1025,31 @@ struct py_context { /** @brief Module cache (Dict: module_name -> PyModule) */ PyObject *module_cache; + + /* ========== Interrupt support ========== */ + + /** + * @brief Protects the exec_* fields below and serialises interrupt delivery + * + * LOCKING INVARIANT: this mutex is only ever taken by a thread that does + * NOT hold the GIL. nif_context_interrupt holds it while blocking on the + * GIL, so py_context_exec_enter() must run BEFORE acquiring the GIL and + * py_context_exec_leave() AFTER releasing it. Taking it with the GIL held + * deadlocks against an in-flight interrupt. + */ + pthread_mutex_t interrupt_mutex; + + /** @brief True once interrupt_mutex has been initialized (destructor guard) */ + bool interrupt_mutex_init; + + /** @brief True while a request is executing Python code in this context */ + _Atomic bool exec_in_flight; + + /** @brief PyThread_get_thread_ident() of the thread executing the request */ + unsigned long exec_thread_id; + + /** @brief True when an async exception was injected and not yet consumed */ + _Atomic bool interrupt_pending; }; /* ============================================================================ @@ -1127,6 +1152,67 @@ typedef struct { * py_context_release(&guard); * @endcode */ +/** + * @brief Mark this thread as the one executing Python for @p ctx + * + * Records the calling thread's Python thread identifier so + * nif_context_interrupt() knows where to deliver an async exception. + * + * @note MUST be called BEFORE acquiring the GIL. See the locking invariant + * on py_context::interrupt_mutex. + */ +static inline void py_context_exec_enter(py_context_t *ctx) { + if (ctx == NULL) { + return; + } + pthread_mutex_lock(&ctx->interrupt_mutex); + ctx->exec_thread_id = PyThread_get_thread_ident(); + atomic_store(&ctx->exec_in_flight, true); + pthread_mutex_unlock(&ctx->interrupt_mutex); +} + +/** + * @brief Clear the executing-thread marker for @p ctx + * + * If an interrupt was injected but never consumed (it landed after the code + * finished, or Python swallowed it), the pending async exception is cleared + * here so it cannot leak into the next request on this context. + * + * @note MUST be called AFTER releasing the GIL. See the locking invariant + * on py_context::interrupt_mutex. + */ +static inline void py_context_exec_leave(py_context_t *ctx) { + if (ctx == NULL) { + return; + } + pthread_mutex_lock(&ctx->interrupt_mutex); + atomic_store(&ctx->exec_in_flight, false); + + if (atomic_load(&ctx->interrupt_pending)) { + /* Drop an async exception that was never raised. Re-acquiring the GIL + * here is safe: we do not hold it, per the invariant above. */ + unsigned long tid = ctx->exec_thread_id; +#ifdef HAVE_SUBINTERPRETERS + if (ctx->uses_own_gil && ctx->own_gil_interp != NULL) { + PyThreadState *tstate = PyThreadState_New(ctx->own_gil_interp); + if (tstate != NULL) { + PyEval_RestoreThread(tstate); + PyThreadState_SetAsyncExc(tid, NULL); + PyThreadState_Clear(tstate); + PyThreadState_DeleteCurrent(); + } + } else +#endif + { + PyGILState_STATE gstate = PyGILState_Ensure(); + PyThreadState_SetAsyncExc(tid, NULL); + PyGILState_Release(gstate); + } + atomic_store(&ctx->interrupt_pending, false); + } + pthread_mutex_unlock(&ctx->interrupt_mutex); +} + static inline py_context_guard_t py_context_acquire(py_context_t *ctx) { py_context_guard_t guard = { .ctx = ctx, @@ -1140,6 +1226,10 @@ static inline py_context_guard_t py_context_acquire(py_context_t *ctx) { return guard; } + /* Register as the executing thread BEFORE taking the GIL, so an interrupt + * blocked on the GIL cannot deadlock against us. */ + py_context_exec_enter(ctx); + /* Acquire the GIL first */ guard.gstate = PyGILState_Ensure(); @@ -1167,6 +1257,9 @@ static inline void py_context_release(py_context_guard_t *guard) { /* Release the GIL */ PyGILState_Release(guard->gstate); guard->acquired = false; + + /* Clear the executing-thread marker AFTER dropping the GIL. */ + py_context_exec_leave(guard->ctx); } /** diff --git a/docs/interrupts.md b/docs/interrupts.md new file mode 100644 index 0000000..63ecb65 --- /dev/null +++ b/docs/interrupts.md @@ -0,0 +1,77 @@ +# Interrupting Python Code + +This guide covers stopping Python code that is already running: a call that +overruns its timeout, or work you want to cancel early. Without this, an +Erlang-side timeout only stops *waiting*: the Python thread keeps running and +the context stays busy. You need it whenever you run code you do not fully +control, such as user-supplied scripts. + +## Timeouts interrupt automatically + +Pass a timeout to any `py_context` call. When it expires the Python code is +interrupted and the context is free again: + +```erlang +{ok, Ctx} = py_context:new(#{mode => owngil}), +{error, timeout} = py_context:eval(Ctx, <<"while True: pass">>, #{}, 500), +%% The context is immediately reusable +{ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000). +``` + +The caller still gets `{error, timeout}`. The Python side sees a +`KeyboardInterrupt` at the point it was executing. + +## Cancelling explicitly + +To stop work before its timeout, or work started with `infinity`, call +`py:interrupt/1` from any process: + +```erlang +Ctx = py:context(), +spawn(fun() -> py_context:eval(Ctx, <<"while True: pass">>, #{}, infinity) end), +%% ... later, from anywhere ... +ok = py:interrupt(Ctx). +``` + +The interrupted call returns `{error, interrupted}`. `py:interrupt/1` returns +`not_running` if the context is idle. + +## Catching it in Python + +The interrupt arrives as `KeyboardInterrupt`, which derives from +`BaseException`, so ordinary handlers do not swallow it: + +```python +try: + do_work() +except Exception: # does NOT catch the interrupt + log_failure() +``` + +Catch it explicitly to clean up, then re-raise: + +```python +try: + do_work() +except KeyboardInterrupt: + release_resources() + raise +``` + +Code that catches `BaseException` and continues will keep running. Destroy the +context to deal with that: + +```erlang +ok = py_context:destroy(Ctx). +``` + +## Limits + +- CPython delivers an async exception at the next bytecode boundary. Code + blocked inside a C call (`time.sleep`, a numpy kernel, a socket read) is + not interrupted until that call returns. The call still times out on the + Erlang side; the context becomes usable once the C call finishes. +- An interrupt targets the context, not an individual request. Interrupting a + context that just finished one call and started another stops the new one. +- `py:call/3,4` and `py:eval/1,2` use `infinity` by default. Pass an explicit + timeout, or use `py:interrupt/1`, if you need a bound. diff --git a/docs/memory.md b/docs/memory.md index baaa07f..92f2bb6 100644 --- a/docs/memory.md +++ b/docs/memory.md @@ -146,6 +146,64 @@ Default thresholds are `{700, 10, 10}`: Objects with circular references and `__del__` methods may be uncollectable. Monitor the `uncollectable` count in gc_stats. +## Per-Context Memory Caps + +Cap how much memory one context may allocate, so a runaway script cannot take +the whole node down. Exceeding the cap raises `MemoryError` inside that +context; every other context is unaffected. + +Caps require `owngil` mode and must be enabled before Python starts, because +they hook the allocator: + +```erlang +%% sys.config +[{erlang_python, [{enable_memory_limits, true}]}]. +``` + +Then set a cap per context: + +```erlang +{ok, Ctx} = py_context:new(#{mode => owngil, memory_limit => 256 * 1024 * 1024}), +{error, {'MemoryError', _}} = + py_context:exec(Ctx, <<"_hog = [[] for _ in range(10000000)]">>), +%% The context stays usable +{ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000). +``` + +Read current usage: + +```erlang +Ref = py_context:get_nif_ref(Ctx), +{ok, UsedBytes, LimitBytes} = py_nif:context_memory_usage(Ref). +``` + +Remove a cap with `py_nif:context_set_memory_limit(Ref, 0)`. Accounting keeps +running, so `context_memory_usage/1` still reports usage. + +### What is counted + +Accounting comes from obmalloc arena traffic, which covers ordinary Python +objects: lists, dicts, tuples, instances, small strings. + +Not counted: + +- Allocations over 512 bytes, which bypass obmalloc: large `bytes`, numpy + buffers, and extensions with their own allocator. +- Memory held by the interpreter before the cap was set. + +Other behaviour to expect: + +- Granularity is one 1 MB arena. +- Enforcement raises `MemoryError` at the next bytecode boundary, so usage can + overshoot the cap slightly before the code stops. +- The cap re-arms once usage drops back below it. +- `worker` mode contexts share the main interpreter, so a per-context cap has + no meaning there. `py_context:new(#{mode => worker, memory_limit => N})` + returns `{error, memory_limit_requires_owngil}` instead of silently applying + a process-wide cap. + +Treat a cap as a guard against runaway object graphs, not as a hard RSS bound. + ## Troubleshooting ### High Memory Usage diff --git a/docs/streaming.md b/docs/streaming.md index 94c2075..a53bb90 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -65,19 +65,35 @@ ok = py:stream_cancel(Ref). ### Async Generators -`stream_start` supports both sync and async generators: +`stream_start` accepts sync and async generators. An async generator is driven +on a private event loop, one value at a time, and delivers the same +`{py_stream, Ref, ...}` events. -```erlang -%% Async generator (e.g., streaming from an async API) -ok = py:exec(<<" -async def async_gen(): - for i in range(5): +Put the generator in an importable module: + +```python +# my_module.py +import asyncio + +async def async_gen(n): + for i in range(n): await asyncio.sleep(0.1) yield i -">>), -{ok, Ref} = py:stream_start('__main__', async_gen, []). ``` +```erlang +{ok, Ref} = py:stream_start(my_module, async_gen, [5]), +receive_loop(Ref). +``` + +Notes: + +- Delivery of each value blocks that private loop, so other coroutines on it do + not progress between yields. Use it for sequential streams. +- `py:stream_cancel/1` works the same for async generators. +- The batch helpers below (`py:stream/4` with kwargs, `py:stream_eval/1,2`) + wrap the call in `list()` and accept sync generators only. + ## Batch Streaming (Collecting All Values) For simpler use cases where you want all values at once: diff --git a/rebar.config b/rebar.config index 2062cff..a93fbe8 100644 --- a/rebar.config +++ b/rebar.config @@ -28,6 +28,10 @@ {apps, [erlang_python]} ]}. +{ct_opts, [ + {sys_config, ["test/test.config"]} +]}. + {project_plugins, [rebar3_ex_doc]}. {hex, [ @@ -51,6 +55,7 @@ <<"docs/channel.md">>, <<"docs/buffer.md">>, <<"docs/streaming.md">>, + <<"docs/interrupts.md">>, <<"docs/memory.md">>, <<"docs/shared-dict.md">>, <<"docs/logging.md">>, @@ -79,6 +84,7 @@ <<"docs/channel.md">>, <<"docs/buffer.md">>, <<"docs/streaming.md">>, + <<"docs/interrupts.md">>, <<"docs/memory.md">>, <<"docs/shared-dict.md">>, <<"docs/logging.md">> diff --git a/src/erlang_python_sup.erl b/src/erlang_python_sup.erl index d6471be..e9e3ebd 100644 --- a/src/erlang_python_sup.erl +++ b/src/erlang_python_sup.erl @@ -36,8 +36,13 @@ init([]) -> erlang:system_info(schedulers)), ContextMode = application:get_env(erlang_python, context_mode, worker), + %% Per-context memory caps hook the obmalloc arena allocator, which has to + %% be installed before Python starts, hence an app env rather than a + %% per-context option. + MemoryLimits = application:get_env(erlang_python, enable_memory_limits, false), + %% Initialize Python runtime first - ok = py_nif:init(#{}), + ok = py_nif:init(#{enable_memory_limits => MemoryLimits}), %% Initialize the semaphore ETS table for rate limiting ok = py_semaphore:init(), @@ -48,6 +53,10 @@ init([]) -> %% Initialize shared state ETS table (owned by supervisor for resilience) ok = py_state:init_tab(), + %% Initialize the context pid -> NIF reference table used by + %% py_context:interrupt/1 (owned by supervisor, must outlive contexts) + ok = py_context:init_ref_tab(), + %% Initialize import/path registry and load config (before contexts start) ok = py_import:init(), diff --git a/src/py.erl b/src/py.erl index 31d88ef..0d16c9a 100644 --- a/src/py.erl +++ b/src/py.erl @@ -111,6 +111,7 @@ %% Process-per-context API (new architecture) context/0, context/1, + interrupt/1, start_contexts/0, start_contexts/1, stop_contexts/0, @@ -490,7 +491,9 @@ stream_eval(Code, Locals) -> %% - `{py_stream, Ref, done}' - Stream completed %% - `{py_stream, Ref, {error, Reason}}' - Stream error %% -%% Supports both sync generators and async generators (coroutines). +%% Accepts sync generators and async generators. An async generator is driven +%% on a private event loop, one value at a time; delivering a value blocks that +%% loop, so other coroutines on it do not progress between yields. %% %% Example: %% ``` @@ -558,13 +561,27 @@ stream_run_python(ModuleBin0, FuncBin0, RefHash) -> <<" _mod = __import__('">>, ModuleBin, <<"')\n">>, <<" _fn = getattr(_mod, '">>, FuncBin, <<"')\n">>, <<" _gen = _fn(*_args) if _args else _fn()\n">>, - <<" for _val in _gen:\n">>, - <<" if erlang.call('_py_stream_cancelled', _rh):\n">>, - <<" erlang.call('_py_stream_send', _rh, 'error', 'cancelled')\n">>, - <<" break\n">>, - <<" erlang.call('_py_stream_send', _rh, 'data', _val)\n">>, + %% Async generators are driven on a private event loop. erlang.call is + %% a blocking pipe read, so it stalls that loop between yields, which + %% is fine for a sequential stream. + <<" if hasattr(_gen, '__anext__'):\n">>, + <<" import asyncio\n">>, + <<" async def _drive():\n">>, + <<" async for _val in _gen:\n">>, + <<" if erlang.call('_py_stream_cancelled', _rh):\n">>, + <<" erlang.call('_py_stream_send', _rh, 'error', 'cancelled')\n">>, + <<" return\n">>, + <<" erlang.call('_py_stream_send', _rh, 'data', _val)\n">>, + <<" erlang.call('_py_stream_send', _rh, 'done', None)\n">>, + <<" asyncio.run(_drive())\n">>, <<" else:\n">>, - <<" erlang.call('_py_stream_send', _rh, 'done', None)\n">>, + <<" for _val in _gen:\n">>, + <<" if erlang.call('_py_stream_cancelled', _rh):\n">>, + <<" erlang.call('_py_stream_send', _rh, 'error', 'cancelled')\n">>, + <<" break\n">>, + <<" erlang.call('_py_stream_send', _rh, 'data', _val)\n">>, + <<" else:\n">>, + <<" erlang.call('_py_stream_send', _rh, 'done', None)\n">>, <<"except Exception as _e:\n">>, <<" erlang.call('_py_stream_send', _rh, 'error', str(_e))\n">>, <<"finally:\n">>, @@ -1453,6 +1470,25 @@ context() -> context(N) -> py_context_router:get_context(N). +%% @doc Interrupt Python code currently running in a context. +%% +%% Raises KeyboardInterrupt in the thread executing the context; the in-flight +%% call returns `{error, interrupted}'. Callable from any process, including +%% while the context process is blocked in a NIF. +%% +%% Calls made with a timeout interrupt themselves when that timeout expires, +%% so this is only needed to cancel work early or to stop a call made with +%% `infinity'. +%% +%% Code blocked inside a C call (`time.sleep', a numpy kernel, a socket read) +%% is only interrupted once that call returns. +%% +%% @param Ctx Context pid +%% @returns ok | not_running +-spec interrupt(pid()) -> ok | not_running. +interrupt(Ctx) when is_pid(Ctx) -> + py_context:interrupt(Ctx). + %%% ============================================================================ %%% py_ref API (Python object references with auto-routing) %%% diff --git a/src/py_context.erl b/src/py_context.erl index a06ab78..b79da74 100644 --- a/src/py_context.erl +++ b/src/py_context.erl @@ -35,6 +35,7 @@ -export([ start_link/2, + start_link/3, new/1, stop/1, destroy/1, @@ -53,15 +54,25 @@ get_interp_id/1, is_subinterp/1, create_local_env/1, - get_nif_ref/1 + get_nif_ref/1, + interrupt/1 ]). %% Internal exports --export([init/3]). +-export([init/3, init/4, init_ref_tab/0]). %% Exported for py_reactor_context -export([extend_erlang_module_in_context/1]). +%% Maps context pid -> NIF context reference. Read by interrupt/1, which must +%% reach the NIF reference while the context process is blocked in a NIF and +%% therefore cannot answer get_nif_ref/1. +-define(REF_TAB, py_context_refs). + +%% How long to wait for an interrupted call to unwind and reply, so the late +%% reply is drained instead of being left in the caller's mailbox. +-define(INTERRUPT_GRACE_MS, 1000). + -type context_mode() :: worker | owngil. -type context() :: pid(). @@ -93,8 +104,21 @@ %% @returns {ok, Pid} | {error, Reason} -spec start_link(pos_integer(), context_mode()) -> {ok, pid()} | {error, term()}. start_link(Id, Mode) -> + start_link(Id, Mode, #{}). + +%% @doc Start a new py_context process with options. +%% +%% See new/1 for the recognised options. +%% +%% @param Id Unique identifier for this context +%% @param Mode Context mode +%% @param Opts Options map +%% @returns {ok, Pid} | {error, Reason} +-spec start_link(pos_integer(), context_mode(), map()) -> + {ok, pid()} | {error, term()}. +start_link(Id, Mode, Opts) when is_map(Opts) -> Parent = self(), - Pid = spawn_link(fun() -> init(Parent, Id, Mode) end), + Pid = spawn_link(fun() -> init(Parent, Id, Mode, Opts) end), receive {Pid, started} -> {ok, Pid}; @@ -126,6 +150,10 @@ stop(Ctx) when is_pid(Ctx) -> %% %% Options: %% - `mode' - Context mode (worker | owngil), default: worker +%% - `memory_limit' - Cap in bytes on memory allocated by this context. +%% Requires `mode => owngil' and the runtime started with +%% `enable_memory_limits'; see py_nif:context_set_memory_limit/2 for what +%% is counted. %% %% @param Opts Options map %% @returns {ok, Pid} | {error, Reason} @@ -133,7 +161,7 @@ stop(Ctx) when is_pid(Ctx) -> new(Opts) when is_map(Opts) -> Mode = maps:get(mode, Opts, worker), Id = erlang:unique_integer([positive]), - start_link(Id, Mode). + start_link(Id, Mode, Opts). %% @doc Alias for stop/1 for API consistency. -spec destroy(context()) -> ok. @@ -175,16 +203,7 @@ call(Ctx, Module, Func, Args, Kwargs, Timeout) when is_pid(Ctx) -> ModuleBin = to_binary(Module), FuncBin = to_binary(Func), Ctx ! {call, self(), MRef, ModuleBin, FuncBin, Args, Kwargs}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after Timeout -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, Timeout). %% @doc Call a Python function with a process-local environment. %% @@ -203,16 +222,7 @@ call(Ctx, Module, Func, Args, Kwargs, Timeout, EnvRef) when is_pid(Ctx), is_refe ModuleBin = to_binary(Module), FuncBin = to_binary(Func), Ctx ! {call, self(), MRef, ModuleBin, FuncBin, Args, Kwargs, EnvRef}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after Timeout -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, Timeout). %% @doc Evaluate a Python expression with empty locals. %% @@ -244,16 +254,7 @@ eval(Ctx, Code, Locals, Timeout) when is_pid(Ctx) -> MRef = erlang:monitor(process, Ctx), CodeBin = to_binary(Code), Ctx ! {eval, self(), MRef, CodeBin, Locals}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after Timeout -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, Timeout). %% @doc Evaluate a Python expression with a process-local environment. %% @@ -269,16 +270,7 @@ eval(Ctx, Code, Locals, Timeout, EnvRef) when is_pid(Ctx), is_reference(EnvRef) MRef = erlang:monitor(process, Ctx), CodeBin = to_binary(Code), Ctx ! {eval, self(), MRef, CodeBin, Locals, EnvRef}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after Timeout -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, Timeout). %% @doc Execute Python statements. %% @@ -290,16 +282,7 @@ exec(Ctx, Code) when is_pid(Ctx) -> MRef = erlang:monitor(process, Ctx), CodeBin = to_binary(Code), Ctx ! {exec, self(), MRef, CodeBin}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after infinity -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, infinity). %% @doc Execute Python statements with a process-local environment. %% @@ -312,16 +295,7 @@ exec(Ctx, Code, EnvRef) when is_pid(Ctx), is_reference(EnvRef) -> MRef = erlang:monitor(process, Ctx), CodeBin = to_binary(Code), Ctx ! {exec, self(), MRef, CodeBin, EnvRef}, - receive - {MRef, Result} -> - erlang:demonitor(MRef, [flush]), - Result; - {'DOWN', MRef, process, Ctx, Reason} -> - {error, {context_died, Reason}} - after infinity -> - erlang:demonitor(MRef, [flush]), - {error, timeout} - end. + await_reply(Ctx, MRef, infinity). %% @doc Call a method on a Python object reference. -spec call_method(context(), reference(), atom() | binary(), list()) -> @@ -408,45 +382,168 @@ get_nif_ref(Ctx) when is_pid(Ctx) -> error({context_died, Reason}) end. +%% @doc Interrupt Python code currently running in this context. +%% +%% Raises KeyboardInterrupt in the thread executing the context; the in-flight +%% call returns `{error, interrupted}'. Callable from any process, including +%% while the context process is blocked in a NIF. +%% +%% Returns `not_running' if the context is idle, unknown, or the exception +%% could not be delivered. Code blocked in a C call (`time.sleep', a numpy +%% kernel, a socket read) is only interrupted once that call returns. +%% +%% @param Ctx Context process +%% @returns ok | not_running +-spec interrupt(context()) -> ok | not_running. +interrupt(Ctx) when is_pid(Ctx) -> + case lookup_nif_ref(Ctx) of + {ok, Ref} -> + try py_nif:context_interrupt(Ref) of + ok -> ok; + _ -> not_running + catch + _:_ -> not_running + end; + error -> + not_running + end. + +%% @private Create the pid -> NIF reference table. Called by the supervisor +%% before any context starts. +-spec init_ref_tab() -> ok. +init_ref_tab() -> + case ets:whereis(?REF_TAB) of + undefined -> + ?REF_TAB = ets:new(?REF_TAB, [ + named_table, public, set, {read_concurrency, true} + ]), + ok; + _ -> + ok + end. + %% ============================================================================ %% Internal functions %% ============================================================================ +%% @private Wait for a context reply, interrupting the running Python code if +%% the timeout expires. +%% +%% On timeout the Python side is interrupted and we wait a bounded grace period +%% for the unwinding call to reply, so the late reply is consumed here rather +%% than left behind in the caller's mailbox. The result is still +%% `{error, timeout}': the caller asked to stop waiting. +await_reply(Ctx, MRef, Timeout) -> + receive + {MRef, Result} -> + erlang:demonitor(MRef, [flush]), + Result; + {'DOWN', MRef, process, Ctx, Reason} -> + {error, {context_died, Reason}} + after Timeout -> + _ = interrupt(Ctx), + receive + {MRef, _Late} -> + erlang:demonitor(MRef, [flush]); + {'DOWN', MRef, process, Ctx, _} -> + ok + after ?INTERRUPT_GRACE_MS -> + erlang:demonitor(MRef, [flush]) + end, + {error, timeout} + end. + +%% @private +register_nif_ref(Ref) -> + try + true = ets:insert(?REF_TAB, {self(), Ref}), + ok + catch + error:badarg -> ok %% table not created (library used without the app) + end. + +%% @private +unregister_nif_ref() -> + try + true = ets:delete(?REF_TAB, self()), + ok + catch + error:badarg -> ok + end. + +%% @private +lookup_nif_ref(Ctx) -> + try ets:lookup(?REF_TAB, Ctx) of + [{Ctx, Ref}] -> {ok, Ref}; + [] -> error + catch + error:badarg -> error + end. + %% @private init(Parent, Id, Mode) -> + init(Parent, Id, Mode, #{}). + +%% @private +init(Parent, Id, Mode, Opts) -> process_flag(trap_exit, true), case create_context(Mode) of {ok, Ref, InterpId} -> - %% Apply all registered imports and paths to this interpreter - apply_registered_imports(Ref), - apply_registered_paths(Ref), - %% Apply preload code (populates globals for process-local envs) - apply_preload(Ref), - %% For subinterpreters, create a dedicated event worker - EventState = setup_event_worker(Ref, InterpId), - %% For thread-model subinterpreters, spawn a dedicated callback handler - %% because the main context process will be blocked in the NIF - CallbackHandler = case maps:get(mode, EventState, normal) of - thread_model -> - Handler = spawn_callback_handler(Ref), - ok = py_nif:context_set_callback_handler(Ref, Handler), - Handler; - _ -> - undefined - end, - Parent ! {self(), started}, - State = #state{ - ref = Ref, - id = Id, - interp_id = InterpId, - event_state = EventState, - callback_handler = CallbackHandler - }, - loop(State); + %% Publish the NIF reference so interrupt/1 can reach it while + %% this process is blocked in a NIF + register_nif_ref(Ref), + case apply_memory_limit(Ref, Opts) of + ok -> + init_started(Parent, Id, Ref, InterpId); + {error, LimitError} -> + unregister_nif_ref(), + try py_nif:context_destroy(Ref) catch _:_ -> ok end, + Parent ! {self(), {error, LimitError}} + end; {error, Reason} -> Parent ! {self(), {error, Reason}} end. +%% @private +apply_memory_limit(Ref, Opts) -> + case maps:get(memory_limit, Opts, undefined) of + undefined -> + ok; + Bytes when is_integer(Bytes), Bytes >= 0 -> + py_nif:context_set_memory_limit(Ref, Bytes); + Other -> + {error, {invalid_memory_limit, Other}} + end. + +%% @private +init_started(Parent, Id, Ref, InterpId) -> + %% Apply all registered imports and paths to this interpreter + apply_registered_imports(Ref), + apply_registered_paths(Ref), + %% Apply preload code (populates globals for process-local envs) + apply_preload(Ref), + %% For subinterpreters, create a dedicated event worker + EventState = setup_event_worker(Ref, InterpId), + %% For thread-model subinterpreters, spawn a dedicated callback handler + %% because the main context process will be blocked in the NIF + CallbackHandler = case maps:get(mode, EventState, normal) of + thread_model -> + Handler = spawn_callback_handler(Ref), + ok = py_nif:context_set_callback_handler(Ref, Handler), + Handler; + _ -> + undefined + end, + Parent ! {self(), started}, + State = #state{ + ref = Ref, + id = Id, + interp_id = InterpId, + event_state = EventState, + callback_handler = CallbackHandler + }, + loop(State). + %% @private Create event worker for subinterpreter contexts setup_event_worker(Ref, InterpId) -> case py_nif:context_get_event_loop(Ref) of @@ -641,6 +738,7 @@ loop(#state{ref = Ref, interp_id = InterpId} = State) -> %% @private Clean up resources on termination terminate(_Reason, #state{ref = Ref, event_state = EventState, callback_handler = CallbackHandler}) -> + unregister_nif_ref(), %% Stop the callback handler if it exists case CallbackHandler of Pid when is_pid(Pid) -> @@ -699,9 +797,8 @@ handle_blocking_callback(Ref, FuncName, Args) -> %% Execute the registered function Response = case py_callback:execute(FuncName, ArgsList) of {ok, Result} -> - %% Format: status_byte (0=ok) + python_repr - ResultStr = term_to_python_repr(Result), - <<0, ResultStr/binary>>; + %% Format: status_byte (2=ok, ETF) + external term format + <<2, (term_to_binary(Result))/binary>>; {error, {not_found, Name}} -> ErrMsg = iolist_to_binary( io_lib:format("Function '~s' not registered", [Name])), @@ -966,8 +1063,7 @@ handle_callback_with_nested_receive(Ref, FuncName, CallbackArgs) -> ArgsList = tuple_to_list(CallbackArgs), case py_callback:execute(FuncName, ArgsList) of {ok, Value} -> - ReprStr = term_to_python_repr(Value), - {ok, <<0, ReprStr/binary>>}; + {ok, <<2, (term_to_binary(Value))/binary>>}; {error, Reason} -> ErrMsg = iolist_to_binary(io_lib:format("~p", [Reason])), {ok, <<1, ErrMsg/binary>>} @@ -1069,79 +1165,10 @@ resume_and_continue(Ref, StateRef, {error, _} = Err) -> %% Utility functions %% ============================================================================ -%% @private -%% Convert Erlang term to Python repr string -term_to_python_repr(Term) when is_integer(Term) -> - integer_to_binary(Term); -term_to_python_repr(Term) when is_float(Term) -> - float_to_binary(Term, [{decimals, 15}, compact]); -term_to_python_repr(true) -> - <<"True">>; -term_to_python_repr(false) -> - <<"False">>; -term_to_python_repr(none) -> - <<"None">>; -term_to_python_repr(nil) -> - <<"None">>; -term_to_python_repr(undefined) -> - <<"None">>; -term_to_python_repr(Term) when is_atom(Term) -> - %% Convert atom to Python string - BinStr = atom_to_binary(Term, utf8), - <<"'", BinStr/binary, "'">>; -term_to_python_repr(Term) when is_binary(Term) -> - %% Escape the binary for Python - Escaped = binary:replace(Term, <<"'">>, <<"\\'">>, [global]), - <<"'", Escaped/binary, "'">>; -term_to_python_repr(Term) when is_list(Term) -> - case io_lib:printable_unicode_list(Term) of - true -> - %% It's a string - Bin = unicode:characters_to_binary(Term), - Escaped = binary:replace(Bin, <<"'">>, <<"\\'">>, [global]), - <<"'", Escaped/binary, "'">>; - false -> - %% It's a list - Items = [term_to_python_repr(E) || E <- Term], - ItemsBin = join_binaries(Items, <<", ">>), - <<"[", ItemsBin/binary, "]">> - end; -term_to_python_repr(Term) when is_tuple(Term) -> - Items = [term_to_python_repr(E) || E <- tuple_to_list(Term)], - ItemsBin = join_binaries(Items, <<", ">>), - case tuple_size(Term) of - 1 -> <<"(", ItemsBin/binary, ",)">>; - _ -> <<"(", ItemsBin/binary, ")">> - end; -term_to_python_repr(Term) when is_map(Term) -> - Items = maps:fold(fun(K, V, Acc) -> - KeyRepr = term_to_python_repr(K), - ValRepr = term_to_python_repr(V), - [<> | Acc] - end, [], Term), - ItemsBin = join_binaries(lists:reverse(Items), <<", ">>), - <<"{", ItemsBin/binary, "}">>; -term_to_python_repr(Term) when is_pid(Term) -> - %% Encode PID using ETF (Erlang Term Format) for exact reconstruction. - %% Format: "__etf__:" - %% The C side will detect this, base64 decode, and use enif_binary_to_term - %% to reconstruct the pid, then convert to ErlangPidObject. - Etf = term_to_binary(Term), - B64 = base64:encode(Etf), - <<"\"__etf__:", B64/binary, "\"">>; -term_to_python_repr(Term) when is_reference(Term) -> - %% References also need ETF encoding for round-trip - Etf = term_to_binary(Term), - B64 = base64:encode(Etf), - <<"\"__etf__:", B64/binary, "\"">>; -term_to_python_repr(_Term) -> - <<"None">>. - -%% @private -join_binaries([], _Sep) -> <<>>; -join_binaries([H], _Sep) -> H; -join_binaries([H|T], Sep) -> - lists:foldl(fun(B, Acc) -> <> end, H, T). +%% Callback results cross to Python as external term format (status byte 2) +%% and are decoded by term_to_py() in c_src/py_convert.c, the same +%% converter used for call arguments. The former Python-repr encoder was +%% removed in favour of it. %% @private to_binary(Atom) when is_atom(Atom) -> diff --git a/src/py_nif.erl b/src/py_nif.erl index 4d07a6d..d59bde9 100644 --- a/src/py_nif.erl +++ b/src/py_nif.erl @@ -156,6 +156,9 @@ %% Process-per-context API (no mutex) context_create/1, context_destroy/1, + context_interrupt/1, + context_set_memory_limit/2, + context_memory_usage/1, context_call/5, context_call/6, context_eval/3, @@ -1174,6 +1177,52 @@ context_create(_Mode) -> context_destroy(_ContextRef) -> ?NIF_STUB. +%% @doc Interrupt Python code currently running in a context. +%% +%% Raises KeyboardInterrupt asynchronously in the thread executing the +%% context. The interrupted call returns `{error, interrupted}'. +%% +%% Safe to call from any process, including while the owning context process +%% is blocked in a NIF. Returns `not_running' if the context is idle or the +%% exception could not be delivered. +%% +%% CPython delivers async exceptions at bytecode boundaries, so code blocked +%% in a C call (`time.sleep', a numpy kernel, a socket read) is not +%% interrupted until that call returns. +%% +%% @param ContextRef Reference returned by context_create/1 +%% @returns ok | not_running | {error, Reason} +-spec context_interrupt(reference()) -> ok | not_running | {error, term()}. +context_interrupt(_ContextRef) -> + ?NIF_STUB. + +%% @doc Set a memory cap for a context, in bytes (0 removes the cap). +%% +%% Requires `owngil' mode and the runtime started with +%% `enable_memory_limits'. Exceeding the cap raises MemoryError in the +%% Python code, which surfaces as `{error, {'MemoryError', _}}'. +%% +%% Only memory routed through obmalloc is counted: allocations over 512 bytes +%% (large binaries, numpy buffers, extensions with their own allocator) bypass +%% it, and the granularity is one 1 MB arena. +%% +%% @param ContextRef Reference returned by context_create/1 +%% @param Bytes Cap in bytes, or 0 for no cap +%% @returns ok | {error, Reason} +-spec context_set_memory_limit(reference(), non_neg_integer()) -> + ok | {error, term()}. +context_set_memory_limit(_ContextRef, _Bytes) -> + ?NIF_STUB. + +%% @doc Report accounted memory usage for a context. +%% +%% @param ContextRef Reference returned by context_create/1 +%% @returns {ok, UsedBytes, LimitBytes} | {error, Reason} +-spec context_memory_usage(reference()) -> + {ok, non_neg_integer(), non_neg_integer()} | {error, term()}. +context_memory_usage(_ContextRef) -> + ?NIF_STUB. + %% @doc Call a Python function in a context. %% %% NO MUTEX - caller must ensure exclusive access (process ownership). diff --git a/src/py_thread_handler.erl b/src/py_thread_handler.erl index 343e4df..3fe00a4 100644 --- a/src/py_thread_handler.erl +++ b/src/py_thread_handler.erl @@ -225,10 +225,8 @@ handle_thread_callback(WriteFd, CallbackId, FuncName, Args) -> %% Execute the registered function Response = case py_callback:execute(FuncName, ArgsList) of {ok, Result} -> - %% Encode result as Python-parseable string - %% Format: status_byte (0=ok) + python_repr - ResultStr = term_to_python_repr(Result), - <<0, ResultStr/binary>>; + %% Format: status_byte (2=ok, ETF) + external term format + <<2, (term_to_binary(Result))/binary>>; {error, {not_found, Name}} -> ErrMsg = iolist_to_binary( io_lib:format("Function '~s' not registered", [Name])), @@ -247,7 +245,7 @@ handle_thread_callback(WriteFd, CallbackId, FuncName, Args) -> end. %% Execute the user's registered function and encode the result as a -%% wire-format response body (`<>`). +%% wire-format response body (`<>`). %% Used by async_writer_loop (and indirectly by handle_thread_callback %% via the same encoding shape — the sync path keeps its own copy %% close to the write call site). @@ -259,8 +257,7 @@ run_async_callback(FuncName, Args) -> end, case py_callback:execute(FuncName, ArgsList) of {ok, Result} -> - ResultStr = term_to_python_repr(Result), - <<0, ResultStr/binary>>; + <<2, (term_to_binary(Result))/binary>>; {error, {not_found, Name}} -> ErrMsg = iolist_to_binary( io_lib:format("Function '~s' not registered", [Name])), @@ -341,88 +338,10 @@ async_writer_loop(WriteFd) -> end. %%% ============================================================================ -%%% Term to Python repr conversion -%%% (Same as py_worker.erl - could be factored out to py_util.erl) +%%% Callback result encoding %%% ============================================================================ - -%% Convert Erlang term to Python-parseable string representation -term_to_python_repr(Term) when is_integer(Term) -> - integer_to_binary(Term); -term_to_python_repr(Term) when is_float(Term) -> - float_to_binary(Term, [{decimals, 17}, compact]); -term_to_python_repr(true) -> - <<"True">>; -term_to_python_repr(false) -> - <<"False">>; -term_to_python_repr(none) -> - <<"None">>; -term_to_python_repr(nil) -> - <<"None">>; -term_to_python_repr(undefined) -> - <<"None">>; -term_to_python_repr(Term) when is_atom(Term) -> - %% Convert atom to Python string - AtomStr = atom_to_binary(Term, utf8), - <<"\"", AtomStr/binary, "\"">>; -term_to_python_repr(Term) when is_binary(Term) -> - %% Escape binary as Python string - Escaped = escape_string(Term), - <<"\"", Escaped/binary, "\"">>; -term_to_python_repr(Term) when is_list(Term) -> - %% Check if it's a string (list of integers) - case io_lib:printable_list(Term) of - true -> - Bin = list_to_binary(Term), - Escaped = escape_string(Bin), - <<"\"", Escaped/binary, "\"">>; - false -> - Items = [term_to_python_repr(E) || E <- Term], - Joined = join_binaries(Items, <<", ">>), - <<"[", Joined/binary, "]">> - end; -term_to_python_repr(Term) when is_tuple(Term) -> - Items = [term_to_python_repr(E) || E <- tuple_to_list(Term)], - Joined = join_binaries(Items, <<", ">>), - case length(Items) of - 1 -> <<"(", Joined/binary, ",)">>; - _ -> <<"(", Joined/binary, ")">> - end; -term_to_python_repr(Term) when is_map(Term) -> - Items = maps:fold(fun(K, V, Acc) -> - KeyRepr = term_to_python_repr(K), - ValRepr = term_to_python_repr(V), - [<> | Acc] - end, [], Term), - Joined = join_binaries(Items, <<", ">>), - <<"{", Joined/binary, "}">>; -term_to_python_repr(Term) when is_pid(Term) -> - %% Encode PID using ETF (Erlang Term Format) for exact reconstruction. - %% Format: "__etf__:" - %% The C side will detect this, base64 decode, and use enif_binary_to_term - %% to reconstruct the pid, then convert to ErlangPidObject. - Etf = term_to_binary(Term), - B64 = base64:encode(Etf), - <<"\"__etf__:", B64/binary, "\"">>; -term_to_python_repr(Term) when is_reference(Term) -> - %% References also need ETF encoding for round-trip - Etf = term_to_binary(Term), - B64 = base64:encode(Etf), - <<"\"__etf__:", B64/binary, "\"">>; -term_to_python_repr(_Term) -> - %% Fallback - return None for unsupported types - <<"None">>. - -escape_string(Bin) -> - %% Escape special characters for Python string - binary:replace( - binary:replace( - binary:replace( - binary:replace(Bin, <<"\\">>, <<"\\\\">>, [global]), - <<"\"">>, <<"\\\"">>, [global]), - <<"\n">>, <<"\\n">>, [global]), - <<"\r">>, <<"\\r">>, [global]). - -join_binaries([], _Sep) -> <<>>; -join_binaries([H], _Sep) -> H; -join_binaries([H|T], Sep) -> - lists:foldl(fun(E, Acc) -> <> end, H, T). +%%% +%%% Results cross to Python as external term format (status byte 2) and are +%%% decoded by term_to_py() in c_src/py_convert.c, the same converter used +%%% for call arguments. The former Python-repr encoder was removed in +%%% favour of it. diff --git a/test/py_callback_encoding_SUITE.erl b/test/py_callback_encoding_SUITE.erl new file mode 100644 index 0000000..103c9d5 --- /dev/null +++ b/test/py_callback_encoding_SUITE.erl @@ -0,0 +1,169 @@ +%%% @doc Common Test suite for the callback result encoding. +%%% +%%% Callback results cross from Erlang to Python as external term format and +%%% are decoded by term_to_py() in c_src/py_convert.c, the same converter that +%%% handles call arguments. These cases pin the resulting Python types and the +%%% round-trip fidelity, including payloads that the previous repr-string +%%% encoder corrupted. +-module(py_callback_encoding_SUITE). + +-include_lib("common_test/include/ct.hrl"). + +-export([ + all/0, + init_per_suite/1, + end_per_suite/1, + end_per_testcase/2 +]). + +-export([ + test_binary_with_escapes/1, + test_binary_non_utf8/1, + test_large_binary/1, + test_atom_becomes_str/1, + test_empty_list/1, + test_erlang_string_is_int_list/1, + test_nested_containers/1, + test_pid_round_trip/1, + test_ref_round_trip/1, + test_float_round_trip/1, + test_booleans_and_none/1, + test_python_types/1 +]). + +all() -> [ + test_binary_with_escapes, + test_binary_non_utf8, + test_large_binary, + test_atom_becomes_str, + test_empty_list, + test_erlang_string_is_int_list, + test_nested_containers, + test_pid_round_trip, + test_ref_round_trip, + test_float_round_trip, + test_booleans_and_none, + test_python_types +]. + +init_per_suite(Config) -> + {ok, _} = application:ensure_all_started(erlang_python), + Config. + +end_per_suite(_Config) -> + ok = application:stop(erlang_python), + ok. + +end_per_testcase(_TestCase, _Config) -> + catch py:unregister_function(cbenc_probe), + ok. + +%%% ============================================================================ +%%% Test Cases +%%% ============================================================================ + +%% @doc Backslashes, both quote styles, newlines and tabs survive intact. The +%% repr encoder escaped only single quotes, so these payloads failed to parse +%% and were silently handed to Python as the raw repr text. +test_binary_with_escapes(_Config) -> + Value = <<"back\\slash \"dq\" 'sq'\nnewline\ttab\r">>, + Value = probe(Value), + <<"str">> = probe_type(Value), + ok. + +%% @doc A binary that is not valid UTF-8 arrives as Python bytes and returns +%% byte-identical. +test_binary_non_utf8(_Config) -> + Value = <<0, 1, 255, 254, 128>>, + Value = probe(Value), + <<"bytes">> = probe_type(Value), + ok. + +test_large_binary(_Config) -> + Value = binary:copy(<<"abcdefghij">>, 20000), + Value = probe(Value), + ok. + +%% @doc Atoms have no Python counterpart and arrive as str. +test_atom_becomes_str(_Config) -> + <<"some_atom">> = probe(some_atom), + <<"str">> = probe_type(some_atom), + ok. + +%% @doc The empty list is a list. The repr encoder classified it as a printable +%% string and produced '' instead. +test_empty_list(_Config) -> + [] = probe([]), + <<"list">> = probe_type([]), + ok. + +%% @doc An Erlang string is a list of integers, matching how call arguments +%% have always been converted. Return a binary for a Python str. +test_erlang_string_is_int_list(_Config) -> + "abc" = probe("abc"), + <<"list">> = probe_type("abc"), + <<"abc">> = probe(<<"abc">>), + <<"str">> = probe_type(<<"abc">>), + ok. + +test_nested_containers(_Config) -> + Value = #{<<"k">> => [1, 2.5, {a, b}, #{<<"inner">> => [[], {}]}]}, + Expected = #{<<"k">> => [1, 2.5, {<<"a">>, <<"b">>}, #{<<"inner">> => [[], {}]}]}, + Expected = probe(Value), + <<"dict">> = probe_type(Value), + ok. + +%% @doc Pids cross as native Pid objects, with no base64 marker round-trip. +test_pid_round_trip(_Config) -> + Pid = self(), + Pid = probe(Pid), + <<"Pid">> = probe_type(Pid), + ok. + +test_ref_round_trip(_Config) -> + Ref = make_ref(), + Ref = probe(Ref), + <<"Ref">> = probe_type(Ref), + ok. + +%% @doc Floats are exact, not routed through a decimal-formatted string. +test_float_round_trip(_Config) -> + lists:foreach(fun(F) -> F = probe(F) end, + [3.14159265358979, 1.0e-300, 1.7976931348623157e308, -0.0]), + ok. + +test_booleans_and_none(_Config) -> + true = probe(true), + false = probe(false), + <<"bool">> = probe_type(true), + lists:foreach(fun(A) -> + none = probe(A), + <<"NoneType">> = probe_type(A) + end, [undefined, nil, none]), + ok. + +%% @doc Tuples stay tuples and integers stay ints on the Python side. +test_python_types(_Config) -> + <<"tuple">> = probe_type({1, 2}), + <<"int">> = probe_type(42), + <<"float">> = probe_type(1.5), + ok. + +%%% ============================================================================ +%%% Helpers +%%% ============================================================================ + +%% Return the value through a callback and back to Erlang. +probe(Value) -> + py:register_function(cbenc_probe, fun(_) -> Value end), + {ok, Got} = py:eval(<<"__import__('erlang').call('cbenc_probe', [])">>), + py:unregister_function(cbenc_probe), + Got. + +%% Return the Python type name the callback result lands as. +probe_type(Value) -> + py:register_function(cbenc_probe, fun(_) -> Value end), + {ok, Type} = py:eval( + <<"type(__import__('erlang').call('cbenc_probe', [])).__name__">>), + py:unregister_function(cbenc_probe), + Type. diff --git a/test/py_interrupt_SUITE.erl b/test/py_interrupt_SUITE.erl new file mode 100644 index 0000000..ecb87aa --- /dev/null +++ b/test/py_interrupt_SUITE.erl @@ -0,0 +1,204 @@ +%%% @doc Common Test suite for interrupting running Python code. +%%% +%%% Covers py_nif:context_interrupt/1 and the automatic interrupt that +%%% py_context issues when a call timeout expires. Every case runs against +%%% both context modes; the owngil group is skipped when the runtime does +%%% not support it. +-module(py_interrupt_SUITE). + +-include_lib("common_test/include/ct.hrl"). + +-export([ + all/0, + groups/0, + init_per_suite/1, + end_per_suite/1, + init_per_group/2, + end_per_group/2, + init_per_testcase/2, + end_per_testcase/2 +]). + +-export([ + test_timeout_interrupts_loop/1, + test_context_reusable_after_interrupt/1, + test_explicit_interrupt_returns_interrupted/1, + test_interrupt_idle_context/1, + test_base_exception_not_swallowed/1, + test_no_stale_reply_in_mailbox/1, + test_interrupt_does_not_leak_to_next_request/1, + test_interrupt_during_callback/1, + test_blocking_c_call_recovers/1 +]). + +%% Python that spins forever without allocating or calling into C. +-define(BUSY, <<"sum(1 for _ in iter(int, 1))">>). + +all() -> + [{group, worker}, {group, owngil}]. + +groups() -> + Cases = [ + test_timeout_interrupts_loop, + test_context_reusable_after_interrupt, + test_explicit_interrupt_returns_interrupted, + test_interrupt_idle_context, + test_base_exception_not_swallowed, + test_no_stale_reply_in_mailbox, + test_interrupt_does_not_leak_to_next_request, + test_interrupt_during_callback, + test_blocking_c_call_recovers + ], + [{worker, [], Cases}, {owngil, [], Cases}]. + +init_per_suite(Config) -> + {ok, _} = application:ensure_all_started(erlang_python), + Config. + +end_per_suite(_Config) -> + ok = application:stop(erlang_python), + ok. + +init_per_group(owngil, Config) -> + case py_nif:owngil_supported() of + true -> [{mode, owngil} | Config]; + false -> {skip, "OWN_GIL requires Python 3.14+"} + end; +init_per_group(worker, Config) -> + [{mode, worker} | Config]. + +end_per_group(_Group, _Config) -> + ok. + +init_per_testcase(_TestCase, Config) -> + {ok, Ctx} = py_context:new(#{mode => ?config(mode, Config)}), + [{ctx, Ctx} | Config]. + +end_per_testcase(_TestCase, Config) -> + catch py_context:stop(?config(ctx, Config)), + flush(), + ok. + +%%% ============================================================================ +%%% Test Cases +%%% ============================================================================ + +%% @doc A timed-out call stops burning CPU instead of running to completion. +test_timeout_interrupts_loop(Config) -> + Ctx = ?config(ctx, Config), + T0 = erlang:monotonic_time(millisecond), + {error, timeout} = py_context:eval(Ctx, ?BUSY, #{}, 500), + Elapsed = erlang:monotonic_time(millisecond) - T0, + %% Returns near the deadline, not after the grace period expires + true = Elapsed < 500 + 900, + ok. + +%% @doc The context is immediately usable again after an interrupt. If the +%% Python loop were still running, this call would queue behind it. +test_context_reusable_after_interrupt(Config) -> + Ctx = ?config(ctx, Config), + {error, timeout} = py_context:eval(Ctx, ?BUSY, #{}, 300), + T0 = erlang:monotonic_time(millisecond), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000), + Elapsed = erlang:monotonic_time(millisecond) - T0, + true = Elapsed < 1000, + ok. + +%% @doc An explicit interrupt surfaces to the caller as {error, interrupted}. +test_explicit_interrupt_returns_interrupted(Config) -> + Ctx = ?config(ctx, Config), + Me = self(), + spawn(fun() -> Me ! {res, py_context:eval(Ctx, ?BUSY, #{}, infinity)} end), + timer:sleep(300), + ok = py:interrupt(Ctx), + receive + {res, Result} -> {error, interrupted} = Result + after 5000 -> + ct:fail(no_reply_after_interrupt) + end, + ok. + +%% @doc Interrupting an idle context is a no-op and does not affect later calls. +test_interrupt_idle_context(Config) -> + Ctx = ?config(ctx, Config), + not_running = py:interrupt(Ctx), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000), + ok. + +%% @doc KeyboardInterrupt is a BaseException, so a bare `except Exception` +%% around the hot loop does not swallow the interrupt. +test_base_exception_not_swallowed(Config) -> + Ctx = ?config(ctx, Config), + Code = <<" +def _guarded(): + while True: + try: + pass + except Exception: + pass +_guarded()">>, + {error, timeout} = py_context:eval(Ctx, <<"exec('''", Code/binary, "''')">>, #{}, 500), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000), + ok. + +%% @doc The reply from the interrupted call is drained, not left in the +%% caller's mailbox. +test_no_stale_reply_in_mailbox(Config) -> + Ctx = ?config(ctx, Config), + {error, timeout} = py_context:eval(Ctx, ?BUSY, #{}, 300), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 5000), + timer:sleep(200), + {messages, []} = process_info(self(), messages), + ok. + +%% @doc An interrupt that lands as a request completes must not be delivered +%% to the next request on the same context. +test_interrupt_does_not_leak_to_next_request(Config) -> + Ctx = ?config(ctx, Config), + Bad = lists:foldl(fun(_, Acc) -> + %% Very short timeout so the interrupt races request completion + _ = py_context:eval(Ctx, ?BUSY, #{}, 20), + case py_context:eval(Ctx, <<"1+1">>, #{}, 5000) of + {ok, 2} -> Acc; + Other -> [Other | Acc] + end + end, [], lists:seq(1, 50)), + [] = Bad, + ok. + +%% @doc Interrupting Python that is suspended in an Erlang callback. The +%% callback path executes on a dirty scheduler via py_context_acquire, a +%% different code path from the context worker thread. +test_interrupt_during_callback(Config) -> + Ctx = ?config(ctx, Config), + ok = py:register_function(<<"interrupt_test_echo">>, fun([X]) -> X end), + Code = <<" +import erlang +_n = 0 +while True: + _n = erlang.call('interrupt_test_echo', _n) + 1 +">>, + {error, timeout} = py_context:eval(Ctx, <<"exec('''", Code/binary, "''')">>, #{}, 500), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 10000), + py:unregister_function(<<"interrupt_test_echo">>), + ok. + +%% @doc Documented limitation: an async exception is delivered at a bytecode +%% boundary, so code blocked in a C call is only interrupted once that call +%% returns. The call still times out and the context recovers afterwards. +test_blocking_c_call_recovers(Config) -> + Ctx = ?config(ctx, Config), + {error, timeout} = py_context:eval( + Ctx, <<"__import__('time').sleep(2) or 1">>, #{}, 200), + %% Once the sleep finishes the context accepts work again + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 10000), + ok. + +%%% ============================================================================ +%%% Helpers +%%% ============================================================================ + +flush() -> + receive _ -> flush() + after 0 -> ok + end. diff --git a/test/py_memory_limit_SUITE.erl b/test/py_memory_limit_SUITE.erl new file mode 100644 index 0000000..d780906 --- /dev/null +++ b/test/py_memory_limit_SUITE.erl @@ -0,0 +1,139 @@ +%%% @doc Common Test suite for per-context memory caps. +%%% +%%% Caps are accounted from obmalloc arena traffic and enforced by raising +%%% MemoryError in the offending context. They require owngil mode and the +%%% runtime started with `enable_memory_limits' (set node-wide in +%%% test/test.config, since the allocator is hooked before Python starts). +-module(py_memory_limit_SUITE). + +-include_lib("common_test/include/ct.hrl"). + +-export([ + all/0, + init_per_suite/1, + end_per_suite/1 +]). + +-export([ + test_cap_raises_memory_error/1, + test_context_survives_cap/1, + test_cap_rearms_after_release/1, + test_usage_is_reported/1, + test_no_cap_is_unlimited/1, + test_worker_mode_rejected/1, + test_invalid_limit_rejected/1 +]). + +%% Allocates well past any cap used here: ~56 bytes per empty list. +-define(GREEDY, <<"_hog = [[] for _ in range(3000000)]">>). + +-define(CAP, (64 * 1024 * 1024)). + +all() -> [ + test_cap_raises_memory_error, + test_context_survives_cap, + test_cap_rearms_after_release, + test_usage_is_reported, + test_no_cap_is_unlimited, + test_worker_mode_rejected, + test_invalid_limit_rejected +]. + +init_per_suite(Config) -> + {ok, _} = application:ensure_all_started(erlang_python), + case py_nif:owngil_supported() of + false -> + {skip, "Memory limits require OWN_GIL (Python 3.14+)"}; + true -> + %% The allocator is hooked before Python starts, so if another + %% suite initialised the runtime without the flag we cannot turn + %% it on here. + case probe_enabled() of + true -> Config; + false -> {skip, "runtime started without enable_memory_limits"} + end + end. + +end_per_suite(_Config) -> + ok = application:stop(erlang_python), + ok. + +%%% ============================================================================ +%%% Test Cases +%%% ============================================================================ + +%% @doc Allocating past the cap raises MemoryError in the Python code. +test_cap_raises_memory_error(_Config) -> + {ok, Ctx} = py_context:new(#{mode => owngil, memory_limit => ?CAP}), + {error, {'MemoryError', _}} = py_context:exec(Ctx, ?GREEDY), + py_context:stop(Ctx), + ok. + +%% @doc The context is still usable after its cap has been hit. +test_context_survives_cap(_Config) -> + {ok, Ctx} = py_context:new(#{mode => owngil, memory_limit => ?CAP}), + {error, {'MemoryError', _}} = py_context:exec(Ctx, ?GREEDY), + {ok, 4} = py_context:eval(Ctx, <<"2+2">>, #{}, 10000), + py_context:stop(Ctx), + ok. + +%% @doc Once the memory is released the cap enforces again, rather than +%% latching after the first breach. +test_cap_rearms_after_release(_Config) -> + {ok, Ctx} = py_context:new(#{mode => owngil, memory_limit => ?CAP}), + {error, {'MemoryError', _}} = py_context:exec(Ctx, ?GREEDY), + %% Whatever survived the unwinding goes away here + _ = py_context:exec(Ctx, <<"_hog = None">>), + {ok, _} = py_context:eval(Ctx, <<"__import__('gc').collect()">>, #{}, 30000), + {error, {'MemoryError', _}} = py_context:exec(Ctx, ?GREEDY), + py_context:stop(Ctx), + ok. + +%% @doc Usage is reported and grows with live objects. +test_usage_is_reported(_Config) -> + {ok, Ctx} = py_context:new(#{mode => owngil}), + Ref = py_context:get_nif_ref(Ctx), + {ok, Before, 0} = py_nif:context_memory_usage(Ref), + ok = py_context:exec(Ctx, <<"_keep = [[] for _ in range(500000)]">>), + {ok, After, 0} = py_nif:context_memory_usage(Ref), + true = After > Before, + py_context:stop(Ctx), + ok. + +%% @doc Without a cap the same allocation succeeds. +test_no_cap_is_unlimited(_Config) -> + {ok, Ctx} = py_context:new(#{mode => owngil}), + ok = py_context:exec(Ctx, ?GREEDY), + {ok, 3000000} = py_context:eval(Ctx, <<"len(_hog)">>, #{}, 10000), + py_context:stop(Ctx), + ok. + +%% @doc Worker-mode contexts share the main interpreter, so a per-context cap +%% is refused rather than silently applied process-wide. +test_worker_mode_rejected(_Config) -> + {error, memory_limit_requires_owngil} = + py_context:new(#{mode => worker, memory_limit => ?CAP}), + ok. + +test_invalid_limit_rejected(_Config) -> + {error, {invalid_memory_limit, -1}} = + py_context:new(#{mode => owngil, memory_limit => -1}), + {error, {invalid_memory_limit, <<"big">>}} = + py_context:new(#{mode => owngil, memory_limit => <<"big">>}), + ok. + +%%% ============================================================================ +%%% Helpers +%%% ============================================================================ + +%% Setting a zero cap is a no-op when limits are on and fails when they are off. +probe_enabled() -> + case py_context:new(#{mode => owngil}) of + {ok, Ctx} -> + Ref = py_context:get_nif_ref(Ctx), + Result = py_nif:context_set_memory_limit(Ref, 0), + py_context:stop(Ctx), + Result =:= ok; + _ -> + false + end. diff --git a/test/py_reentrant_SUITE.erl b/test/py_reentrant_SUITE.erl index 64666f6..4eb6f18 100644 --- a/test/py_reentrant_SUITE.erl +++ b/test/py_reentrant_SUITE.erl @@ -116,21 +116,23 @@ test_reentrant_resume_stress(_Config) -> ok. %% @doc Regression for the binary_to_term SAFE-flag hardening (atom exhaustion). -%% A `__etf__:` callback result that encodes a brand-new atom must be rejected (not -%% decoded) so it cannot mint non-GC'd atoms, while a valid existing-atom payload -%% still round-trips through the same enif_binary_to_term path. +%% +%% Callback results cross as external term format built by the Erlang side, so +%% pids and references round-trip natively. A binary that merely looks like the +%% old `__etf__:` marker is data: it must be passed through untouched and must +%% never be re-interpreted as a term, which is what could mint non-GC'd atoms. test_etf_decode_safe(_Config) -> - %% Positive: an EXISTING atom encoded via __etf__: still decodes under SAFE, - %% proving the decode path runs and the change is non-breaking. - OkMarker = etf_marker(term_to_binary(ok)), - py:register_function(etf_probe_ok, fun(_) -> OkMarker end), - {ok, GotOk} = py:eval(<<"__import__('erlang').call('etf_probe_ok', [])">>), - true = (GotOk =/= OkMarker), %% decoded to the term, not passed through verbatim + %% Positive: pids and refs cross natively, no marker encoding involved. + %% (The atom comes back as a binary: atoms have no Python counterpart.) + Pid = self(), + Ref = make_ref(), + py:register_function(etf_probe_ok, fun(_) -> {Pid, Ref, ok} end), + {ok, {Pid, Ref, <<"ok">>}} = + py:eval(<<"__import__('erlang').call('etf_probe_ok', [])">>), py:unregister_function(etf_probe_ok), - %% Negative: many DISTINCT brand-new atoms encoded via __etf__: must all be - %% rejected. Pre-fix (flags=0) each would create a permanent atom and come back - %% as that atom; with ERL_NIF_BIN2TERM_SAFE the raw marker is returned unchanged. + %% Negative: many DISTINCT brand-new atoms wrapped in marker-shaped binaries + %% must all come back verbatim, never decoded into atoms. Before = erlang:system_info(atom_count), N = 50, lists:foreach( @@ -139,7 +141,7 @@ test_etf_decode_safe(_Config) -> Marker = etf_marker(novel_atom_etf(Name)), py:register_function(etf_probe_novel, fun(_) -> Marker end), {ok, Got} = py:eval(<<"__import__('erlang').call('etf_probe_novel', [])">>), - Marker = Got, %% rejected: returned verbatim, never decoded to the atom + Marker = Got, %% passed through as data, never decoded to the atom assert_atom_absent(Name) end, lists:seq(1, N) @@ -150,8 +152,8 @@ test_etf_decode_safe(_Config) -> true = (After - Before) < (N div 2), ok. -%% @private Build a "__etf__:" marker the C side base64-decodes and feeds to -%% enif_binary_to_term (see term_to_python_repr/1 in py_context for the real encoder). +%% @private Build a "__etf__:" marker, the shape the legacy repr encoder +%% used for pids and refs, now just an ordinary binary as far as callbacks go. etf_marker(Etf) -> <<"__etf__:", (base64:encode(Etf))/binary>>. diff --git a/test/py_stream_SUITE.erl b/test/py_stream_SUITE.erl index ff4dfb1..57447b9 100644 --- a/test/py_stream_SUITE.erl +++ b/test/py_stream_SUITE.erl @@ -18,7 +18,12 @@ test_stream_error/1, test_stream_empty/1, test_stream_large/1, - test_stream_rejects_injection/1 + test_stream_rejects_injection/1, + test_stream_async_generator/1, + test_stream_async_generator_args/1, + test_stream_async_generator_empty/1, + test_stream_async_generator_error/1, + test_stream_async_cancel/1 ]). all() -> @@ -31,11 +36,20 @@ all() -> test_stream_error, test_stream_empty, test_stream_large, - test_stream_rejects_injection + test_stream_rejects_injection, + test_stream_async_generator, + test_stream_async_generator_args, + test_stream_async_generator_empty, + test_stream_async_generator_error, + test_stream_async_cancel ]. init_per_suite(Config) -> {ok, _} = application:ensure_all_started(erlang_python), + %% Make test/py_test_agen.py importable for the async generator cases + TestDir = filename:join(code:lib_dir(erlang_python), "test"), + ok = py:exec(iolist_to_binary(io_lib:format( + "import sys; sys.path.insert(0, '~s')", [TestDir]))), Config. end_per_suite(_Config) -> @@ -175,3 +189,54 @@ test_stream_large(_Config) -> 0 = hd(Values), 999 = lists:last(Values), ok. + +%%% ============================================================================ +%%% Async generators +%%% +%%% stream_start/3,4 drives an async generator on a private event loop. The +%%% collect-style py:stream/4-with-kwargs and py:stream_eval/1,2 wrap the call +%%% in list() and remain sync-only. +%%% ============================================================================ + +%% Test streaming from an async generator +test_stream_async_generator(_Config) -> + {ok, Ref} = py:stream_start(<<"py_test_agen">>, <<"counter">>, [4]), + {ok, Values} = collect_stream(Ref), + [0, 1, 2, 3] = Values, + ok. + +%% Test an async generator taking several arguments +test_stream_async_generator_args(_Config) -> + {ok, Ref} = py:stream_start(<<"py_test_agen">>, <<"scaled">>, [3, 10]), + {ok, Values} = collect_stream(Ref), + [0, 10, 20] = Values, + ok. + +%% An async generator that yields nothing still completes +test_stream_async_generator_empty(_Config) -> + {ok, Ref} = py:stream_start(<<"py_test_agen">>, <<"empty">>, []), + {ok, Values} = collect_stream(Ref), + [] = Values, + ok. + +%% An exception raised mid-iteration is reported as a stream error, after the +%% values yielded before it +test_stream_async_generator_error(_Config) -> + {ok, Ref} = py:stream_start(<<"py_test_agen">>, <<"failing">>, [5]), + {error, Reason} = collect_stream(Ref), + true = is_binary(Reason), + {_, _} = binary:match(Reason, <<"agen boom">>), + ok. + +%% Cancelling an async stream stops it with {error, cancelled} +test_stream_async_cancel(_Config) -> + {ok, Ref} = py:stream_start(<<"py_test_agen">>, <<"slow">>, [50]), + receive + {py_stream, Ref, {data, _}} -> ok + after 5000 -> + ct:fail(no_first_value) + end, + ok = py:stream_cancel(Ref), + {error, <<"cancelled">>} = collect_stream(Ref), + drain_stream(Ref), + ok. diff --git a/test/py_test_agen.py b/test/py_test_agen.py new file mode 100644 index 0000000..205f913 --- /dev/null +++ b/test/py_test_agen.py @@ -0,0 +1,36 @@ +# Async generators used by py_stream_SUITE to exercise py:stream_start/3,4. +import asyncio + + +async def counter(n): + """Yield 0..n-1, awaiting between values so the loop really suspends.""" + for i in range(n): + await asyncio.sleep(0) + yield i + + +async def scaled(n, factor): + for i in range(n): + await asyncio.sleep(0) + yield i * factor + + +async def empty(): + """An async generator that yields nothing.""" + return + yield # pragma: no cover - makes this an async generator + + +async def failing(n): + for i in range(n): + await asyncio.sleep(0) + if i == 2: + raise ValueError("agen boom") + yield i + + +async def slow(n): + """Yield slowly enough that a cancel can land mid-stream.""" + for i in range(n): + await asyncio.sleep(0.05) + yield i diff --git a/test/test.config b/test/test.config new file mode 100644 index 0000000..ccec148 --- /dev/null +++ b/test/test.config @@ -0,0 +1,10 @@ +%% Node-wide settings for the Common Test run. +%% +%% Memory limits hook the obmalloc arena allocator before Python starts, so the +%% flag has to be set for whichever suite initialises the runtime first; +%% py_memory_limit_SUITE cannot turn it on by itself mid-run. +[ + {erlang_python, [ + {enable_memory_limits, true} + ]} +].