Add a safe ensure() variant for SemiFuture

Reviewed By: yfeldblum, larabear

Differential Revision: D14692263

fbshipit-source-id: fbdab61363ea83667832dbc56a883d069a1bd597

folly/futures/Future-inl.h

@@ -2433,6 +2433,16 @@ mapTry(Executor& exec, It first, It last, F func, int) {
   return results;
 }
 
+template <typename F, class Ensure>
+auto ensure(F&& f, Ensure&& ensure) {
+  return makeSemiFuture()
+      .deferValue([f = std::forward<F>(f)](auto) mutable { return f(); })
+      .defer([ensure = std::forward<Ensure>(ensure)](auto resultTry) mutable {
+        ensure();
+        return std::move(resultTry).value();
+      });
+}
+
 } // namespace futures
 
 template <class Clock>

folly/futures/helpers.h

@@ -152,6 +152,15 @@ SemiFuture<Unit> wait(std::unique_ptr<fibers::Baton> baton);
 
 #endif
 
+/**
+ * Returns a lazy SemiFuture constructed by f, which also ensures that ensure is
+ * called before completion.
+ * f doesn't get called until the SemiFuture is activated (e.g. through a .get()
+ * or .via() call). If f gets called, ensure is guaranteed to be called as well.
+ */
+template <typename F, class Ensure>
+auto ensure(F&& f, Ensure&& ensure);
+
 } // namespace futures
 
 /**

folly/futures/test/SemiFutureTest.cpp

@@ -1098,3 +1098,46 @@ TEST(SemiFuture, within) {
     p.setValue();
   }
 }
+
+TEST(SemiFuture, ensure) {
+  {
+    bool fCalled{false};
+    bool ensureCalled{false};
+    auto sf = futures::ensure(
+        [&] {
+          fCalled = true;
+          return 42;
+        },
+        [&] { ensureCalled = true; });
+    EXPECT_EQ(42, std::move(sf).get());
+    EXPECT_TRUE(fCalled);
+    EXPECT_TRUE(ensureCalled);
+  }
+  {
+    bool fCalled{false};
+    bool ensureCalled{false};
+    futures::ensure(
+        [&] {
+          fCalled = true;
+          return 42;
+        },
+        [&] { ensureCalled = true; });
+    EXPECT_FALSE(fCalled);
+    EXPECT_FALSE(ensureCalled);
+  }
+  {
+    struct ExpectedException {};
+
+    bool fCalled{false};
+    bool ensureCalled{false};
+    auto sf = futures::ensure(
+        [&] {
+          fCalled = true;
+          throw ExpectedException();
+        },
+        [&] { ensureCalled = true; });
+    EXPECT_THROW(std::move(sf).get(), ExpectedException);
+    EXPECT_TRUE(fCalled);
+    EXPECT_TRUE(ensureCalled);
+  }
+}