Commit 136979f1 authored by Nick Wolchko's avatar Nick Wolchko Committed by Facebook Github Bot

fix blockingWait interaction with RequestContext

Summary:
We need to wrap coro_.resume() with a RequestContextScopeGuard so that
if the task changes the request context, we restore it after the coroutine
suspends itself.

Reviewed By: lewissbaker

Differential Revision: D16497375

fbshipit-source-id: 439a6f41b2294e56d2ec1b0013a6013524b5ae6a
parent db4dba1b
...@@ -196,7 +196,10 @@ class BlockingWaitTask { ...@@ -196,7 +196,10 @@ class BlockingWaitTask {
folly::Try<detail::lift_lvalue_reference_t<T>> result; folly::Try<detail::lift_lvalue_reference_t<T>> result;
auto& promise = coro_.promise(); auto& promise = coro_.promise();
promise.setTry(&result); promise.setTry(&result);
coro_.resume(); {
RequestContextScopeGuard guard{RequestContext::saveContext()};
coro_.resume();
}
promise.wait(); promise.wait();
return result; return result;
} }
...@@ -209,7 +212,10 @@ class BlockingWaitTask { ...@@ -209,7 +212,10 @@ class BlockingWaitTask {
folly::Try<detail::lift_lvalue_reference_t<T>> result; folly::Try<detail::lift_lvalue_reference_t<T>> result;
auto& promise = coro_.promise(); auto& promise = coro_.promise();
promise.setTry(&result); promise.setTry(&result);
coro_.resume(); {
RequestContextScopeGuard guard{RequestContext::saveContext()};
coro_.resume();
}
while (!promise.done()) { while (!promise.done()) {
executor->drive(); executor->drive();
} }
......
...@@ -249,4 +249,20 @@ TEST(BlockingWait, WaitOnSemiFuture) { ...@@ -249,4 +249,20 @@ TEST(BlockingWait, WaitOnSemiFuture) {
CHECK_EQ(result, 123); CHECK_EQ(result, 123);
} }
TEST(BlockingWait, RequestContext) {
folly::RequestContext::create();
std::shared_ptr<folly::RequestContext> ctx1, ctx2;
ctx1 = folly::RequestContext::saveContext();
folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
EXPECT_EQ(ctx1.get(), folly::RequestContext::get());
folly::RequestContextScopeGuard guard;
ctx2 = folly::RequestContext::saveContext();
EXPECT_NE(ctx1, ctx2);
co_await folly::coro::co_reschedule_on_current_executor;
EXPECT_EQ(ctx2.get(), folly::RequestContext::get());
co_return;
}());
EXPECT_EQ(ctx1.get(), folly::RequestContext::get());
}
#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