1 //
2 // Copyright (c) 2016-2019 Vinnie Falco (vinnie dot falco at gmail dot com)
3 //
4 // Distributed under the Boost Software License, Version 1.0. (See accompanying
5 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6 //
7 // Official repository: https://github.com/boostorg/beast
8 //
9
10 #ifndef BOOST_BEAST_TEST_YIELD_TO_HPP
11 #define BOOST_BEAST_TEST_YIELD_TO_HPP
12
13 #include <boost/asio/executor_work_guard.hpp>
14 #include <boost/asio/io_context.hpp>
15 #include <boost/asio/spawn.hpp>
16 #include <condition_variable>
17 #include <functional>
18 #include <mutex>
19 #include <thread>
20 #include <vector>
21
22 namespace boost {
23 namespace beast {
24 namespace test {
25
26 /** Mix-in to support tests using asio coroutines.
27
28 Derive from this class and use yield_to to launch test
29 functions inside coroutines. This is handy for testing
30 asynchronous asio code.
31 */
32 class enable_yield_to
33 {
34 protected:
35 net::io_context ioc_;
36
37 private:
38 beast::detail::select_work_guard_t<
39 net::io_context::executor_type>
40 work_;
41 std::vector<std::thread> threads_;
42 std::mutex m_;
43 std::condition_variable cv_;
44 std::size_t running_ = 0;
45
46 public:
47 /// The type of yield context passed to functions.
48 using yield_context =
49 net::yield_context;
50
51 explicit
enable_yield_to(std::size_t concurrency=1)52 enable_yield_to(std::size_t concurrency = 1)
53 : work_(beast::detail::make_work_guard(
54 ioc_.get_executor()))
55 {
56 threads_.reserve(concurrency);
57 while(concurrency--)
58 threads_.emplace_back(
59 [&]{ ioc_.run(); });
60 }
61
~enable_yield_to()62 ~enable_yield_to()
63 {
64 work_.reset();
65 for(auto& t : threads_)
66 t.join();
67 }
68
69 /// Return the `io_context` associated with the object
70 net::io_context&
get_io_service()71 get_io_service()
72 {
73 return ioc_;
74 }
75
76 /** Run one or more functions, each in a coroutine.
77
78 This call will block until all coroutines terminate.
79
80 Each functions should have this signature:
81 @code
82 void f(yield_context);
83 @endcode
84
85 @param fn... One or more functions to invoke.
86 */
87 #if BOOST_BEAST_DOXYGEN
88 template<class... FN>
89 void
90 yield_to(FN&&... fn)
91 #else
92 template<class F0, class... FN>
93 void
94 yield_to(F0&& f0, FN&&... fn);
95 #endif
96
97 private:
98 void
spawn()99 spawn()
100 {
101 }
102
103 template<class F0, class... FN>
104 void
105 spawn(F0&& f, FN&&... fn);
106 };
107
108 template<class F0, class... FN>
109 void
110 enable_yield_to::
yield_to(F0 && f0,FN &&...fn)111 yield_to(F0&& f0, FN&&... fn)
112 {
113 running_ = 1 + sizeof...(FN);
114 spawn(f0, fn...);
115 std::unique_lock<std::mutex> lock{m_};
116 cv_.wait(lock, [&]{ return running_ == 0; });
117 }
118
119 template<class F0, class... FN>
120 inline
121 void
122 enable_yield_to::
spawn(F0 && f,FN &&...fn)123 spawn(F0&& f, FN&&... fn)
124 {
125 asio::spawn(ioc_,
126 [&](yield_context yield)
127 {
128 f(yield);
129 std::lock_guard<std::mutex> lock{m_};
130 if(--running_ == 0)
131 cv_.notify_all();
132 }
133 , boost::coroutines::attributes(2 * 1024 * 1024));
134 spawn(fn...);
135 }
136
137 } // test
138 } // beast
139 } // boost
140
141
142 #endif
143