1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
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
17 #pragma once
18
19 #include <atomic>
20
21 #include <folly/Portability.h>
22 #include <folly/detail/Futex.h>
23 #include <folly/experimental/coro/Coroutine.h>
24 #include <folly/io/async/HHWheelTimer.h>
25
26 namespace folly {
27 namespace fibers {
28
29 class Fiber;
30 class FiberManager;
31
32 /**
33 * @class Baton
34 *
35 * Primitive which allows one to put current Fiber to sleep and wake it from
36 * another Fiber/thread.
37 */
38 class Baton {
39 public:
40 class TimeoutHandler;
41
42 class Waiter {
43 public:
44 virtual void post() = 0;
45
~Waiter()46 virtual ~Waiter() {}
47 };
48
49 Baton() noexcept;
50
51 ~Baton() noexcept = default;
52
ready()53 bool ready() const {
54 auto state = waiter_.load();
55 return state == POSTED;
56 }
57
58 /**
59 * Registers a waiter for the baton. The waiter will be notified when
60 * the baton is posted.
61 */
62 void setWaiter(Waiter& waiter);
63
64 /**
65 * Puts active fiber to sleep. Returns when post is called.
66 */
67 void wait();
68
69 /**
70 * Put active fiber to sleep indefinitely. However, timeoutHandler may
71 * be used elsewhere on the same thread in order to schedule a wakeup
72 * for the active fiber. Users of timeoutHandler must be on the same thread
73 * as the active fiber and may only schedule one timeout, which must occur
74 * after the active fiber calls wait.
75 */
76 void wait(TimeoutHandler& timeoutHandler);
77
78 /**
79 * Puts active fiber to sleep. Returns when post is called.
80 *
81 * @param mainContextFunc this function is immediately executed on the main
82 * context.
83 */
84 template <typename F>
85 void wait(F&& mainContextFunc);
86
87 /**
88 * Checks if the baton has been posted without blocking.
89 *
90 * @return true iff the baton has been posted.
91 */
92 bool try_wait();
93
94 /**
95 * Puts active fiber to sleep. Returns when post is called or the timeout
96 * expires.
97 *
98 * @param timeout Baton will be automatically awaken if timeout expires
99 *
100 * @return true if was posted, false if timeout expired
101 */
102 template <typename Rep, typename Period>
try_wait_for(const std::chrono::duration<Rep,Period> & timeout)103 bool try_wait_for(const std::chrono::duration<Rep, Period>& timeout) {
104 return try_wait_for(timeout, [] {});
105 }
106
107 /**
108 * Puts active fiber to sleep. Returns when post is called or the timeout
109 * expires.
110 *
111 * @param timeout Baton will be automatically awaken if timeout expires
112 * @param mainContextFunc this function is immediately executed on the main
113 * context.
114 *
115 * @return true if was posted, false if timeout expired
116 */
117 template <typename Rep, typename Period, typename F>
118 bool try_wait_for(
119 const std::chrono::duration<Rep, Period>& timeout, F&& mainContextFunc);
120
121 /**
122 * Puts active fiber to sleep. Returns when post is called or the deadline
123 * expires.
124 *
125 * @param deadline Baton will be automatically awaken if deadline expires
126 *
127 * @return true if was posted, false if timeout expired
128 */
129 template <typename Clock, typename Duration>
try_wait_until(const std::chrono::time_point<Clock,Duration> & deadline)130 bool try_wait_until(
131 const std::chrono::time_point<Clock, Duration>& deadline) {
132 return try_wait_until(deadline, [] {});
133 }
134
135 /**
136 * Puts active fiber to sleep. Returns when post is called or the deadline
137 * expires.
138 *
139 * @param deadline Baton will be automatically awaken if deadline expires
140 * @param mainContextFunc this function is immediately executed on the main
141 * context.
142 *
143 * @return true if was posted, false if timeout expired
144 */
145 template <typename Clock, typename Duration, typename F>
146 bool try_wait_until(
147 const std::chrono::time_point<Clock, Duration>& deadline,
148 F&& mainContextFunc);
149
150 /**
151 * Puts active fiber to sleep. Returns when post is called or the deadline
152 * expires.
153 *
154 * @param deadline Baton will be automatically awaken if deadline expires
155 * @param mainContextFunc this function is immediately executed on the main
156 * context.
157 *
158 * @return true if was posted, false if timeout expired
159 */
160 template <typename Clock, typename Duration, typename F>
161 bool try_wait_for(
162 const std::chrono::time_point<Clock, Duration>& deadline,
163 F&& mainContextFunc);
164
165 /// Alias to try_wait_for. Deprecated.
166 template <typename Rep, typename Period>
timed_wait(const std::chrono::duration<Rep,Period> & timeout)167 bool timed_wait(const std::chrono::duration<Rep, Period>& timeout) {
168 return try_wait_for(timeout);
169 }
170
171 /// Alias to try_wait_for. Deprecated.
172 template <typename Rep, typename Period, typename F>
timed_wait(const std::chrono::duration<Rep,Period> & timeout,F && mainContextFunc)173 bool timed_wait(
174 const std::chrono::duration<Rep, Period>& timeout, F&& mainContextFunc) {
175 return try_wait_for(timeout, static_cast<F&&>(mainContextFunc));
176 }
177
178 /// Alias to try_wait_until. Deprecated.
179 template <typename Clock, typename Duration>
timed_wait(const std::chrono::time_point<Clock,Duration> & deadline)180 bool timed_wait(const std::chrono::time_point<Clock, Duration>& deadline) {
181 return try_wait_until(deadline);
182 }
183
184 /// Alias to try_wait_until. Deprecated.
185 template <typename Clock, typename Duration, typename F>
timed_wait(const std::chrono::time_point<Clock,Duration> & deadline,F && mainContextFunc)186 bool timed_wait(
187 const std::chrono::time_point<Clock, Duration>& deadline,
188 F&& mainContextFunc) {
189 return try_wait_until(deadline, static_cast<F&&>(mainContextFunc));
190 }
191
192 /**
193 * Wakes up Fiber which was waiting on this Baton (or if no Fiber is waiting,
194 * next wait() call will return immediately).
195 */
196 void post();
197
198 /**
199 * Reset's the baton (equivalent to destroying the object and constructing
200 * another one in place).
201 * Caller is responsible for making sure no one is waiting on/posting the
202 * baton when reset() is called.
203 */
204 void reset();
205
206 /**
207 * Provides a way to schedule a wakeup for a wait()ing fiber.
208 * A TimeoutHandler must be passed to Baton::wait(TimeoutHandler&)
209 * before a timeout is scheduled. It is only safe to use the
210 * TimeoutHandler on the same thread as the wait()ing fiber.
211 * scheduleTimeout() may only be called once prior to the end of the
212 * associated Baton's life.
213 */
214 class TimeoutHandler final : public HHWheelTimer::Callback {
215 public:
216 void scheduleTimeout(std::chrono::milliseconds timeout);
217
218 private:
219 friend class Baton;
220
221 std::function<void()> timeoutFunc_{nullptr};
222 FiberManager* fiberManager_{nullptr};
223
timeoutExpired()224 void timeoutExpired() noexcept override {
225 assert(timeoutFunc_ != nullptr);
226 timeoutFunc_();
227 }
228
callbackCanceled()229 void callbackCanceled() noexcept override {}
230 };
231
232 private:
233 class FiberWaiter;
234
Baton(intptr_t state)235 explicit Baton(intptr_t state) : waiter_(state) {}
236
237 void postHelper(intptr_t new_value);
238 void postThread();
239 void waitThread();
240
241 template <typename F>
242 inline void waitFiber(FiberManager& fm, F&& mainContextFunc);
243
244 template <typename Clock, typename Duration>
245 bool timedWaitThread(
246 const std::chrono::time_point<Clock, Duration>& deadline);
247
248 static constexpr intptr_t NO_WAITER = 0;
249 static constexpr intptr_t POSTED = -1;
250 static constexpr intptr_t TIMEOUT = -2;
251 static constexpr intptr_t THREAD_WAITING = -3;
252
253 struct _futex_wrapper {
254 folly::detail::Futex<> futex{};
255 int32_t _unused_packing;
256 };
257
258 union {
259 std::atomic<intptr_t> waiter_;
260 struct _futex_wrapper futex_;
261 };
262 };
263
264 #if FOLLY_HAS_COROUTINES
265 namespace detail {
266 class BatonAwaitableWaiter : public Baton::Waiter {
267 public:
BatonAwaitableWaiter(Baton & baton)268 explicit BatonAwaitableWaiter(Baton& baton) : baton_(baton) {}
269
post()270 void post() override {
271 assert(h_);
272 h_();
273 }
274
await_ready()275 bool await_ready() const { return baton_.ready(); }
276
await_resume()277 void await_resume() {}
278
await_suspend(coro::coroutine_handle<> h)279 void await_suspend(coro::coroutine_handle<> h) {
280 assert(!h_);
281 h_ = std::move(h);
282 baton_.setWaiter(*this);
283 }
284
285 private:
286 coro::coroutine_handle<> h_;
287 Baton& baton_;
288 };
289 } // namespace detail
290
co_await(Baton & baton)291 inline detail::BatonAwaitableWaiter /* implicit */ operator co_await(
292 Baton& baton) {
293 return detail::BatonAwaitableWaiter(baton);
294 }
295 #endif
296 } // namespace fibers
297 } // namespace folly
298
299 #include <folly/fibers/Baton-inl.h>
300