fix(io_uring): balance io_uring_inflight_ in drain_cqes_for (#299)

include/boost/corosio/native/detail/io_uring/io_uring_scheduler.hpp

@@ -159,6 +159,23 @@ class BOOST_COROSIO_DECL io_uring_scheduler final
         io_uring_inflight_.fetch_add(1, std::memory_order_release);
     }
 
+    /** Return the current io_uring in-flight counter.
+
+        Test-only helper: `io_uring_inflight_` is an internal accounting
+        counter (it gates the `do_one` ring pump), with no bearing on the
+        public API. It is exposed solely so tests can assert the counter
+        stays balanced across op submission and teardown — in particular
+        that `drain_cqes_for` does not leak counts. Not thread-safe with
+        respect to a concurrently running scheduler; call it from a quiesced
+        context.
+
+        @return The number of SQEs currently counted as in flight.
+    */
+    std::int64_t inflight() const noexcept
+    {
+        return io_uring_inflight_.load(std::memory_order_acquire);
+    }
+
     /// Initialize the io_uring ring on first access. Idempotent.
     void lazy_init_ring() const;
 
@@ -1386,28 +1403,80 @@ io_uring_scheduler::drain_cqes_for(io_uring_op* target) noexcept
         ::io_uring_cqe* cqe;
         unsigned        consumed = 0;
         bool            saw_target = false;
+        std::int64_t    inflight_dec = 0;
 
         io_uring_for_each_cqe(&ring_, head, cqe)
         {
+            // Mirror process_completions' io_uring_inflight_ accounting.
+            // That counter gates the do_one ring pump, so every CQE we
+            // advance past here must adjust it exactly as the normal
+            // drain would — otherwise it drifts upward (each teardown
+            // leaks the counts of the CQEs it swallows), defeating the
+            // idle-skip optimisation for the lifetime of the io_context.
+            // We do NOT dispatch real ops — the target is being
+            // destructed and siblings may already be freed — but we still
+            // account for and house-keep each CQE we consume.
             void* ud = io_uring_cqe_get_data(cqe);
-            if (ud == target)
+            if (ud == nullptr)
+            {
+                // Wakeup eventfd CQE — our own interrupt_reactor() above
+                // very likely produced one. Drain the byte and re-arm if
+                // the multishot terminated, exactly as process_completions
+                // does. Never incremented, so never decremented.
+                drain_wakeup_eventfd();
+                if ((cqe->flags & IORING_CQE_F_MORE) == 0)
+                    prep_multishot_poll(wakeup_eventfd_, nullptr);
+            }
+            else if (ud == &signal_pipe_sentinel_)
+            {
+                // Signal self-pipe readiness. Re-arm if the multishot
+                // terminated; the still-readable pipe re-fires on the next
+                // kernel enter so process_completions delivers the signal —
+                // we deliberately do NOT enqueue signal_drain_op_ from this
+                // teardown path. Not counted by io_uring_inflight_ (the poll
+                // was armed via prep_multishot_poll, which never increments),
+                // so it must NOT be decremented.
+                if ((cqe->flags & IORING_CQE_F_MORE) == 0)
+                    prep_multishot_poll(
+                        signal_pipe_read_fd_, &signal_pipe_sentinel_);
+            }
+            else if (ud == &cancel_sentinel_)
+            {
+                // ASYNC_CANCEL CQE (one-shot, no F_MORE), including the
+                // cancel SQE we submitted just above. Decrement inflight.
+                ++inflight_dec;
+            }
+            else if (ud == target)
             {
                 saw_target = true;
-                // Don't dispatch — caller is destructing target;
-                // just consume so the CQE doesn't dangle.
+                // Don't dispatch — caller is destructing target; just
+                // consume so the CQE doesn't dangle. Decrement inflight on
+                // the terminal CQE only: the target is a multishot op
+                // whose intermediate CQEs carry F_MORE.
+                if ((cqe->flags & IORING_CQE_F_MORE) == 0)
+                    ++inflight_dec;
+            }
+            else
+            {
+                // Some other op's CQE. Intentionally NOT dispatched: it
+                // may belong to an op freed by a sibling teardown (other
+                // acceptors / sockets), and dispatching would UAF. We
+                // still account for its terminal CQE so inflight stays
+                // balanced — the submit that produced it incremented the
+                // counter. The io_context's destructor sequence runs
+                // services' shutdowns before ~scheduler, so any still-live
+                // ops drain through their own paths first.
+                if ((cqe->flags & IORING_CQE_F_MORE) == 0)
+                    ++inflight_dec;
             }
-            // Other CQEs are intentionally NOT dispatched here. They
-            // may belong to ops freed by sibling teardowns (other
-            // acceptors / sockets), and dispatching would UAF. The
-            // next normal run-loop iteration will handle them; the
-            // io_context's destructor sequence runs services'
-            // shutdowns before ~scheduler so any still-live ops get
-            // a chance to drain through their own paths first.
             ++consumed;
         }
         if (consumed)
         {
             io_uring_cq_advance(&ring_, consumed);
+            if (inflight_dec)
+                io_uring_inflight_.fetch_sub(
+                    inflight_dec, std::memory_order_acq_rel);
             if (saw_target)
                 break;
             continue;

test/unit/native/native_io_uring_specific.cpp

@@ -1,5 +1,6 @@
 //
 // Copyright (c) 2026 Steve Gerbino
+// Copyright (c) 2026 Michael Vandeberg
 //
 // Distributed under the Boost Software License, Version 1.0. (See accompanying
 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
@@ -15,6 +16,12 @@
 
 #include <boost/corosio/backend.hpp>
 #include <boost/corosio/io_context.hpp>
+#include <boost/corosio/native/native_io_context.hpp>
+#include <boost/corosio/socket_option.hpp>
+#include <boost/corosio/tcp.hpp>
+#include <boost/corosio/tcp_acceptor.hpp>
+
+#include <cstdint>
 
 namespace boost::corosio {
 
@@ -23,8 +30,7 @@ namespace boost::corosio {
    Most io_uring behaviors (multishot accept queueing, cancel-by-fd, op
    lifecycle) are exercised by the existing backend-templated test suites
    (tcp_acceptor.io_uring, tcp_socket.io_uring, cancel.io_uring, etc.).
-   This file is the slot for io_uring-only tests when they're needed —
-   currently just a smoke test.
+   This file is the slot for io_uring-only tests when they're needed.
 
    Future additions when there's a specific behavior to pin:
    - SQ ring backpressure (>256 in-flight ops): current behavior surfaces
@@ -33,6 +39,17 @@ namespace boost::corosio {
    - Probe-and-fall-back: requires loading a seccomp filter at process
      start; deferred to test infrastructure work.
 */
+
+// Exposes the io_uring scheduler's internal in-flight counter so the
+// teardown-accounting test can assert it stays balanced.
+struct io_uring_test_context : native_io_context<io_uring>
+{
+    std::int64_t inflight()
+    {
+        return static_cast<detail::io_uring_scheduler*>(sched_)->inflight();
+    }
+};
+
 struct native_io_uring_specific_test
 {
     void testTagAvailable()
@@ -43,9 +60,48 @@ struct native_io_uring_specific_test
         BOOST_TEST(!ioc.stopped());
     }
 
+    // Regression: drain_cqes_for (run from the multishot acceptor's
+    // destructor) used to advance CQEs without decrementing
+    // io_uring_inflight_, leaking the counts of the multishot accept SQE
+    // and the cancel SQEs it submits. The counter gates the do_one ring
+    // pump, so the leak accumulated across every acceptor teardown for the
+    // lifetime of the io_context. After the fix the counter returns to
+    // zero once teardown settles.
+    void testDrainCqesBalancesInflight()
+    {
+        io_uring_test_context ctx;
+
+        // No io_uring SQEs are counted before any op is submitted.
+        BOOST_TEST_EQ(ctx.inflight(), 0);
+
+        {
+            tcp_acceptor acc(ctx);
+            acc.open();
+            acc.set_option(socket_option::reuse_address(true));
+            BOOST_TEST(!acc.bind(endpoint(0)));
+            // listen() calls start_multishot(), submitting a multishot
+            // accept SQE (counted once).
+            BOOST_TEST(!acc.listen());
+
+            // Flush the SQE to the kernel so it is genuinely in flight.
+            ctx.poll();
+            BOOST_TEST(ctx.inflight() >= 1);
+        }
+        // acc destroyed: the impl destructor runs cancel-by-fd and
+        // drain_cqes_for(multi_op_). Drain the CQEs those produce that the
+        // destructor's bounded loop did not itself consume.
+        for (int i = 0; i < 64 && ctx.inflight() != 0; ++i)
+            ctx.poll();
+
+        // Every submitted SQE (the multishot accept plus the teardown
+        // cancels) is now accounted for. Before the fix this stayed > 0.
+        BOOST_TEST_EQ(ctx.inflight(), 0);
+    }
+
     void run()
     {
         testTagAvailable();
+        testDrainCqesBalancesInflight();
     }
 };