proper handling of immediate completion in IoCP

Author: vinniefalco

src/corosio/src/detail/iocp/sockets.cpp

@@ -363,6 +363,18 @@ win_socket_impl_internal(win_sockets& svc) noexcept
 win_socket_impl_internal::
 ~win_socket_impl_internal()
 {
+    // Destroy any active initiator coroutines
+    if (read_initiator_handle_)
+        read_initiator_handle_.destroy();
+    if (write_initiator_handle_)
+        write_initiator_handle_.destroy();
+
+    // Free cached frame storage (operator delete in promise_type is no-op)
+    if (read_initiator_frame_)
+        ::operator delete(read_initiator_frame_);
+    if (write_initiator_frame_)
+        ::operator delete(write_initiator_frame_);
+
     svc_.unregister_impl(*this);
 }
 
@@ -469,46 +481,31 @@ connect(
     }
 }
 
-std::coroutine_handle<>
-win_socket_impl_internal::
-read_some(
-    capy::coro h,
-    capy::executor_ref d,
-    io_buffer_param param,
-    std::stop_token token,
-    std::error_code* ec,
-    std::size_t* bytes_out)
-{
-    // Keep internal alive during I/O
-    rd_.internal_ptr = shared_from_this();
+//------------------------------------------------------------------------------
+// Initiator coroutines - receive control via symmetric transfer after caller
+// suspends, then initiate the actual I/O.
 
-    auto& op = rd_;
-    op.reset();
-    op.h = h;
-    op.d = d;
-    op.ec_out = ec;
-    op.bytes_out = bytes_out;
-    op.start(token);
+read_initiator
+make_read_initiator(void*& cached, win_socket_impl_internal* impl)
+{
+    impl->do_read_io();
+    co_return;
+}
 
-    capy::mutable_buffer bufs[read_op::max_buffers];
-    op.wsabuf_count = static_cast<DWORD>(
-        param.copy_to(bufs, read_op::max_buffers));
+write_initiator
+make_write_initiator(void*& cached, win_socket_impl_internal* impl)
+{
+    impl->do_write_io();
+    co_return;
+}
 
-    // Handle empty buffer: complete with 0 bytes via post for consistency
-    if (op.wsabuf_count == 0)
-    {
-        op.bytes_transferred = 0;
-        op.dwError = 0;
-        op.empty_buffer = true;
-        svc_.post(&op);
-        return std::noop_coroutine();
-    }
+//------------------------------------------------------------------------------
 
-    for (DWORD i = 0; i < op.wsabuf_count; ++i)
-    {
-        op.wsabufs[i].buf = static_cast<char*>(bufs[i].data());
-        op.wsabufs[i].len = static_cast<ULONG>(bufs[i].size());
-    }
+void
+win_socket_impl_internal::
+do_read_io()
+{
+    auto& op = rd_;
 
     op.flags = 0;
 
@@ -532,7 +529,7 @@ read_some(
             svc_.work_finished();
             op.dwError = err;
             op.complete_immediate();
-            return std::noop_coroutine();
+            return;
         }
     }
     else
@@ -555,7 +552,109 @@ read_some(
             svc_.post(&op);
         }
     }
-    return std::noop_coroutine();
+}
+
+void
+win_socket_impl_internal::
+do_write_io()
+{
+    auto& op = wr_;
+
+    svc_.work_started();
+
+    int result = ::WSASend(
+        socket_,
+        op.wsabufs,
+        op.wsabuf_count,
+        nullptr,
+        0,
+        &op,
+        nullptr);
+
+    if (result == SOCKET_ERROR)
+    {
+        DWORD err = ::WSAGetLastError();
+        if (err != WSA_IO_PENDING)
+        {
+            // Immediate error - must use post(). See do_read_io for explanation.
+            svc_.work_finished();
+            op.dwError = err;
+            svc_.post(&op);
+            return;
+        }
+    }
+    else
+    {
+        // Synchronous completion - use CAS to race with IOCP.
+        // See do_read_io for detailed explanation.
+        //
+        // CRITICAL: Must call work_finished() ONLY if we win the CAS, and must
+        // not access op after CAS fails. If IOCP wins, it processes the op
+        // (which may destroy it), so any access to op is use-after-free.
+        // The IOCP handler calls work_finished() via its work_guard.
+        if (::InterlockedCompareExchange(&op.ready_, 1, 0) == 0)
+        {
+            svc_.work_finished();
+            op.bytes_transferred = static_cast<DWORD>(op.InternalHigh);
+            op.dwError = 0;
+            svc_.post(&op);
+        }
+    }
+}
+
+//------------------------------------------------------------------------------
+
+std::coroutine_handle<>
+win_socket_impl_internal::
+read_some(
+    capy::coro h,
+    capy::executor_ref d,
+    io_buffer_param param,
+    std::stop_token token,
+    std::error_code* ec,
+    std::size_t* bytes_out)
+{
+    // Keep internal alive during I/O
+    rd_.internal_ptr = shared_from_this();
+
+    auto& op = rd_;
+    op.reset();
+    op.h = h;
+    op.d = d;
+    op.ec_out = ec;
+    op.bytes_out = bytes_out;
+    op.start(token);
+
+    // Prepare buffers (must happen before initiator runs)
+    capy::mutable_buffer bufs[read_op::max_buffers];
+    op.wsabuf_count = static_cast<DWORD>(
+        param.copy_to(bufs, read_op::max_buffers));
+
+    // Handle empty buffer: complete with 0 bytes via post for consistency
+    if (op.wsabuf_count == 0)
+    {
+        op.bytes_transferred = 0;
+        op.dwError = 0;
+        op.empty_buffer = true;
+        svc_.post(&op);
+        return std::noop_coroutine();
+    }
+
+    for (DWORD i = 0; i < op.wsabuf_count; ++i)
+    {
+        op.wsabufs[i].buf = static_cast<char*>(bufs[i].data());
+        op.wsabufs[i].len = static_cast<ULONG>(bufs[i].size());
+    }
+
+    // Destroy previous initiator if any, construct new one into cached frame
+    if (read_initiator_handle_)
+        read_initiator_handle_.destroy();
+
+    auto initiator = make_read_initiator(read_initiator_frame_, this);
+    read_initiator_handle_ = initiator.h;
+
+    // Symmetric transfer to initiator - I/O starts after caller is suspended
+    return initiator.h;
 }
 
 std::coroutine_handle<>
@@ -579,6 +678,7 @@ write_some(
     op.bytes_out = bytes_out;
     op.start(token);
 
+    // Prepare buffers (must happen before initiator runs)
     capy::mutable_buffer bufs[write_op::max_buffers];
     op.wsabuf_count = static_cast<DWORD>(
         param.copy_to(bufs, write_op::max_buffers));
@@ -598,47 +698,15 @@ write_some(
         op.wsabufs[i].len = static_cast<ULONG>(bufs[i].size());
     }
 
-    svc_.work_started();
+    // Destroy previous initiator if any, construct new one into cached frame
+    if (write_initiator_handle_)
+        write_initiator_handle_.destroy();
 
-    int result = ::WSASend(
-        socket_,
-        op.wsabufs,
-        op.wsabuf_count,
-        nullptr,
-        0,
-        &op,
-        nullptr);
+    auto initiator = make_write_initiator(write_initiator_frame_, this);
+    write_initiator_handle_ = initiator.h;
 
-    if (result == SOCKET_ERROR)
-    {
-        DWORD err = ::WSAGetLastError();
-        if (err != WSA_IO_PENDING)
-        {
-            // Immediate error - must use post(). See read_some for explanation.
-            svc_.work_finished();
-            op.dwError = err;
-            svc_.post(&op);
-            return std::noop_coroutine();
-        }
-    }
-    else
-    {
-        // Synchronous completion - use CAS to race with IOCP.
-        // See read_some for detailed explanation.
-        //
-        // CRITICAL: Must call work_finished() ONLY if we win the CAS, and must
-        // not access op after CAS fails. If IOCP wins, it processes the op
-        // (which may destroy it), so any access to op is use-after-free.
-        // The IOCP handler calls work_finished() via its work_guard.
-        if (::InterlockedCompareExchange(&op.ready_, 1, 0) == 0)
-        {
-            svc_.work_finished();
-            op.bytes_transferred = static_cast<DWORD>(op.InternalHigh);
-            op.dwError = 0;
-            svc_.post(&op);
-        }
-    }
-    return std::noop_coroutine();
+    // Symmetric transfer to initiator - I/O starts after caller is suspended
+    return initiator.h;
 }
 
 void

src/corosio/src/detail/iocp/sockets.hpp

@@ -27,6 +27,7 @@
 #include "src/detail/iocp/mutex.hpp"
 #include "src/detail/iocp/wsa_init.hpp"
 
+#include <coroutine>
 #include <memory>
 
 #include <MSWSock.h>
@@ -108,6 +109,90 @@ struct accept_op : overlapped_op
 
 //------------------------------------------------------------------------------
 
+/** Initiator coroutine for read operations.
+
+    This coroutine receives control via symmetric transfer after the caller
+    has fully suspended, then initiates the actual I/O. Uses cached frame
+    allocation to avoid per-operation heap allocations.
+*/
+struct read_initiator
+{
+    struct promise_type
+    {
+        win_socket_impl_internal* impl;
+
+        /** Cached allocation - first call allocates, subsequent calls reuse. */
+        static void* operator new(std::size_t n, void*& cached, win_socket_impl_internal*)
+        {
+            if (!cached)
+                cached = ::operator new(n);
+            return cached;
+        }
+
+        /** No-op - frame memory freed in socket destructor. */
+        static void operator delete(void*) noexcept {}
+
+        std::suspend_always initial_suspend() noexcept { return {}; }
+        std::suspend_always final_suspend() noexcept { return {}; }
+
+        read_initiator get_return_object()
+        {
+            return {std::coroutine_handle<promise_type>::from_promise(*this)};
+        }
+
+        void return_void() {}
+        void unhandled_exception() { std::terminate(); }
+    };
+
+    using handle_type = std::coroutine_handle<promise_type>;
+    handle_type h;
+};
+
+/** Initiator coroutine for write operations.
+
+    This coroutine receives control via symmetric transfer after the caller
+    has fully suspended, then initiates the actual I/O. Uses cached frame
+    allocation to avoid per-operation heap allocations.
+*/
+struct write_initiator
+{
+    struct promise_type
+    {
+        win_socket_impl_internal* impl;
+
+        /** Cached allocation - first call allocates, subsequent calls reuse. */
+        static void* operator new(std::size_t n, void*& cached, win_socket_impl_internal*)
+        {
+            if (!cached)
+                cached = ::operator new(n);
+            return cached;
+        }
+
+        /** No-op - frame memory freed in socket destructor. */
+        static void operator delete(void*) noexcept {}
+
+        std::suspend_always initial_suspend() noexcept { return {}; }
+        std::suspend_always final_suspend() noexcept { return {}; }
+
+        write_initiator get_return_object()
+        {
+            return {std::coroutine_handle<promise_type>::from_promise(*this)};
+        }
+
+        void return_void() {}
+        void unhandled_exception() { std::terminate(); }
+    };
+
+    using handle_type = std::coroutine_handle<promise_type>;
+    handle_type h;
+};
+
+// Coroutine factory functions (defined in sockets.cpp)
+read_initiator make_read_initiator(void*& cached, win_socket_impl_internal* impl);
+write_initiator make_write_initiator(void*& cached, win_socket_impl_internal* impl);
+
+//------------------------------------------------------------------------------
+
 /** Internal socket state for IOCP-based I/O.
 
     This class contains the actual state for a single socket, including
@@ -132,6 +217,12 @@ class win_socket_impl_internal
     write_op wr_;
     SOCKET socket_ = INVALID_SOCKET;
 
+    // Cached initiator coroutine frames (allocated on first use)
+    void* read_initiator_frame_ = nullptr;
+    void* write_initiator_frame_ = nullptr;
+    read_initiator::handle_type read_initiator_handle_;
+    write_initiator::handle_type write_initiator_handle_;
+
 public:
     explicit win_socket_impl_internal(win_sockets& svc) noexcept;
     ~win_socket_impl_internal();
@@ -174,6 +265,12 @@ class win_socket_impl_internal
         remote_endpoint_ = remote;
     }
 
+    /** Execute the read I/O operation (called by initiator coroutine). */
+    void do_read_io();
+
+    /** Execute the write I/O operation (called by initiator coroutine). */
+    void do_write_io();
+
 private:
     endpoint local_endpoint_;
     endpoint remote_endpoint_;