Notify the completion handler through the handler's associated executor

azmq/detail/receive_op.hpp

@@ -14,11 +14,19 @@
 #include "socket_ops.hpp"
 #include "reactor_op.hpp"
 
+#include <boost/version.hpp>
 #include <boost/asio/io_service.hpp>
+#include <boost/asio/dispatch.hpp>
+#include <boost/asio/executor_work_guard.hpp>
+#if BOOST_VERSION >= 107900
+#include <boost/asio/recycling_allocator.hpp>
+#include <boost/asio/bind_allocator.hpp>
+#endif
 
 #include <zmq.h>
 
 #include <iterator>
+#include <type_traits>
 
 namespace azmq {
 namespace detail {
@@ -31,6 +39,23 @@ class receive_buffer_op_base : public reactor_op {
         { }
 
     virtual bool do_perform(socket_type& socket) override {
+        return do_perform_impl(socket);
+    }
+
+private:
+    template<typename Buff = MutableBufferSequence>
+    typename std::enable_if<std::is_same<Buff, azmq::message>::value, bool>::type do_perform_impl(socket_type& socket)
+    {
+        ec_ = boost::system::error_code();
+        bytes_transferred_ += socket_ops::receive(const_cast<azmq::message&>(buffers_), socket, flags_ | ZMQ_DONTWAIT, ec_);
+        if (ec_)
+            return !try_again();
+        return true;
+    }
+
+    template<typename Buff = MutableBufferSequence>
+    typename std::enable_if<!std::is_same<Buff, azmq::message>::value, bool>::type do_perform_impl(socket_type& socket)
+    {
         ec_ = boost::system::error_code();
         bytes_transferred_ += socket_ops::receive(buffers_, socket, flags_ | ZMQ_DONTWAIT, ec_);
         if (ec_)
@@ -44,7 +69,7 @@ class receive_buffer_op_base : public reactor_op {
     }
 
 private:
-    MutableBufferSequence buffers_;
+    typename std::conditional<std::is_same<MutableBufferSequence, azmq::message>::value, MutableBufferSequence const&, MutableBufferSequence>::type buffers_;
     flags_type flags_;
 };
 
@@ -57,14 +82,30 @@ class receive_buffer_op : public receive_buffer_op_base<MutableBufferSequence> {
                       socket_ops::flags_type flags)
         : receive_buffer_op_base<MutableBufferSequence>(buffers, flags)
         , handler_(std::move(handler))
+        , work_guard(boost::asio::make_work_guard(handler_))
         { }
 
     virtual void do_complete() override {
-        handler_(this->ec_, this->bytes_transferred_);
+#if BOOST_VERSION >= 107900
+        auto alloc = boost::asio::get_associated_allocator(
+            handler_, boost::asio::recycling_allocator<void>());
+#endif
+        boost::asio::dispatch(work_guard.get_executor(),
+#if BOOST_VERSION >= 107900
+            boost::asio::bind_allocator(alloc,
+#endif
+                [ec_ = this->ec_, handler_ = std::move(handler_), bytes_transferred_ = this->bytes_transferred_]() mutable {
+            handler_(ec_, bytes_transferred_);
+        })
+#if BOOST_VERSION >= 107900
+        )
+#endif
+            ;
     }
 
 private:
     Handler handler_;
+    boost::asio::executor_work_guard<typename boost::asio::associated_executor<Handler>::type> work_guard;
 };
 
 template<typename MutableBufferSequence,
@@ -76,14 +117,30 @@ class receive_more_buffer_op : public receive_buffer_op_base<MutableBufferSequen
                            socket_ops::flags_type flags)
         : receive_buffer_op_base<MutableBufferSequence>(buffers, flags)
         , handler_(std::move(handler))
+        , work_guard(boost::asio::make_work_guard(handler_))
         { }
 
     virtual void do_complete() override {
-        handler_(this->ec_, std::make_pair(this->bytes_transferred_, this->more()));
+#if BOOST_VERSION >= 107900
+        auto alloc = boost::asio::get_associated_allocator(
+            handler_, boost::asio::recycling_allocator<void>());
+#endif
+        boost::asio::dispatch(work_guard.get_executor(),
+#if BOOST_VERSION >= 107900
+            boost::asio::bind_allocator(alloc,
+#endif
+                [ec_ = this->ec_, handler_ = std::move(handler_), bytes_transferred_ = this->bytes_transferred_, more = this->more()]() mutable {
+            handler_(ec_, std::make_pair(bytes_transferred_, more));
+        })
+#if BOOST_VERSION >= 107900
+        )
+#endif
+            ;
     }
 
 private:
     Handler handler_;
+    boost::asio::executor_work_guard<typename boost::asio::associated_executor<Handler>::type> work_guard;
 };
 
 class receive_op_base : public reactor_op {
@@ -112,14 +169,30 @@ class receive_op : public receive_op_base {
                socket_ops::flags_type flags)
         : receive_op_base(flags)
         , handler_(std::move(handler))
+        , work_guard(boost::asio::make_work_guard(handler_))
         { }
 
     virtual void do_complete() override {
-        handler_(ec_, msg_, bytes_transferred_);
+#if BOOST_VERSION >= 107900
+        auto alloc = boost::asio::get_associated_allocator(
+            handler_, boost::asio::recycling_allocator<void>());
+#endif
+        boost::asio::dispatch(work_guard.get_executor(),
+#if BOOST_VERSION >= 107900
+            boost::asio::bind_allocator(alloc,
+#endif
+                [ec_ = this->ec_, handler_ = std::move(handler_), msg_ = std::move(msg_), bytes_transferred_ = this->bytes_transferred_]() mutable {
+            handler_(ec_, msg_, bytes_transferred_);
+        })
+#if BOOST_VERSION >= 107900
+        )
+#endif
+            ;
     }
 
 private:
     Handler handler_;
+    boost::asio::executor_work_guard<typename boost::asio::associated_executor<Handler>::type> work_guard;
 };
 } // namespace detail
 } // namespace azmq

azmq/detail/send_op.hpp

@@ -13,7 +13,14 @@
 #include "socket_ops.hpp"
 #include "reactor_op.hpp"
 
+#include <boost/version.hpp>
 #include <boost/asio/io_service.hpp>
+#include <boost/asio/dispatch.hpp>
+#include <boost/asio/executor_work_guard.hpp>
+#if BOOST_VERSION >= 107900
+#include <boost/asio/recycling_allocator.hpp>
+#include <boost/asio/bind_allocator.hpp>
+#endif
 
 #include <zmq.h>
 #include <iterator>
@@ -52,14 +59,30 @@ class send_buffer_op : public send_buffer_op_base<ConstBufferSequence> {
                    reactor_op::flags_type flags)
         : send_buffer_op_base<ConstBufferSequence>(buffers, flags)
         , handler_(std::move(handler))
+        , work_guard(boost::asio::make_work_guard(handler_))
     { }
 
     virtual void do_complete() override {
-        handler_(this->ec_, this->bytes_transferred_);
+#if BOOST_VERSION >= 107900
+        auto alloc = boost::asio::get_associated_allocator(
+            handler_, boost::asio::recycling_allocator<void>());
+#endif
+        boost::asio::dispatch(work_guard.get_executor(),
+#if BOOST_VERSION >= 107900
+            boost::asio::bind_allocator(alloc,
+#endif
+                [ec_ = this->ec_, handler_ = std::move(handler_), bytes_transferred_ = this->bytes_transferred_]() mutable {
+            handler_(ec_, bytes_transferred_);
+        })
+#if BOOST_VERSION >= 107900
+        )
+#endif
+            ;
     }
 
 private:
     Handler handler_;
+    boost::asio::executor_work_guard<typename boost::asio::associated_executor<Handler>::type> work_guard;
 };
 
 class send_op_base : public reactor_op {
@@ -90,14 +113,31 @@ class send_op : public send_op_base {
             flags_type flags)
         : send_op_base(std::move(msg), flags)
         , handler_(std::move(handler))
+        , work_guard(boost::asio::make_work_guard(handler_))
     { }
 
     virtual void do_complete() override {
-        handler_(ec_, bytes_transferred_);
+#if BOOST_VERSION >= 107900
+        auto alloc = boost::asio::get_associated_allocator(
+                handler_, boost::asio::recycling_allocator<void>());
+#endif
+        boost::asio::dispatch(work_guard.get_executor(),
+#if BOOST_VERSION >= 107900
+            boost::asio::bind_allocator(alloc,
+#endif
+                [ec_ = this->ec_, handler_ = std::move(handler_), bytes_transferred_ = this->bytes_transferred_]() mutable {
+            handler_(ec_, bytes_transferred_);
+        })
+#if BOOST_VERSION >= 107900
+        )
+#endif
+        ;
+
     }
 
 private:
     Handler handler_;
+    boost::asio::executor_work_guard<typename boost::asio::associated_executor<Handler>::type> work_guard;
 };
 
 } // namespace detail

test/socket/main.cpp

@@ -9,6 +9,7 @@
 #include <azmq/socket.hpp>
 #include <azmq/util/scope_guard.hpp>
 
+#include <boost/current_function.hpp>
 #include <boost/utility/string_ref.hpp>
 #include <boost/algorithm/string/split.hpp>
 #include <boost/algorithm/string/classification.hpp>
@@ -18,6 +19,9 @@
 #include <boost/asio/use_future.hpp>
 #include <boost/asio/spawn.hpp>
 #include <boost/optional.hpp>
+#if BOOST_VERSION >= 107400
+#include <boost/asio/any_io_executor.hpp>
+#endif
 #endif
 
 #include <array>
@@ -879,19 +883,74 @@ TEST_CASE("Async Operation Send/Receive single message, stackful coroutine, one
       auto frame1 = azmq::message{};
       auto const btb1 = azmq::async_receive(sb, frame1, yield);
       REQUIRE(btb1 == 5);
+      REQUIRE(frame1.more());
 
       auto frame2 = azmq::message{};
       auto const btb2 = azmq::async_receive(sb, frame2, yield);
       REQUIRE(btb2 == 2);
+      REQUIRE(frame2.more());
+      REQUIRE(message_ref(snd_bufs.at(0)) == message_ref(frame2));
 
       auto frame3 = azmq::message{};
       auto const btb3 = azmq::async_receive(sb, frame3, yield);
       REQUIRE(btb3 == 2);
+      REQUIRE(!frame3.more());
+      REQUIRE(message_ref(snd_bufs.at(1)) == message_ref(frame3));
 
     });
 
     ios.run();
 }
 
+
+TEST_CASE("Async Operation Send/Receive single message, check thread safety", "[socket_ops]") {
+	boost::asio::io_service ios;
+#if BOOST_VERSION >= 107400
+	boost::asio::strand<boost::asio::any_io_executor> strand{ios.get_executor()};
+#else
+	boost::asio::strand<boost::asio::executor> strand{ios.get_executor()};
+#endif
+
+	azmq::socket sb(ios, ZMQ_ROUTER);
+	sb.bind(subj(BOOST_CURRENT_FUNCTION));
+
+	azmq::socket sc(ios, ZMQ_DEALER);
+	sc.connect(subj(BOOST_CURRENT_FUNCTION));
+
+	//send coroutine task
+	boost::asio::spawn(strand, [&](boost::asio::yield_context yield) {
+		REQUIRE(strand.running_in_this_thread());
+        boost::system::error_code ecc;
+		auto const btc = azmq::async_send(sc, snd_bufs, yield[ecc]);
+		REQUIRE(strand.running_in_this_thread());
+        REQUIRE(!ecc);
+		REQUIRE(btc == 4);
+	});
+
+	//receive coroutine task
+	boost::asio::spawn(strand, [&](boost::asio::yield_context yield) {
+		std::array<char, 5> ident;
+		std::array<char, 2> a;
+		std::array<char, 2> b;
+
+		std::array<boost::asio::mutable_buffer, 3> rcv_bufs = { {boost::asio::buffer(ident),
+						boost::asio::buffer(a),
+						boost::asio::buffer(b)}};
+
+		boost::system::error_code ecc;
+
+		REQUIRE(strand.running_in_this_thread());
+		auto const btb = azmq::async_receive(sb, rcv_bufs, yield[ecc]);
+		REQUIRE(strand.running_in_this_thread());
+		REQUIRE(!ecc);
+		REQUIRE(btb == 9);
+
+		REQUIRE(message_ref(snd_bufs.at(0)) == boost::string_ref(a.data(), 2));
+		REQUIRE(message_ref(snd_bufs.at(1)) == boost::string_ref(b.data(), 2));
+	});
+
+	ios.run();
+}
+
 #endif // BOOST_VERSION >= 107000
 

test/socket_ops/main.cpp

@@ -2,6 +2,7 @@
 #include <azmq/detail/socket_ops.hpp>
 
 #include <boost/asio/buffer.hpp>
+#include <boost/current_function.hpp>
 
 #include <array>
 #include <chrono>