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 <folly/ScopeGuard.h>
20 #include <folly/Try.h>
21 #include <folly/experimental/coro/Coroutine.h>
22 #include <folly/experimental/coro/WithAsyncStack.h>
23 #include <folly/experimental/coro/detail/Malloc.h>
24 #include <folly/lang/Assume.h>
25 #include <folly/tracing/AsyncStack.h>
26
27 #include <cassert>
28 #include <utility>
29
30 #if FOLLY_HAS_COROUTINES
31
32 namespace folly {
33 namespace coro {
34 namespace detail {
35
36 /// InlineTask<T> is a coroutine-return type where the coroutine is launched
37 /// inline in the current execution context when it is co_awaited and the
38 /// task's continuation is launched inline in the execution context that the
39 /// task completed on.
40 ///
41 /// This task type is primarily intended as a building block for certain
42 /// coroutine operators. It is not intended for general use in application
43 /// code or in library interfaces exposed to library code as it can easily be
44 /// abused to accidentally run logic on the wrong execution context.
45 ///
46 /// For this reason, the InlineTask<T> type has been placed inside the
47 /// folly::coro::detail namespace to discourage general usage.
48 template <typename T>
49 class InlineTask;
50
51 class InlineTaskPromiseBase {
52 struct FinalAwaiter {
await_readyFinalAwaiter53 bool await_ready() noexcept { return false; }
54
55 template <typename Promise>
await_suspendFinalAwaiter56 coroutine_handle<> await_suspend(coroutine_handle<Promise> h) noexcept {
57 InlineTaskPromiseBase& promise = h.promise();
58 return promise.continuation_;
59 }
60
await_resumeFinalAwaiter61 void await_resume() noexcept {}
62 };
63
64 protected:
65 InlineTaskPromiseBase() noexcept = default;
66
67 InlineTaskPromiseBase(const InlineTaskPromiseBase&) = delete;
68 InlineTaskPromiseBase(InlineTaskPromiseBase&&) = delete;
69 InlineTaskPromiseBase& operator=(const InlineTaskPromiseBase&) = delete;
70 InlineTaskPromiseBase& operator=(InlineTaskPromiseBase&&) = delete;
71
72 public:
new(std::size_t size)73 static void* operator new(std::size_t size) {
74 return ::folly_coro_async_malloc(size);
75 }
76
delete(void * ptr,std::size_t size)77 static void operator delete(void* ptr, std::size_t size) {
78 ::folly_coro_async_free(ptr, size);
79 }
80
initial_suspend()81 suspend_always initial_suspend() noexcept { return {}; }
82
final_suspend()83 auto final_suspend() noexcept { return FinalAwaiter{}; }
84
set_continuation(coroutine_handle<> continuation)85 void set_continuation(coroutine_handle<> continuation) noexcept {
86 assert(!continuation_);
87 continuation_ = continuation;
88 }
89
90 private:
91 coroutine_handle<> continuation_;
92 };
93
94 template <typename T>
95 class InlineTaskPromise : public InlineTaskPromiseBase {
96 public:
97 static_assert(
98 std::is_move_constructible<T>::value,
99 "InlineTask<T> only supports types that are move-constructible.");
100 static_assert(
101 !std::is_rvalue_reference<T>::value, "InlineTask<T&&> is not supported");
102
103 InlineTaskPromise() noexcept = default;
104
105 ~InlineTaskPromise() = default;
106
107 InlineTask<T> get_return_object() noexcept;
108
109 template <
110 typename Value = T,
111 std::enable_if_t<std::is_convertible<Value&&, T>::value, int> = 0>
return_value(Value && value)112 void return_value(Value&& value) noexcept(
113 std::is_nothrow_constructible<T, Value&&>::value) {
114 result_.emplace(static_cast<Value&&>(value));
115 }
116
unhandled_exception()117 void unhandled_exception() noexcept {
118 result_.emplaceException(
119 folly::exception_wrapper{std::current_exception()});
120 }
121
result()122 T result() { return std::move(result_).value(); }
123
124 private:
125 // folly::Try<T> doesn't support storing reference types so we store a
126 // std::reference_wrapper instead.
127 using StorageType = std::conditional_t<
128 std::is_lvalue_reference<T>::value,
129 std::reference_wrapper<std::remove_reference_t<T>>,
130 T>;
131
132 folly::Try<StorageType> result_;
133 };
134
135 template <>
136 class InlineTaskPromise<void> : public InlineTaskPromiseBase {
137 public:
138 InlineTaskPromise() noexcept = default;
139
140 InlineTask<void> get_return_object() noexcept;
141
return_void()142 void return_void() noexcept {}
143
unhandled_exception()144 void unhandled_exception() noexcept {
145 result_.emplaceException(
146 folly::exception_wrapper{std::current_exception()});
147 }
148
result()149 void result() { return result_.value(); }
150
151 private:
152 folly::Try<void> result_;
153 };
154
155 template <typename T>
156 class InlineTask {
157 public:
158 using promise_type = detail::InlineTaskPromise<T>;
159
160 private:
161 using handle_t = coroutine_handle<promise_type>;
162
163 public:
InlineTask(InlineTask && other)164 InlineTask(InlineTask&& other) noexcept
165 : coro_(std::exchange(other.coro_, {})) {}
166
~InlineTask()167 ~InlineTask() {
168 if (coro_) {
169 coro_.destroy();
170 }
171 }
172
173 class Awaiter {
174 public:
~Awaiter()175 ~Awaiter() {
176 if (coro_) {
177 coro_.destroy();
178 }
179 }
180
await_ready()181 bool await_ready() noexcept { return false; }
182
await_suspend(coroutine_handle<> awaitingCoroutine)183 handle_t await_suspend(coroutine_handle<> awaitingCoroutine) noexcept {
184 assert(coro_ && !coro_.done());
185 coro_.promise().set_continuation(awaitingCoroutine);
186 return coro_;
187 }
188
await_resume()189 T await_resume() {
190 auto destroyOnExit =
191 folly::makeGuard([this] { std::exchange(coro_, {}).destroy(); });
192 return coro_.promise().result();
193 }
194
195 private:
196 friend class InlineTask<T>;
Awaiter(handle_t coro)197 explicit Awaiter(handle_t coro) noexcept : coro_(coro) {}
198 handle_t coro_;
199 };
200
co_await()201 Awaiter operator co_await() && {
202 assert(coro_ && !coro_.done());
203 return Awaiter{std::exchange(coro_, {})};
204 }
205
206 private:
207 friend class InlineTaskPromise<T>;
InlineTask(handle_t coro)208 explicit InlineTask(handle_t coro) noexcept : coro_(coro) {}
209 handle_t coro_;
210 };
211
212 template <typename T>
get_return_object()213 inline InlineTask<T> InlineTaskPromise<T>::get_return_object() noexcept {
214 return InlineTask<T>{
215 coroutine_handle<InlineTaskPromise<T>>::from_promise(*this)};
216 }
217
get_return_object()218 inline InlineTask<void> InlineTaskPromise<void>::get_return_object() noexcept {
219 return InlineTask<void>{
220 coroutine_handle<InlineTaskPromise<void>>::from_promise(*this)};
221 }
222
223 /// InlineTaskDetached is a coroutine-return type where the coroutine is
224 /// launched in the current execution context when it is created and the
225 /// task's continuation is launched inline in the execution context that the
226 /// task completed on.
227 ///
228 /// This task type is primarily intended as a building block for certain
229 /// coroutine operators. It is not intended for general use in application
230 /// code or in library interfaces exposed to library code as it can easily be
231 /// abused to accidentally run logic on the wrong execution context.
232 ///
233 /// For this reason, the InlineTaskDetached type has been placed inside the
234 /// folly::coro::detail namespace to discourage general usage.
235 struct InlineTaskDetached {
236 class promise_type {
237 struct FinalAwaiter {
await_readyInlineTaskDetached::FinalAwaiter238 bool await_ready() noexcept { return false; }
await_suspendInlineTaskDetached::FinalAwaiter239 void await_suspend(coroutine_handle<promise_type> h) noexcept {
240 folly::deactivateAsyncStackFrame(h.promise().getAsyncFrame());
241 h.destroy();
242 }
await_resumeInlineTaskDetached::FinalAwaiter243 [[noreturn]] void await_resume() noexcept { folly::assume_unreachable(); }
244 };
245
246 public:
newInlineTaskDetached247 static void* operator new(std::size_t size) {
248 return ::folly_coro_async_malloc(size);
249 }
250
deleteInlineTaskDetached251 static void operator delete(void* ptr, std::size_t size) {
252 ::folly_coro_async_free(ptr, size);
253 }
254
promise_typeInlineTaskDetached255 promise_type() noexcept {
256 asyncFrame_.setParentFrame(folly::getDetachedRootAsyncStackFrame());
257 }
258
get_return_objectInlineTaskDetached259 InlineTaskDetached get_return_object() noexcept {
260 return InlineTaskDetached{
261 coroutine_handle<promise_type>::from_promise(*this)};
262 }
263
initial_suspendInlineTaskDetached264 suspend_always initial_suspend() noexcept { return {}; }
265
final_suspendInlineTaskDetached266 FinalAwaiter final_suspend() noexcept { return {}; }
267
return_voidInlineTaskDetached268 void return_void() noexcept {}
269
unhandled_exceptionInlineTaskDetached270 [[noreturn]] void unhandled_exception() noexcept { std::terminate(); }
271
272 template <typename Awaitable>
decltypeInlineTaskDetached273 decltype(auto) await_transform(Awaitable&& awaitable) {
274 return folly::coro::co_withAsyncStack(
275 static_cast<Awaitable&&>(awaitable));
276 }
277
getAsyncFrameInlineTaskDetached278 folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; }
279
280 private:
281 folly::AsyncStackFrame asyncFrame_;
282 };
283
InlineTaskDetachedInlineTaskDetached284 InlineTaskDetached(InlineTaskDetached&& other) noexcept
285 : coro_(std::exchange(other.coro_, {})) {}
286
~InlineTaskDetachedInlineTaskDetached287 ~InlineTaskDetached() {
288 if (coro_) {
289 coro_.destroy();
290 }
291 }
292
startInlineTaskDetached293 FOLLY_NOINLINE void start() noexcept {
294 start(FOLLY_ASYNC_STACK_RETURN_ADDRESS());
295 }
296
startInlineTaskDetached297 void start(void* returnAddress) noexcept {
298 coro_.promise().getAsyncFrame().setReturnAddress(returnAddress);
299 folly::resumeCoroutineWithNewAsyncStackRoot(std::exchange(coro_, {}));
300 }
301
302 private:
InlineTaskDetachedInlineTaskDetached303 explicit InlineTaskDetached(coroutine_handle<promise_type> h) noexcept
304 : coro_(h) {}
305
306 coroutine_handle<promise_type> coro_;
307 };
308
309 } // namespace detail
310 } // namespace coro
311 } // namespace folly
312
313 #endif // FOLLY_HAS_COROUTINES
314