Add support for rate limiting to transform and resumableTransform

Summary:
The channels framework provides most of its memory savings by allowing channels to be used without long-lived coroutine frames consuming them. However, coroutines are still used for short-lived transformation functions. While these coroutines should not increase steady state memory (as they are short-lived), it is possible that a large number of transformations occur at exactly the same time. This can cause a large memory spike (due to the size of the coroutine frames), and lead to heap fragmentation.

This diff solves this problem by adding support for rate limiting. An optional RateLimiter object can be provided to any transform or resumableTransform. A RateLimiter has a maximum concurrency specified on construction. For any transforms using the same rate limiter, the channels framework will ensure that the concurrency constraint is not violated. This limits the number of simultaneous coroutine frames (and therefore allows one to eliminate these spikes).

Reviewed By: aary

Differential Revision: D33037310

fbshipit-source-id: f7c5d30e08d155db7825d01f1b4c4b7354b14def

folly/experimental/channels/RateLimiter.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its 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 <folly/experimental/channels/RateLimiter.h>
+
+namespace folly {
+namespace channels {
+
+std::shared_ptr<RateLimiter> RateLimiter::create(size_t maxConcurrent) {
+  return std::shared_ptr<RateLimiter>(new RateLimiter(maxConcurrent));
+}
+
+RateLimiter::RateLimiter(size_t maxConcurrent)
+    : maxConcurrent_(maxConcurrent) {}
+
+void RateLimiter::executeWhenReady(
+    RateLimiter::QueuedFunc func,
+    Executor::KeepAlive<SequencedExecutor> executor) {
+  auto state = state_.wlock();
+  if (state->running < maxConcurrent_) {
+    CHECK(state->queue.empty());
+    state->running++;
+    executor->add(
+        [func = std::move(func), token = Token(shared_from_this())]() mutable {
+          func(std::move(token));
+        });
+  } else {
+    state->queue.enqueue(QueueItem{std::move(func), std::move(executor)});
+  }
+}
+
+RateLimiter::Token::Token(std::shared_ptr<RateLimiter> rateLimiter)
+    : rateLimiter_(std::move(rateLimiter)) {}
+
+RateLimiter::Token::~Token() {
+  if (!rateLimiter_) {
+    return;
+  }
+  auto state = rateLimiter_->state_.wlock();
+  if (!state->queue.empty()) {
+    auto queueItem = state->queue.dequeue();
+    queueItem.executor->add(
+        [func = std::move(queueItem.func),
+         token = Token(rateLimiter_->shared_from_this())]() mutable {
+          func(std::move(token));
+        });
+  } else {
+    state->running--;
+  }
+}
+} // namespace channels
+} // namespace folly

folly/experimental/channels/RateLimiter.h

+/*
+ * Copyright (c) Facebook, Inc. and its 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.
+ */
+
+#pragma once
+
+#include <folly/Synchronized.h>
+#include <folly/concurrency/UnboundedQueue.h>
+#include <folly/executors/SequencedExecutor.h>
+
+namespace folly {
+namespace channels {
+
+class RateLimiter : public std::enable_shared_from_this<RateLimiter> {
+ public:
+  static std::shared_ptr<RateLimiter> create(size_t maxConcurrent);
+
+  class Token {
+   public:
+    explicit Token(std::shared_ptr<RateLimiter> rateLimiter);
+    ~Token();
+
+    Token(const Token&) = delete;
+    Token& operator=(const Token&) = delete;
+    Token(Token&&) = default;
+    Token& operator=(Token&&) = default;
+
+   private:
+    std::shared_ptr<RateLimiter> rateLimiter_;
+  };
+
+  using QueuedFunc = folly::Function<void(Token)>;
+
+  void executeWhenReady(
+      QueuedFunc func, Executor::KeepAlive<SequencedExecutor> executor);
+
+ private:
+  explicit RateLimiter(size_t maxConcurrent);
+
+  struct QueueItem {
+    QueuedFunc func;
+    Executor::KeepAlive<SequencedExecutor> executor;
+  };
+
+  struct State {
+    USPSCQueue<QueueItem, false /* MayBlock */, 6 /* LgSegmentSize */> queue;
+    size_t running{0};
+  };
+
+  const size_t maxConcurrent_;
+  folly::Synchronized<State> state_;
+};
+} // namespace channels
+} // namespace folly

folly/experimental/channels/Transform.h

@@ -18,6 +18,7 @@
 
 #include <folly/executors/SequencedExecutor.h>
 #include <folly/experimental/channels/Channel.h>
+#include <folly/experimental/channels/RateLimiter.h>
 
 namespace folly {
 namespace channels {
@@ -52,6 +53,10 @@ namespace channels {
  *
  * @param transformValue: A function as described above.
  *
+ * @param rateLimiter: An optional rate limiter. If specified, the given rate
+ *     limiter will limit the number of transformation functions that are
+ *     simultaneously running.
+ *
  * Example:
  *
  *  // Function that returns a receiver
@@ -77,7 +82,8 @@ template <
 Receiver<OutputValueType> transform(
     ReceiverType inputReceiver,
     folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
-    TransformValueFunc transformValue);
+    TransformValueFunc transformValue,
+    std::shared_ptr<RateLimiter> rateLimiter = nullptr);
 
 /**
  * This overload accepts arguments in the form of a transformer object. The
@@ -87,6 +93,8 @@ Receiver<OutputValueType> transform(
  *
  * folly::coro::AsyncGenerator<OutputValueType&&> transformValue(
  *     folly::Try<InputValueType> inputValue);
+ *
+ * std::shared_ptr<RateLimiter> getRateLimiter(); // Can return nullptr
  */
 template <
     typename ReceiverType,
@@ -128,6 +136,10 @@ Receiver<OutputValueType> transform(
  *
  * @param transformValue: The TransformValue function as described above.
  *
+ * @param rateLimiter: An optional rate limiter. If specified, the given rate
+ *     limiter will limit the number of transformation functions that are
+ *     simultaneously running.
+ *
  * Example:
  *
  *  struct InitializeArg {
@@ -173,7 +185,8 @@ Receiver<OutputValueType> resumableTransform(
     folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
     InitializeArg initializeArg,
     InitializeTransformFunc initializeTransform,
-    TransformValueFunc transformValue);
+    TransformValueFunc transformValue,
+    std::shared_ptr<RateLimiter> rateLimiter = nullptr);
 
 /**
  * This overload accepts arguments in the form of a transformer object. The
@@ -186,6 +199,8 @@ Receiver<OutputValueType> resumableTransform(
  *
  * folly::coro::AsyncGenerator<OutputValueType&&> transformValue(
  *     folly::Try<InputValueType> inputValue);
+ *
+ * std::shared_ptr<RateLimiter> getRateLimiter(); // Can return nullptr
  */
 template <
     typename InitializeArg,

folly/experimental/channels/detail/Utility.h

@@ -22,6 +22,7 @@
 #include <folly/ScopeGuard.h>
 #include <folly/executors/SequencedExecutor.h>
 #include <folly/experimental/channels/Channel.h>
+#include <folly/experimental/channels/RateLimiter.h>
 #include <folly/experimental/coro/Task.h>
 
 namespace folly {
@@ -225,6 +226,8 @@ class SenderCancellationCallback : public IChannelCallback {
  *  coroutine operation is complete.
  *
  * @param operation: The operation to run.
+ *
+ * @param token: The rate limiter token for this operation.
  */
 template <typename TSender>
 void runOperationWithSenderCancellation(
@@ -232,7 +235,8 @@ void runOperationWithSenderCancellation(
     TSender& sender,
     bool alreadyStartedWaiting,
     IChannelCallback* channelCallbackToRestore,
-    folly::coro::Task<void> operation) noexcept {
+    folly::coro::Task<void> operation,
+    RateLimiter::Token token) noexcept {
   if (alreadyStartedWaiting && (!sender || !sender->cancelSenderWait())) {
     // The output receiver was cancelled before starting this operation
     // (indicating that the channel callback already ran).
@@ -242,6 +246,7 @@ void runOperationWithSenderCancellation(
       [&sender,
        executor,
        channelCallbackToRestore,
+       token = std::move(token),
        operation = std::move(operation)]() mutable -> folly::coro::Task<void> {
         auto senderCancellationCallback = SenderCancellationCallback(
             sender, executor, channelCallbackToRestore);

folly/experimental/channels/test/TransformTest.cpp

@@ -345,6 +345,81 @@ TEST_F(SimpleTransformFixture, Chained) {
   executor_.drain();
 }
 
+TEST_F(SimpleTransformFixture, MultipleTransformsWithRateLimiter) {
+  auto rateLimiter = RateLimiter::create(1 /* maxConcurrent */);
+
+  auto [untransformedReceiver1, sender1] = Channel<int>::create();
+  auto [controlReceiver1, controlSender1] = Channel<folly::Unit>::create();
+  int transform1Executions = 0;
+  auto transformedReceiver1 = transform(
+      std::move(untransformedReceiver1),
+      &executor_,
+      [&, &controlReceiver1 = controlReceiver1](folly::Try<int> result)
+          -> folly::coro::AsyncGenerator<std::string&&> {
+        transform1Executions++;
+        co_await controlReceiver1.next();
+        co_yield folly::to<std::string>(result.value());
+      },
+      rateLimiter);
+
+  auto [untransformedReceiver2, sender2] = Channel<int>::create();
+  auto [controlReceiver2, controlSender2] = Channel<folly::Unit>::create();
+  int transform2Executions = 0;
+  auto transformedReceiver2 = transform(
+      std::move(untransformedReceiver2),
+      &executor_,
+      [&, &controlReceiver2 = controlReceiver2](folly::Try<int> result)
+          -> folly::coro::AsyncGenerator<std::string&&> {
+        transform2Executions++;
+        co_await controlReceiver2.next();
+        co_yield folly::to<std::string>(result.value());
+      },
+      rateLimiter);
+
+  auto callbackHandle1 = processValues(std::move(transformedReceiver1));
+  auto callbackHandle2 = processValues(std::move(transformedReceiver2));
+
+  EXPECT_CALL(onNext_, onValue("1"));
+  EXPECT_CALL(onNext_, onValue("1000"));
+  EXPECT_CALL(onNext_, onValue("2"));
+  EXPECT_CALL(onNext_, onValue("2000"));
+  EXPECT_CALL(onNext_, onClosed()).Times(2);
+
+  sender1.write(1);
+  sender2.write(1000);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 1);
+  EXPECT_EQ(transform2Executions, 0);
+  controlSender1.write(folly::unit);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 1);
+  EXPECT_EQ(transform2Executions, 1);
+  controlSender2.write(folly::unit);
+  executor_.drain();
+
+  sender2.write(2000);
+  sender1.write(2);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 1);
+  EXPECT_EQ(transform2Executions, 2);
+  controlSender2.write(folly::unit);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 2);
+  EXPECT_EQ(transform2Executions, 2);
+  controlSender1.write(folly::unit);
+  executor_.drain();
+
+  std::move(sender1).close();
+  std::move(sender2).close();
+  std::move(controlSender1).close();
+  std::move(controlSender2).close();
+  executor_.drain();
+}
+
 class TransformFixtureStress : public Test {
  protected:
   TransformFixtureStress()
@@ -722,6 +797,112 @@ TEST_F(ResumableTransformFixture, TransformThrows_NoReinitialization_Rethrows) {
   executor_.drain();
 }
 
+TEST_F(ResumableTransformFixture, MultipleResumableTransformsWithRateLimiter) {
+  auto rateLimiter = RateLimiter::create(1 /* maxConcurrent */);
+
+  auto [untransformedReceiver1, sender1] = Channel<int>::create();
+  auto [controlReceiver1, controlSender1] = Channel<folly::Unit>::create();
+  int transform1Executions = 0;
+  auto transformedReceiver1 = resumableTransform(
+      &executor_,
+      toVector("init1"s),
+      [&,
+       receiver = std::move(untransformedReceiver1),
+       &controlReceiver1 =
+           controlReceiver1](std::vector<std::string> initializeArg) mutable
+      -> folly::coro::Task<std::pair<std::vector<std::string>, Receiver<int>>> {
+        transform1Executions++;
+        co_await controlReceiver1.next();
+        co_return std::make_pair(std::move(initializeArg), std::move(receiver));
+      },
+      [&, &controlReceiver1 = controlReceiver1](folly::Try<int> result)
+          -> folly::coro::AsyncGenerator<std::string&&> {
+        transform1Executions++;
+        co_await controlReceiver1.next();
+        co_yield folly::to<std::string>(result.value());
+      },
+      rateLimiter);
+
+  auto [untransformedReceiver2, sender2] = Channel<int>::create();
+  auto [controlReceiver2, controlSender2] = Channel<folly::Unit>::create();
+  int transform2Executions = 0;
+  auto transformedReceiver2 = resumableTransform(
+      &executor_,
+      toVector("init2"s),
+      [&,
+       receiver = std::move(untransformedReceiver2),
+       &controlReceiver2 =
+           controlReceiver2](std::vector<std::string> initializeArg) mutable
+      -> folly::coro::Task<std::pair<std::vector<std::string>, Receiver<int>>> {
+        transform2Executions++;
+        co_await controlReceiver2.next();
+        co_return std::make_pair(std::move(initializeArg), std::move(receiver));
+      },
+      [&, &controlReceiver2 = controlReceiver2](folly::Try<int> result)
+          -> folly::coro::AsyncGenerator<std::string&&> {
+        transform2Executions++;
+        co_await controlReceiver2.next();
+        co_yield folly::to<std::string>(result.value());
+      },
+      rateLimiter);
+
+  auto callbackHandle1 = processValues(std::move(transformedReceiver1));
+  auto callbackHandle2 = processValues(std::move(transformedReceiver2));
+
+  EXPECT_CALL(onNext_, onValue("init1"));
+  EXPECT_CALL(onNext_, onValue("init2"));
+  EXPECT_CALL(onNext_, onValue("1"));
+  EXPECT_CALL(onNext_, onValue("1000"));
+  EXPECT_CALL(onNext_, onValue("2"));
+  EXPECT_CALL(onNext_, onValue("2000"));
+  EXPECT_CALL(onNext_, onClosed()).Times(2);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 1);
+  EXPECT_EQ(transform2Executions, 0);
+  controlSender1.write(folly::unit);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 1);
+  EXPECT_EQ(transform2Executions, 1);
+  controlSender2.write(folly::unit);
+  executor_.drain();
+
+  sender1.write(1);
+  sender2.write(1000);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 2);
+  EXPECT_EQ(transform2Executions, 1);
+  controlSender1.write(folly::unit);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 2);
+  EXPECT_EQ(transform2Executions, 2);
+  controlSender2.write(folly::unit);
+  executor_.drain();
+
+  sender2.write(2000);
+  sender1.write(2);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 2);
+  EXPECT_EQ(transform2Executions, 3);
+  controlSender2.write(folly::unit);
+  executor_.drain();
+
+  EXPECT_EQ(transform1Executions, 3);
+  EXPECT_EQ(transform2Executions, 3);
+  controlSender1.write(folly::unit);
+  executor_.drain();
+
+  std::move(sender1).close();
+  std::move(sender2).close();
+  std::move(controlSender1).close();
+  std::move(controlSender2).close();
+  executor_.drain();
+}
+
 class ResumableTransformFixtureStress : public Test {
  protected:
   ResumableTransformFixtureStress()