feat(worker): bounded queue, drop-with-error stop, off-scheduler join

RELEASE.md

@@ -1,5 +1,29 @@
+### Added
+
+- `Emily.Stream.close/1` stops a stream's worker thread deterministically
+  instead of waiting for garbage collection: queued operations are
+  cancelled (their callers get a `RuntimeError`), the in-flight op
+  finishes, and the OS thread is joined off the BEAM schedulers.
+- `config :emily, worker_queue_limit: N` (default `8192`) bounds the
+  per-worker async queue, and `config :emily, await_timeout: ms` (default
+  `:infinity`) sets an optional timeout for awaiting native results.
+
 ### Security
 
+- Worker-thread teardown no longer blocks a BEAM scheduler. The resource
+  destructor previously drained the worker's entire queue and joined the
+  OS thread inline, so collecting a busy stream during GC could stall a
+  scheduler. Workers are now joined off-scheduler by a dedicated reaper
+  (itself joined at NIF unload), and on stop the worker cancels its
+  queued tasks — replying `{:error, :stopped}` — instead of running them.
+
+- The async NIF worker queue is now bounded (`worker_queue_limit`, reject
+  when full) so a flood of operations can't grow it without limit and pin
+  host/GPU memory, and a stopped or dropped worker now replies
+  `{:error, :stopped}` to every queued caller instead of leaving it
+  blocked forever. `Emily.Native.worker_queue_depth/1` exposes the depth
+  for observability.
+
 - The dev/CI source-build path now refuses to trust an MLX install
   directory it doesn't own and keeps the build cache `0700`, so a shared
   or attacker-controlled `EMILY_CACHE` can't plant a `libmlx.a` that is

c_src/emily/async.hpp

@@ -110,19 +110,29 @@ fine::Term async_reply(ErlNifEnv *env,
   try {
     w->run_async([msg_env, ref_in_msg, caller,
                   build_payload = std::forward<BuildPayload>(build_payload)]
-                 (mx::Stream &s) mutable {
+                 (mx::Stream &s, bool cancelled) mutable {
       ERL_NIF_TERM reply;
-      try {
-        ERL_NIF_TERM payload = build_payload(s, msg_env);
-        ERL_NIF_TERM ok_tuple = enif_make_tuple2(
-            msg_env, enif_make_atom(msg_env, "ok"), payload);
-        reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple);
-      } catch (...) {
+      if (cancelled) {
+        // The worker was stopped before this task ran. Report
+        // {:error, :stopped} so the awaiting process unblocks instead
+        // of hanging on a reply that will never come.
         reply = enif_make_tuple2(
             msg_env, ref_in_msg,
-            enif_make_tuple2(
-                msg_env, enif_make_atom(msg_env, "error"),
-                __async::error_reason_from_current_exception(msg_env)));
+            enif_make_tuple2(msg_env, enif_make_atom(msg_env, "error"),
+                             enif_make_atom(msg_env, "stopped")));
+      } else {
+        try {
+          ERL_NIF_TERM payload = build_payload(s, msg_env);
+          ERL_NIF_TERM ok_tuple = enif_make_tuple2(
+              msg_env, enif_make_atom(msg_env, "ok"), payload);
+          reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple);
+        } catch (...) {
+          reply = enif_make_tuple2(
+              msg_env, ref_in_msg,
+              enif_make_tuple2(
+                  msg_env, enif_make_atom(msg_env, "error"),
+                  __async::error_reason_from_current_exception(msg_env)));
+        }
       }
 
       // enif_send invalidates msg_env on success but does not take

c_src/emily/worker.hpp

@@ -1,132 +1,237 @@
-// WorkerThread: dedicated OS thread owning an MLX stream.
+// WorkerThread: a dedicated OS thread owning an MLX stream, plus the
+// Reaper that joins finished worker threads off the BEAM schedulers.
 //
-// MLX uses thread-local CommandEncoders — a stream's encoder only
-// exists on the thread that created it. BEAM processes migrate
-// between OS threads, so we pin each MLX stream to a dedicated
-// thread and dispatch work to it via run_sync (promise/future) or
-// run_async (fire-and-forget, with the task posting its own reply
-// via enif_send — see emily/async.hpp).
+// MLX uses thread-local CommandEncoders — a stream's encoder only exists
+// on the thread that created it. BEAM processes migrate between OS
+// threads, so we pin each MLX stream to a dedicated thread and dispatch
+// work to it via run_async (fire-and-forget, with the task posting its
+// own reply via enif_send — see emily/async.hpp).
+//
+// Lifetime / teardown:
+//   - Per-worker mutable state lives in a heap `State` (shared_ptr), so a
+//     worker thread safely outlives the BEAM resource object.
+//   - On stop the worker *cancels* its queued tasks (each posts an
+//     {:error, :stopped} reply) instead of running them, then exits — so
+//     a join waits for at most the one in-flight op, never the backlog.
+//   - A resource destructor runs during BEAM GC and must never join() on
+//     a scheduler. Instead it signals stop and hands the thread to the
+//     Reaper: one long-lived thread that joins finished workers
+//     off-scheduler and is itself joined in the NIF unload callback
+//     (`Reaper::shutdown`, wired in worker_nif.cpp). The Reaper also
+//     tracks live workers so unload can stop+join stragglers.
 
 #pragma once
 
 #include <mlx/mlx.h>
 
 #include <condition_variable>
+#include <cstddef>
 #include <functional>
-#include <future>
+#include <map>
+#include <memory>
 #include <mutex>
 #include <queue>
 #include <stdexcept>
+#include <string>
 #include <thread>
-#include <type_traits>
+#include <utility>
+#include <vector>
 
 namespace emily {
 
 namespace mx = mlx::core;
 
-class WorkerThread {
+// A queued unit of work. `cancelled == true` means the worker is
+// stopping: the task must post its {:error, :stopped} reply (and free its
+// env) instead of running the MLX op. Invoked exactly once per task.
+using Task = std::function<void(mx::Stream &, bool /*cancelled*/)>;
+
+// Per-worker mutable state. Heap-allocated and shared between the BEAM
+// resource and the worker thread so it outlives whichever drops last.
+struct State {
+  explicit State(std::size_t queue_limit) : cap(queue_limit) {}
+
+  std::mutex mtx;
+  std::condition_variable cv;
+  std::queue<Task> queue;
+  bool stop = false;
+  bool ready = false;
+  std::size_t cap;
+  mx::Stream stream{0, mx::Device(mx::Device::DeviceType::gpu)};
+};
+
+inline void signal_stop(State &st) {
+  {
+    std::lock_guard<std::mutex> lock(st.mtx);
+    st.stop = true;
+  }
+  st.cv.notify_all();
+}
+
+// Joins worker threads off the BEAM schedulers. Singleton; its thread is
+// created lazily on first use and joined in the NIF unload callback.
+class Reaper {
 public:
-  WorkerThread() {
-    thread_ = std::thread(&WorkerThread::run, this);
-    std::unique_lock<std::mutex> lock(mtx_);
-    cv_.wait(lock, [this] { return ready_; });
+  static Reaper &instance() {
+    static Reaper reaper;
+    return reaper;
   }
 
-  ~WorkerThread() { stop(); }
+  // Track a newly-spawned worker as live. `key` identifies the worker
+  // (its State pointer); `st` is a non-owning handle used to signal stop
+  // at unload.
+  void track(State *key, std::thread thread, std::weak_ptr<State> st) {
+    std::lock_guard<std::mutex> lock(mtx_);
+    live_.emplace(key, Live{std::move(thread), std::move(st)});
+  }
 
-  // Enqueue a task without blocking. The caller is responsible for
-  // whatever side-effect the task performs (typically enif_send back
-  // to a caller PID captured at NIF entry; see emily/async.hpp).
-  // Exceptions thrown by the task are swallowed — the task owns
-  // error propagation because there is no future to carry an
-  // exception through.
-  template <typename F>
-  void run_async(F &&f) {
+  // The worker is going away (resource GC or explicit close): stop it and
+  // hand its thread off to be joined. Non-blocking — no join on the
+  // calling (scheduler) thread.
+  void retire(State *key) {
     {
       std::lock_guard<std::mutex> lock(mtx_);
-      if (stop_)
-        throw std::runtime_error("worker thread has been stopped");
-      queue_.push([f = std::forward<F>(f), this]() mutable {
-        try {
-          f(stream_);
-        } catch (...) {
-          // Swallow — the task owns error propagation.
-        }
-      });
+      auto it = live_.find(key);
+      if (it == live_.end()) {
+        return;  // already retired (e.g. double close)
+      }
+      if (auto st = it->second.st.lock()) {
+        signal_stop(*st);
+      }
+      dying_.push_back(std::move(it->second.thread));
+      live_.erase(it);
     }
-    cv_.notify_one();
+    cv_.notify_all();
   }
 
-  template <typename F>
-  auto run_sync(F &&f) -> decltype(f(std::declval<mx::Stream &>())) {
-    using R = decltype(f(std::declval<mx::Stream &>()));
-    auto p = std::make_shared<std::promise<R>>();
-    auto fut = p->get_future();
+  // NIF unload: stop and join every worker (live + dying), then stop the
+  // reaper thread. Blocking is fine here — the VM is tearing down.
+  void shutdown() {
     {
       std::lock_guard<std::mutex> lock(mtx_);
-      if (stop_)
-        throw std::runtime_error("worker thread has been stopped");
-      queue_.push([p, f = std::forward<F>(f), this]() mutable {
-        try {
-          if constexpr (std::is_void_v<R>) {
-            f(stream_);
-            p->set_value();
-          } else {
-            p->set_value(f(stream_));
-          }
-        } catch (...) {
-          p->set_exception(std::current_exception());
+      for (auto &entry : live_) {
+        if (auto st = entry.second.st.lock()) {
+          signal_stop(*st);
         }
-      });
+        dying_.push_back(std::move(entry.second.thread));
+      }
+      live_.clear();
+      shutdown_ = true;
+    }
+    cv_.notify_all();
+    if (thread_.joinable()) {
+      thread_.join();
     }
-    cv_.notify_one();
-    return fut.get();
   }
 
-  void stop() {
+  ~Reaper() { shutdown(); }
+
+private:
+  struct Live {
+    std::thread thread;
+    std::weak_ptr<State> st;
+  };
+
+  Reaper() : thread_([this] { run(); }) {}
+
+  void run() {
+    while (true) {
+      std::vector<std::thread> batch;
+      {
+        std::unique_lock<std::mutex> lock(mtx_);
+        cv_.wait(lock, [this] { return shutdown_ || !dying_.empty(); });
+        batch.swap(dying_);
+        if (batch.empty() && shutdown_) {
+          break;
+        }
+      }
+      for (auto &t : batch) {
+        if (t.joinable()) {
+          t.join();
+        }
+      }
+    }
+  }
+
+  std::mutex mtx_;
+  std::condition_variable cv_;
+  std::map<State *, Live> live_;
+  std::vector<std::thread> dying_;
+  bool shutdown_ = false;
+  std::thread thread_;
+};
+
+class WorkerThread {
+public:
+  explicit WorkerThread(std::size_t queue_limit)
+      : state_(std::make_shared<State>(queue_limit)) {
+    std::thread thread([st = state_] { run(st); });
+    Reaper::instance().track(state_.get(), std::move(thread), state_);
+
+    std::unique_lock<std::mutex> lock(state_->mtx);
+    state_->cv.wait(lock, [this] { return state_->ready; });
+  }
+
+  // Non-blocking: signal stop and hand the thread to the Reaper to join
+  // off-scheduler. Pending tasks are cancelled with {:error, :stopped}.
+  ~WorkerThread() { Reaper::instance().retire(state_.get()); }
+
+  // Enqueue a task. Throws if the worker has been stopped or the queue is
+  // at capacity (back-pressure) — both surface synchronously as an
+  // exception from the calling NIF.
+  void run_async(Task task) {
     {
-      std::lock_guard<std::mutex> lock(mtx_);
-      if (stop_)
-        return;
-      stop_ = true;
+      std::lock_guard<std::mutex> lock(state_->mtx);
+      if (state_->stop) {
+        throw std::runtime_error("worker thread has been stopped");
+      }
+      if (state_->queue.size() >= state_->cap) {
+        throw std::runtime_error(
+            "worker queue is full (limit " + std::to_string(state_->cap) +
+            "); too many operations are queued on this stream");
+      }
+      state_->queue.push(std::move(task));
     }
-    cv_.notify_one();
-    if (thread_.joinable())
-      thread_.join();
+    state_->cv.notify_one();
   }
 
-  mx::Stream stream() const { return stream_; }
+  // Signal stop without blocking. The Reaper joins the thread when the
+  // resource is collected; queued tasks are cancelled with :stopped.
+  void stop() { signal_stop(*state_); }
+
+  std::size_t queue_depth() {
+    std::lock_guard<std::mutex> lock(state_->mtx);
+    return state_->queue.size();
+  }
 
 private:
-  void run() {
-    stream_ = mx::new_stream(mx::Device(mx::Device::DeviceType::gpu));
+  static void run(std::shared_ptr<State> st) {
+    st->stream = mx::new_stream(mx::Device(mx::Device::DeviceType::gpu));
     {
-      std::lock_guard<std::mutex> lock(mtx_);
-      ready_ = true;
+      std::lock_guard<std::mutex> lock(st->mtx);
+      st->ready = true;
     }
-    cv_.notify_one();
+    st->cv.notify_all();
 
     while (true) {
-      std::function<void()> task;
+      Task task;
+      bool cancelled = false;
       {
-        std::unique_lock<std::mutex> lock(mtx_);
-        cv_.wait(lock, [this] { return stop_ || !queue_.empty(); });
-        if (queue_.empty())
-          break;
-        task = std::move(queue_.front());
-        queue_.pop();
+        std::unique_lock<std::mutex> lock(st->mtx);
+        st->cv.wait(lock, [&] { return st->stop || !st->queue.empty(); });
+        if (st->queue.empty()) {
+          break;  // stop_ set and the queue is drained
+        }
+        task = std::move(st->queue.front());
+        st->queue.pop();
+        cancelled = st->stop;
       }
-      task();
+      // Exceptions are owned by the task (it posts its own error reply).
+      task(st->stream, cancelled);
     }
   }
 
-  std::thread thread_;
-  std::queue<std::function<void()>> queue_;
-  std::mutex mtx_;
-  std::condition_variable cv_;
-  bool stop_ = false;
-  bool ready_ = false;
-  mx::Stream stream_{0, mx::Device(mx::Device::DeviceType::gpu)};
+  std::shared_ptr<State> state_;
 };
 
 } // namespace emily