diff --git a/folly/coro/BUCK b/folly/coro/BUCK index b5a9fbd76b8..fb0269c6093 100644 --- a/folly/coro/BUCK +++ b/folly/coro/BUCK @@ -283,10 +283,7 @@ cpp_library( name = "future_util", headers = ["FutureUtil.h"], exported_deps = [ - "//folly:cancellation_token", - "//folly/experimental/coro:baton", "//folly/experimental/coro:coroutine", - "//folly/experimental/coro:current_executor", "//folly/experimental/coro:detail_traits", "//folly/experimental/coro:invoke", "//folly/experimental/coro:task", @@ -489,8 +486,10 @@ cpp_library( "Sleep-inl.h", ], exported_deps = [ + "//folly:cancellation_token", + "//folly/experimental/coro:baton", "//folly/experimental/coro:coroutine", - "//folly/experimental/coro:future_util", + "//folly/experimental/coro:current_executor", "//folly/experimental/coro:task", "//folly/futures:core", ], diff --git a/folly/coro/FutureUtil.h b/folly/coro/FutureUtil.h index 98339b9fe61..82508b0b028 100644 --- a/folly/coro/FutureUtil.h +++ b/folly/coro/FutureUtil.h @@ -16,10 +16,7 @@ #pragma once -#include -#include #include -#include #include #include #include @@ -51,40 +48,6 @@ inline Task toTask(folly::SemiFuture a) { co_yield co_result(co_await co_awaitTry(std::move(a))); } -template -Task> toTaskInterruptOnCancel(folly::SemiFuture sf) { - bool cancelled{false}; - Baton baton; - Try result; - auto f = std::move(sf).via(co_await co_current_executor); - f.setCallback_( - [&result, &baton](Executor::KeepAlive<>&&, Try&& t) { - result = std::move(t); - baton.post(); - }, - // No user logic runs in the callback, we can avoid the cost of switching - // the context. - /* context */ nullptr); - - { - CancellationCallback cancelCallback( - co_await co_current_cancellation_token, [&]() noexcept { - cancelled = true; - f.cancel(); - }); - co_await baton; - } - if (cancelled) { - co_yield co_cancelled; - } - co_yield co_result(std::move(result)); -} - -template -Task> toTaskInterruptOnCancel(folly::Future f) { - return toTaskInterruptOnCancel(std::move(f).semi()); -} - // Converts the given SemiAwaitable to a SemiFuture (without starting it) template folly::SemiFuture< diff --git a/folly/coro/Sleep-inl.h b/folly/coro/Sleep-inl.h index 267301d5fad..baa18d1bd85 100644 --- a/folly/coro/Sleep-inl.h +++ b/folly/coro/Sleep-inl.h @@ -16,7 +16,9 @@ #pragma once -#include +#include +#include +#include #if FOLLY_HAS_COROUTINES @@ -24,7 +26,31 @@ namespace folly { namespace coro { inline Task sleep(HighResDuration d, Timekeeper* tk) { - co_await co_nothrow(toTaskInterruptOnCancel(folly::futures::sleep(d, tk))); + bool cancelled{false}; + folly::coro::Baton baton; + Try result; + auto future = folly::futures::sleep(d, tk).toUnsafeFuture(); + future.setCallback_( + [&result, &baton](Executor::KeepAlive<>&&, Try&& t) { + result = std::move(t); + baton.post(); + }, + // No user logic runs in the callback, we can avoid the cost of switching + // the context. + /* context */ nullptr); + + { + CancellationCallback cancelCallback( + co_await co_current_cancellation_token, [&]() noexcept { + cancelled = true; + future.cancel(); + }); + co_await baton; + } + if (cancelled) { + co_yield co_cancelled; + } + co_yield co_result(std::move(result)); } inline Task sleepReturnEarlyOnCancel(HighResDuration d, Timekeeper* tk) { diff --git a/folly/coro/test/FutureUtilTest.cpp b/folly/coro/test/FutureUtilTest.cpp index e71ad241d38..1b30f638f4b 100644 --- a/folly/coro/test/FutureUtilTest.cpp +++ b/folly/coro/test/FutureUtilTest.cpp @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -114,53 +113,4 @@ TEST(FutureUtilTest, VoidRoundtrip) { task = folly::coro::toTask(std::move(semi)); folly::coro::blockingWait(std::move(task)); } - -TEST(FutureUtilTest, ToTaskInterruptOnCancelFutureWithCancellation) { - auto [p, f] = folly::makePromiseContract(); - - // to verify that cancellation propagates into the future interrupt-handler - folly::exception_wrapper interrupt; - p.setInterruptHandler([&, p_ = &p](auto&& ew) { - interrupt = ew; - p_->setException(std::move(ew)); - }); - - // to verify that deferred work runs - folly::Try touched; - auto f1 = std::move(f).defer([&](folly::Try t) { touched = t; }); - ASSERT_FALSE(touched.tryGetExceptionObject()); // sanity check - - // run the scenario within blocking-wait - auto result = folly::coro::blocking_wait( - std::invoke([&, f_ = &f1]() -> folly::coro::Task> { - folly::CancellationSource csource; - - co_return std::get<0>(co_await folly::coro::collectAllTry( - - folly::coro::co_withCancellation( - csource.getToken(), - // a task that will be cancelled, wrapping a future - std::invoke([&]() -> folly::coro::Task<> { - EXPECT_FALSE(touched.tryGetExceptionObject()); // sanity check - co_await folly::coro::toTaskInterruptOnCancel(std::move(*f_)); - })), - - // a task that will do the cancelling, after waiting a bit - std::invoke([&]() -> folly::coro::Task<> { - EXPECT_FALSE(touched.tryGetExceptionObject()); // sanity check - co_await folly::coro::co_reschedule_on_current_executor; - EXPECT_FALSE(touched.tryGetExceptionObject()); // sanity check - csource.requestCancellation(); - EXPECT_FALSE(touched.tryGetExceptionObject()); // sanity check - co_await folly::coro::co_reschedule_on_current_executor; - EXPECT_TRUE( // sanity check: events happen here - touched.tryGetExceptionObject()); - }))); - })); - - EXPECT_TRUE(touched.tryGetExceptionObject()); - EXPECT_TRUE(result.tryGetExceptionObject()); - EXPECT_TRUE(interrupt.get_exception()); -} - #endif