Commit 4d2f1db9 authored by Lewis Baker's avatar Lewis Baker Committed by Facebook GitHub Bot

Allow the CancellationToken to propagate through the co_awaitTry() algorithm

Summary:
Modify the TrySemiAwaitable adapter so that it customises the co_withCancellation() customisation-point and passes through the CancellationToken to the underlying awaitable (eg. a Task).

This fixes code that tries to do the following so that the cancellation token propagates to a Task adapted with co_awaitTry().
```
co_await co_withCancellation(cancelToken, co_awaitTry(co_foo()));
```

Reviewed By: andriigrynenko

Differential Revision: D20443798

fbshipit-source-id: 945e0f98abad4cdef0be29365642006c5687ad5f
parent e38fc326
......@@ -22,6 +22,7 @@
#include <folly/Executor.h>
#include <folly/Traits.h>
#include <folly/experimental/coro/Traits.h>
#include <folly/experimental/coro/WithCancellation.h>
#include <folly/io/async/Request.h>
#include <folly/lang/CustomizationPoint.h>
......@@ -431,6 +432,16 @@ class TrySemiAwaitable {
co_viaIfAsync(std::move(executor), std::move(self.semiAwaitable_))));
}
friend auto co_withCancellation(
const CancellationToken& cancelToken,
TrySemiAwaitable&& awaitable) {
auto cancelAwaitable = folly::coro::co_withCancellation(
std::move(cancelToken),
static_cast<SemiAwaitable&&>(awaitable.semiAwaitable_));
return TrySemiAwaitable<decltype(cancelAwaitable)>(
std::move(cancelAwaitable));
}
private:
SemiAwaitable semiAwaitable_;
};
......
......@@ -409,6 +409,20 @@ TEST_F(TaskTest, CancellationPropagation) {
}());
}
TEST_F(TaskTest, CancellationPropagatesThroughCoAwaitTry) {
folly::CancellationSource source;
folly::Try<int> result =
folly::coro::blockingWait(folly::coro::co_withCancellation(
source.getToken(),
folly::coro::co_awaitTry([&]() -> folly::coro::Task<int> {
auto cancelToken =
co_await folly::coro::co_current_cancellation_token;
EXPECT_TRUE(cancelToken == source.getToken());
co_return 42;
}())));
EXPECT_EQ(42, result.value());
}
TEST_F(TaskTest, StartInlineUnsafe) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto executor = co_await folly::coro::co_current_executor;
......
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