Do not capture RCTX in HHWheelTimer's underlying AsyncTimeout

Summary:
We should only be capturing `folly::RequestContext` in HHWheelTimer's callback
objects (already  happens). And avoid capturing it in underlying AsyncTimeout.

Notice that `AsyncTimeout::scheduleTimeout`/`scheduleTimeoutHighRes` APIs capture
current `folly::RequestContext` and `AsyncTimeout::cancelTimeout`
releases the captured `folly::RequestContext`

Now, the way HHWheelTimer works, it re-uses one instance of AsyncTimeout for
multiple timeouts. This opens up possibility for leaking lifetime of RCTX.
Imagine first  request A and then B schedule timeouts using same HHWheelTimer
instnace. A's RCTX is captured by underlying AsyncTimeout. Next, request A is
complete and we need to cancel its timeout. Notice, however that A's RCTX will
not be destroyed as long as B is still running. That's the lifetime management
leak.

Reviewed By: jordalgo, andriigrynenko

Differential Revision: D23260733

fbshipit-source-id: 0bdf79db40d34d6682e89e87ce30811ffccc8cbf

folly/io/async/AsyncTimeout.cpp

@@ -17,7 +17,6 @@
 #include <folly/io/async/AsyncTimeout.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/io/async/EventUtil.h>
-#include <folly/io/async/Request.h>
 #include <folly/net/NetworkSocket.h>
 
 #include <glog/logging.h>
@@ -87,21 +86,27 @@ AsyncTimeout::~AsyncTimeout() {
   cancelTimeout();
 }
 
-bool AsyncTimeout::scheduleTimeout(TimeoutManager::timeout_type timeout) {
+bool AsyncTimeout::scheduleTimeout(
+    TimeoutManager::timeout_type timeout,
+    std::shared_ptr<RequestContext>&& rctx) {
   assert(timeoutManager_ != nullptr);
-  context_ = RequestContext::saveContext();
+  context_ = std::move(rctx);
   return timeoutManager_->scheduleTimeout(this, timeout);
 }
 
 bool AsyncTimeout::scheduleTimeoutHighRes(
-    TimeoutManager::timeout_type_high_res timeout) {
+    TimeoutManager::timeout_type_high_res timeout,
+    std::shared_ptr<RequestContext>&& rctx) {
   assert(timeoutManager_ != nullptr);
-  context_ = RequestContext::saveContext();
+  context_ = std::move(rctx);
   return timeoutManager_->scheduleTimeoutHighRes(this, timeout);
 }
 
-bool AsyncTimeout::scheduleTimeout(uint32_t milliseconds) {
-  return scheduleTimeout(TimeoutManager::timeout_type(milliseconds));
+bool AsyncTimeout::scheduleTimeout(
+    uint32_t milliseconds,
+    std::shared_ptr<RequestContext>&& rctx) {
+  return scheduleTimeout(
+      TimeoutManager::timeout_type(milliseconds), std::move(rctx));
 }
 
 void AsyncTimeout::cancelTimeout() {

folly/io/async/AsyncTimeout.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <folly/io/async/EventBaseBackendBase.h>
+#include <folly/io/async/Request.h>
 #include <folly/io/async/TimeoutManager.h>
 
 #include <folly/portability/Event.h>
@@ -27,7 +28,6 @@
 namespace folly {
 
 class EventBase;
-class RequestContext;
 
 /**
  * AsyncTimeout is used to asynchronously wait for a timeout to occur.
@@ -91,15 +91,23 @@ class AsyncTimeout {
    * new timeout value.
    *
    * @param milliseconds  The timeout duration, in milliseconds.
+   * @param rctx request context to be captured by the callback
+   *             set to empty if current context should not be saved.
    *
    * @return Returns true if the timeout was successfully scheduled,
    *         and false if an error occurred.  After an error, the timeout is
    *         always unscheduled, even if scheduleTimeout() was just
    *         rescheduling an existing timeout.
    */
-  bool scheduleTimeout(uint32_t milliseconds);
-  bool scheduleTimeout(TimeoutManager::timeout_type timeout);
-  bool scheduleTimeoutHighRes(TimeoutManager::timeout_type_high_res timeout);
+  bool scheduleTimeout(
+      uint32_t milliseconds,
+      std::shared_ptr<RequestContext>&& rctx = RequestContext::saveContext());
+  bool scheduleTimeout(
+      TimeoutManager::timeout_type timeout,
+      std::shared_ptr<RequestContext>&& rctx = RequestContext::saveContext());
+  bool scheduleTimeoutHighRes(
+      TimeoutManager::timeout_type_high_res timeout,
+      std::shared_ptr<RequestContext>&& rctx = RequestContext::saveContext());
 
   /**
    * Cancel the timeout, if it is running.

folly/io/async/HHWheelTimer.cpp

@@ -228,7 +228,7 @@ bool HHWheelTimerBase<Duration>::cascadeTimers(
 
 template <class Duration>
 void HHWheelTimerBase<Duration>::scheduleTimeoutInternal(Duration timeout) {
-  this->AsyncTimeout::scheduleTimeout(timeout);
+  this->AsyncTimeout::scheduleTimeout(timeout, {});
 }
 
 template <class Duration>
@@ -390,7 +390,7 @@ int64_t HHWheelTimerBase<Duration>::calcNextTick(
 template <>
 void HHWheelTimerBase<std::chrono::microseconds>::scheduleTimeoutInternal(
     std::chrono::microseconds timeout) {
-  this->AsyncTimeout::scheduleTimeoutHighRes(timeout);
+  this->AsyncTimeout::scheduleTimeoutHighRes(timeout, {});
 }
 
 // std::chrono::milliseconds

folly/io/async/test/EventBaseTestLib.h

@@ -1076,13 +1076,10 @@ TYPED_TEST_P(EventBaseTest, RescheduleTimeout) {
   t2.scheduleTimeout(30);
   t3.scheduleTimeout(30);
 
-  auto f = static_cast<bool (AsyncTimeout::*)(uint32_t)>(
-      &AsyncTimeout::scheduleTimeout);
-
   // after 10ms, reschedule t2 to run sooner than originally scheduled
-  eb.tryRunAfterDelay(std::bind(f, &t2, 10), 10);
+  eb.tryRunAfterDelay([&] { t2.scheduleTimeout(10); }, 10);
   // after 10ms, reschedule t3 to run later than originally scheduled
-  eb.tryRunAfterDelay(std::bind(f, &t3, 40), 10);
+  eb.tryRunAfterDelay([&] { t3.scheduleTimeout(40); }, 10);
 
   eb.loop();
   TimePoint end;

folly/io/async/test/HHWheelTimerTest.cpp

@@ -92,6 +92,50 @@ TEST_F(HHWheelTimerTest, FireOnce) {
   T_CHECK_TIMEOUT(start, end, milliseconds(10));
 }
 
+TEST_F(HHWheelTimerTest, NoRequestContextLeak) {
+  StackWheelTimer t(&eventBase, milliseconds(1));
+  std::set<int> destructed;
+
+  class TestData : public RequestData {
+   public:
+    TestData(int data, std::set<int>& destructed)
+        : data_(data), destructed_(destructed) {}
+    ~TestData() override {
+      destructed_.insert(data_);
+    }
+
+    bool hasCallback() override {
+      return false;
+    }
+
+   private:
+    int data_;
+    std::set<int>& destructed_;
+  };
+
+  folly::Optional<TestTimeout> t1 = TestTimeout{};
+  folly::Optional<TestTimeout> t2 = TestTimeout{};
+
+  {
+    RequestContextScopeGuard g;
+    RequestContext::get()->setContextData(
+        "k", std::make_unique<TestData>(1, destructed));
+    t.scheduleTimeout(&*t1, milliseconds(5));
+  }
+
+  {
+    RequestContextScopeGuard g;
+    RequestContext::get()->setContextData(
+        "k", std::make_unique<TestData>(2, destructed));
+    t.scheduleTimeout(&*t2, milliseconds(5));
+  }
+
+  EXPECT_EQ(0, destructed.size());
+  t1.reset();
+  EXPECT_EQ(1, destructed.count(1));
+  EXPECT_EQ(0, destructed.count(2));
+}
+
 /*
  * Test scheduling a timeout from another timeout callback.
  */