Add ensureCleanupAfterTask

Summary:
adds ensureCleanupAfterTask()

takes a SemiFuture<T> task and a SemiFuture<Unit> cleanup.

returns a SemiFuture<T>

always starts cleanup after task completes. stores the result of the task and replays it after the cleanup completes.

cleanup is not allowed to fail.

Reviewed By: yfeldblum, andriigrynenko

Differential Revision: D21481825

fbshipit-source-id: 901a91c25f7b2b4fa794a590a5433e0fad27d168

folly/experimental/Cleanup.h

@@ -50,6 +50,36 @@ constexpr bool is_cleanup_v = folly::is_invocable_v<detail::cleanup_fn, T>;
 template <typename T>
 using is_cleanup = std::bool_constant<is_cleanup_v<T>>;
 
+// Structured Async Cleanup
+//
+// This helps compose a task with async cleanup
+// The task result is stored until cleanup completes and is then produced
+// The cleanup task is not allowed to fail.
+//
+// This can be used with collectAll to combine multiple async resources
+//
+// ensureCleanupAfterTask(collectAll(a.run(), b.run()), collectAll(a.cleanup(),
+// b.cleanup())).wait();
+//
+template <typename T>
+folly::SemiFuture<T> ensureCleanupAfterTask(
+    folly::SemiFuture<T> task,
+    folly::SemiFuture<folly::Unit> cleanup) {
+  return folly::makeSemiFuture()
+      .deferValue([task_ = std::move(task)](folly::Unit) mutable {
+        return std::move(task_);
+      })
+      .defer([cleanup_ = std::move(cleanup)](folly::Try<T> taskResult) mutable {
+        return std::move(cleanup_).defer(
+            [taskResult_ = std::move(taskResult)](folly::Try<folly::Unit> t) {
+              if (t.hasException()) {
+                terminate_with<std::logic_error>("cleanup must not throw");
+              }
+              return std::move(taskResult_).value();
+            });
+      });
+}
+
 // Structured Async Cleanup
 //
 // This implementation is a base class that collects a set of cleanup tasks

folly/experimental/test/CleanupTest.cpp

@@ -62,6 +62,34 @@ TEST(CleanupTest, Basic) {
   EXPECT_EQ(index, 0);
 }
 
+TEST(CleanupTest, EnsureCleanupAfterTaskBasic) {
+  Cleaned cleaned;
+  int phase = 0;
+  int index = 0;
+
+  cleaned.addCleanup(
+      folly::makeSemiFuture().deferValue([&, expected = index++](folly::Unit) {
+        EXPECT_EQ(phase, 1);
+        EXPECT_EQ(--index, expected);
+      }));
+
+  auto task =
+      folly::makeSemiFuture().deferValue([&, expected = index++](folly::Unit) {
+        EXPECT_EQ(phase, 1);
+        EXPECT_EQ(--index, expected);
+      });
+  EXPECT_EQ(index, 2);
+
+  folly::ManualExecutor exec;
+  phase = 1;
+  folly::ensureCleanupAfterTask(std::move(task), cleaned.cleanup())
+      .within(1s)
+      .via(folly::getKeepAliveToken(exec))
+      .getVia(&exec);
+  phase = 2;
+  EXPECT_EQ(index, 0);
+}
+
 TEST(CleanupTest, Errors) {
   auto cleaned = std::make_unique<Cleaned>();