ConsumerBase sends excessive numbers of requestN messages

Fuat Geleri · Fuat Geleri · commit b37d1dc6dbab · 2017-12-07T12:14:42.000-08:00
diff --git a/rsocket/statemachine/ConsumerBase.cpp b/rsocket/statemachine/ConsumerBase.cpp
@@ -35,6 +35,11 @@ void ConsumerBase::cancelConsumer() {
   consumingSubscriber_ = nullptr;
 }
 
+void ConsumerBase::addImplicitAllowance(size_t n) {
+  allowance_.add(n);
+  activeRequests_.add(n);
+}
+
 void ConsumerBase::generateRequest(size_t n) {
   allowance_.add(n);
   pendingAllowance_.add(n);
@@ -64,14 +69,14 @@ void ConsumerBase::processPayload(Payload&& payload, bool onNext) {
   if (payload || onNext) {
     // Frames carry application-level payloads are taken into account when
     // figuring out flow control allowance.
-    if (allowance_.tryConsume(1)) {
+    if (allowance_.tryConsume(1) && activeRequests_.tryConsume(1)) {
       sendRequests();
       if (consumingSubscriber_) {
         consumingSubscriber_->onNext(std::move(payload));
       } else {
         LOG(ERROR)
             << "consuming subscriber is missing, might be a race condition on "
-                " cancel/onNext.";
+               " cancel/onNext.";
       }
     } else {
       handleFlowControlError();
@@ -97,12 +102,16 @@ void ConsumerBase::errorConsumer(folly::exception_wrapper ex) {
 }
 
 void ConsumerBase::sendRequests() {
-  // TODO(stupaq): batch if remote end has some spare allowance
-  // TODO(stupaq): limit how much is synced to the other end
-  size_t toSync = Frame_REQUEST_N::kMaxRequestN;
-  toSync = pendingAllowance_.consumeUpTo(toSync);
-  if (toSync > 0) {
-    writeRequestN(static_cast<uint32_t>(toSync));
+  auto toSync =
+      std::min<size_t>(pendingAllowance_.get(), Frame_REQUEST_N::kMaxRequestN);
+  auto actives = activeRequests_.get();
+  if (actives < (toSync + 1) / 2) {
+    toSync = toSync - actives;
+    toSync = pendingAllowance_.consumeUpTo(toSync);
+    if (toSync > 0) {
+      writeRequestN(static_cast<uint32_t>(toSync));
+      activeRequests_.add(toSync);
+    }
   }
 }
 
diff --git a/rsocket/statemachine/ConsumerBase.h b/rsocket/statemachine/ConsumerBase.h
@@ -27,9 +27,7 @@ class ConsumerBase : public StreamStateMachineBase,
   ///
   /// This portion of allowance will not be synced to the remote end, but will
   /// count towards the limit of allowance the remote PublisherBase may use.
-  void addImplicitAllowance(size_t n) {
-    allowance_.add(n);
-  }
+  void addImplicitAllowance(size_t n);
 
   void subscribe(
       yarpl::Reference<yarpl::flowable::Subscriber<Payload>> subscriber);
@@ -68,6 +66,10 @@ class ConsumerBase : public StreamStateMachineBase,
   /// REQUEST_N frames.
   Allowance pendingAllowance_;
 
+  /// The number of already requested payload count.
+  /// Prevent excessive requestN calls.
+  Allowance activeRequests_;
+
   enum class State : uint8_t {
     RESPONDING,
     CLOSED,
diff --git a/test/RSocketTests.cpp b/test/RSocketTests.cpp
@@ -17,14 +17,16 @@ std::unique_ptr<TcpConnectionFactory> getConnFactory(
 }
 
 std::unique_ptr<RSocketServer> makeServer(
-    std::shared_ptr<rsocket::RSocketResponder> responder) {
+    std::shared_ptr<rsocket::RSocketResponder> responder,
+    std::shared_ptr<RSocketStats> stats) {
   TcpConnectionAcceptor::Options opts;
   opts.threads = 2;
   opts.address = folly::SocketAddress("0.0.0.0", 0);
 
   // RSocket server accepting on TCP.
   auto rs = RSocket::createServer(
-      std::make_unique<TcpConnectionAcceptor>(std::move(opts)));
+      std::make_unique<TcpConnectionAcceptor>(std::move(opts)),
+      std::move(stats));
 
   rs->start([r = std::move(responder)](const SetupParameters&) { return r; });
   return rs;
@@ -45,14 +47,15 @@ std::unique_ptr<RSocketServer> makeResumableServer(
 folly::Future<std::unique_ptr<RSocketClient>> makeClientAsync(
     folly::EventBase* eventBase,
     uint16_t port,
-    folly::EventBase* stateMachineEvb) {
+    folly::EventBase* stateMachineEvb,
+    std::shared_ptr<RSocketStats> stats) {
   CHECK(eventBase);
   return RSocket::createConnectedClient(
       getConnFactory(eventBase, port),
       SetupParameters(),
       std::make_shared<RSocketResponder>(),
       kDefaultKeepaliveInterval,
-      RSocketStats::noop(),
+      std::move(stats),
       std::shared_ptr<RSocketConnectionEvents>(),
       std::shared_ptr<ResumeManager>(),
       std::shared_ptr<ColdResumeHandler>(),
@@ -62,8 +65,10 @@ folly::Future<std::unique_ptr<RSocketClient>> makeClientAsync(
 std::unique_ptr<RSocketClient> makeClient(
     folly::EventBase* eventBase,
     uint16_t port,
-    folly::EventBase* stateMachineEvb) {
-  return makeClientAsync(eventBase, port, stateMachineEvb).get();
+    folly::EventBase* stateMachineEvb,
+    std::shared_ptr<RSocketStats> stats) {
+  return makeClientAsync(
+      eventBase, port, stateMachineEvb, std::move(stats)).get();
 }
 
 namespace {
diff --git a/test/RSocketTests.h b/test/RSocketTests.h
@@ -13,28 +13,50 @@ namespace rsocket {
 namespace tests {
 namespace client_server {
 
+class RSocketStatsFlowControl : public RSocketStats {
+public:
+  void frameWritten(FrameType frameType) {
+    if (frameType == FrameType::REQUEST_N) {
+      ++writeRequestN_;
+    }
+  }
+
+  void frameRead(FrameType frameType) {
+    if (frameType == FrameType::REQUEST_N) {
+      ++readRequestN_;
+    }
+  }
+
+public:
+  int writeRequestN_{0};
+  int readRequestN_{0};
+};
+
 std::unique_ptr<TcpConnectionFactory> getConnFactory(
     folly::EventBase* eventBase,
     uint16_t port);
 
 std::unique_ptr<RSocketServer> makeServer(
-    std::shared_ptr<rsocket::RSocketResponder> responder);
+    std::shared_ptr<rsocket::RSocketResponder> responder,
+    std::shared_ptr<RSocketStats> stats = RSocketStats::noop());
 
 std::unique_ptr<RSocketServer> makeResumableServer(
     std::shared_ptr<RSocketServiceHandler> serviceHandler);
 
 std::unique_ptr<RSocketClient> makeClient(
     folly::EventBase* eventBase,
     uint16_t port,
-    folly::EventBase* stateMachineEvb = nullptr);
+    folly::EventBase* stateMachineEvb = nullptr,
+    std::shared_ptr<RSocketStats> stats = RSocketStats::noop());
 
 std::unique_ptr<RSocketClient> makeDisconnectedClient(
     folly::EventBase* eventBase);
 
 folly::Future<std::unique_ptr<RSocketClient>> makeClientAsync(
     folly::EventBase* eventBase,
     uint16_t port,
-    folly::EventBase* stateMachineEvb = nullptr);
+    folly::EventBase* stateMachineEvb = nullptr,
+    std::shared_ptr<RSocketStats> stats = RSocketStats::noop());
 
 std::unique_ptr<RSocketClient> makeWarmResumableClient(
     folly::EventBase* eventBase,
diff --git a/test/RequestChannelTest.cpp b/test/RequestChannelTest.cpp
@@ -61,6 +61,36 @@ TEST(RequestChannelTest, Hello) {
   ts->assertValueAt(1, "[/hello] Hello Jane!");
 }
 
+TEST(RequestChannelTest, HelloNoFlowControl) {
+  folly::ScopedEventBaseThread worker;
+  auto server = makeServer(std::make_shared<TestHandlerHello>());
+  auto stats = std::make_shared<RSocketStatsFlowControl>();
+  auto client = makeClient(
+      worker.getEventBase(), *server->listeningPort(), nullptr, stats);
+  auto requester = client->getRequester();
+
+  auto ts = TestSubscriber<std::string>::create();
+  requester
+      ->requestChannel(
+          Flowables::justN({"/hello", "Bob", "Jane"})->map([](std::string v) {
+            return Payload(v);
+          }))
+      ->map([](auto p) { return p.moveDataToString(); })
+      ->subscribe(ts);
+
+  ts->awaitTerminalEvent();
+  ts->assertSuccess();
+  ts->assertValueCount(2);
+  // assert that we echo back the 2nd and 3rd request values
+  // with the 1st initial payload prepended to each
+  ts->assertValueAt(0, "[/hello] Hello Bob!");
+  ts->assertValueAt(1, "[/hello] Hello Jane!");
+
+  // Make sure that the initial requestN in the Stream Request Frame
+  // is already enough and no other requestN messages are sent.
+  EXPECT_EQ(stats->writeRequestN_, 0);
+}
+
 TEST(RequestChannelTest, RequestOnDisconnectedClient) {
   folly::ScopedEventBaseThread worker;
   auto client = makeDisconnectedClient(worker.getEventBase());
diff --git a/test/RequestStreamTest.cpp b/test/RequestStreamTest.cpp
@@ -76,6 +76,28 @@ TEST(RequestStreamTest, HelloFlowControl) {
   ts->assertSuccess();
 }
 
+TEST(RequestStreamTest, HelloNoFlowControl) {
+  folly::ScopedEventBaseThread worker;
+  auto server = makeServer(std::make_shared<TestHandlerSync>());
+  auto stats = std::make_shared<RSocketStatsFlowControl>();
+  auto client = makeClient(
+      worker.getEventBase(), *server->listeningPort(), nullptr, stats);
+  auto requester = client->getRequester();
+  auto ts = TestSubscriber<std::string>::create();
+  requester->requestStream(Payload("Bob"))
+      ->map([](auto p) { return p.moveDataToString(); })
+      ->subscribe(ts);
+  ts->awaitTerminalEvent();
+  ts->assertSuccess();
+  ts->assertValueCount(10);
+  ts->assertValueAt(0, "Hello Bob 1!");
+  ts->assertValueAt(9, "Hello Bob 10!");
+
+  // Make sure that the initial requestN in the Stream Request Frame
+  // is already enough and no other requestN messages are sent.
+  EXPECT_EQ(stats->writeRequestN_, 0);
+}
+
 class TestHandlerAsync : public rsocket::RSocketResponder {
  public:
   Reference<Flowable<Payload>> handleRequestStream(Payload request, StreamId)
@@ -98,7 +120,8 @@ class TestHandlerAsync : public rsocket::RSocketResponder {
               return Payload(s, "metadata");
             })
             ->subscribe(subscriber);
-      }).detach();
+      })
+          .detach();
     });
   }
 };
