include/kpn/pool_node.hpp

@@ -117,9 +117,7 @@ public:
     void stop() override {
         stop_flag_.store(true, std::memory_order_seq_cst);
         disable_inputs(std::make_index_sequence<input_count>{});
-        // fire_once() observes stop_flag_ and will not resubmit.
+        await_quiescence();
-        // We do not wait for an in-flight fire_once() to complete here;
-        // callers that need that guarantee should call scheduler_->drain() first.
     }
 
     bool running() const override {
@@ -418,7 +416,55 @@ private:
         else if (want_more)     try_submit(prio);
     }
 
+    /// Block until no firing of this node is in flight or queued.
+    ///
+    /// stop() used to set the flag and return, leaving an executing fire_once
+    /// touching input_channels_, stats_ and pending_ while the caller went on
+    /// to destroy them. For a node with a private pool that was survivable by
+    /// accident — Node::stop() calls pool->stop(), which joins — but a node
+    /// sharing a pool had nothing joining it at all, so ~PoolNode raced its own
+    /// members. The old comment said callers wanting the guarantee should call
+    /// scheduler_->drain() first; a destructor cannot, and the default should
+    /// not be a use-after-free.
+    ///
+    /// The gate is exactly the right thing to wait on: it is claimed for the
+    /// whole of a firing and released as the last act of one. A queued but
+    /// unstarted firing also holds it, and will run, observe stop_flag_ and
+    /// release — which is why the pool must still be running when this is
+    /// called. That is already the documented order (stop nodes, then the
+    /// pool), and Node/ObjectNode do it that way.
+    ///
+    /// Bounded, because a node function that never returns must not turn
+    /// teardown into a hang; and skipped entirely when called from the firing
+    /// thread itself, since an error handler that stops its own node would
+    /// otherwise wait for a firing that is waiting for it.
+    void await_quiescence() {
+        if (firing_thread_.load(std::memory_order_acquire) == std::this_thread::get_id())
+            return;
+        const auto deadline = clock_t::now() + std::chrono::seconds(5);
+        while (gate_.queued()) {
+            if (clock_t::now() >= deadline) {
+                std::cerr << "[kpn] stop: node '" << name_
+                          << "' still had work in flight after 5 s; "
+                             "continuing without it\n";
+                return;
+            }
+            std::this_thread::sleep_for(std::chrono::microseconds(50));
+        }
+    }
+
+    /// Marks fire_once's thread for the duration of a firing, so await_quiescence
+    /// can tell a re-entrant stop() from an external one.
+    struct FiringMark {
+        std::atomic<std::thread::id>& slot;
+        explicit FiringMark(std::atomic<std::thread::id>& s) : slot(s) {
+            slot.store(std::this_thread::get_id(), std::memory_order_release);
+        }
+        ~FiringMark() { slot.store(std::thread::id{}, std::memory_order_release); }
+    };
+
     void fire_once() {
+        FiringMark mark(firing_thread_);
         if (stop_flag_.load(std::memory_order_relaxed)) {
             gate_.force_idle();
             return;
@@ -624,6 +670,9 @@ private:
     /// cleared: the callbacks capture `this` and stay valid across a restart, so
     /// re-registering them would be a pointless write to a live channel.
     bool                                     prepared_{false};
+    /// Thread currently inside fire_once, or a default id when none is.
+    /// See await_quiescence.
+    std::atomic<std::thread::id>             firing_thread_{};
 
     /// The hidden one-slot output buffer (see push_outputs). Holding the value
     /// here is what lets a node stop running without dropping it or occupying a
@@ -707,6 +756,7 @@ public:
     void stop() override {
         stop_flag_.store(true, std::memory_order_seq_cst);
         disable_inputs(std::make_index_sequence<input_count>{});
+        await_quiescence();
     }
 
     bool running() const override { return !stop_flag_.load(std::memory_order_relaxed); }
@@ -960,7 +1010,55 @@ private:
         else if (want_more)     try_submit(prio);
     }
 
+    /// Block until no firing of this node is in flight or queued.
+    ///
+    /// stop() used to set the flag and return, leaving an executing fire_once
+    /// touching input_channels_, stats_ and pending_ while the caller went on
+    /// to destroy them. For a node with a private pool that was survivable by
+    /// accident — Node::stop() calls pool->stop(), which joins — but a node
+    /// sharing a pool had nothing joining it at all, so ~PoolNode raced its own
+    /// members. The old comment said callers wanting the guarantee should call
+    /// scheduler_->drain() first; a destructor cannot, and the default should
+    /// not be a use-after-free.
+    ///
+    /// The gate is exactly the right thing to wait on: it is claimed for the
+    /// whole of a firing and released as the last act of one. A queued but
+    /// unstarted firing also holds it, and will run, observe stop_flag_ and
+    /// release — which is why the pool must still be running when this is
+    /// called. That is already the documented order (stop nodes, then the
+    /// pool), and Node/ObjectNode do it that way.
+    ///
+    /// Bounded, because a node function that never returns must not turn
+    /// teardown into a hang; and skipped entirely when called from the firing
+    /// thread itself, since an error handler that stops its own node would
+    /// otherwise wait for a firing that is waiting for it.
+    void await_quiescence() {
+        if (firing_thread_.load(std::memory_order_acquire) == std::this_thread::get_id())
+            return;
+        const auto deadline = clock_t::now() + std::chrono::seconds(5);
+        while (gate_.queued()) {
+            if (clock_t::now() >= deadline) {
+                std::cerr << "[kpn] stop: node '" << name_
+                          << "' still had work in flight after 5 s; "
+                             "continuing without it\n";
+                return;
+            }
+            std::this_thread::sleep_for(std::chrono::microseconds(50));
+        }
+    }
+
+    /// Marks fire_once's thread for the duration of a firing, so await_quiescence
+    /// can tell a re-entrant stop() from an external one.
+    struct FiringMark {
+        std::atomic<std::thread::id>& slot;
+        explicit FiringMark(std::atomic<std::thread::id>& s) : slot(s) {
+            slot.store(std::this_thread::get_id(), std::memory_order_release);
+        }
+        ~FiringMark() { slot.store(std::thread::id{}, std::memory_order_release); }
+    };
+
     void fire_once() {
+        FiringMark mark(firing_thread_);
         if (stop_flag_.load(std::memory_order_relaxed)) {
             gate_.force_idle();
             return;
@@ -1135,6 +1233,9 @@ private:
     /// cleared: the callbacks capture `this` and stay valid across a restart, so
     /// re-registering them would be a pointless write to a live channel.
     bool                                     prepared_{false};
+    /// Thread currently inside fire_once, or a default id when none is.
+    /// See await_quiescence.
+    std::atomic<std::thread::id>             firing_thread_{};
 
     /// The hidden one-slot output buffer (see push_outputs). Holding the value
     /// here is what lets a node stop running without dropping it or occupying a

tests/test_pool_node.cpp

@@ -668,3 +668,55 @@ TEST_CASE("a node woken with empty inputs does not stop itself", "[pool_node]")
     node.stop();
     pool->stop();
 }
+
+// Regression: stop() must not return while a firing is still running.
+//
+// stop() set the flag, disabled the inputs and returned, leaving an executing
+// fire_once touching input_channels_, stats_ and pending_ while the caller went
+// on to destroy them. The old comment was explicit that callers wanting the
+// guarantee should call scheduler_->drain() first — but ~PoolNode calls stop(),
+// and a destructor cannot ask its caller to have done that.
+//
+// A node with a private pool survived by accident: Node::stop() calls
+// pool->stop(), which joins the worker. A node sharing a pool, which
+// make_pool_node exists to create, had nothing joining it at all, so its own
+// destructor raced the firing.
+//
+// Asserted through an observable side effect rather than by trying to catch the
+// use-after-free: if stop() returns before the node function has finished, the
+// flag it sets on the way out is still false.
+namespace {
+
+struct SlowFiring {
+    static constexpr std::string_view label() { return "slow_firing"; }
+    std::atomic<bool>* entered;
+    std::atomic<bool>* finished;
+    void operator()(int) {
+        entered->store(true, std::memory_order_release);
+        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        finished->store(true, std::memory_order_release);
+    }
+};
+
+}  // namespace
+
+TEST_CASE("stop waits for a firing already in flight", "[pool_node]") {
+    std::atomic<bool> entered{false}, finished{false};
+
+    auto pool = std::make_shared<ThreadPool>(2);
+    pool->start();
+
+    SlowFiring fn{&entered, &finished};
+    auto node = make_pool_node(fn, pool, 4);
+    node.start();
+    node.input_channel<0>().push(1);
+
+    // Stop only once the node is demonstrably inside its function.
+    while (!entered.load(std::memory_order_acquire))
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+
+    node.stop();
+    CHECK(finished.load(std::memory_order_acquire));
+
+    pool->stop();
+}