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

folly/experimental/coro/BlockingWait.h

@@ -196,7 +196,10 @@ class BlockingWaitTask {
     folly::Try<detail::lift_lvalue_reference_t<T>> result;
     auto& promise = coro_.promise();
     promise.setTry(&result);
-    coro_.resume();
+    {
+      RequestContextScopeGuard guard{RequestContext::saveContext()};
+      coro_.resume();
+    }
     promise.wait();
     return result;
   }
@@ -209,7 +212,10 @@ class BlockingWaitTask {
     folly::Try<detail::lift_lvalue_reference_t<T>> result;
     auto& promise = coro_.promise();
     promise.setTry(&result);
-    coro_.resume();
+    {
+      RequestContextScopeGuard guard{RequestContext::saveContext()};
+      coro_.resume();
+    }
     while (!promise.done()) {
       executor->drive();
     }

folly/experimental/coro/test/BlockingWaitTest.cpp

@@ -249,4 +249,20 @@ TEST(BlockingWait, WaitOnSemiFuture) {
   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