From bba44789baaa394a2afc052f3ba6fb9126ee060f Mon Sep 17 00:00:00 2001 From: Benoit Chesneau Date: Sat, 15 Aug 2026 11:56:23 +0200 Subject: [PATCH] Add interrupts, memory caps, async streaming, ETF callback results Close four gaps that blocked running untrusted Python: no way to stop running code, no per-context memory bound, no async-generator streaming, and a lossy callback result encoding. py:interrupt/1 raises KeyboardInterrupt in the thread executing a context, and py_context:call/eval/exec now interrupt on timeout instead of only abandoning the reply while the thread kept running. py_context:new(#{mode => owngil, memory_limit => Bytes}) caps memory per context, accounted from obmalloc arenas and enforced with MemoryError. Opt-in via enable_memory_limits, since the allocator is hooked before Python starts. py:stream_start/3,4 drives async generators again, which the docs had claimed since 3.0.0 without the code behind it. Callback results now cross as external term format instead of Python repr strings parsed with ast.literal_eval, fixing binaries with backslashes, quotes or newlines, empty lists, float precision, and the base64 pid/ref round-trip. Breaking: an Erlang string returned from a callback arrives as a list of integers, matching call arguments. --- CHANGELOG.md | 44 ++++ c_src/py_callback.c | 68 ++++-- c_src/py_convert.c | 10 + c_src/py_mem_limit.c | 280 +++++++++++++++++++++ c_src/py_nif.c | 202 +++++++++++++++- c_src/py_nif.h | 93 +++++++ docs/interrupts.md | 77 ++++++ docs/memory.md | 58 +++++ docs/streaming.md | 32 ++- rebar.config | 6 + src/erlang_python_sup.erl | 11 +- src/py.erl | 50 +++- src/py_context.erl | 363 +++++++++++++++------------- src/py_nif.erl | 49 ++++ src/py_thread_handler.erl | 101 +------- test/py_callback_encoding_SUITE.erl | 169 +++++++++++++ test/py_interrupt_SUITE.erl | 204 ++++++++++++++++ test/py_memory_limit_SUITE.erl | 139 +++++++++++ test/py_reentrant_SUITE.erl | 32 +-- test/py_stream_SUITE.erl | 69 +++++- test/py_test_agen.py | 36 +++ test/test.config | 10 + 22 files changed, 1785 insertions(+), 318 deletions(-) create mode 100644 c_src/py_mem_limit.c create mode 100644 docs/interrupts.md create mode 100644 test/py_callback_encoding_SUITE.erl create mode 100644 test/py_interrupt_SUITE.erl create mode 100644 test/py_memory_limit_SUITE.erl create mode 100644 test/py_test_agen.py create mode 100644 test/test.config 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} + ]} +].