1/*
2 * Copyright 2014-present Facebook, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16#pragma once
17
18#include <atomic>
19
20#include <folly/Portability.h>
21#include <folly/detail/Futex.h>
22#include <folly/io/async/HHWheelTimer.h>
23
24#if FOLLY_HAS_COROUTINES
25#include <experimental/coroutine>
26#endif
27
28namespace folly {
29namespace fibers {
30
31class Fiber;
32class FiberManager;
33
34/**
35 * @class Baton
36 *
37 * Primitive which allows one to put current Fiber to sleep and wake it from
38 * another Fiber/thread.
39 */
40class Baton {
41 public:
42 class TimeoutHandler;
43
44 class Waiter {
45 public:
46 virtual void post() = 0;
47
48 virtual ~Waiter() {}
49 };
50
51 Baton() noexcept;
52
53 ~Baton() noexcept = default;
54
55 bool ready() const {
56 auto state = waiter_.load();
57 return state == POSTED;
58 }
59
60 /**
61 * Registers a waiter for the baton. The waiter will be notified when
62 * the baton is posted.
63 */
64 void setWaiter(Waiter& waiter);
65
66 /**
67 * Puts active fiber to sleep. Returns when post is called.
68 */
69 void wait();
70
71 /**
72 * Put active fiber to sleep indefinitely. However, timeoutHandler may
73 * be used elsewhere on the same thread in order to schedule a wakeup
74 * for the active fiber. Users of timeoutHandler must be on the same thread
75 * as the active fiber and may only schedule one timeout, which must occur
76 * after the active fiber calls wait.
77 */
78 void wait(TimeoutHandler& timeoutHandler);
79
80 /**
81 * Puts active fiber to sleep. Returns when post is called.
82 *
83 * @param mainContextFunc this function is immediately executed on the main
84 * context.
85 */
86 template <typename F>
87 void wait(F&& mainContextFunc);
88
89 /**
90 * Checks if the baton has been posted without blocking.
91 *
92 * @return true iff the baton has been posted.
93 */
94 bool try_wait();
95
96 /**
97 * Puts active fiber to sleep. Returns when post is called or the timeout
98 * expires.
99 *
100 * @param timeout Baton will be automatically awaken if timeout expires
101 *
102 * @return true if was posted, false if timeout expired
103 */
104 template <typename Rep, typename Period>
105 bool try_wait_for(const std::chrono::duration<Rep, Period>& timeout) {
106 return try_wait_for(timeout, [] {});
107 }
108
109 /**
110 * Puts active fiber to sleep. Returns when post is called or the timeout
111 * expires.
112 *
113 * @param timeout Baton will be automatically awaken if timeout expires
114 * @param mainContextFunc this function is immediately executed on the main
115 * context.
116 *
117 * @return true if was posted, false if timeout expired
118 */
119 template <typename Rep, typename Period, typename F>
120 bool try_wait_for(
121 const std::chrono::duration<Rep, Period>& timeout,
122 F&& mainContextFunc);
123
124 /**
125 * Puts active fiber to sleep. Returns when post is called or the deadline
126 * expires.
127 *
128 * @param timeout Baton will be automatically awaken if deadline expires
129 *
130 * @return true if was posted, false if timeout expired
131 */
132 template <typename Clock, typename Duration>
133 bool try_wait_until(
134 const std::chrono::time_point<Clock, Duration>& deadline) {
135 return try_wait_until(deadline, [] {});
136 }
137
138 /**
139 * Puts active fiber to sleep. Returns when post is called or the deadline
140 * expires.
141 *
142 * @param timeout Baton will be automatically awaken if deadline expires
143 * @param mainContextFunc this function is immediately executed on the main
144 * context.
145 *
146 * @return true if was posted, false if timeout expired
147 */
148 template <typename Clock, typename Duration, typename F>
149 bool try_wait_until(
150 const std::chrono::time_point<Clock, Duration>& deadline,
151 F&& mainContextFunc);
152
153 /**
154 * Puts active fiber to sleep. Returns when post is called or the deadline
155 * expires.
156 *
157 * @param timeout Baton will be automatically awaken if deadline expires
158 * @param mainContextFunc this function is immediately executed on the main
159 * context.
160 *
161 * @return true if was posted, false if timeout expired
162 */
163 template <typename Clock, typename Duration, typename F>
164 bool try_wait_for(
165 const std::chrono::time_point<Clock, Duration>& deadline,
166 F&& mainContextFunc);
167
168 /// Alias to try_wait_for. Deprecated.
169 template <typename Rep, typename Period>
170 bool timed_wait(const std::chrono::duration<Rep, Period>& timeout) {
171 return try_wait_for(timeout);
172 }
173
174 /// Alias to try_wait_for. Deprecated.
175 template <typename Rep, typename Period, typename F>
176 bool timed_wait(
177 const std::chrono::duration<Rep, Period>& timeout,
178 F&& mainContextFunc) {
179 return try_wait_for(timeout, static_cast<F&&>(mainContextFunc));
180 }
181
182 /// Alias to try_wait_until. Deprecated.
183 template <typename Clock, typename Duration>
184 bool timed_wait(const std::chrono::time_point<Clock, Duration>& deadline) {
185 return try_wait_until(deadline);
186 }
187
188 /// Alias to try_wait_until. Deprecated.
189 template <typename Clock, typename Duration, typename F>
190 bool timed_wait(
191 const std::chrono::time_point<Clock, Duration>& deadline,
192 F&& mainContextFunc) {
193 return try_wait_until(deadline, static_cast<F&&>(mainContextFunc));
194 }
195
196 /**
197 * Wakes up Fiber which was waiting on this Baton (or if no Fiber is waiting,
198 * next wait() call will return immediately).
199 */
200 void post();
201
202 /**
203 * Reset's the baton (equivalent to destroying the object and constructing
204 * another one in place).
205 * Caller is responsible for making sure no one is waiting on/posting the
206 * baton when reset() is called.
207 */
208 void reset();
209
210 /**
211 * Provides a way to schedule a wakeup for a wait()ing fiber.
212 * A TimeoutHandler must be passed to Baton::wait(TimeoutHandler&)
213 * before a timeout is scheduled. It is only safe to use the
214 * TimeoutHandler on the same thread as the wait()ing fiber.
215 * scheduleTimeout() may only be called once prior to the end of the
216 * associated Baton's life.
217 */
218 class TimeoutHandler final : private HHWheelTimer::Callback {
219 public:
220 void scheduleTimeout(std::chrono::milliseconds timeout);
221
222 private:
223 friend class Baton;
224
225 std::function<void()> timeoutFunc_{nullptr};
226 FiberManager* fiberManager_{nullptr};
227
228 void timeoutExpired() noexcept override {
229 assert(timeoutFunc_ != nullptr);
230 timeoutFunc_();
231 }
232
233 void callbackCanceled() noexcept override {}
234 };
235
236 private:
237 class FiberWaiter;
238
239 enum {
240 /**
241 * Must be positive. If multiple threads are actively using a
242 * higher-level data structure that uses batons internally, it is
243 * likely that the post() and wait() calls happen almost at the same
244 * time. In this state, we lose big 50% of the time if the wait goes
245 * to sleep immediately. On circa-2013 devbox hardware it costs about
246 * 7 usec to FUTEX_WAIT and then be awoken (half the t/iter as the
247 * posix_sem_pingpong test in BatonTests). We can improve our chances
248 * of early post by spinning for a bit, although we have to balance
249 * this against the loss if we end up sleeping any way. Spins on this
250 * hw take about 7 nanos (all but 0.5 nanos is the pause instruction).
251 * We give ourself 300 spins, which is about 2 usec of waiting. As a
252 * partial consolation, since we are using the pause instruction we
253 * are giving a speed boost to the colocated hyperthread.
254 */
255 PreBlockAttempts = 300,
256 };
257
258 explicit Baton(intptr_t state) : waiter_(state) {}
259
260 void postHelper(intptr_t new_value);
261 void postThread();
262 void waitThread();
263
264 template <typename F>
265 inline void waitFiber(FiberManager& fm, F&& mainContextFunc);
266 /**
267 * Spin for "some time" (see discussion on PreBlockAttempts) waiting
268 * for a post.
269 * @return true if we received a post the spin wait, false otherwise. If the
270 * function returns true then Baton state is guaranteed to be POSTED
271 */
272 bool spinWaitForEarlyPost();
273
274 bool timedWaitThread(std::chrono::milliseconds timeout);
275
276 static constexpr intptr_t NO_WAITER = 0;
277 static constexpr intptr_t POSTED = -1;
278 static constexpr intptr_t TIMEOUT = -2;
279 static constexpr intptr_t THREAD_WAITING = -3;
280
281 union {
282 std::atomic<intptr_t> waiter_;
283 struct {
284 folly::detail::Futex<> futex{};
285 int32_t _unused_packing;
286 } futex_;
287 };
288};
289
290#if FOLLY_HAS_COROUTINES
291namespace detail {
292class BatonAwaitableWaiter : public Baton::Waiter {
293 public:
294 explicit BatonAwaitableWaiter(Baton& baton) : baton_(baton) {}
295
296 void post() override {
297 assert(h_);
298 h_();
299 }
300
301 bool await_ready() const {
302 return baton_.ready();
303 }
304
305 void await_resume() {}
306
307 void await_suspend(std::experimental::coroutine_handle<> h) {
308 assert(!h_);
309 h_ = std::move(h);
310 baton_.setWaiter(*this);
311 }
312
313 private:
314 std::experimental::coroutine_handle<> h_;
315 Baton& baton_;
316};
317} // namespace detail
318
319inline detail::BatonAwaitableWaiter /* implicit */ operator co_await(
320 Baton& baton) {
321 return detail::BatonAwaitableWaiter(baton);
322}
323#endif
324} // namespace fibers
325} // namespace folly
326
327#include <folly/fibers/Baton-inl.h>
328