From: Richard Meng <richmeng@fb.com>
Date: Tue, 29 Aug 2017 17:18:29 +0000 (-0700)
Subject: Use weak_ptr to hold future context in timekeeper to allow clean up when future complete
X-Git-Tag: v2017.09.04.00~12
X-Git-Url: http://demsky.eecs.uci.edu/git/?a=commitdiff_plain;h=c02d603390a8ece8e5378d92f3971955667426db;p=folly.git

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 e682dcab..fa39dc84 100644
--- 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 a8f3eb8d..49ce1cac 100644
--- 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());
+      });
+}