Commit 3c2162df authored by Andrii Grynenko's avatar Andrii Grynenko Committed by Facebook Github Bot

Fix SemiAwaitable coro::blockingWait to work on fibers

Summary: We can't use futures::WaitExecutor here, because it may drop some work if detached, which is not supported by coro::Task.

Reviewed By: lewissbaker

Differential Revision: D16664921

fbshipit-source-id: 9797b2c59af156ec96e64244d5810a5924cfb3cc
parent bd918f3f
...@@ -282,6 +282,59 @@ auto makeRefBlockingWaitTask(Awaitable&& awaitable) ...@@ -282,6 +282,59 @@ auto makeRefBlockingWaitTask(Awaitable&& awaitable)
co_yield co_await static_cast<Awaitable&&>(awaitable); co_yield co_await static_cast<Awaitable&&>(awaitable);
} }
class BlockingWaitExecutor final : public folly::DrivableExecutor {
public:
~BlockingWaitExecutor() {
while (keepAliveCount_.load() > 0) {
drive();
}
}
void add(Func func) override {
auto wQueue = queue_.wlock();
bool empty = wQueue->empty();
wQueue->push_back(std::move(func));
if (empty) {
baton_.post();
}
}
void drive() override {
baton_.wait();
baton_.reset();
folly::fibers::runInMainContext([&]() {
std::vector<Func> funcs;
queue_.swap(funcs);
for (auto& func : funcs) {
std::exchange(func, nullptr)();
}
});
}
private:
bool keepAliveAcquire() override {
auto keepAliveCount =
keepAliveCount_.fetch_add(1, std::memory_order_relaxed);
DCHECK(keepAliveCount >= 0);
return true;
}
void keepAliveRelease() override {
auto keepAliveCount =
keepAliveCount_.fetch_sub(1, std::memory_order_acq_rel);
DCHECK(keepAliveCount > 0);
if (keepAliveCount == 1) {
add([] {});
}
}
folly::Synchronized<std::vector<Func>> queue_;
fibers::Baton baton_;
std::atomic<ssize_t> keepAliveCount_{0};
};
} // namespace detail } // namespace detail
/// blockingWait(Awaitable&&) -> await_result_t<Awaitable> /// blockingWait(Awaitable&&) -> await_result_t<Awaitable>
...@@ -310,7 +363,7 @@ template < ...@@ -310,7 +363,7 @@ template <
std::enable_if_t<!is_awaitable_v<SemiAwaitable>, int> = 0> std::enable_if_t<!is_awaitable_v<SemiAwaitable>, int> = 0>
auto blockingWait(SemiAwaitable&& awaitable) auto blockingWait(SemiAwaitable&& awaitable)
-> detail::decay_rvalue_reference_t<semi_await_result_t<SemiAwaitable>> { -> detail::decay_rvalue_reference_t<semi_await_result_t<SemiAwaitable>> {
folly::ManualExecutor executor; detail::BlockingWaitExecutor executor;
return static_cast< return static_cast<
std::add_rvalue_reference_t<semi_await_result_t<SemiAwaitable>>>( std::add_rvalue_reference_t<semi_await_result_t<SemiAwaitable>>>(
detail::makeRefBlockingWaitTask( detail::makeRefBlockingWaitTask(
......
...@@ -244,6 +244,26 @@ TEST(BlockingWait, WaitInFiber) { ...@@ -244,6 +244,26 @@ TEST(BlockingWait, WaitInFiber) {
EXPECT_EQ(42, std::move(future).get()); EXPECT_EQ(42, std::move(future).get());
} }
TEST(BlockingWait, WaitTaskInFiber) {
SimplePromise<int> promise;
folly::EventBase evb;
auto& fm = folly::fibers::getFiberManager(evb);
auto future = fm.addTaskFuture([&] {
return folly::coro::blockingWait(folly::coro::co_invoke(
[&]() -> folly::coro::Task<int> { co_return co_await promise; }));
});
evb.loopOnce();
EXPECT_FALSE(future.isReady());
promise.emplace(42);
evb.loopOnce();
EXPECT_TRUE(future.isReady());
EXPECT_EQ(42, std::move(future).get());
}
TEST(BlockingWait, WaitOnSemiFuture) { TEST(BlockingWait, WaitOnSemiFuture) {
int result = folly::coro::blockingWait(folly::makeSemiFuture(123)); int result = folly::coro::blockingWait(folly::makeSemiFuture(123));
CHECK_EQ(result, 123); CHECK_EQ(result, 123);
......
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