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

Make folly::coro::collectAll() isolate RequestContext state of child tasks.

Summary:
No longer allow a child task's modifications to the RequestContext made before it's first suspension point to bleed into the next child task's context.

Now always restore the parent task's RequestContext before starting subsequent tasks.

Reviewed By: andriigrynenko

Differential Revision: D17846070

fbshipit-source-id: 44e206d4c833759fc36c7d1bd0f3545ca6bb9ab1
parent 368f00f0
......@@ -93,7 +93,16 @@ auto collectAllTryImpl(
// Use std::initializer_list to ensure that the sub-tasks are launched
// in the order they appear in the parameter pack.
(void)std::initializer_list<int>{(tasks[Indices].start(&barrier), 0)...};
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
(void)std::initializer_list<int>{
(tasks[Indices].start(&barrier),
RequestContext::setContext(context),
0)...};
// Wait for all of the sub-tasks to finish execution.
// Should be safe to avoid an executor transition here even if the
......@@ -137,9 +146,18 @@ auto collectAllImpl(
folly::coro::detail::Barrier barrier{sizeof...(SemiAwaitables) + 1};
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
// Use std::initializer_list to ensure that the sub-tasks are launched
// in the order they appear in the parameter pack.
(void)std::initializer_list<int>{(tasks[Indices].start(&barrier), 0)...};
(void)std::initializer_list<int>{
(tasks[Indices].start(&barrier),
RequestContext::setContext(context),
0)...};
// Wait for all of the sub-tasks to finish execution.
// Should be safe to avoid an executor transition here even if the
......@@ -213,11 +231,18 @@ auto collectAllTryRangeImpl(
// executing the tasks.
results.resize(tasks.size());
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
// Launch the tasks and wait for them all to finish.
{
detail::Barrier barrier{tasks.size() + 1};
for (auto&& task : tasks) {
task.start(&barrier);
RequestContext::setContext(context);
}
co_await detail::UnsafeResumeInlineSemiAwaitable{barrier.arriveAndWait()};
}
......@@ -360,6 +385,12 @@ auto collectAllTryWindowedImpl(
exception_wrapper workerCreationException;
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
try {
auto lock = co_await mutex.co_scoped_lock();
while (iter != iterEnd && workerTasks.size() < maxConcurrency) {
......@@ -374,6 +405,8 @@ auto collectAllTryWindowedImpl(
barrier.add(1);
workerTasks.back().start(&barrier);
RequestContext::setContext(context);
lock = co_await mutex.co_scoped_lock();
}
} catch (const std::exception& ex) {
......@@ -506,11 +539,18 @@ auto collectAllRange(InputRange awaitables) -> folly::coro::Task<void> {
makeTask(static_cast<decltype(semiAwaitable)&&>(semiAwaitable)));
}
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
// Launch the tasks and wait for them all to finish.
{
detail::Barrier barrier{tasks.size() + 1};
for (auto&& task : tasks) {
task.start(&barrier);
RequestContext::setContext(context);
}
co_await detail::UnsafeResumeInlineSemiAwaitable{barrier.arriveAndWait()};
}
......@@ -613,6 +653,12 @@ auto collectAllWindowed(InputRange awaitables, std::size_t maxConcurrency)
exception_wrapper workerCreationException;
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
try {
auto lock = co_await mutex.co_scoped_lock();
......@@ -628,6 +674,8 @@ auto collectAllWindowed(InputRange awaitables, std::size_t maxConcurrency)
barrier.add(1);
workerTasks.back().start(&barrier);
RequestContext::setContext(context);
lock = co_await mutex.co_scoped_lock();
}
} catch (const std::exception& ex) {
......
......@@ -28,6 +28,7 @@
#include <folly/experimental/coro/Mutex.h>
#include <folly/experimental/coro/Sleep.h>
#include <folly/experimental/coro/Task.h>
#include <folly/io/async/Request.h>
#include <folly/portability/GTest.h>
#include <numeric>
......@@ -307,6 +308,55 @@ TEST(CollectAll, CollectAllCancelsSubtasksWhenParentTaskCancelled) {
}());
}
namespace {
class TestRequestData : public folly::RequestData {
public:
explicit TestRequestData() noexcept {}
bool hasCallback() override {
return false;
}
};
} // namespace
TEST(CollectAll, CollectAllKeepsRequestContextOfChildTasksIndependent) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::RequestContextScopeGuard requestScope;
auto getContextData = []() {
return folly::RequestContext::get()->getContextData("test");
};
auto setContextData = []() {
folly::RequestContext::get()->setContextData(
"test", std::make_unique<TestRequestData>());
};
setContextData();
auto initialContextData = getContextData();
auto makeChildTask = [&]() -> folly::coro::Task<void> {
CHECK(getContextData() == initialContextData);
folly::RequestContextScopeGuard childScope;
CHECK(getContextData() == nullptr);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == nullptr);
setContextData();
auto newContextData = getContextData();
CHECK(newContextData != nullptr);
CHECK(newContextData != initialContextData);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == newContextData);
};
co_await folly::coro::collectAll(makeChildTask(), makeChildTask());
CHECK(getContextData() == initialContextData);
}());
}
/////////////////////////////////////////////////////////
// folly::coro::collectAllTry() tests
......@@ -464,6 +514,42 @@ TEST(CollectAllTry, CollectAllCancelsSubtasksWhenParentTaskCancelled) {
}());
}
TEST(CollectAllTry, CollectAllTryKeepsRequestContextOfChildTasksIndependent) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::RequestContextScopeGuard requestScope;
auto getContextData = []() {
return folly::RequestContext::get()->getContextData("test");
};
auto setContextData = []() {
folly::RequestContext::get()->setContextData(
"test", std::make_unique<TestRequestData>());
};
setContextData();
auto initialContextData = getContextData();
auto makeChildTask = [&]() -> folly::coro::Task<void> {
CHECK(getContextData() == initialContextData);
folly::RequestContextScopeGuard childScope;
CHECK(getContextData() == nullptr);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == nullptr);
setContextData();
auto newContextData = getContextData();
CHECK(newContextData != nullptr);
CHECK(newContextData != initialContextData);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == newContextData);
};
co_await folly::coro::collectAllTry(makeChildTask(), makeChildTask());
CHECK(getContextData() == initialContextData);
}());
}
/////////////////////////////////////////////////////////////
// collectAllRange() tests
......@@ -633,6 +719,49 @@ TEST(CollectAllRange, SubtasksCancelledWhenParentTaskCancelled) {
}());
}
TEST(
CollectAllRange,
CollectAllRangeKeepsRequestContextOfChildTasksIndependent) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::RequestContextScopeGuard requestScope;
auto getContextData = []() {
return folly::RequestContext::get()->getContextData("test");
};
auto setContextData = []() {
folly::RequestContext::get()->setContextData(
"test", std::make_unique<TestRequestData>());
};
setContextData();
auto initialContextData = getContextData();
auto makeChildTask = [&]() -> folly::coro::Task<void> {
CHECK(getContextData() == initialContextData);
folly::RequestContextScopeGuard childScope;
CHECK(getContextData() == nullptr);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == nullptr);
setContextData();
auto newContextData = getContextData();
CHECK(newContextData != nullptr);
CHECK(newContextData != initialContextData);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == newContextData);
};
std::vector<folly::coro::Task<void>> tasks;
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
co_await folly::coro::collectAllRange(std::move(tasks));
CHECK(getContextData() == initialContextData);
}());
}
////////////////////////////////////////////////////////////////////
// folly::coro::collectAllTryRange() tests
......@@ -776,6 +905,49 @@ TEST(CollectAllTryRange, SubtasksCancelledWhenParentTaskCancelled) {
}());
}
TEST(
CollectAllTryRange,
CollectAllTryRangeKeepsRequestContextOfChildTasksIndependent) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::RequestContextScopeGuard requestScope;
auto getContextData = []() {
return folly::RequestContext::get()->getContextData("test");
};
auto setContextData = []() {
folly::RequestContext::get()->setContextData(
"test", std::make_unique<TestRequestData>());
};
setContextData();
auto initialContextData = getContextData();
auto makeChildTask = [&]() -> folly::coro::Task<void> {
CHECK(getContextData() == initialContextData);
folly::RequestContextScopeGuard childScope;
CHECK(getContextData() == nullptr);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == nullptr);
setContextData();
auto newContextData = getContextData();
CHECK(newContextData != nullptr);
CHECK(newContextData != initialContextData);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == newContextData);
};
std::vector<folly::coro::Task<void>> tasks;
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
co_await folly::coro::collectAllTryRange(std::move(tasks));
CHECK(getContextData() == initialContextData);
}());
}
////////////////////////////////////////////////////////////////////
// folly::coro::collectAllWindowed() tests
......@@ -1024,6 +1196,49 @@ TEST(CollectAllWindowed, SubtasksCancelledWhenParentTaskCancelled) {
}());
}
TEST(
CollectAllWindowed,
CollectAllWindowedKeepsRequestContextOfChildTasksIndependent) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::RequestContextScopeGuard requestScope;
auto getContextData = []() {
return folly::RequestContext::get()->getContextData("test");
};
auto setContextData = []() {
folly::RequestContext::get()->setContextData(
"test", std::make_unique<TestRequestData>());
};
setContextData();
auto initialContextData = getContextData();
auto makeChildTask = [&]() -> folly::coro::Task<void> {
CHECK(getContextData() == initialContextData);
folly::RequestContextScopeGuard childScope;
CHECK(getContextData() == nullptr);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == nullptr);
setContextData();
auto newContextData = getContextData();
CHECK(newContextData != nullptr);
CHECK(newContextData != initialContextData);
co_await folly::coro::co_reschedule_on_current_executor;
CHECK(getContextData() == newContextData);
};
std::vector<folly::coro::Task<void>> tasks;
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
tasks.emplace_back(makeChildTask());
co_await folly::coro::collectAllWindowed(std::move(tasks), 2);
CHECK(getContextData() == initialContextData);
}());
}
////////////////////////////////////////////////////////////////////
// folly::coro::collectAllTryWindowed() tests
......
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