From f62ac819f506ad8013200930d282ac84bde0bc61 Mon Sep 17 00:00:00 2001
From: Vinnie Falco <vinnie.falco@gmail.com>
Date: Wed, 24 May 2017 19:42:24 -0700
Subject: [PATCH] Better test::enable_yield_to

---
 CHANGELOG.md                         |  1 +
 extras/beast/test/string_istream.hpp |  3 +-
 extras/beast/test/yield_to.hpp       | 78 +++++++++++++++-------------
 test/core/buffered_read_stream.cpp   |  3 +-
 test/http/read.cpp                   |  9 ++--
 test/http/write.cpp                  |  9 ++--
 6 files changed, 57 insertions(+), 46 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 735dbeb9..9c9fbe68 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@ Version 45
 
 * Workaround for boost::asio::basic_streambuf type check
 * Fix message doc image
+* Better test::enable_yield_to
 
 --------------------------------------------------------------------------------
 
diff --git a/extras/beast/test/string_istream.hpp b/extras/beast/test/string_istream.hpp
index 059ad3eb..ff4d7421 100644
--- a/extras/beast/test/string_istream.hpp
+++ b/extras/beast/test/string_istream.hpp
@@ -10,7 +10,6 @@
 
 #include <beast/core/async_result.hpp>
 #include <beast/core/bind_handler.hpp>
-#include <beast/core/buffer_prefix.hpp>
 #include <beast/core/error.hpp>
 #include <beast/websocket/teardown.hpp>
 #include <boost/asio/buffer.hpp>
@@ -68,7 +67,7 @@ public:
         error_code& ec)
     {
         auto const n = boost::asio::buffer_copy(
-            buffers, buffer_prefix(read_max_, cb_));
+            buffers, cb_, read_max_);
         if(n > 0)
             cb_ = cb_ + n;
         else
diff --git a/extras/beast/test/yield_to.hpp b/extras/beast/test/yield_to.hpp
index 16dea02d..689c165f 100644
--- a/extras/beast/test/yield_to.hpp
+++ b/extras/beast/test/yield_to.hpp
@@ -35,7 +35,7 @@ private:
     std::thread thread_;
     std::mutex m_;
     std::condition_variable cv_;
-    bool running_ = false;
+    std::size_t running_ = 0;
 
 public:
     /// The type of yield context passed to functions.
@@ -65,61 +65,65 @@ public:
         return ios_;
     }
 
-    /** Run a function in a coroutine.
+    /** Run one or more functions, each in a coroutine.
 
-        This call will block until the coroutine terminates.
-
-        Function will be called with this signature:
+        This call will block until all coroutines terminate.
 
+        Each functions should have this signature:
         @code
-            void f(args..., yield_context);
+            void f(yield_context);
         @endcode
 
-        @param f The Callable object to invoke.
-
-        @param args Optional arguments forwarded to the callable object.
+        @param fn... One or more functions to invoke.
     */
 #if BEAST_DOXYGEN
-    template<class F, class... Args>
+    template<class... FN>
     void
-    yield_to(F&& f, Args&&... args);
+    yield_to(FN&&... fn)
 #else
-    template<class F>
+    template<class F0, class... FN>
     void
-    yield_to(F&& f);
-
-    template<class Function, class Arg, class... Args>
-    void
-    yield_to(Function&& f, Arg&& arg, Args&&... args)
-    {
-        yield_to(std::bind(f,
-            std::forward<Arg>(arg),
-                std::forward<Args>(args)...,
-                    std::placeholders::_1));
-    }
+    yield_to(F0&& f0, FN&&... fn);
 #endif
+
+private:
+    void
+    spawn()
+    {
+    }
+
+    template<class F0, class... FN>
+    void
+    spawn(F0&& f, FN&&... fn);
 };
 
-template<class Function>
+template<class F0, class... FN>
 void
-enable_yield_to::yield_to(Function&& f)
+enable_yield_to::
+yield_to(F0&& f0, FN&&... fn)
+{
+    running_ = 1 + sizeof...(FN);
+    spawn(f0, fn...);
+    std::unique_lock<std::mutex> lock{m_};
+    cv_.wait(lock, [&]{ return running_ == 0; });
+}
+
+template<class F0, class... FN>
+inline
+void
+enable_yield_to::
+spawn(F0&& f, FN&&... fn)
 {
-    {
-        std::lock_guard<std::mutex> lock(m_);
-        running_ = true;
-    }
     boost::asio::spawn(ios_,
-        [&](boost::asio::yield_context do_yield)
+        [&](yield_context yield)
         {
-            f(do_yield);
-            std::lock_guard<std::mutex> lock(m_);
-            running_ = false;
-            cv_.notify_all();
+            f(yield);
+            std::lock_guard<std::mutex> lock{m_};
+            if(--running_ == 0)
+                cv_.notify_all();
         }
         , boost::coroutines::attributes(2 * 1024 * 1024));
-
-    std::unique_lock<std::mutex> lock(m_);
-    cv_.wait(lock, [&]{ return ! running_; });
+    spawn(fn...);
 }
 
 } // test
diff --git a/test/core/buffered_read_stream.cpp b/test/core/buffered_read_stream.cpp
index e36b73a7..e4338814 100644
--- a/test/core/buffered_read_stream.cpp
+++ b/test/core/buffered_read_stream.cpp
@@ -132,7 +132,8 @@ public:
     {
         testSpecialMembers();
 
-        yield_to(&self::testRead, this);
+        yield_to([&](yield_context yield){
+            testRead(yield);});
     }
 };
 
diff --git a/test/http/read.cpp b/test/http/read.cpp
index 2d9ff567..a90f33b7 100644
--- a/test/http/read.cpp
+++ b/test/http/read.cpp
@@ -319,9 +319,12 @@ public:
     {
         testThrow();
 
-        yield_to(&read_test::testFailures, this);
-        yield_to(&read_test::testRead, this);
-        yield_to(&read_test::testEof, this);
+        yield_to([&](yield_context yield){
+            testFailures(yield); });
+        yield_to([&](yield_context yield){
+            testRead(yield); });
+        yield_to([&](yield_context yield){
+            testEof(yield); });
 
         testIoService();
     }
diff --git a/test/http/write.cpp b/test/http/write.cpp
index 9cb37f8f..2fce413f 100644
--- a/test/http/write.cpp
+++ b/test/http/write.cpp
@@ -666,9 +666,12 @@ public:
 
     void run() override
     {
-        yield_to(&write_test::testAsyncWriteHeaders, this);
-        yield_to(&write_test::testAsyncWrite, this);
-        yield_to(&write_test::testFailures, this);
+        yield_to([&](yield_context yield){
+            testAsyncWriteHeaders(yield); });
+        yield_to([&](yield_context yield){
+            testAsyncWrite(yield); });
+        yield_to([&](yield_context yield){
+            testFailures(yield); });
         testOutput();
         test_std_ostream();
         testOstream();