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

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));
       }
     }
   });

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());
+      });
+}