Commit f01cb070 authored by Lewis Baker's avatar Lewis Baker Committed by Facebook Github Bot

Add .startInlineUnsafe() method to folly::coro::Task

Summary: This allows the caller to guarantee that the coroutine starts executing on the current thread. This can be necessary for correctness in some situations, where having the coroutine enqueued to the executor, as `.start()` does, and potentially allowing other code to run in the meantime could invalidate some assumptions.

Reviewed By: andriigrynenko

Differential Revision: D18076197

fbshipit-source-id: a6c5c226f2d7189f37e8f9713a30896767e9b8e9
parent 1b99423a
...@@ -259,6 +259,45 @@ class FOLLY_NODISCARD TaskWithExecutor { ...@@ -259,6 +259,45 @@ class FOLLY_NODISCARD TaskWithExecutor {
}(std::move(*this), std::forward<F>(tryCallback)); }(std::move(*this), std::forward<F>(tryCallback));
} }
// Start execution of this task eagerly, inline on the current thread.
// Assumes that the current thread is already on the associated execution
// context.
template <typename F>
void startInlineUnsafe(
F&& tryCallback,
folly::CancellationToken cancelToken = {}) && {
coro_.promise().setCancelToken(std::move(cancelToken));
RequestContextScopeGuard contextScope{RequestContext::saveContext()};
[](TaskWithExecutor task,
std::decay_t<F> cb) -> detail::InlineTaskDetached {
try {
cb(co_await InlineAwaiter{std::exchange(task.coro_, {})});
} catch (const std::exception& e) {
cb(Try<StorageType>(exception_wrapper(std::current_exception(), e)));
} catch (...) {
cb(Try<StorageType>(exception_wrapper(std::current_exception())));
}
}(std::move(*this), std::forward<F>(tryCallback));
}
// Start execution of this task eagerly inline on the current thread,
// assuming the current thread is already on the associated executor,
// and return a folly::SemiFuture<T> that will complete with the result.
auto startInlineUnsafe() && {
Promise<lift_unit_t<StorageType>> p;
auto sf = p.getSemiFuture();
std::move(*this).startInlineUnsafe(
[promise = std::move(p)](Try<StorageType>&& result) mutable {
promise.setTry(std::move(result));
});
return sf;
}
template <typename ResultCreator> template <typename ResultCreator>
class Awaiter { class Awaiter {
public: public:
...@@ -301,6 +340,41 @@ class FOLLY_NODISCARD TaskWithExecutor { ...@@ -301,6 +340,41 @@ class FOLLY_NODISCARD TaskWithExecutor {
handle_t coro_; handle_t coro_;
}; };
class InlineAwaiter {
public:
InlineAwaiter(handle_t coro) noexcept : coro_(coro) {}
~InlineAwaiter() {
if (coro_) {
coro_.destroy();
}
}
bool await_ready() {
return false;
}
auto await_suspend(std::experimental::coroutine_handle<> continuation) {
auto& promise = coro_.promise();
DCHECK(!promise.continuation_);
DCHECK(promise.executor_);
promise.continuation_ = continuation;
return coro_;
}
folly::Try<StorageType> await_resume() {
// Eagerly destroy the coroutine-frame once we have retrieved the result.
SCOPE_EXIT {
std::exchange(coro_, {}).destroy();
};
return std::move(coro_.promise().result());
}
private:
handle_t coro_;
};
struct ValueCreator { struct ValueCreator {
T operator()(Try<StorageType>&& t) const { T operator()(Try<StorageType>&& t) const {
return std::move(t).value(); return std::move(t).value();
......
...@@ -406,4 +406,30 @@ TEST(Task, CancellationPropagation) { ...@@ -406,4 +406,30 @@ TEST(Task, CancellationPropagation) {
}()); }());
} }
TEST(Task, StartInlineUnsafe) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto executor = co_await folly::coro::co_current_executor;
bool hasStarted = false;
bool hasFinished = false;
auto sf = [&]() -> folly::coro::Task<void> {
hasStarted = true;
co_await folly::coro::co_reschedule_on_current_executor;
hasFinished = true;
}()
.scheduleOn(executor)
.startInlineUnsafe();
// Check that the task started inline on the current thread.
// It should not yet have completed, however, since the rest
// of the coroutine needs this coroutine to suspend so the
// executor can schedule the rest of it.
CHECK(hasStarted);
CHECK(!hasFinished);
co_await std::move(sf);
CHECK(hasFinished);
}());
}
#endif #endif
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment