Merge pull request #55 from benoitc/feature/owngil-imports-paths

Apply imports/paths to OWN_GIL sessions

Author: benoitc

c_src/py_nif.c

@@ -6881,6 +6881,132 @@ static ERL_NIF_TERM nif_owngil_destroy_session(ErlNifEnv *env, int argc,
     return ATOM_OK;
 }
 
+/**
+ * @brief NIF: Apply imports to OWN_GIL session
+ *
+ * Imports modules into the worker's sys.modules.
+ * Args: WorkerId, HandleId, Imports (list of {ModuleBin, FuncBin | all})
+ */
+static ERL_NIF_TERM nif_owngil_apply_imports(ErlNifEnv *env, int argc,
+                                              const ERL_NIF_TERM argv[]) {
+    if (argc != 3) {
+        return enif_make_badarg(env);
+    }
+
+    if (!subinterp_thread_pool_is_ready()) {
+        return ATOM_OK;  /* Silently succeed if pool not ready */
+    }
+
+    unsigned int worker_id;
+    ErlNifUInt64 handle_id;
+
+    if (!enif_get_uint(env, argv[0], &worker_id) ||
+        !enif_get_uint64(env, argv[1], &handle_id)) {
+        return enif_make_badarg(env);
+    }
+
+    if (worker_id >= (unsigned int)g_thread_pool.num_workers) {
+        return ATOM_OK;  /* Invalid worker, silently succeed */
+    }
+
+    /* Serialize imports list to ETF */
+    ErlNifBinary payload_bin;
+    if (!enif_term_to_binary(env, argv[2], &payload_bin)) {
+        return ATOM_OK;  /* Serialization failed, silently succeed */
+    }
+
+    subinterp_thread_worker_t *w = &g_thread_pool.workers[worker_id];
+
+    pthread_mutex_lock(&w->dispatch_mutex);
+
+    uint64_t request_id = atomic_fetch_add(&g_thread_pool.next_request_id, 1);
+    owngil_header_t header = {
+        .magic = OWNGIL_MAGIC,
+        .version = OWNGIL_PROTOCOL_VERSION,
+        .msg_type = MSG_REQUEST,
+        .req_type = REQ_APPLY_IMPORTS,
+        .request_id = request_id,
+        .handle_id = handle_id,
+        .payload_len = payload_bin.size,
+    };
+
+    /* Write header and payload */
+    if (write(w->cmd_pipe[1], &header, sizeof(header)) == sizeof(header)) {
+        write(w->cmd_pipe[1], payload_bin.data, payload_bin.size);
+        /* Wait for response */
+        owngil_header_t resp;
+        read(w->result_pipe[0], &resp, sizeof(resp));
+    }
+
+    enif_release_binary(&payload_bin);
+    pthread_mutex_unlock(&w->dispatch_mutex);
+
+    return ATOM_OK;
+}
+
+/**
+ * @brief NIF: Apply paths to OWN_GIL session
+ *
+ * Adds paths to the worker's sys.path.
+ * Args: WorkerId, HandleId, Paths (list of path binaries)
+ */
+static ERL_NIF_TERM nif_owngil_apply_paths(ErlNifEnv *env, int argc,
+                                            const ERL_NIF_TERM argv[]) {
+    if (argc != 3) {
+        return enif_make_badarg(env);
+    }
+
+    if (!subinterp_thread_pool_is_ready()) {
+        return ATOM_OK;  /* Silently succeed if pool not ready */
+    }
+
+    unsigned int worker_id;
+    ErlNifUInt64 handle_id;
+
+    if (!enif_get_uint(env, argv[0], &worker_id) ||
+        !enif_get_uint64(env, argv[1], &handle_id)) {
+        return enif_make_badarg(env);
+    }
+
+    if (worker_id >= (unsigned int)g_thread_pool.num_workers) {
+        return ATOM_OK;  /* Invalid worker, silently succeed */
+    }
+
+    /* Serialize paths list to ETF */
+    ErlNifBinary payload_bin;
+    if (!enif_term_to_binary(env, argv[2], &payload_bin)) {
+        return ATOM_OK;  /* Serialization failed, silently succeed */
+    }
+
+    subinterp_thread_worker_t *w = &g_thread_pool.workers[worker_id];
+
+    pthread_mutex_lock(&w->dispatch_mutex);
+
+    uint64_t request_id = atomic_fetch_add(&g_thread_pool.next_request_id, 1);
+    owngil_header_t header = {
+        .magic = OWNGIL_MAGIC,
+        .version = OWNGIL_PROTOCOL_VERSION,
+        .msg_type = MSG_REQUEST,
+        .req_type = REQ_APPLY_PATHS,
+        .request_id = request_id,
+        .handle_id = handle_id,
+        .payload_len = payload_bin.size,
+    };
+
+    /* Write header and payload */
+    if (write(w->cmd_pipe[1], &header, sizeof(header)) == sizeof(header)) {
+        write(w->cmd_pipe[1], payload_bin.data, payload_bin.size);
+        /* Wait for response */
+        owngil_header_t resp;
+        read(w->result_pipe[0], &resp, sizeof(resp));
+    }
+
+    enif_release_binary(&payload_bin);
+    pthread_mutex_unlock(&w->dispatch_mutex);
+
+    return ATOM_OK;
+}
+
 #else /* !HAVE_SUBINTERPRETERS */
 
 /* Stub implementations for Python < 3.12 */
@@ -6980,6 +7106,18 @@ static ERL_NIF_TERM nif_owngil_destroy_session(ErlNifEnv *env, int argc,
     return ATOM_OK;
 }
 
+static ERL_NIF_TERM nif_owngil_apply_imports(ErlNifEnv *env, int argc,
+                                              const ERL_NIF_TERM argv[]) {
+    (void)argc; (void)argv;
+    return ATOM_OK;
+}
+
+static ERL_NIF_TERM nif_owngil_apply_paths(ErlNifEnv *env, int argc,
+                                            const ERL_NIF_TERM argv[]) {
+    (void)argc; (void)argv;
+    return ATOM_OK;
+}
+
 #endif /* HAVE_SUBINTERPRETERS */
 
 /* ============================================================================
@@ -7239,6 +7377,8 @@ static ErlNifFunc nif_funcs[] = {
     {"owngil_create_session", 1, nif_owngil_create_session, 0},
     {"owngil_submit_task", 7, nif_owngil_submit_task, 0},
     {"owngil_destroy_session", 2, nif_owngil_destroy_session, 0},
+    {"owngil_apply_imports", 3, nif_owngil_apply_imports, 0},
+    {"owngil_apply_paths", 3, nif_owngil_apply_paths, 0},
 
     /* Execution mode info */
     {"execution_mode", 0, nif_execution_mode, 0},

c_src/py_subinterp_thread.c

@@ -425,6 +425,105 @@ static void *worker_thread_main(void *arg) {
             continue;
         }
 
+        /* Handle apply imports - imports modules into sys.modules */
+        if (header.req_type == REQ_APPLY_IMPORTS) {
+            /* Payload is ETF list of {ModuleBin, FuncBin | all} tuples */
+            if (payload != NULL && header.payload_len > 0) {
+                ErlNifEnv *tmp_env = enif_alloc_env();
+                if (tmp_env != NULL) {
+                    ERL_NIF_TERM imports_list;
+                    if (enif_binary_to_term(tmp_env, payload, header.payload_len,
+                                            &imports_list, 0) != 0) {
+                        ERL_NIF_TERM head, tail = imports_list;
+                        int arity;
+                        const ERL_NIF_TERM *tuple;
+                        while (enif_get_list_cell(tmp_env, tail, &head, &tail)) {
+                            if (enif_get_tuple(tmp_env, head, &arity, &tuple) && arity == 2) {
+                                ErlNifBinary module_bin;
+                                if (enif_inspect_binary(tmp_env, tuple[0], &module_bin)) {
+                                    char *module_name = enif_alloc(module_bin.size + 1);
+                                    if (module_name != NULL) {
+                                        memcpy(module_name, module_bin.data, module_bin.size);
+                                        module_name[module_bin.size] = '\0';
+                                        /* Skip __main__ */
+                                        if (strcmp(module_name, "__main__") != 0) {
+                                            PyObject *mod = PyImport_ImportModule(module_name);
+                                            if (mod != NULL) {
+                                                Py_DECREF(mod);
+                                            } else {
+                                                PyErr_Clear();
+                                            }
+                                        }
+                                        enif_free(module_name);
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    enif_free_env(tmp_env);
+                }
+            }
+            PyEval_SaveThread();
+            /* Send OK response */
+            owngil_header_t resp = {
+                .magic = OWNGIL_MAGIC,
+                .version = OWNGIL_PROTOCOL_VERSION,
+                .msg_type = MSG_RESPONSE,
+                .request_id = header.request_id,
+                .payload_len = 0,
+            };
+            write_full(w->result_pipe[1], &resp, sizeof(resp));
+            free(payload);
+            continue;
+        }
+
+        /* Handle apply paths - add paths to sys.path */
+        if (header.req_type == REQ_APPLY_PATHS) {
+            /* Payload is ETF list of path binaries */
+            if (payload != NULL && header.payload_len > 0) {
+                ErlNifEnv *tmp_env = enif_alloc_env();
+                if (tmp_env != NULL) {
+                    ERL_NIF_TERM paths_list;
+                    if (enif_binary_to_term(tmp_env, payload, header.payload_len,
+                                            &paths_list, 0) != 0) {
+                        PyObject *sys_path = PySys_GetObject("path");
+                        if (sys_path != NULL && PyList_Check(sys_path)) {
+                            ERL_NIF_TERM head, tail = paths_list;
+                            /* Insert in reverse order so first path ends up first */
+                            while (enif_get_list_cell(tmp_env, tail, &head, &tail)) {
+                                ErlNifBinary path_bin;
+                                if (enif_inspect_binary(tmp_env, head, &path_bin)) {
+                                    PyObject *path_str = PyUnicode_FromStringAndSize(
+                                        (const char *)path_bin.data, path_bin.size);
+                                    if (path_str != NULL) {
+                                        /* Check if path already in sys.path */
+                                        int contains = PySequence_Contains(sys_path, path_str);
+                                        if (contains == 0) {
+                                            PyList_Insert(sys_path, 0, path_str);
+                                        }
+                                        Py_DECREF(path_str);
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    enif_free_env(tmp_env);
+                }
+            }
+            PyEval_SaveThread();
+            /* Send OK response */
+            owngil_header_t resp = {
+                .magic = OWNGIL_MAGIC,
+                .version = OWNGIL_PROTOCOL_VERSION,
+                .msg_type = MSG_RESPONSE,
+                .request_id = header.request_id,
+                .payload_len = 0,
+            };
+            write_full(w->result_pipe[1], &resp, sizeof(resp));
+            free(payload);
+            continue;
+        }
+
         /* Find namespace for this handle */
         subinterp_namespace_t *ns = worker_find_namespace(w, header.handle_id);
         if (ns == NULL) {

c_src/py_subinterp_thread.h

@@ -92,6 +92,8 @@ typedef enum {
     REQ_ASYNC_CALL      = 5,   /**< Async - response via erlang.send() */
     REQ_CREATE_NS       = 10,  /**< Create namespace for handle */
     REQ_DESTROY_NS      = 11,  /**< Destroy namespace for handle */
+    REQ_APPLY_IMPORTS   = 12,  /**< Apply imports to namespace's sys.modules */
+    REQ_APPLY_PATHS     = 13,  /**< Apply paths to namespace's sys.path */
     REQ_SHUTDOWN        = 99,  /**< Shutdown the worker */
 } owngil_req_type_t;
 

src/py_event_loop_pool.erl

@@ -44,7 +44,9 @@
     %% Per-process namespace API
     exec/1, exec/2,
     eval/1, eval/2,
-    get_all_loops/0
+    get_all_loops/0,
+    get_all_sessions/0,
+    is_owngil_enabled/0
 ]).
 
 %% Legacy API
@@ -349,6 +351,24 @@ get_all_loops() ->
             {ok, [element(Idx, Loops) || Idx <- lists:seq(1, N)]}
     end.
 
+%% @doc Get all OWN_GIL sessions in the pool.
+%%
+%% Returns a list of {WorkerId, HandleId} tuples for all active sessions.
+%% Used by py_import to apply imports/paths to running OWN_GIL sessions.
+-spec get_all_sessions() -> [{non_neg_integer(), non_neg_integer()}].
+get_all_sessions() ->
+    case get_sessions_table() of
+        undefined -> [];
+        Tid ->
+            ets:foldl(
+                fun(#owngil_session{worker_id = W, handle_id = H}, Acc) ->
+                    [{W, H} | Acc]
+                end,
+                [],
+                Tid
+            )
+    end.
+
 %%% ============================================================================
 %%% Legacy API
 %%% ============================================================================

src/py_import.erl

@@ -387,7 +387,7 @@ load_config() ->
     ),
     ok.
 
-%% @private Apply import to all running interpreters (contexts + event loops)
+%% @private Apply import to all running interpreters (contexts + event loops + OWN_GIL sessions)
 apply_import_to_interpreters(ModuleBin) ->
     Imports = [{ModuleBin, all}],
     %% Apply to all contexts
@@ -418,9 +418,21 @@ apply_import_to_interpreters(ModuleBin) ->
             );
         _ -> ok
     end,
+    %% Apply to all OWN_GIL sessions (if enabled)
+    case py_event_loop_pool:is_owngil_enabled() of
+        true ->
+            lists:foreach(
+                fun({WorkerId, HandleId}) ->
+                    catch py_nif:owngil_apply_imports(WorkerId, HandleId, Imports)
+                end,
+                py_event_loop_pool:get_all_sessions()
+            );
+        false ->
+            ok
+    end,
     ok.
 
-%% @private Apply path to all running interpreters (contexts + event loops)
+%% @private Apply path to all running interpreters (contexts + event loops + OWN_GIL sessions)
 apply_path_to_interpreters(PathBin) ->
     Paths = [PathBin],
     %% Apply to all contexts
@@ -451,6 +463,18 @@ apply_path_to_interpreters(PathBin) ->
             );
         _ -> ok
     end,
+    %% Apply to all OWN_GIL sessions (if enabled)
+    case py_event_loop_pool:is_owngil_enabled() of
+        true ->
+            lists:foreach(
+                fun({WorkerId, HandleId}) ->
+                    catch py_nif:owngil_apply_paths(WorkerId, HandleId, Paths)
+                end,
+                py_event_loop_pool:get_all_sessions()
+            );
+        false ->
+            ok
+    end,
     ok.
 
 %% @private Get all context pids from all pools

src/py_nif.erl

@@ -79,6 +79,8 @@
     owngil_create_session/1,
     owngil_submit_task/7,
     owngil_destroy_session/2,
+    owngil_apply_imports/3,
+    owngil_apply_paths/3,
     %% Execution mode info
     execution_mode/0,
     num_executors/0,
@@ -645,6 +647,18 @@ owngil_submit_task(_WorkerId, _HandleId, _CallerPid, _Ref, _Module, _Func, _Args
 owngil_destroy_session(_WorkerId, _HandleId) ->
     ?NIF_STUB.
 
+%% @doc Apply imports to an OWN_GIL session.
+%% Imports modules into the worker's sys.modules.
+-spec owngil_apply_imports(non_neg_integer(), non_neg_integer(), [{binary(), binary() | all}]) -> ok | {error, term()}.
+owngil_apply_imports(_WorkerId, _HandleId, _Imports) ->
+    ?NIF_STUB.
+
+%% @doc Apply paths to an OWN_GIL session.
+%% Adds paths to the worker's sys.path.
+-spec owngil_apply_paths(non_neg_integer(), non_neg_integer(), [binary()]) -> ok | {error, term()}.
+owngil_apply_paths(_WorkerId, _HandleId, _Paths) ->
+    ?NIF_STUB.
+
 %%% ============================================================================
 %%% Execution Mode Info
 %%% ============================================================================