%% 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. %%% @doc Python worker process. %%% %%% Each worker maintains its own Python execution context. Workers %%% receive requests from the pool and execute Python code, sending %%% results back to callers. %%% %%% The NIF functions use ERL_NIF_DIRTY_JOB_IO_BOUND to run on dirty %%% I/O schedulers, so the worker process itself runs on a normal scheduler. %%% %%% @private -module(py_worker). -export([ start_link/0, init/1 ]). %% Timeout for checking shutdown (ms) -define(RECV_TIMEOUT, 1000). %%% ============================================================================ %%% API %%% ============================================================================ -spec start_link() -> {ok, pid()}. start_link() -> Pid = spawn_link(?MODULE, init, [self()]), receive {Pid, ready} -> {ok, Pid}; {Pid, {error, Reason}} -> {error, Reason} after 10000 -> exit(Pid, kill), {error, timeout} end. %%% ============================================================================ %%% Worker Process %%% ============================================================================ init(Parent) -> %% Create worker context case py_nif:worker_new() of {ok, WorkerRef} -> %% Spawn a separate callback handler process CallbackHandler = spawn_link(fun() -> callback_handler_loop() end), %% Set up callback handler with the separate process case py_nif:set_callback_handler(WorkerRef, CallbackHandler) of {ok, CallbackFd} -> CallbackHandler ! {set_fd, CallbackFd}, Parent ! {self(), ready}, loop(WorkerRef, Parent, CallbackFd); {error, Reason} -> exit(CallbackHandler, kill), Parent ! {self(), {error, Reason}} end; {error, Reason} -> Parent ! {self(), {error, Reason}} end. %% Separate process that handles callbacks from Python callback_handler_loop() -> receive {set_fd, CallbackFd} -> callback_handler_loop(CallbackFd) end. callback_handler_loop(CallbackFd) -> receive {erlang_callback, _CallbackId, FuncName, Args} -> handle_callback(CallbackFd, FuncName, Args), callback_handler_loop(CallbackFd); shutdown -> ok; _Other -> callback_handler_loop(CallbackFd) end. loop(WorkerRef, _Parent, _CallbackFd) -> receive {py_request, Request} -> handle_request(WorkerRef, Request), loop(WorkerRef, _Parent, _CallbackFd); shutdown -> py_nif:worker_destroy(WorkerRef), ok; _Other -> loop(WorkerRef, _Parent, _CallbackFd) end. %%% ============================================================================ %%% Request Handling %%% ============================================================================ %% Call with timeout handle_request(WorkerRef, {call, Ref, Caller, Module, Func, Args, Kwargs, TimeoutMs}) -> ModuleBin = to_binary(Module), FuncBin = to_binary(Func), Result = py_nif:worker_call(WorkerRef, ModuleBin, FuncBin, Args, Kwargs, TimeoutMs), handle_call_result(Result, Ref, Caller); %% Call without timeout (backward compatible) handle_request(WorkerRef, {call, Ref, Caller, Module, Func, Args, Kwargs}) -> ModuleBin = to_binary(Module), FuncBin = to_binary(Func), Result = py_nif:worker_call(WorkerRef, ModuleBin, FuncBin, Args, Kwargs), handle_call_result(Result, Ref, Caller); %% Eval with timeout handle_request(WorkerRef, {eval, Ref, Caller, Code, Locals, TimeoutMs}) -> CodeBin = to_binary(Code), Result = py_nif:worker_eval(WorkerRef, CodeBin, Locals, TimeoutMs), handle_call_result(Result, Ref, Caller); %% Eval without timeout (backward compatible) handle_request(WorkerRef, {eval, Ref, Caller, Code, Locals}) -> CodeBin = to_binary(Code), Result = py_nif:worker_eval(WorkerRef, CodeBin, Locals), handle_call_result(Result, Ref, Caller); handle_request(WorkerRef, {exec, Ref, Caller, Code}) -> CodeBin = to_binary(Code), Result = py_nif:worker_exec(WorkerRef, CodeBin), send_response(Caller, Ref, Result); handle_request(WorkerRef, {stream, Ref, Caller, Module, Func, Args, Kwargs}) -> ModuleBin = to_binary(Module), FuncBin = to_binary(Func), %% For streaming, we call a special function that yields chunks case py_nif:worker_call(WorkerRef, ModuleBin, FuncBin, Args, Kwargs) of {ok, {generator, GenRef}} -> stream_chunks(WorkerRef, GenRef, Ref, Caller); {ok, Value} -> %% Not a generator, send as single chunk Caller ! {py_chunk, Ref, Value}, Caller ! {py_end, Ref}; {error, _} = Error -> Caller ! {py_error, Ref, Error} end; handle_request(WorkerRef, {stream_eval, Ref, Caller, Code, Locals}) -> %% Evaluate expression and stream if result is a generator CodeBin = to_binary(Code), case py_nif:worker_eval(WorkerRef, CodeBin, Locals) of {ok, {generator, GenRef}} -> stream_chunks(WorkerRef, GenRef, Ref, Caller); {ok, Value} -> %% Not a generator, send as single value Caller ! {py_chunk, Ref, Value}, Caller ! {py_end, Ref}; {error, _} = Error -> Caller ! {py_error, Ref, Error} end. stream_chunks(WorkerRef, GenRef, Ref, Caller) -> case py_nif:worker_next(WorkerRef, GenRef) of {ok, {generator, NestedGen}} -> %% Nested generator - stream it inline stream_chunks(WorkerRef, NestedGen, Ref, Caller); {ok, Chunk} -> Caller ! {py_chunk, Ref, Chunk}, stream_chunks(WorkerRef, GenRef, Ref, Caller); {error, stop_iteration} -> Caller ! {py_end, Ref}; {error, Error} -> Caller ! {py_error, Ref, Error} end. %%% ============================================================================ %%% Suspended Callback Handling %%% %%% When Python code calls erlang.call(), it may return a suspension marker %%% instead of blocking. This allows the dirty scheduler to be freed while %%% the Erlang callback is executed. %%% ============================================================================ %% Handle the result of a worker_call - either normal result or suspended callback handle_call_result({suspended, _CallbackId, StateRef, {FuncName, CallArgs}}, Ref, Caller) -> %% Python code called erlang.call() - spawn a process to handle the callback. %% This prevents deadlock when the callback itself calls py:eval, which would %% otherwise block this worker while waiting for another worker (or even this %% same worker via round-robin). spawn_link(fun() -> handle_suspended_callback(StateRef, FuncName, CallArgs, Ref, Caller) end), ok; %% Don't block the worker - it can process other requests handle_call_result(Result, Ref, Caller) -> %% Normal result - send directly to caller send_response(Caller, Ref, Result). %% Execute a suspended callback and resume Python execution. %% This runs in a separate process to avoid blocking the worker. handle_suspended_callback(StateRef, FuncName, CallArgs, Ref, Caller) -> %% Convert Args from tuple/list to list ArgsList = case CallArgs of T when is_tuple(T) -> tuple_to_list(T); L when is_list(L) -> L; _ -> [CallArgs] end, %% Execute the registered Erlang function %% This can call py:eval, py:call, etc. without deadlocking CallbackResult = py_callback:execute(FuncName, ArgsList), %% Encode result as binary (status byte + python repr) ResultBinary = case CallbackResult of {ok, Value} -> ValueRepr = term_to_python_repr(Value), <<0, ValueRepr/binary>>; {error, {not_found, Name}} -> ErrMsg = iolist_to_binary(io_lib:format("Function '~s' not registered", [Name])), <<1, ErrMsg/binary>>; {error, {Class, Reason, _Stack}} -> ErrMsg = iolist_to_binary(io_lib:format("~p: ~p", [Class, Reason])), <<1, ErrMsg/binary>> end, %% Resume Python execution with the callback result %% The NIF parses the result using Python's ast.literal_eval and returns %% {ok, ParsedTerm} or {error, Reason} FinalResult = py_nif:resume_callback(StateRef, ResultBinary), %% Handle the final result (could be another suspension for nested callbacks) %% Note: for nested callbacks, this will spawn another process recursively forward_final_result(FinalResult, Ref, Caller). %% Forward the final result to the caller, handling nested suspensions forward_final_result({suspended, _CallbackId, StateRef, {FuncName, CallArgs}}, Ref, Caller) -> %% Another suspension - handle it recursively handle_suspended_callback(StateRef, FuncName, CallArgs, Ref, Caller); forward_final_result(Result, Ref, Caller) -> %% Final result - send to the original caller send_response(Caller, Ref, Result). %%% ============================================================================ %%% Callback Handling %%% ============================================================================ handle_callback(CallbackFd, FuncName, Args) -> %% Convert Args from tuple to list if needed ArgsList = case Args of T when is_tuple(T) -> tuple_to_list(T); L when is_list(L) -> L; _ -> [Args] end, %% Execute the registered function 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), Response = <<0, ResultStr/binary>>, py_nif:send_callback_response(CallbackFd, Response); {error, {not_found, Name}} -> ErrMsg = iolist_to_binary(io_lib:format("Function '~s' not registered", [Name])), Response = <<1, ErrMsg/binary>>, py_nif:send_callback_response(CallbackFd, Response); {error, {Class, Reason, _Stack}} -> ErrMsg = iolist_to_binary(io_lib:format("~p: ~p", [Class, Reason])), Response = <<1, ErrMsg/binary>>, py_nif:send_callback_response(CallbackFd, Response) end. %% 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) -> %% 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). %%% ============================================================================ %%% Internal Functions %%% ============================================================================ send_response(Caller, Ref, Result) -> py_util:send_response(Caller, Ref, Result). to_binary(Term) -> py_util:to_binary(Term).