AtScopeExit for coroutine cleanups

Reviewed By: iahs

Differential Revision: D31271338

fbshipit-source-id: 10cbd083c4da7461a43d4d0a9d4ab5455dcee270

folly/experimental/coro/Task.h

@@ -32,6 +32,7 @@
 #include <folly/experimental/coro/CurrentExecutor.h>
 #include <folly/experimental/coro/Invoke.h>
 #include <folly/experimental/coro/Result.h>
+#include <folly/experimental/coro/ScopeExit.h>
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithAsyncStack.h>
@@ -66,7 +67,15 @@ class TaskPromiseBase {
     FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES coroutine_handle<>
     await_suspend(coroutine_handle<Promise> coro) noexcept {
       TaskPromiseBase& promise = coro.promise();
+      // If the continuation has been exchanged, then we expect that the
+      // exchanger will handle the lifetime of the async stack. See
+      // ScopeExitTaskPromise's FinalAwaiter for more details.
+      //
+      // This is a bit untidy, and hopefully something we can replace with
+      // a virtual wrapper over coroutine_handle that handles the pop for us.
+      if (promise.ownsAsyncFrame_) {
         folly::popAsyncStackFrameCallee(promise.asyncFrame_);
+      }
       return promise.continuation_;
     }
 
@@ -125,6 +134,10 @@ class TaskPromiseBase {
 
   folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; }
 
+  folly::Executor::KeepAlive<> getExecutor() const noexcept {
+    return executor_;
+  }
+
  private:
   template <typename T>
   friend class folly::coro::TaskWithExecutor;
@@ -132,11 +145,22 @@ class TaskPromiseBase {
   template <typename T>
   friend class folly::coro::Task;
 
+  friend std::tuple<bool, coroutine_handle<>> tag_invoke(
+      cpo_t<co_attachScopeExit>,
+      TaskPromiseBase& p,
+      coroutine_handle<> continuation) noexcept {
+    return {
+        std::exchange(p.ownsAsyncFrame_, false),
+        std::exchange(p.continuation_, continuation),
+    };
+  }
+
   coroutine_handle<> continuation_;
   folly::AsyncStackFrame asyncFrame_;
   folly::Executor::KeepAlive<> executor_;
   folly::CancellationToken cancelToken_;
   bool hasCancelTokenOverride_ = false;
+  bool ownsAsyncFrame_ = true;
 };
 
 template <typename T>

folly/experimental/coro/test/ScopeExitTest.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+
+#include <folly/experimental/coro/AsyncScope.h>
+#include <folly/experimental/coro/Baton.h>
+#include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Task.h>
+#include <folly/portability/GTest.h>
+
+#include <stdexcept>
+#include <type_traits>
+
+#if FOLLY_HAS_COROUTINES
+
+using namespace folly::coro;
+
+namespace {
+class ScopeExitTest : public testing::Test {
+ protected:
+  int count = 0;
+};
+} // namespace
+
+TEST_F(ScopeExitTest, OneExitAction) {
+  folly::coro::blockingWait([this]() -> Task<> {
+    ++count;
+    co_await co_scope_exit([this]() -> Task<> {
+      count *= 2;
+      co_return;
+    });
+    ++count;
+  }());
+  EXPECT_EQ(count, 4);
+}
+
+TEST_F(ScopeExitTest, TwoExitActions) {
+  folly::coro::blockingWait([this]() -> Task<> {
+    ++count;
+    co_await co_scope_exit([this]() -> Task<> {
+      count *= count;
+      co_return;
+    });
+    co_await co_scope_exit([this]() -> Task<> {
+      count *= 2;
+      co_return;
+    });
+    ++count;
+  }());
+  EXPECT_EQ(count, 16);
+}
+
+TEST_F(ScopeExitTest, OneExitActionWithException) {
+  EXPECT_THROW(
+      folly::coro::blockingWait([this]() -> Task<> {
+        ++count;
+        co_await co_scope_exit([this]() -> Task<> {
+          count *= 2;
+          co_return;
+        });
+        throw std::runtime_error("Something bad happened!");
+      }()),
+      std::runtime_error);
+  EXPECT_EQ(count, 2);
+}
+
+TEST_F(ScopeExitTest, ExceptionInExitActionCausesTermination) {
+  ASSERT_DEATH(
+      folly::coro::blockingWait([]() -> Task<> {
+        co_await co_scope_exit([]() -> Task<> {
+          throw std::runtime_error("Something bad happened!");
+        });
+      }()),
+      "");
+}
+
+TEST_F(ScopeExitTest, ExceptionInExitActionDuringExceptionCausesTermination) {
+  ASSERT_DEATH(
+      folly::coro::blockingWait([]() -> Task<> {
+        co_await co_scope_exit([]() -> Task<> {
+          throw std::runtime_error("Something bad happened!");
+        });
+        throw std::runtime_error("Throwing from parent");
+      }()),
+      "Something bad happened!");
+}
+
+TEST_F(ScopeExitTest, StatefulExitAction) {
+  folly::coro::blockingWait([this]() -> Task<> {
+    auto&& [i] = co_await co_scope_exit(
+        [this](int&& i) -> Task<void> {
+          count += i;
+          co_return;
+        },
+        3);
+    ++count;
+    i *= i;
+  }());
+  EXPECT_EQ(count, 10);
+}
+
+TEST_F(ScopeExitTest, NonMoveableState) {
+  folly::coro::blockingWait([this]() -> Task<> {
+    auto&& [asyncScope] = co_await co_scope_exit(
+        [this](auto&& scope) -> Task<void> {
+          co_await scope->joinAsync();
+          count *= 2;
+        },
+        std::make_unique<AsyncScope>());
+
+    auto ex = co_await co_current_executor;
+    asyncScope->add(co_invoke([this]() -> Task<> {
+                      ++count;
+                      co_return;
+                    }).scheduleOn(ex));
+  }());
+  EXPECT_EQ(count, 2);
+}
+
+#endif // FOLLY_HAS_COROUTINES