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

Integrate CancellationToken support into Task

Summary:
The `folly::coro::Task` coroutine type now has an associated `CancellationToken` that is implicitly passed down to child operations that it awaits.

This is a first step towards supporting cancellation of persistent async streams.

This adds a new `co_withCancellation()` customisation point that awaitable types can customise to allow them to opt-in to cancellation. Currently only `Task` customises this operation.

Also provided a new `co_current_cancellation_token` object that can be awaited within a `Task` to retrieve the current `CancellationToken`.

Note that I have not yet hooked up the `Future` or `SemiFuture` to integrate with this cancellation mechanism (most `Future`-based code is not cancellation-aware). So a coroutine that is currently suspended awaiting for a `Future` may not currently respond to a cancellation request.

Reviewed By: andriigrynenko

Differential Revision: D16610810

fbshipit-source-id: 72a31d7a3ba4c281db54c1942ab41d4ea2f34d21
parent fbb2e01e
......@@ -106,5 +106,17 @@ using co_reschedule_on_current_executor_t =
inline constexpr co_reschedule_on_current_executor_t
co_reschedule_on_current_executor;
namespace detail {
struct co_current_cancellation_token_ {
enum class secret_ { token_ };
explicit constexpr co_current_cancellation_token_(secret_) {}
};
} // namespace detail
using co_current_cancellation_token_t = detail::co_current_cancellation_token_;
inline constexpr co_current_cancellation_token_t co_current_cancellation_token{
co_current_cancellation_token_t::secret_::token_};
} // namespace coro
} // namespace folly
......@@ -20,6 +20,7 @@
#include <glog/logging.h>
#include <folly/CancellationToken.h>
#include <folly/Executor.h>
#include <folly/Portability.h>
#include <folly/ScopeGuard.h>
......@@ -28,6 +29,7 @@
#include <folly/experimental/coro/Traits.h>
#include <folly/experimental/coro/Utils.h>
#include <folly/experimental/coro/ViaIfAsync.h>
#include <folly/experimental/coro/WithCancellation.h>
#include <folly/experimental/coro/detail/InlineTask.h>
#include <folly/experimental/coro/detail/Traits.h>
#include <folly/futures/Future.h>
......@@ -78,13 +80,26 @@ class TaskPromiseBase {
template <typename Awaitable>
auto await_transform(Awaitable&& awaitable) noexcept {
return folly::coro::co_viaIfAsync(
executor_.get_alias(), static_cast<Awaitable&&>(awaitable));
executor_.get_alias(),
folly::coro::co_withCancellation(
cancelToken_, static_cast<Awaitable&&>(awaitable)));
}
auto await_transform(co_current_executor_t) noexcept {
return AwaitableReady<folly::Executor*>{executor_.get()};
}
auto await_transform(co_current_cancellation_token_t) noexcept {
return AwaitableReady<const folly::CancellationToken&>{cancelToken_};
}
void setCancelToken(const folly::CancellationToken& cancelToken) {
if (!hasCancelTokenOverride_) {
cancelToken_ = cancelToken;
hasCancelTokenOverride_ = true;
}
}
private:
template <typename T>
friend class folly::coro::TaskWithExecutor;
......@@ -94,6 +109,8 @@ class TaskPromiseBase {
std::experimental::coroutine_handle<> continuation_;
folly::Executor::KeepAlive<> executor_;
folly::CancellationToken cancelToken_;
bool hasCancelTokenOverride_ = false;
};
template <typename T>
......@@ -212,6 +229,7 @@ class FOLLY_NODISCARD TaskWithExecutor {
// that will complete with the result.
auto start() && {
Promise<lift_unit_t<StorageType>> p;
auto sf = p.getSemiFuture();
std::move(*this).start(
......@@ -224,7 +242,9 @@ class FOLLY_NODISCARD TaskWithExecutor {
// Start execution of this task eagerly and call the callback when complete.
template <typename F>
void start(F&& tryCallback) && {
void start(F&& tryCallback, folly::CancellationToken cancelToken = {}) && {
coro_.promise().setCancelToken(std::move(cancelToken));
[](TaskWithExecutor task,
std::decay_t<F> cb) -> detail::InlineTaskDetached {
try {
......@@ -299,6 +319,13 @@ class FOLLY_NODISCARD TaskWithExecutor {
return Awaiter<TryCreator>{std::exchange(coro_, {})};
}
friend TaskWithExecutor co_withCancellation(
const folly::CancellationToken& cancelToken,
TaskWithExecutor&& task) noexcept {
task.coro_.promise().setCancelToken(cancelToken);
return std::move(task);
}
private:
friend class Task<T>;
......@@ -385,6 +412,13 @@ class FOLLY_NODISCARD Task {
return Awaiter{std::exchange(t.coro_, {})};
}
friend Task co_withCancellation(
const folly::CancellationToken& cancelToken,
Task&& task) noexcept {
task.coro_.promise().setCancelToken(cancelToken);
return std::move(task);
}
private:
friend class detail::TaskPromiseBase;
friend class detail::TaskPromise<T>;
......
/*
* Copyright 2017-present Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include <folly/CancellationToken.h>
#include <folly/lang/CustomizationPoint.h>
namespace folly {
namespace coro {
namespace detail {
namespace adl {
// Default implementation that does not hook the cancellation token.
// Types must opt-in to hooking cancellation by customising this function.
template <typename Awaitable>
Awaitable&& co_withCancellation(
const folly::CancellationToken&,
Awaitable&& awaitable) noexcept {
return (Awaitable &&) awaitable;
}
struct WithCancellationFunction {
template <typename Awaitable>
auto operator()(
const folly::CancellationToken& cancelToken,
Awaitable&& awaitable) const
noexcept(
noexcept(co_withCancellation(cancelToken, (Awaitable &&) awaitable)))
-> decltype(
co_withCancellation(cancelToken, (Awaitable &&) awaitable)) {
return co_withCancellation(cancelToken, (Awaitable &&) awaitable);
}
template <typename Awaitable>
auto operator()(folly::CancellationToken&& cancelToken, Awaitable&& awaitable)
const noexcept(noexcept(co_withCancellation(
std::move(cancelToken),
(Awaitable &&) awaitable)))
-> decltype(co_withCancellation(
std::move(cancelToken),
(Awaitable &&) awaitable)) {
return co_withCancellation(
std::move(cancelToken), (Awaitable &&) awaitable);
}
};
} // namespace adl
} // namespace detail
FOLLY_DEFINE_CPO(detail::adl::WithCancellationFunction, co_withCancellation)
} // namespace coro
} // namespace folly
......@@ -18,13 +18,19 @@
#if FOLLY_HAS_COROUTINES
#include <folly/CancellationToken.h>
#include <folly/Chrono.h>
#include <folly/executors/ManualExecutor.h>
#include <folly/experimental/coro/Baton.h>
#include <folly/experimental/coro/BlockingWait.h>
#include <folly/experimental/coro/Collect.h>
#include <folly/experimental/coro/CurrentExecutor.h>
#include <folly/experimental/coro/Task.h>
#include <folly/experimental/coro/TimedWait.h>
#include <folly/experimental/coro/Utils.h>
#include <folly/experimental/coro/WithCancellation.h>
#include <folly/fibers/Semaphore.h>
#include <folly/futures/Future.h>
#include <folly/io/async/ScopedEventBaseThread.h>
#include <folly/portability/GTest.h>
......@@ -539,29 +545,51 @@ TEST(Coro, FutureTry) {
}());
}
TEST(Coro, CoReturnTry) {
EXPECT_EQ(42, folly::coro::blockingWait([]() -> folly::coro::Task<int> {
co_return folly::Try<int>(42);
}()));
template <typename T>
folly::coro::Task<T> cancellableFuture(folly::SemiFuture<T> future) {
folly::coro::Baton baton;
struct ExpectedException : public std::runtime_error {
ExpectedException() : std::runtime_error("ExpectedException") {}
};
EXPECT_THROW(
folly::coro::blockingWait([]() -> folly::coro::Task<int> {
co_return folly::Try<int>(ExpectedException());
}()),
ExpectedException);
// Attach an executor to ensure that the operation has been started.
auto future2 = std::move(future)
.via(co_await folly::coro::co_current_executor)
.ensure([&]() { baton.post(); });
folly::Try<int> t(42);
EXPECT_EQ(42, folly::coro::blockingWait([&]() -> folly::coro::Task<int> {
co_return t;
}()));
{
folly::CancellationCallback cb{
co_await folly::coro::co_current_cancellation_token,
[&] { future2.cancel(); }};
co_await baton;
}
const folly::Try<int> tConst(42);
EXPECT_EQ(42, folly::coro::blockingWait([&]() -> folly::coro::Task<int> {
co_return tConst;
}()));
co_return co_await std::move(future2);
}
TEST(Coro, CancellableSleep) {
using namespace std::chrono;
using namespace std::chrono_literals;
CancellationSource cancelSrc;
auto start = steady_clock::now();
coro::blockingWait([&]() -> coro::Task<void> {
co_await coro::collectAll(
[&]() -> coro::Task<void> {
try {
co_await coro::co_withCancellation(
cancelSrc.getToken(),
cancellableFuture(folly::futures::sleep(10s)));
} catch (const folly::FutureCancellation&) {
}
}(),
[&]() -> coro::Task<void> {
co_await coro::co_reschedule_on_current_executor;
co_await coro::co_reschedule_on_current_executor;
co_await coro::co_reschedule_on_current_executor;
cancelSrc.requestCancellation();
}());
}());
auto end = steady_clock::now();
CHECK((end - start) < 1s);
}
#endif
......@@ -379,4 +379,31 @@ TEST(Task, TaskOfLvalueReferenceAsTry) {
}());
}
TEST(Task, CancellationPropagation) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto token = co_await folly::coro::co_current_cancellation_token;
CHECK(!token.canBeCancelled());
folly::CancellationSource cancelSource;
co_await folly::coro::co_withCancellation(
cancelSource.getToken(), [&]() -> folly::coro::Task<void> {
auto token2 = co_await folly::coro::co_current_cancellation_token;
CHECK(token2.canBeCancelled());
CHECK(!token2.isCancellationRequested());
// The cancellation token should implicitly propagate into the
//
co_await[&]()->folly::coro::Task<void> {
auto token3 = co_await folly::coro::co_current_cancellation_token;
CHECK(token3 == token2);
cancelSource.requestCancellation();
CHECK(token3.isCancellationRequested());
}
();
CHECK(token2.isCancellationRequested());
}());
}());
}
#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