Use weak_ptr to hold future context in timekeeper to allow clean up when future complete

Summary:
Before this change, future context will not be cleaned up until timekeeper times out. These objects has been occupying memory when more shorter future tasks are registered.
Switch to use weak ptr to hold context, so that context objects are deallocated as soon as the future completes (or times out)

Reviewed By: yfeldblum

Differential Revision: D5692040

fbshipit-source-id: b3b74a29b2ccafef6c4a06011699b069feb3a847

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index e682dcab47c5f8f5ce5a13eaf6e12a9e584cd11c..fa39dc84987a81a3bac8156719c9285b1716065a 100644 (file)
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -1076,20 +1076,26 @@ Future<T> Future<T>::within(Duration dur, E e, Timekeeper* tk) {
   auto ctx = std::make_shared<Context>(std::move(e));
 
   ctx->thisFuture = this->then([ctx](Try<T>&& t) mutable {
-    // TODO: "this" completed first, cancel "after"
     if (ctx->token.exchange(true) == false) {
       ctx->promise.setTry(std::move(t));
     }
   });
 
-  tk->after(dur).then([ctx](Try<Unit> const& t) mutable {
+  // Have time keeper use a weak ptr to hold ctx,
+  // so that ctx can be deallocated as soon as the future job finished.
+  tk->after(dur).then([weakCtx = to_weak_ptr(ctx)](Try<Unit> const& t) mutable {
+    auto lockedCtx = weakCtx.lock();
+    if (!lockedCtx) {
+      // ctx already released. "this" completed first, cancel "after"
+      return;
+    }
     // "after" completed first, cancel "this"
-    ctx->thisFuture.raise(TimedOut());
-    if (ctx->token.exchange(true) == false) {
+    lockedCtx->thisFuture.raise(TimedOut());
+    if (lockedCtx->token.exchange(true) == false) {
       if (t.hasException()) {
-        ctx->promise.setException(std::move(t.exception()));
+        lockedCtx->promise.setException(std::move(t.exception()));
       } else {
-        ctx->promise.setException(std::move(ctx->exception));
+        lockedCtx->promise.setException(std::move(lockedCtx->exception));
       }
     }
   });

diff --git a/folly/futures/test/FutureTest.cpp b/folly/futures/test/FutureTest.cpp
index a8f3eb8df0ccff95359083b1d4628310da6d5bc8..49ce1cac3d9dcea9f1b35ae538ef65a9d37f7373 100644 (file)
--- a/folly/futures/test/FutureTest.cpp
+++ b/folly/futures/test/FutureTest.cpp
@@ -936,3 +936,35 @@ TEST(Future, invokeCallbackReturningFutureAsRvalue) {
   EXPECT_EQ(203, makeFuture<int>(200).then(cfoo).value());
   EXPECT_EQ(303, makeFuture<int>(300).then(Foo()).value());
 }
+
+TEST(Future, futureWithinCtxCleanedUpWhenTaskFinishedInTime) {
+  // Used to track the use_count of callbackInput even outside of its scope
+  std::weak_ptr<int> target;
+  {
+    Promise<std::shared_ptr<int>> promise;
+    auto input = std::make_shared<int>(1);
+    auto longEnough = std::chrono::milliseconds(1000);
+
+    promise.getFuture()
+        .within(longEnough)
+        .then([&target](
+                  folly::Try<std::shared_ptr<int>>&& callbackInput) mutable {
+          target = callbackInput.value();
+        });
+    promise.setValue(input);
+  }
+  // After promise's life cycle is finished, make sure no one is holding the
+  // input anymore, in other words, ctx should have been cleaned up.
+  EXPECT_EQ(0, target.use_count());
+}
+
+TEST(Future, futureWithinNoValueReferenceWhenTimeOut) {
+  Promise<std::shared_ptr<int>> promise;
+  auto veryShort = std::chrono::milliseconds(1);
+
+  promise.getFuture().within(veryShort).then(
+      [](folly::Try<std::shared_ptr<int>>&& callbackInput) {
+        // Timeout is fired. Verify callbackInput is not referenced
+        EXPECT_EQ(0, callbackInput.value().use_count());
+      });
+}