Commit eadf2192 authored by Andrew Smith's avatar Andrew Smith Committed by Facebook GitHub Bot

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
parent 05e0144b
/*
* 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
/*
* 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
...@@ -65,14 +65,15 @@ class TransformProcessorBase : public IChannelCallback { ...@@ -65,14 +65,15 @@ class TransformProcessorBase : public IChannelCallback {
template <typename ReceiverType> template <typename ReceiverType>
void startTransform(ReceiverType receiver) { void startTransform(ReceiverType receiver) {
transformer_.getExecutor()->add( executeWhenReady(
[=, receiver = std::move(receiver)]() mutable { [=, receiver = std::move(receiver)](RateLimiter::Token token) mutable {
runOperationWithSenderCancellation( runOperationWithSenderCancellation(
transformer_.getExecutor(), transformer_.getExecutor(),
this->sender_, this->sender_,
false /* alreadyStartedWaiting */, false /* alreadyStartedWaiting */,
this /* channelCallbackToRestore */, this /* channelCallbackToRestore */,
startTransformImpl(std::move(receiver))); startTransformImpl(std::move(receiver)),
std::move(token));
}); });
} }
...@@ -97,7 +98,7 @@ class TransformProcessorBase : public IChannelCallback { ...@@ -97,7 +98,7 @@ class TransformProcessorBase : public IChannelCallback {
* sender). * sender).
*/ */
void consume(ChannelBridgeBase* bridge) override { void consume(ChannelBridgeBase* bridge) override {
transformer_.getExecutor()->add([=]() { executeWhenReady([=](RateLimiter::Token token) {
if (bridge == receiver_.get()) { if (bridge == receiver_.get()) {
// We have received new values from the input receiver. // We have received new values from the input receiver.
CHECK_NE(getReceiverState(), ChannelState::CancellationProcessed); CHECK_NE(getReceiverState(), ChannelState::CancellationProcessed);
...@@ -106,7 +107,8 @@ class TransformProcessorBase : public IChannelCallback { ...@@ -106,7 +107,8 @@ class TransformProcessorBase : public IChannelCallback {
this->sender_, this->sender_,
true /* alreadyStartedWaiting */, true /* alreadyStartedWaiting */,
this /* channelCallbackToRestore */, this /* channelCallbackToRestore */,
processAllAvailableValues()); processAllAvailableValues(),
std::move(token));
} else { } else {
CHECK_NE(getSenderState(), ChannelState::CancellationProcessed); CHECK_NE(getSenderState(), ChannelState::CancellationProcessed);
// The consumer of the output receiver has stopped consuming. // The consumer of the output receiver has stopped consuming.
...@@ -123,15 +125,19 @@ class TransformProcessorBase : public IChannelCallback { ...@@ -123,15 +125,19 @@ class TransformProcessorBase : public IChannelCallback {
* listening to. * listening to.
*/ */
void canceled(ChannelBridgeBase* bridge) override { void canceled(ChannelBridgeBase* bridge) override {
transformer_.getExecutor()->add([=]() { executeWhenReady([=](RateLimiter::Token token) {
if (bridge == receiver_.get()) { if (bridge == receiver_.get()) {
// We previously cancelled the input receiver (because the consumer of // We previously cancelled the input receiver (because the consumer of
// the output receiver stopped consuming). Process the cancellation for // the output receiver stopped consuming). Process the cancellation for
// the input receiver. // the input receiver.
CHECK_EQ(getReceiverState(), ChannelState::CancellationTriggered); CHECK_EQ(getReceiverState(), ChannelState::CancellationTriggered);
processReceiverCancelled(CloseResult()) runOperationWithSenderCancellation(
.scheduleOn(transformer_.getExecutor()) transformer_.getExecutor(),
.start(); this->sender_,
true /* alreadyStartedWaiting */,
this /* channelCallbackToRestore */,
processReceiverCancelled(CloseResult()),
std::move(token));
} else { } else {
// We previously cancelled the sender due to the closure of the input // We previously cancelled the sender due to the closure of the input
// receiver. Process the cancellation for the sender. // receiver. Process the cancellation for the sender.
...@@ -268,6 +274,18 @@ class TransformProcessorBase : public IChannelCallback { ...@@ -268,6 +274,18 @@ class TransformProcessorBase : public IChannelCallback {
return detail::getSenderState(sender_.get()); return detail::getSenderState(sender_.get());
} }
void executeWhenReady(folly::Function<void(RateLimiter::Token)> func) {
auto rateLimiter = transformer_.getRateLimiter();
if (rateLimiter != nullptr) {
rateLimiter->executeWhenReady(
std::move(func), transformer_.getExecutor());
} else {
transformer_.getExecutor()->add([func = std::move(func)]() mutable {
func(RateLimiter::Token(nullptr));
});
}
}
ChannelBridgePtr<InputValueType> receiver_; ChannelBridgePtr<InputValueType> receiver_;
ChannelBridgePtr<OutputValueType> sender_; ChannelBridgePtr<OutputValueType> sender_;
TransformerType transformer_; TransformerType transformer_;
...@@ -330,14 +348,15 @@ class ResumableTransformProcessor : public TransformProcessorBase< ...@@ -330,14 +348,15 @@ class ResumableTransformProcessor : public TransformProcessorBase<
using Base::Base; using Base::Base;
void initialize(InitializeArg initializeArg) { void initialize(InitializeArg initializeArg) {
this->transformer_.getExecutor()->add( this->executeWhenReady([=, initializeArg = std::move(initializeArg)](
[=, initializeArg = std::move(initializeArg)]() mutable { RateLimiter::Token token) mutable {
runOperationWithSenderCancellation( runOperationWithSenderCancellation(
this->transformer_.getExecutor(), this->transformer_.getExecutor(),
this->sender_, this->sender_,
false /* currentlyWaiting */, false /* currentlyWaiting */,
this /* channelCallbackToRestore */, this /* channelCallbackToRestore */,
initializeImpl(std::move(initializeArg))); initializeImpl(std::move(initializeArg)),
std::move(token));
}); });
} }
...@@ -416,16 +435,44 @@ class ResumableTransformProcessor : public TransformProcessorBase< ...@@ -416,16 +435,44 @@ class ResumableTransformProcessor : public TransformProcessorBase<
} }
}; };
template <bool Enabled>
class RateLimiterHolder;
template <>
class RateLimiterHolder<true> {
public:
explicit RateLimiterHolder(std::shared_ptr<RateLimiter> rateLimiter)
: rateLimiter_(std::move(rateLimiter)) {}
std::shared_ptr<RateLimiter> getRateLimiter() { return rateLimiter_; }
private:
std::shared_ptr<RateLimiter> rateLimiter_;
};
template <>
class RateLimiterHolder<false> {
public:
explicit RateLimiterHolder(std::shared_ptr<RateLimiter> rateLimiter) {
CHECK_NULL(rateLimiter.get());
}
std::shared_ptr<RateLimiter> getRateLimiter() { return nullptr; }
};
template < template <
typename InputValueType, typename InputValueType,
typename OutputValueType, typename OutputValueType,
typename TransformValueFunc> typename TransformValueFunc,
class DefaultTransformer { bool RateLimiterEnabled>
class DefaultTransformer : public RateLimiterHolder<RateLimiterEnabled> {
public: public:
DefaultTransformer( DefaultTransformer(
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
TransformValueFunc transformValue) TransformValueFunc transformValue,
: executor_(std::move(executor)), std::shared_ptr<RateLimiter> rateLimiter)
: RateLimiterHolder<RateLimiterEnabled>(std::move(rateLimiter)),
executor_(std::move(executor)),
transformValue_(std::move(transformValue)) {} transformValue_(std::move(transformValue)) {}
folly::Executor::KeepAlive<folly::SequencedExecutor> getExecutor() { folly::Executor::KeepAlive<folly::SequencedExecutor> getExecutor() {
...@@ -446,20 +493,29 @@ template < ...@@ -446,20 +493,29 @@ template <
typename InputValueType, typename InputValueType,
typename OutputValueType, typename OutputValueType,
typename InitializeTransformFunc, typename InitializeTransformFunc,
typename TransformValueFunc> typename TransformValueFunc,
bool RateLimiterEnabled>
class DefaultResumableTransformer : public DefaultTransformer< class DefaultResumableTransformer : public DefaultTransformer<
InputValueType, InputValueType,
OutputValueType, OutputValueType,
TransformValueFunc> { TransformValueFunc,
RateLimiterEnabled> {
public: public:
using Base = using Base = DefaultTransformer<
DefaultTransformer<InputValueType, OutputValueType, TransformValueFunc>; InputValueType,
OutputValueType,
TransformValueFunc,
RateLimiterEnabled>;
DefaultResumableTransformer( DefaultResumableTransformer(
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
InitializeTransformFunc initializeTransform, InitializeTransformFunc initializeTransform,
TransformValueFunc transformValue) TransformValueFunc transformValue,
: Base(std::move(executor), std::move(transformValue)), std::shared_ptr<RateLimiter> rateLimiter)
: Base(
std::move(executor),
std::move(transformValue),
std::move(rateLimiter)),
initializeTransform_(std::move(initializeTransform)) {} initializeTransform_(std::move(initializeTransform)) {}
auto initializeTransform(InitializeArg initializeArg) { auto initializeTransform(InitializeArg initializeArg) {
...@@ -479,13 +535,34 @@ template < ...@@ -479,13 +535,34 @@ template <
Receiver<OutputValueType> transform( Receiver<OutputValueType> transform(
ReceiverType inputReceiver, ReceiverType inputReceiver,
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
TransformValueFunc transformValue) { TransformValueFunc transformValue,
std::shared_ptr<RateLimiter> rateLimiter) {
if (rateLimiter != nullptr) {
using TransformerType = detail::DefaultTransformer<
InputValueType,
OutputValueType,
TransformValueFunc,
true /* RateLimiterEnabled */>;
return transform( return transform(
std::move(inputReceiver), std::move(inputReceiver),
detail::DefaultTransformer< TransformerType(
std::move(executor),
std::move(transformValue),
std::move(rateLimiter)));
} else {
using TransformerType = detail::DefaultTransformer<
InputValueType, InputValueType,
OutputValueType, OutputValueType,
TransformValueFunc>(std::move(executor), std::move(transformValue))); TransformValueFunc,
false /* RateLimiterEnabled */>;
return transform(
std::move(inputReceiver),
TransformerType(
std::move(executor),
std::move(transformValue),
nullptr /* rateLimiter */));
}
} }
template < template <
...@@ -515,18 +592,39 @@ Receiver<OutputValueType> resumableTransform( ...@@ -515,18 +592,39 @@ Receiver<OutputValueType> resumableTransform(
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
InitializeArg initializeArg, InitializeArg initializeArg,
InitializeTransformFunc initializeTransform, InitializeTransformFunc initializeTransform,
TransformValueFunc transformValue) { TransformValueFunc transformValue,
std::shared_ptr<RateLimiter> rateLimiter) {
if (rateLimiter != nullptr) {
using TransformerType = detail::DefaultResumableTransformer<
InitializeArg,
InputValueType,
OutputValueType,
InitializeTransformFunc,
TransformValueFunc,
true /* RateLimiterEnabled */>;
return resumableTransform( return resumableTransform(
std::move(initializeArg), std::move(initializeArg),
detail::DefaultResumableTransformer< TransformerType(
std::move(executor),
std::move(initializeTransform),
std::move(transformValue),
std::move(rateLimiter)));
} else {
using TransformerType = detail::DefaultResumableTransformer<
InitializeArg, InitializeArg,
InputValueType, InputValueType,
OutputValueType, OutputValueType,
InitializeTransformFunc, InitializeTransformFunc,
TransformValueFunc>( TransformValueFunc,
false /* RateLimiterEnabled */>;
return resumableTransform(
std::move(initializeArg),
TransformerType(
std::move(executor), std::move(executor),
std::move(initializeTransform), std::move(initializeTransform),
std::move(transformValue))); std::move(transformValue),
nullptr /* rateLimiter */));
}
} }
template < template <
......
...@@ -18,6 +18,7 @@ ...@@ -18,6 +18,7 @@
#include <folly/executors/SequencedExecutor.h> #include <folly/executors/SequencedExecutor.h>
#include <folly/experimental/channels/Channel.h> #include <folly/experimental/channels/Channel.h>
#include <folly/experimental/channels/RateLimiter.h>
namespace folly { namespace folly {
namespace channels { namespace channels {
...@@ -52,6 +53,10 @@ namespace channels { ...@@ -52,6 +53,10 @@ namespace channels {
* *
* @param transformValue: A function as described above. * @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: * Example:
* *
* // Function that returns a receiver * // Function that returns a receiver
...@@ -77,7 +82,8 @@ template < ...@@ -77,7 +82,8 @@ template <
Receiver<OutputValueType> transform( Receiver<OutputValueType> transform(
ReceiverType inputReceiver, ReceiverType inputReceiver,
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, 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 * This overload accepts arguments in the form of a transformer object. The
...@@ -87,6 +93,8 @@ Receiver<OutputValueType> transform( ...@@ -87,6 +93,8 @@ Receiver<OutputValueType> transform(
* *
* folly::coro::AsyncGenerator<OutputValueType&&> transformValue( * folly::coro::AsyncGenerator<OutputValueType&&> transformValue(
* folly::Try<InputValueType> inputValue); * folly::Try<InputValueType> inputValue);
*
* std::shared_ptr<RateLimiter> getRateLimiter(); // Can return nullptr
*/ */
template < template <
typename ReceiverType, typename ReceiverType,
...@@ -128,6 +136,10 @@ Receiver<OutputValueType> transform( ...@@ -128,6 +136,10 @@ Receiver<OutputValueType> transform(
* *
* @param transformValue: The TransformValue function as described above. * @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: * Example:
* *
* struct InitializeArg { * struct InitializeArg {
...@@ -173,7 +185,8 @@ Receiver<OutputValueType> resumableTransform( ...@@ -173,7 +185,8 @@ Receiver<OutputValueType> resumableTransform(
folly::Executor::KeepAlive<folly::SequencedExecutor> executor, folly::Executor::KeepAlive<folly::SequencedExecutor> executor,
InitializeArg initializeArg, InitializeArg initializeArg,
InitializeTransformFunc initializeTransform, InitializeTransformFunc initializeTransform,
TransformValueFunc transformValue); TransformValueFunc transformValue,
std::shared_ptr<RateLimiter> rateLimiter = nullptr);
/** /**
* This overload accepts arguments in the form of a transformer object. The * This overload accepts arguments in the form of a transformer object. The
...@@ -186,6 +199,8 @@ Receiver<OutputValueType> resumableTransform( ...@@ -186,6 +199,8 @@ Receiver<OutputValueType> resumableTransform(
* *
* folly::coro::AsyncGenerator<OutputValueType&&> transformValue( * folly::coro::AsyncGenerator<OutputValueType&&> transformValue(
* folly::Try<InputValueType> inputValue); * folly::Try<InputValueType> inputValue);
*
* std::shared_ptr<RateLimiter> getRateLimiter(); // Can return nullptr
*/ */
template < template <
typename InitializeArg, typename InitializeArg,
......
...@@ -22,6 +22,7 @@ ...@@ -22,6 +22,7 @@
#include <folly/ScopeGuard.h> #include <folly/ScopeGuard.h>
#include <folly/executors/SequencedExecutor.h> #include <folly/executors/SequencedExecutor.h>
#include <folly/experimental/channels/Channel.h> #include <folly/experimental/channels/Channel.h>
#include <folly/experimental/channels/RateLimiter.h>
#include <folly/experimental/coro/Task.h> #include <folly/experimental/coro/Task.h>
namespace folly { namespace folly {
...@@ -225,6 +226,8 @@ class SenderCancellationCallback : public IChannelCallback { ...@@ -225,6 +226,8 @@ class SenderCancellationCallback : public IChannelCallback {
* coroutine operation is complete. * coroutine operation is complete.
* *
* @param operation: The operation to run. * @param operation: The operation to run.
*
* @param token: The rate limiter token for this operation.
*/ */
template <typename TSender> template <typename TSender>
void runOperationWithSenderCancellation( void runOperationWithSenderCancellation(
...@@ -232,7 +235,8 @@ void runOperationWithSenderCancellation( ...@@ -232,7 +235,8 @@ void runOperationWithSenderCancellation(
TSender& sender, TSender& sender,
bool alreadyStartedWaiting, bool alreadyStartedWaiting,
IChannelCallback* channelCallbackToRestore, IChannelCallback* channelCallbackToRestore,
folly::coro::Task<void> operation) noexcept { folly::coro::Task<void> operation,
RateLimiter::Token token) noexcept {
if (alreadyStartedWaiting && (!sender || !sender->cancelSenderWait())) { if (alreadyStartedWaiting && (!sender || !sender->cancelSenderWait())) {
// The output receiver was cancelled before starting this operation // The output receiver was cancelled before starting this operation
// (indicating that the channel callback already ran). // (indicating that the channel callback already ran).
...@@ -242,6 +246,7 @@ void runOperationWithSenderCancellation( ...@@ -242,6 +246,7 @@ void runOperationWithSenderCancellation(
[&sender, [&sender,
executor, executor,
channelCallbackToRestore, channelCallbackToRestore,
token = std::move(token),
operation = std::move(operation)]() mutable -> folly::coro::Task<void> { operation = std::move(operation)]() mutable -> folly::coro::Task<void> {
auto senderCancellationCallback = SenderCancellationCallback( auto senderCancellationCallback = SenderCancellationCallback(
sender, executor, channelCallbackToRestore); sender, executor, channelCallbackToRestore);
......
...@@ -345,6 +345,81 @@ TEST_F(SimpleTransformFixture, Chained) { ...@@ -345,6 +345,81 @@ TEST_F(SimpleTransformFixture, Chained) {
executor_.drain(); 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 { class TransformFixtureStress : public Test {
protected: protected:
TransformFixtureStress() TransformFixtureStress()
...@@ -722,6 +797,112 @@ TEST_F(ResumableTransformFixture, TransformThrows_NoReinitialization_Rethrows) { ...@@ -722,6 +797,112 @@ TEST_F(ResumableTransformFixture, TransformThrows_NoReinitialization_Rethrows) {
executor_.drain(); 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 { class ResumableTransformFixtureStress : public Test {
protected: protected:
ResumableTransformFixtureStress() ResumableTransformFixtureStress()
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment