Add automatic env reuse for py_event_loop:run and spawn_task

When a process has a local Python environment (from py:context/py:exec),
py_event_loop functions now automatically pass that env to the NIF. This
allows functions defined via py:exec to be called directly without
manual env parameter passing.

Changes:
- Add submit_task_with_env/8 NIF that accepts an env reference
- Add get_env_resource_type() getter for accessing env resource type
- Modify process_ready_tasks to handle 7-tuple tasks with env ref
- Update py_event_loop:create_task to auto-detect and pass process env
- Update py_event_loop:spawn_task to use caller's env (not receiver's)
- Add get_process_env() helper to retrieve env from process dictionary
- Add tests for env reuse with sync, async, and spawn_task functions

Author: benoitc

c_src/py_event_loop.c

@@ -2501,6 +2501,76 @@ ERL_NIF_TERM nif_submit_task(ErlNifEnv *env, int argc,
     return ATOM_OK;
 }
 
+/**
+ * submit_task_with_env(LoopRef, CallerPid, Ref, Module, Func, Args, Kwargs, EnvRef) -> ok | {error, Reason}
+ *
+ * Like submit_task but includes a process-local env resource reference.
+ * The env's globals dict is used for function lookup, allowing functions
+ * defined via py:exec() to be called from the event loop.
+ */
+ERL_NIF_TERM nif_submit_task_with_env(ErlNifEnv *env, int argc,
+                                       const ERL_NIF_TERM argv[]) {
+    (void)argc;
+
+    erlang_event_loop_t *loop;
+    if (!enif_get_resource(env, argv[0], EVENT_LOOP_RESOURCE_TYPE,
+                           (void **)&loop)) {
+        return make_error(env, "invalid_loop");
+    }
+
+    if (!loop->task_queue_initialized) {
+        return make_error(env, "task_queue_not_initialized");
+    }
+
+    /* Validate caller_pid */
+    ErlNifPid caller_pid;
+    if (!enif_get_local_pid(env, argv[1], &caller_pid)) {
+        return make_error(env, "invalid_caller_pid");
+    }
+
+    /* Create task tuple: {CallerPid, Ref, Module, Func, Args, Kwargs, EnvRef} */
+    /* argv[1] = CallerPid, argv[2] = Ref, argv[3] = Module,
+     * argv[4] = Func, argv[5] = Args, argv[6] = Kwargs, argv[7] = EnvRef */
+    ERL_NIF_TERM task_tuple = enif_make_tuple7(env,
+        argv[1], argv[2], argv[3], argv[4], argv[5], argv[6], argv[7]);
+
+    /* Serialize to binary */
+    ErlNifBinary task_bin;
+    if (!enif_term_to_binary(env, task_tuple, &task_bin)) {
+        return make_error(env, "serialization_failed");
+    }
+
+    /* Thread-safe enqueue */
+    pthread_mutex_lock(&loop->task_queue_mutex);
+    int enq_result = enif_ioq_enq_binary(loop->task_queue, &task_bin, 0);
+    pthread_mutex_unlock(&loop->task_queue_mutex);
+
+    if (enq_result != 1) {
+        enif_release_binary(&task_bin);
+        return make_error(env, "enqueue_failed");
+    }
+
+    /* Increment task count */
+    atomic_fetch_add(&loop->task_count, 1);
+
+    /* Coalesced wakeup (uvloop-style) */
+    if (loop->has_worker) {
+        if (!atomic_exchange(&loop->task_wake_pending, true)) {
+            ErlNifEnv *msg_env = enif_alloc_env();
+            if (msg_env != NULL) {
+                if (ATOM_TASK_READY == 0) {
+                    ATOM_TASK_READY = enif_make_atom(msg_env, "task_ready");
+                }
+                ERL_NIF_TERM msg = enif_make_atom(msg_env, "task_ready");
+                enif_send(NULL, &loop->worker_pid, msg_env, msg);
+                enif_free_env(msg_env);
+            }
+        }
+    }
+
+    return ATOM_OK;
+}
+
 /**
  * Maximum tasks to dequeue in one batch before acquiring GIL.
  * This bounds memory usage while still amortizing GIL acquisition cost.
@@ -2789,10 +2859,12 @@ ERL_NIF_TERM nif_process_ready_tasks(ErlNifEnv *env, int argc,
         ErlNifEnv *term_env = tasks[task_idx].term_env;
         ERL_NIF_TERM task_term = tasks[task_idx].task_term;
 
-        /* Extract: {CallerPid, Ref, Module, Func, Args, Kwargs} */
+        /* Extract: {CallerPid, Ref, Module, Func, Args, Kwargs} or
+         *          {CallerPid, Ref, Module, Func, Args, Kwargs, EnvRef} */
         int arity;
         const ERL_NIF_TERM *tuple_elems;
-        if (!enif_get_tuple(term_env, task_term, &arity, &tuple_elems) || arity != 6) {
+        if (!enif_get_tuple(term_env, task_term, &arity, &tuple_elems) ||
+            (arity != 6 && arity != 7)) {
             enif_free_env(term_env);
             continue;
         }
@@ -2810,6 +2882,16 @@ ERL_NIF_TERM nif_process_ready_tasks(ErlNifEnv *env, int argc,
             continue;
         }
 
+        /* Check for env resource (7-tuple) */
+        py_env_resource_t *task_env = NULL;
+        if (arity == 7) {
+            /* Try to get env resource from tuple_elems[6] */
+            if (!enif_get_resource(term_env, tuple_elems[6],
+                                   get_env_resource_type(), (void **)&task_env)) {
+                task_env = NULL;  /* Invalid env ref, continue without it */
+            }
+        }
+
         /* Convert module/func to C strings */
         char *module_name = enif_alloc(module_bin.size + 1);
         char *func_name = enif_alloc(func_bin.size + 1);
@@ -2824,11 +2906,27 @@ ERL_NIF_TERM nif_process_ready_tasks(ErlNifEnv *env, int argc,
         memcpy(func_name, func_bin.data, func_bin.size);
         func_name[func_bin.size] = '\0';
 
-        /* Look up namespace for caller process (only exists if they called exec/eval) */
+        /* Look up namespace for caller process (used for reentrant calls) */
         process_namespace_t *ns = lookup_process_namespace(loop, &caller_pid);
 
-        /* Look up function (checks process namespace for __main__, then cache/import) */
-        PyObject *func = get_function_for_task(loop, ns, module_name, func_name);
+        /* Look up function - check task_env first, then process namespace, then import */
+        PyObject *func = NULL;
+
+        /* First, check the passed env's globals (from py:exec) */
+        if (task_env != NULL && task_env->globals != NULL) {
+            if (strcmp(module_name, "__main__") == 0 ||
+                strcmp(module_name, "_process_") == 0) {
+                func = PyDict_GetItemString(task_env->globals, func_name);
+                if (func != NULL) {
+                    Py_INCREF(func);
+                }
+            }
+        }
+
+        /* Fallback to process namespace and cache/import */
+        if (func == NULL) {
+            func = get_function_for_task(loop, ns, module_name, func_name);
+        }
 
         enif_free(module_name);
         enif_free(func_name);

c_src/py_event_loop.h

@@ -624,6 +624,18 @@ ERL_NIF_TERM nif_dispatch_sleep_complete(ErlNifEnv *env, int argc,
 ERL_NIF_TERM nif_submit_task(ErlNifEnv *env, int argc,
                               const ERL_NIF_TERM argv[]);
 
+/**
+ * @brief Submit an async task with process-local env (thread-safe)
+ *
+ * Like submit_task but includes an env resource reference. The env's globals
+ * dict is used for function lookup, allowing functions defined via py:exec()
+ * to be called from the event loop.
+ *
+ * NIF: submit_task_with_env(LoopRef, CallerPid, Ref, Module, Func, Args, Kwargs, EnvRef) -> ok | {error, Reason}
+ */
+ERL_NIF_TERM nif_submit_task_with_env(ErlNifEnv *env, int argc,
+                                       const ERL_NIF_TERM argv[]);
+
 /**
  * @brief Process all pending tasks from the task queue
  *

c_src/py_nif.c

@@ -70,6 +70,11 @@ ErlNifResourceType *INLINE_CONTINUATION_RESOURCE_TYPE = NULL;
 /* Process-local Python environment resource type */
 ErlNifResourceType *PY_ENV_RESOURCE_TYPE = NULL;
 
+/* Getter for PY_ENV_RESOURCE_TYPE (used by py_event_loop.c) */
+ErlNifResourceType *get_env_resource_type(void) {
+    return PY_ENV_RESOURCE_TYPE;
+}
+
 _Atomic uint32_t g_context_id_counter = 1;
 
 /* ============================================================================
@@ -6534,6 +6539,7 @@ static ErlNifFunc nif_funcs[] = {
     {"event_loop_run_async", 7, nif_event_loop_run_async, ERL_NIF_DIRTY_JOB_IO_BOUND},
     /* Async task queue NIFs (uvloop-inspired) */
     {"submit_task", 7, nif_submit_task, 0},  /* Thread-safe, no GIL needed */
+    {"submit_task_with_env", 8, nif_submit_task_with_env, 0},  /* With process-local env */
     {"process_ready_tasks", 1, nif_process_ready_tasks, ERL_NIF_DIRTY_JOB_CPU_BOUND},
     {"event_loop_set_py_loop", 2, nif_event_loop_set_py_loop, 0},
     /* Per-process namespace NIFs */

c_src/py_nif.h

@@ -1313,6 +1313,12 @@ extern ErlNifResourceType *PY_CONTEXT_SUSPENDED_RESOURCE_TYPE;
 /** @brief Resource type for inline_continuation_t (inline scheduler continuation) */
 extern ErlNifResourceType *INLINE_CONTINUATION_RESOURCE_TYPE;
 
+/** @brief Resource type for py_env_resource_t (process-local Python environment) */
+extern ErlNifResourceType *PY_ENV_RESOURCE_TYPE;
+
+/** @brief Get the PY_ENV_RESOURCE_TYPE (for use by other modules) */
+ErlNifResourceType *get_env_resource_type(void);
+
 /** @brief Atomic counter for unique interpreter IDs */
 extern _Atomic uint32_t g_context_id_counter;
 

src/py_event_loop.erl

@@ -166,9 +166,32 @@ create_task(Module, Func, Args, Kwargs) ->
     Caller = self(),
     ModuleBin = py_util:to_binary(Module),
     FuncBin = py_util:to_binary(Func),
-    ok = py_nif:submit_task(LoopRef, Caller, Ref, ModuleBin, FuncBin, Args, Kwargs),
+    %% Check if there's a process-local env (from py:exec) and use it
+    ok = case get_process_env() of
+        undefined ->
+            py_nif:submit_task(LoopRef, Caller, Ref, ModuleBin, FuncBin, Args, Kwargs);
+        EnvRef ->
+            py_nif:submit_task_with_env(LoopRef, Caller, Ref, ModuleBin, FuncBin, Args, Kwargs, EnvRef)
+    end,
     Ref.
 
+%% @doc Get the process-local env reference if available.
+%% Returns the first env found (typically there's only one per process).
+-spec get_process_env() -> reference() | undefined.
+get_process_env() ->
+    case get(py_local_env) of
+        undefined -> undefined;
+        Envs when is_map(Envs) ->
+            %% Return any env (for __main__ lookup, any env works)
+            case maps:values(Envs) of
+                [EnvRef | _] -> EnvRef;
+                [] -> undefined
+            end;
+        EnvRef when is_reference(EnvRef) ->
+            %% Legacy single-env format
+            EnvRef
+    end.
+
 %% @doc Wait for an async task result.
 %%
 %% Blocks until the result is received or timeout is reached.
@@ -207,6 +230,8 @@ spawn_task(Module, Func, Args) ->
 spawn_task(Module, Func, Args, Kwargs) ->
     {ok, LoopRef} = get_loop(),
     Ref = make_ref(),
+    %% Get env from caller's process BEFORE spawning receiver
+    CallerEnv = get_process_env(),
     %% Spawn a process that will receive and discard the result
     Receiver = erlang:spawn(fun() ->
         receive
@@ -218,7 +243,13 @@ spawn_task(Module, Func, Args, Kwargs) ->
     end),
     ModuleBin = py_util:to_binary(Module),
     FuncBin = py_util:to_binary(Func),
-    ok = py_nif:submit_task(LoopRef, Receiver, Ref, ModuleBin, FuncBin, Args, Kwargs),
+    %% Submit task with caller's env if available
+    ok = case CallerEnv of
+        undefined ->
+            py_nif:submit_task(LoopRef, Receiver, Ref, ModuleBin, FuncBin, Args, Kwargs);
+        EnvRef ->
+            py_nif:submit_task_with_env(LoopRef, Receiver, Ref, ModuleBin, FuncBin, Args, Kwargs, EnvRef)
+    end,
     ok.
 
 %% ============================================================================

src/py_nif.erl

@@ -102,6 +102,7 @@
     event_loop_run_async/7,
     %% Async task queue NIFs (uvloop-inspired)
     submit_task/7,
+    submit_task_with_env/8,
     process_ready_tasks/1,
     event_loop_set_py_loop/2,
     %% Per-process namespace NIFs
@@ -769,6 +770,16 @@ event_loop_run_async(_LoopRef, _CallerPid, _Ref, _Module, _Func, _Args, _Kwargs)
 submit_task(_LoopRef, _CallerPid, _Ref, _Module, _Func, _Args, _Kwargs) ->
     ?NIF_STUB.
 
+%% @doc Submit an async task with process-local env.
+%%
+%% Like submit_task but includes an env resource reference. The env's globals
+%% dict is used for function lookup, allowing functions defined via py:exec()
+%% to be called from the event loop.
+-spec submit_task_with_env(reference(), pid(), reference(), binary(), binary(), list(), map(), reference()) ->
+    ok | {error, term()}.
+submit_task_with_env(_LoopRef, _CallerPid, _Ref, _Module, _Func, _Args, _Kwargs, _EnvRef) ->
+    ?NIF_STUB.
+
 %% @doc Process all pending tasks from the task queue.
 %%
 %% Called by the event worker when it receives 'task_ready' message.