connection: adding watermarks to the read buffer.

source/common/buffer/watermark_buffer.h

@@ -38,6 +38,7 @@ class WatermarkBuffer : public OwnedImpl {
   void setWatermarks(uint32_t watermark) { setWatermarks(watermark / 2, watermark); }
   void setWatermarks(uint32_t low_watermark, uint32_t high_watermark);
   uint32_t highWatermark() const { return high_watermark_; }
+  bool aboveHighWatermark() const { return above_high_watermark_called_; }
 
 private:
   void checkHighWatermark();

source/common/network/connection_impl.cc

@@ -48,6 +48,8 @@ ConnectionImpl::ConnectionImpl(Event::Dispatcher& dispatcher, ConnectionSocketPt
     : ConnectionImplBase(dispatcher, next_global_id_++),
       transport_socket_(std::move(transport_socket)), socket_(std::move(socket)),
       stream_info_(stream_info), filter_manager_(*this),
+      read_buffer_([this]() -> void { this->onReadBufferLowWatermark(); },
+                   [this]() -> void { this->onReadBufferHighWatermark(); }),
       write_buffer_(dispatcher.getWatermarkFactory().create(
           [this]() -> void { this->onWriteBufferLowWatermark(); },
           [this]() -> void { this->onWriteBufferHighWatermark(); })),
@@ -186,6 +188,11 @@ Connection::State ConnectionImpl::state() const {
 
 void ConnectionImpl::closeConnectionImmediately() { closeSocket(ConnectionEvent::LocalClose); }
 
+bool ConnectionImpl::consumerWantsToRead() {
+  return read_disable_count_ == 0 ||
+         (read_disable_count_ == 1 && read_buffer_.aboveHighWatermark());
+}
+
 void ConnectionImpl::closeSocket(ConnectionEvent close_type) {
   if (!ioHandle().isOpen()) {
     return;
@@ -268,7 +275,7 @@ void ConnectionImpl::noDelay(bool enable) {
 }
 
 void ConnectionImpl::onRead(uint64_t read_buffer_size) {
-  if (read_disable_count_ != 0 || inDelayedClose()) {
+  if (inDelayedClose() || !consumerWantsToRead()) {
     return;
   }
   ASSERT(ioHandle().isOpen());
@@ -342,24 +349,23 @@ void ConnectionImpl::readDisable(bool disable) {
     }
   } else {
     --read_disable_count_;
-    if (read_disable_count_ != 0) {
-      // The socket should stay disabled.
-      return;
-    }
     if (state() != State::Open || file_event_ == nullptr) {
       // If readDisable is called on a closed connection, do not crash.
       return;
     }
 
-    // We never ask for both early close and read at the same time. If we are reading, we want to
-    // consume all available data.
-    file_event_->setEnabled(Event::FileReadyType::Read | Event::FileReadyType::Write);
-    // If the connection has data buffered there's no guarantee there's also data in the kernel
-    // which will kick off the filter chain. Instead fake an event to make sure the buffered data
-    // gets processed regardless and ensure that we dispatch it via onRead.
-    if (read_buffer_.length() > 0) {
+    if (read_disable_count_ == 0) {
+      // We never ask for both early close and read at the same time. If we are reading, we want to
+      // consume all available data.
+      file_event_->setEnabled(Event::FileReadyType::Read | Event::FileReadyType::Write);
+    }
+
+    if (consumerWantsToRead() && read_buffer_.length() > 0) {
+      // If the connection has data buffered there's no guarantee there's also data in the kernel
+      // which will kick off the filter chain. Instead fake an event to make sure the buffered data
+      // gets processed regardless and ensure that we dispatch it via onRead.
       dispatch_buffered_data_ = true;
-      file_event_->activate(Event::FileReadyType::Read);
+      setReadBufferReady();
     }
   }
 }
@@ -465,9 +471,22 @@ void ConnectionImpl::setBufferLimits(uint32_t limit) {
   // would result in respecting the exact buffer limit.
   if (limit > 0) {
     static_cast<Buffer::WatermarkBuffer*>(write_buffer_.get())->setWatermarks(limit + 1);
+    read_buffer_.setWatermarks(limit + 1);
   }
 }
 
+void ConnectionImpl::onReadBufferLowWatermark() {
+  ENVOY_CONN_LOG(debug, "onBelowReadBufferLowWatermark", *this);
+  if (state() == State::Open) {
+    readDisable(false);
+  }
+}
+
+void ConnectionImpl::onReadBufferHighWatermark() {
+  ENVOY_CONN_LOG(debug, "onAboveReadBufferHighWatermark", *this);
+  readDisable(true);
+}
+
 void ConnectionImpl::onWriteBufferLowWatermark() {
   ENVOY_CONN_LOG(debug, "onBelowWriteBufferLowWatermark", *this);
   ASSERT(write_buffer_above_high_watermark_);
@@ -529,6 +548,19 @@ void ConnectionImpl::onReadReady() {
 
   ASSERT(!connecting_);
 
+  // We get here while read disabled in two ways.
+  // 1) There was a call to setReadBufferReady(), for example if a raw buffer socket ceded due to
+  //    shouldDrainReadBuffer(). In this case we defer the event until the socket is read enabled.
+  // 2) The consumer of connection data called readDisable(true), and instead of reading from the
+  //    socket we simply need to dispatch already read data.
+  if (read_disable_count_ != 0) {
+    if (dispatch_buffered_data_ && consumerWantsToRead()) {
+      onRead(read_buffer_.length());
+      dispatch_buffered_data_ = false;
+    }
+    return;
+  }
+
   IoResult result = transport_socket_->doRead(read_buffer_);
   uint64_t new_buffer_size = read_buffer_.length();
   updateReadBufferStats(result.bytes_processed_, new_buffer_size);

source/common/network/connection_impl.h

@@ -122,11 +122,19 @@ class ConnectionImpl : public ConnectionImplBase, public TransportSocketCallback
   static uint64_t nextGlobalIdForTest() { return next_global_id_; }
 
 protected:
+  // A convenience function which returns true if
+  // 1) The read disable count is zero or
+  // 2) The read disable count is one, due to the read buffer being overrun.
+  // In either case the consumer of the data would like to read from the buffer.
+  bool consumerWantsToRead();
+
   // Network::ConnectionImplBase
   void closeConnectionImmediately() override;
 
   void closeSocket(ConnectionEvent close_type);
 
+  void onReadBufferLowWatermark();
+  void onReadBufferHighWatermark();
   void onWriteBufferLowWatermark();
   void onWriteBufferHighWatermark();
 
@@ -135,7 +143,9 @@ class ConnectionImpl : public ConnectionImplBase, public TransportSocketCallback
   StreamInfo::StreamInfo& stream_info_;
   FilterManagerImpl filter_manager_;
 
-  Buffer::OwnedImpl read_buffer_;
+  // Ensure that if the consumer of the data from this connection isn't
+  // consuming, that the connection eventually stops reading from the wire.
+  Buffer::WatermarkBuffer read_buffer_;
   // This must be a WatermarkBuffer, but as it is created by a factory the ConnectionImpl only has
   // a generic pointer.
   // It MUST be defined after the filter_manager_ as some filters may have callbacks that

test/common/network/connection_impl_test.cc

@@ -93,6 +93,12 @@ TEST_P(ConnectionImplDeathTest, BadFd) {
       ".*assert failure: SOCKET_VALID\\(ConnectionImpl::ioHandle\\(\\)\\.fd\\(\\)\\).*");
 }
 
+class TestClientConnectionImpl : public Network::ClientConnectionImpl {
+public:
+  using ClientConnectionImpl::ClientConnectionImpl;
+  Buffer::WatermarkBuffer& readBuffer() { return read_buffer_; }
+};
+
 class ConnectionImplTest : public testing::TestWithParam<Address::IpVersion> {
 protected:
   ConnectionImplTest() : api_(Api::createApiForTest(time_system_)), stream_info_(time_system_) {}
@@ -104,9 +110,9 @@ class ConnectionImplTest : public testing::TestWithParam<Address::IpVersion> {
     socket_ = std::make_shared<Network::TcpListenSocket>(Network::Test::getAnyAddress(GetParam()),
                                                          nullptr, true);
     listener_ = dispatcher_->createListener(socket_, listener_callbacks_, true);
-    client_connection_ = dispatcher_->createClientConnection(
-        socket_->localAddress(), source_address_, Network::Test::createRawBufferSocket(),
-        socket_options_);
+    client_connection_ = std::make_unique<Network::TestClientConnectionImpl>(
+        *dispatcher_, socket_->localAddress(), source_address_,
+        Network::Test::createRawBufferSocket(), socket_options_);
     client_connection_->addConnectionCallbacks(client_callbacks_);
     EXPECT_EQ(nullptr, client_connection_->ssl());
     const Network::ClientConnection& const_connection = *client_connection_;
@@ -215,6 +221,9 @@ class ConnectionImplTest : public testing::TestWithParam<Address::IpVersion> {
     return ConnectionMocks{std::move(dispatcher), timer, std::move(transport_socket), file_event,
                            &file_ready_cb_};
   }
+  Network::TestClientConnectionImpl* testClientConnection() {
+    return dynamic_cast<Network::TestClientConnectionImpl*>(client_connection_.get());
+  }
 
   Event::FileReadyCb file_ready_cb_;
   Event::SimulatedTimeSystem time_system_;
@@ -742,7 +751,7 @@ TEST_P(ConnectionImplTest, HalfCloseNoEarlyCloseDetection) {
 }
 
 // Test that as watermark levels are changed, the appropriate callbacks are triggered.
-TEST_P(ConnectionImplTest, Watermarks) {
+TEST_P(ConnectionImplTest, WriteWatermarks) {
   useMockBuffer();
 
   setUpBasicConnection();
@@ -791,6 +800,75 @@ TEST_P(ConnectionImplTest, Watermarks) {
   disconnect(false);
 }
 
+// Test that as watermark levels are changed, the appropriate callbacks are triggered.
+TEST_P(ConnectionImplTest, ReadWatermarks) {
+
+  setUpBasicConnection();
+  client_connection_->setBufferLimits(2);
+  std::shared_ptr<MockReadFilter> client_read_filter(new NiceMock<MockReadFilter>());
+  client_connection_->addReadFilter(client_read_filter);
+  connect();
+
+  EXPECT_FALSE(testClientConnection()->readBuffer().aboveHighWatermark());
+  EXPECT_TRUE(client_connection_->readEnabled());
+  // Add 4 bytes to the buffer and verify the connection becomes read disabled.
+  {
+    Buffer::OwnedImpl buffer("data");
+    server_connection_->write(buffer, false);
+    EXPECT_CALL(*client_read_filter, onData(_, false))
+        .WillRepeatedly(Invoke([&](Buffer::Instance&, bool) -> FilterStatus {
+          dispatcher_->exit();
+          return FilterStatus::StopIteration;
+        }));
+    dispatcher_->run(Event::Dispatcher::RunType::Block);
+
+    EXPECT_TRUE(testClientConnection()->readBuffer().aboveHighWatermark());
+    EXPECT_FALSE(client_connection_->readEnabled());
+  }
+
+  // Drain 3 bytes from the buffer. This bring sit below the low watermark, and
+  // read enables, as well as triggering a kick for the remaining byte.
+  {
+    testClientConnection()->readBuffer().drain(3);
+    EXPECT_FALSE(testClientConnection()->readBuffer().aboveHighWatermark());
+    EXPECT_TRUE(client_connection_->readEnabled());
+
+    EXPECT_CALL(*client_read_filter, onData(_, false));
+    dispatcher_->run(Event::Dispatcher::RunType::NonBlock);
+  }
+
+  // Add 3 bytes to the buffer and verify the connection becomes read disabled
+  // again.
+  {
+    Buffer::OwnedImpl buffer("bye");
+    server_connection_->write(buffer, false);
+    EXPECT_CALL(*client_read_filter, onData(_, false))
+        .WillRepeatedly(Invoke([&](Buffer::Instance&, bool) -> FilterStatus {
+          dispatcher_->exit();
+          return FilterStatus::StopIteration;
+        }));
+    dispatcher_->run(Event::Dispatcher::RunType::Block);
+
+    EXPECT_TRUE(testClientConnection()->readBuffer().aboveHighWatermark());
+    EXPECT_FALSE(client_connection_->readEnabled());
+  }
+
+  // Now have the consumer read disable.
+  // This time when the buffer is drained, there will be no kick as the consumer
+  // does not want to read.
+  {
+    client_connection_->readDisable(true);
+    testClientConnection()->readBuffer().drain(3);
+    EXPECT_FALSE(testClientConnection()->readBuffer().aboveHighWatermark());
+    EXPECT_FALSE(client_connection_->readEnabled());
+
+    EXPECT_CALL(*client_read_filter, onData(_, false)).Times(0);
+    dispatcher_->run(Event::Dispatcher::RunType::NonBlock);
+  }
+
+  disconnect(true);
+}
+
 // Write some data to the connection. It will automatically attempt to flush
 // it to the upstream file descriptor via a write() call to buffer_, which is
 // configured to succeed and accept all bytes read.