prevent TLStatsAsyncAggregator from saving RequestContext to AsyncTim…

…eout

Summary:
Currently, TLStatsAsyncAggregator::scheduleAggregation will call AsyncTimeout::scheduleTimeout which saves as a default parameter the current RequestContext. The possible issue with this is that RequestContext can possibly persist indefinitely until program end since on the timer expiration, the class calls scheduleTimeout, this cycle will continue.

This issue has happened in other areas too, such as HHWheelTimer. Since this class's counters operate on a global scope, it does not make sense to carry any request specific information into the timeout. Thus, for the invocations of scheduleTimeout, we pass in nullptr for the RequestContext.

The test case now shows that the RequestContext object is no longer persisted into the timeout.

Update: introduced member function scheduleAggregationTimeout to handle the calling of scheduleTimeout. Name helps avoid name clashing with AsyncTimeout and is a bit more descriptive.

Reviewed By: ot

Differential Revision: D65911229

fbshipit-source-id: 552493f7957425a9a0cfeb2661a53d619871999f

fb303/TLStatsAsyncAggregator.cpp

@@ -40,12 +40,20 @@ void TLStatsAsyncAggregator::scheduleAggregation(
   // that can generate stats.
   attachEventBase(eventBase, folly::AsyncTimeout::InternalEnum::INTERNAL);
   intervalMS_ = intervalMS;
-  scheduleTimeout(intervalMS_);
+
+  scheduleAggregationTimeout();
+}
+
+void TLStatsAsyncAggregator::scheduleAggregationTimeout() {
+  // Pass in nullptr for RequestContext so it is not persisted.
+  // These counters are global and should not carry request specific
+  // context.
+  scheduleTimeout(intervalMS_, nullptr);
 }
 
 void TLStatsAsyncAggregator::timeoutExpired() noexcept {
   stats_->aggregate();
-  scheduleTimeout(intervalMS_);
+  scheduleAggregationTimeout();
 }
 
 } // namespace facebook::fb303

fb303/TLStatsAsyncAggregator.h

@@ -73,6 +73,10 @@ class TLStatsAsyncAggregator : private folly::AsyncTimeout {
   // AsyncTimeout methods
   void timeoutExpired() noexcept override;
 
+  // Wrapper around AsyncTimeout::scheduleTimeout to ensure properly handling
+  // of timeout scheduling.
+  void scheduleAggregationTimeout();
+
   uint32_t intervalMS_;
   ThreadLocalStats* stats_;
 };

fb303/test/BUCK

@@ -287,6 +287,17 @@ cpp_unittest(
     ],
 )
 
+cpp_unittest(
+    name = "tl_stats_async_aggregator_test",
+    srcs = [
+        "TLStatsAsyncAggregatorTest.cpp",
+    ],
+    deps = [
+        "//fb303:tl_stats_async_aggregator",
+        "//folly/io/async:async_base",
+    ],
+)
+
 cpp_unittest(
     name = "timeseries_histogram_test",
     srcs = [

fb303/test/TLStatsAsyncAggregatorTest.cpp

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <fb303/TLStatsAsyncAggregator.h>
+#include <folly/io/async/EventBase.h>
+
+#include <gtest/gtest.h>
+
+using namespace std;
+using namespace facebook;
+using namespace facebook::fb303;
+using namespace folly;
+
+TEST(TLStatsAsyncAggregatorTest, NoRequestContextLeak) {
+  // Setup fields
+  ThreadLocalStatsT<TLStatsNoLocking> tlstats{};
+  TLStatsAsyncAggregator aggregator{&tlstats};
+  EventBase eventBase;
+
+  // Weak pointer to attach to request context
+  std::weak_ptr<RequestContext> weakContext;
+  {
+    // Guard to clean up RequestContext in this scope.
+    RequestContextScopeGuard guard;
+    weakContext = RequestContext::saveContext();
+    aggregator.scheduleAggregation(&eventBase);
+  }
+
+  // Since we are not carrying request context anymore,
+  // the RequestContext gets cleared up by guard,
+  // and no other reference is alive, so nullptr is expected.
+  EXPECT_TRUE(weakContext.expired());
+}