|
| 1 | +// Async NIF machinery. See docs/planning/async-worker-exploration.md |
| 2 | +// for the design rationale. |
| 3 | +// |
| 4 | +// The pattern: a NIF captures the caller PID via enif_self, mints a |
| 5 | +// fresh ref in a process-independent env, enqueues a task onto the |
| 6 | +// target WorkerThread, and returns the ref synchronously. The worker |
| 7 | +// thread runs the task, builds the reply term in the msg_env, and |
| 8 | +// posts {ref, {:ok, payload}} or {ref, {:error, reason}} back to |
| 9 | +// the caller via enif_send. The Elixir side awaits the reply with a |
| 10 | +// pattern-match receive in `Emily.Native.Async.call/1`. |
| 11 | +// |
| 12 | +// Three invariants kept by this code: |
| 13 | +// |
| 14 | +// 1. enif_self must be called on the scheduler thread. The NIF |
| 15 | +// captures the ErlNifPid by value into the lambda; the worker |
| 16 | +// thread (non-scheduler) must never call enif_self itself. |
| 17 | +// |
| 18 | +// 2. enif_send with a non-NULL msg_env does not transfer env |
| 19 | +// ownership. The worker calls enif_free_env after the send |
| 20 | +// unconditionally — success invalidates the env but the caller |
| 21 | +// still owns the env object. |
| 22 | +// |
| 23 | +// 3. Resource terms built via enif_make_resource(msg_env, ptr) |
| 24 | +// internally bump the resource's refcount. ResourcePtrs captured |
| 25 | +// by the lambda release their refs on lambda exit, but the term |
| 26 | +// held by msg_env keeps the resource alive until the term is |
| 27 | +// delivered + GC'd on the receiver or msg_env is freed (if the |
| 28 | +// receiver PID is dead). |
| 29 | + |
| 30 | +#pragma once |
| 31 | + |
| 32 | +#include "worker.hpp" |
| 33 | + |
| 34 | +#include <fine.hpp> |
| 35 | +#include <mlx/mlx.h> |
| 36 | + |
| 37 | +#include <cstring> |
| 38 | +#include <exception> |
| 39 | +#include <stdexcept> |
| 40 | +#include <utility> |
| 41 | + |
| 42 | +namespace emily { |
| 43 | + |
| 44 | +namespace mx = mlx::core; |
| 45 | + |
| 46 | +namespace __async { |
| 47 | + |
| 48 | +// Build a binary term in msg_env from a null-terminated C string. |
| 49 | +inline ERL_NIF_TERM make_binary_from_cstr(ErlNifEnv *msg_env, const char *s) { |
| 50 | + size_t len = std::strlen(s); |
| 51 | + ERL_NIF_TERM term; |
| 52 | + unsigned char *data = enif_make_new_binary(msg_env, len, &term); |
| 53 | + std::memcpy(data, s, len); |
| 54 | + return term; |
| 55 | +} |
| 56 | + |
| 57 | +// Build an error reply term in msg_env, classifying the exception |
| 58 | +// type. Matches fine::nif_impl's sync catch ladder so the Elixir |
| 59 | +// side can raise the same exception classes: |
| 60 | +// std::invalid_argument -> {:argument, message} |
| 61 | +// std::runtime_error -> {:runtime, message} |
| 62 | +// std::exception -> {:runtime, message} |
| 63 | +// ... (any) -> :unknown |
| 64 | +inline ERL_NIF_TERM |
| 65 | +error_reason_from_current_exception(ErlNifEnv *msg_env) { |
| 66 | + try { |
| 67 | + throw; // re-raise the current exception to classify it |
| 68 | + } catch (const std::invalid_argument &e) { |
| 69 | + return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "argument"), |
| 70 | + make_binary_from_cstr(msg_env, e.what())); |
| 71 | + } catch (const std::runtime_error &e) { |
| 72 | + return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"), |
| 73 | + make_binary_from_cstr(msg_env, e.what())); |
| 74 | + } catch (const std::exception &e) { |
| 75 | + return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"), |
| 76 | + make_binary_from_cstr(msg_env, e.what())); |
| 77 | + } catch (...) { |
| 78 | + return enif_make_atom(msg_env, "unknown"); |
| 79 | + } |
| 80 | +} |
| 81 | + |
| 82 | +} // namespace __async |
| 83 | + |
| 84 | +// Run `build_payload` on the worker thread of `w` and post the |
| 85 | +// result back to the caller PID as a message. Returns a fresh ref |
| 86 | +// synchronously; the caller awaits the reply via |
| 87 | +// `Emily.Native.Async.call/1`. |
| 88 | +// |
| 89 | +// `build_payload` signature: |
| 90 | +// (mx::Stream &stream, ErlNifEnv *msg_env) -> ERL_NIF_TERM |
| 91 | +// |
| 92 | +// The returned ERL_NIF_TERM is the *payload* term; this helper wraps |
| 93 | +// it as `{ref, {:ok, payload}}`. Exceptions thrown by the lambda are |
| 94 | +// caught and posted as `{ref, {:error, reason}}` where `reason` is |
| 95 | +// `{:argument | :runtime, binary}` or `:unknown`. |
| 96 | +template <typename BuildPayload> |
| 97 | +fine::Term async_reply(ErlNifEnv *env, |
| 98 | + fine::ResourcePtr<WorkerThread> w, |
| 99 | + BuildPayload &&build_payload) { |
| 100 | + ErlNifPid caller; |
| 101 | + enif_self(env, &caller); |
| 102 | + |
| 103 | + // Mint the ref in a durable (process-independent) env so it |
| 104 | + // survives the NIF return. Copy into the caller's env for the |
| 105 | + // synchronous return value. |
| 106 | + ErlNifEnv *msg_env = enif_alloc_env(); |
| 107 | + ERL_NIF_TERM ref_in_msg = enif_make_ref(msg_env); |
| 108 | + ERL_NIF_TERM ref_to_return = enif_make_copy(env, ref_in_msg); |
| 109 | + |
| 110 | + try { |
| 111 | + w->run_async([msg_env, ref_in_msg, caller, |
| 112 | + build_payload = std::forward<BuildPayload>(build_payload)] |
| 113 | + (mx::Stream &s) mutable { |
| 114 | + ERL_NIF_TERM reply; |
| 115 | + try { |
| 116 | + ERL_NIF_TERM payload = build_payload(s, msg_env); |
| 117 | + ERL_NIF_TERM ok_tuple = enif_make_tuple2( |
| 118 | + msg_env, enif_make_atom(msg_env, "ok"), payload); |
| 119 | + reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple); |
| 120 | + } catch (...) { |
| 121 | + reply = enif_make_tuple2( |
| 122 | + msg_env, ref_in_msg, |
| 123 | + enif_make_tuple2( |
| 124 | + msg_env, enif_make_atom(msg_env, "error"), |
| 125 | + __async::error_reason_from_current_exception(msg_env))); |
| 126 | + } |
| 127 | + |
| 128 | + // enif_send invalidates msg_env on success but does not take |
| 129 | + // ownership of the env object itself; free it unconditionally. |
| 130 | + enif_send(nullptr, &caller, msg_env, reply); |
| 131 | + enif_free_env(msg_env); |
| 132 | + }); |
| 133 | + } catch (...) { |
| 134 | + // Enqueue failed (worker stopped). Reclaim the env and rethrow |
| 135 | + // so fine::nif_impl surfaces the exception synchronously. |
| 136 | + enif_free_env(msg_env); |
| 137 | + throw; |
| 138 | + } |
| 139 | + |
| 140 | + return fine::Term(ref_to_return); |
| 141 | +} |
| 142 | + |
| 143 | +} // namespace emily |
0 commit comments