Tcp flow control

include/envoy/network/connection.h

@@ -45,6 +45,17 @@ class ConnectionCallbacks {
    * @param events supplies the ConnectionEvent events that occurred as a bitmask.
    */
   virtual void onEvent(uint32_t events) PURE;
+
+  /**
+   * Called when the write buffer for a connection goes over its high watermark.
+   */
+  virtual void onAboveWriteBufferHighWatermark() PURE;
+
+  /**
+   * Called when the write buffer for a connection goes from over its high
+   * watermark to under its low watermark.
+   */
+  virtual void onBelowWriteBufferLowWatermark() PURE;
 };
 
 /**
@@ -162,6 +173,20 @@ class Connection : public Event::DeferredDeletable, public FilterManager {
    * Get the value set with setReadBufferLimit.
    */
   virtual uint32_t readBufferLimit() const PURE;
+
+  /**
+   * Sets the high and low watermarks which trigger onAboveWriteBufferHighWatermark
+   * and onBelowWriteBufferHighWatermark callbacks.
+   * The connection is assumed to start out with less than high_watermark
+   * worth of data buffered, so onAboveWriteBufferHighWatermark will always be
+   * called before onAboveWriteBufferHighWatermark
+   * @param low_watermark if the connection was above the high watermark and the
+   * connection buffer is drained below this many bytes, onBelowWriteBufferHighWatermark will be
+   * called.
+   * @param high_watermark if the connection has more bytes than this buffered,
+   * onAboveWriteBufferHighWatermark will be called.
+   */
+  virtual void setWriteBufferWatermarks(uint32_t low_watermark, uint32_t high_watermark) PURE;
 };
 
 typedef std::unique_ptr<Connection> ConnectionPtr;

source/common/filter/auth/client_ssl.h

@@ -119,6 +119,8 @@ class Instance : public Network::ReadFilter, public Network::ConnectionCallbacks
 
   // Network::ConnectionCallbacks
   void onEvent(uint32_t events) override;
+  void onAboveWriteBufferHighWatermark() override {}
+  void onBelowWriteBufferLowWatermark() override {}
 
 private:
   ConfigSharedPtr config_;

source/common/filter/ratelimit.h

@@ -82,6 +82,8 @@ class Instance : public Network::ReadFilter,
 
   // Network::ConnectionCallbacks
   void onEvent(uint32_t events) override;
+  void onAboveWriteBufferHighWatermark() override {}
+  void onBelowWriteBufferLowWatermark() override {}
 
   // RateLimit::RequestCallbacks
   void complete(LimitStatus status) override;

source/common/filter/tcp_proxy.cc

@@ -123,6 +123,40 @@ void TcpProxy::initializeReadFilterCallbacks(Network::ReadFilterCallbacks& callb
                                                 config_->stats().downstream_cx_tx_bytes_buffered_});
 }
 
+void TcpProxy::readDisableUpstream(bool disable) { upstream_connection_->readDisable(disable); }
+
+void TcpProxy::readDisableDownstream(bool disable) {
+  read_callbacks_->connection().readDisable(disable);
+}
+
+void TcpProxy::DownstreamCallbacks::onAboveWriteBufferHighWatermark() {
+  ASSERT(!on_high_watermark_called_);
+  on_high_watermark_called_ = true;
+  // If downstream has too much data buffered, stop reading on the upstream connection.
+  parent_.readDisableUpstream(true);
+}
+
+void TcpProxy::DownstreamCallbacks::onBelowWriteBufferLowWatermark() {
+  ASSERT(on_high_watermark_called_);
+  on_high_watermark_called_ = false;
+  // The downstream buffer has been drained.  Resume reading from upstream.
+  parent_.readDisableUpstream(false);
+}
+
+void TcpProxy::UpstreamCallbacks::onAboveWriteBufferHighWatermark() {
+  ASSERT(!on_high_watermark_called_);
+  on_high_watermark_called_ = true;
+  // There's too much data buffered in the upstream write buffer, so stop reading.
+  parent_.readDisableDownstream(true);
+}
+
+void TcpProxy::UpstreamCallbacks::onBelowWriteBufferLowWatermark() {
+  ASSERT(on_high_watermark_called_);
+  on_high_watermark_called_ = false;
+  // The upstream write buffer is drained.  Resume reading.
+  parent_.readDisableDownstream(false);
+}
+
 Network::FilterStatus TcpProxy::initializeUpstreamConnection() {
   const std::string& cluster_name = config_->getRouteFromEntries(read_callbacks_->connection());
 

source/common/filter/tcp_proxy.h

@@ -94,14 +94,22 @@ class TcpProxy : public Network::ReadFilter, Logger::Loggable<Logger::Id::filter
   Network::FilterStatus onNewConnection() override { return initializeUpstreamConnection(); }
   void initializeReadFilterCallbacks(Network::ReadFilterCallbacks& callbacks) override;
 
+  // These two functions allow enabling/disabling reads on the upstream and downstream connections.
+  // They are called by the Downstream/Upstream Watermark callbacks to limit buffering.
+  void readDisableUpstream(bool disable);
+  void readDisableDownstream(bool disable);
+
 private:
   struct DownstreamCallbacks : public Network::ConnectionCallbacks {
     DownstreamCallbacks(TcpProxy& parent) : parent_(parent) {}
 
     // Network::ConnectionCallbacks
     void onEvent(uint32_t event) override { parent_.onDownstreamEvent(event); }
+    void onAboveWriteBufferHighWatermark() override;
+    void onBelowWriteBufferLowWatermark() override;
 
     TcpProxy& parent_;
+    bool on_high_watermark_called_{false};
   };
 
   struct UpstreamCallbacks : public Network::ConnectionCallbacks,
@@ -110,6 +118,8 @@ class TcpProxy : public Network::ReadFilter, Logger::Loggable<Logger::Id::filter
 
     // Network::ConnectionCallbacks
     void onEvent(uint32_t event) override { parent_.onUpstreamEvent(event); }
+    void onAboveWriteBufferHighWatermark() override;
+    void onBelowWriteBufferLowWatermark() override;
 
     // Network::ReadFilter
     Network::FilterStatus onData(Buffer::Instance& data) override {
@@ -118,6 +128,7 @@ class TcpProxy : public Network::ReadFilter, Logger::Loggable<Logger::Id::filter
     }
 
     TcpProxy& parent_;
+    bool on_high_watermark_called_{false};
   };
 
   Network::FilterStatus initializeUpstreamConnection();

source/common/http/codec_client.h

@@ -180,6 +180,8 @@ class CodecClient : Logger::Loggable<Logger::Id::client>,
 
   // Network::ConnectionCallbacks
   void onEvent(uint32_t events) override;
+  void onAboveWriteBufferHighWatermark() override {}
+  void onBelowWriteBufferLowWatermark() override {}
 
   std::list<ActiveRequestPtr> active_requests_;
   Http::ConnectionCallbacks* codec_callbacks_{};

source/common/http/conn_manager_impl.h

@@ -279,6 +279,8 @@ class ConnectionManagerImpl : Logger::Loggable<Logger::Id::http>,
 
   // Network::ConnectionCallbacks
   void onEvent(uint32_t events) override;
+  void onAboveWriteBufferHighWatermark() override {}
+  void onBelowWriteBufferLowWatermark() override {}
 
 private:
   struct ActiveStream;

source/common/http/http1/conn_pool.h

@@ -76,6 +76,8 @@ class ConnPoolImpl : Logger::Loggable<Logger::Id::pool>, public ConnectionPool::
 
     // Network::ConnectionCallbacks
     void onEvent(uint32_t events) override { parent_.onConnectionEvent(*this, events); }
+    void onAboveWriteBufferHighWatermark() override {}
+    void onBelowWriteBufferLowWatermark() override {}
 
     ConnPoolImpl& parent_;
     CodecClientPtr codec_client_;

source/common/http/http2/conn_pool.h

@@ -43,6 +43,8 @@ class ConnPoolImpl : Logger::Loggable<Logger::Id::pool>, public ConnectionPool::
 
     // Network::ConnectionCallbacks
     void onEvent(uint32_t events) override { parent_.onConnectionEvent(*this, events); }
+    void onAboveWriteBufferHighWatermark() override {}
+    void onBelowWriteBufferLowWatermark() override {}
 
     // CodecClientCallbacks
     void onStreamDestroy() override { parent_.onStreamDestroy(*this); }

source/common/mongo/proxy.h

@@ -107,6 +107,8 @@ class ProxyFilter : public Network::Filter,
 
   // Network::ConnectionCallbacks
   void onEvent(uint32_t event) override;
+  void onAboveWriteBufferHighWatermark() override {}
+  void onBelowWriteBufferLowWatermark() override {}
 
 private:
   struct ActiveQuery {

source/common/network/connection_impl.cc

@@ -98,7 +98,7 @@ void ConnectionImpl::close(ConnectionCloseType type) {
     return;
   }
 
-  uint64_t data_to_write = write_buffer_.length();
+  uint64_t data_to_write = write_buffer_->length();
   ENVOY_CONN_LOG(debug, "closing data_to_write={} type={}", *this, data_to_write, enumToInt(type));
   if (data_to_write == 0 || type == ConnectionCloseType::NoFlush) {
     if (data_to_write > 0) {
@@ -207,15 +207,26 @@ void ConnectionImpl::readDisable(bool disable) {
   // TODO(mattklein123): Potentially support half-closed TCP connections. It's unclear if this is
   // required for any scenarios in which Envoy will be used (I don't know of any).
   if (disable) {
+    if (!read_enabled) {
+      ++read_disable_count_;
+      return;
+    }
     ASSERT(read_enabled);
     state_ &= ~InternalState::ReadEnabled;
     file_event_->setEnabled(Event::FileReadyType::Write | Event::FileReadyType::Closed);
   } else {
+    if (read_disable_count_ > 0) {
+      --read_disable_count_;
+      return;
+    }
     ASSERT(!read_enabled);
     state_ |= InternalState::ReadEnabled;
     // 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.
     if (read_buffer_.length() > 0) {
       file_event_->activate(Event::FileReadyType::Read);
     }
@@ -254,13 +265,52 @@ void ConnectionImpl::write(Buffer::Instance& data) {
     // ever changed, read the comment in Ssl::ConnectionImpl::doWriteToSocket() VERY carefully.
     // That code assumes that we never change existing write_buffer_ chain elements between calls
     // to SSL_write(). That code will have to change if we ever copy here.
-    write_buffer_.move(data);
+    write_buffer_->move(data);
+    checkForHighWatermark();
+
     if (!(state_ & InternalState::Connecting)) {
       file_event_->activate(Event::FileReadyType::Write);
     }
   }
 }
 
+void ConnectionImpl::setWriteBufferWatermarks(uint32_t new_low_watermark,
+                                              uint32_t new_high_watermark) {
+  ENVOY_CONN_LOG(debug, "Setting watermarks: {} {}", *this, new_low_watermark, new_high_watermark);
+  ASSERT(new_low_watermark < new_high_watermark);
+
+  high_watermark_ = new_high_watermark;
+  low_watermark_ = new_low_watermark;
+
+  checkForLowWatermark();
+  checkForHighWatermark();
+}
+
+void ConnectionImpl::checkForLowWatermark() {
+  if (!above_high_watermark_called_ || write_buffer_->length() >= low_watermark_) {
+    return;
+  }
+  ENVOY_CONN_LOG(debug, "onBelowWriteBufferLowWatermark", *this);
+
+  above_high_watermark_called_ = false;
+  for (ConnectionCallbacks* callback : callbacks_) {
+    callback->onBelowWriteBufferLowWatermark();
+  }
+}
+
+void ConnectionImpl::checkForHighWatermark() {
+  if (above_high_watermark_called_ || high_watermark_ == 0 ||
+      write_buffer_->length() <= high_watermark_) {
+    return;
+  }
+  ENVOY_CONN_LOG(debug, "onAboveWriteBufferHighWatermark", *this);
+
+  above_high_watermark_called_ = true;
+  for (ConnectionCallbacks* callback : callbacks_) {
+    callback->onAboveWriteBufferHighWatermark();
+  }
+}
+
 void ConnectionImpl::onFileEvent(uint32_t events) {
   ENVOY_CONN_LOG(trace, "socket event: {}", *this, events);
 
@@ -347,12 +397,11 @@ ConnectionImpl::IoResult ConnectionImpl::doWriteToSocket() {
   PostIoAction action;
   uint64_t bytes_written = 0;
   do {
-    if (write_buffer_.length() == 0) {
+    if (write_buffer_->length() == 0) {
       action = PostIoAction::KeepOpen;
       break;
     }
-
-    int rc = write_buffer_.write(fd_);
+    int rc = write_buffer_->write(fd_);
     ENVOY_CONN_LOG(trace, "write returns: {}", *this, rc);
     if (rc == -1) {
       ENVOY_CONN_LOG(trace, "write error: {}", *this, errno);
@@ -364,6 +413,7 @@ ConnectionImpl::IoResult ConnectionImpl::doWriteToSocket() {
 
       break;
     } else {
+      checkForLowWatermark();
       bytes_written += rc;
     }
   } while (true);
@@ -400,7 +450,7 @@ void ConnectionImpl::onWriteReady() {
   }
 
   IoResult result = doWriteToSocket();
-  uint64_t new_buffer_size = write_buffer_.length();
+  uint64_t new_buffer_size = write_buffer_->length();
   updateWriteBufferStats(result.bytes_processed_, new_buffer_size);
 
   if (result.action_ == PostIoAction::Close) {

source/common/network/connection_impl.h

@@ -72,11 +72,16 @@ class ConnectionImpl : public virtual Connection,
   void write(Buffer::Instance& data) override;
   void setReadBufferLimit(uint32_t limit) override { read_buffer_limit_ = limit; }
   uint32_t readBufferLimit() const override { return read_buffer_limit_; }
+  void setWriteBufferWatermarks(uint32_t low_watermark, uint32_t high_watermark) override;
 
   // Network::BufferSource
   Buffer::Instance& getReadBuffer() override { return read_buffer_; }
   Buffer::Instance& getWriteBuffer() override { return *current_write_buffer_; }
 
+  void replaceWriteBufferForTest(std::unique_ptr<Buffer::OwnedImpl> new_buffer) {
+    write_buffer_ = std::move(new_buffer);
+  }
+
 protected:
   enum class PostIoAction { Close, KeepOpen };
 
@@ -99,12 +104,29 @@ class ConnectionImpl : public virtual Connection,
   // Reconsider how to make fairness happen.
   void setReadBufferReady() { file_event_->activate(Event::FileReadyType::Read); }
 
+  // Called when data is drained from the write buffer, to see if onBelowWriteBufferLowWatermark
+  // should be called.
+  void checkForLowWatermark();
+  // Called when data is added to the write buffer, to see if onAboveWriteBufferHighWatermark should
+  // be called.
+  void checkForHighWatermark();
+
   FilterManagerImpl filter_manager_;
   Address::InstanceConstSharedPtr remote_address_;
   Address::InstanceConstSharedPtr local_address_;
   Buffer::OwnedImpl read_buffer_;
-  Buffer::OwnedImpl write_buffer_;
+  Buffer::InstancePtr write_buffer_{new Buffer::OwnedImpl};
   uint32_t read_buffer_limit_ = 0;
+  // Used for network level buffer limits (off by default).  If these are non-zero, when the write
+  // buffer passes |high_watermark_|, onAboveWriteBufferHighWatermark will be called to disable
+  // reading further data.  When the buffer drains below |low_watermark_|,
+  // onBelowWriteBufferLowWatermark will be called to resume reads.
+  uint32_t high_watermark_{0};
+  uint32_t low_watermark_{0};
+  // Tracks the latest state of watermark callbacks.
+  // True between the time onAboveWriteBufferHighWatermark is called until the next call to
+  // onBelowLowWatermark.
+  bool above_high_watermark_called_{false};
 
 private:
   // clang-format off
@@ -138,6 +160,10 @@ class ConnectionImpl : public virtual Connection,
   uint64_t last_read_buffer_size_{};
   uint64_t last_write_buffer_size_{};
   std::unique_ptr<BufferStats> buffer_stats_;
+  // Tracks the number of times reads have been disabled.  If N different components call
+  // readDisabled(true) this allows the connection to only resume reads when readDisabled(false)
+  // has been called N times.
+  uint32_t read_disable_count_{0};
 };
 
 /**

source/common/network/listener_impl.cc

@@ -102,6 +102,17 @@ void ListenerImpl::newConnection(int fd, Address::InstanceConstSharedPtr remote_
                                  Address::InstanceConstSharedPtr local_address) {
   ConnectionPtr new_connection(new ConnectionImpl(dispatcher_, fd, remote_address, local_address));
   new_connection->setReadBufferLimit(options_.per_connection_buffer_limit_bytes_);
+  // Due to the fact that writes to the connection and flushing data from the connection are done
+  // asynchronously, we have the option of either setting the watermarks aggressively, and regularly
+  // enabling/disabling reads from the socket, or allowing more data, but then not triggering
+  // based on watermarks until 2x the data is buffered in the common case.  Given these are all soft
+  // limits we err on the side of buffeing more and having better performace.
+  // If the connection class is changed to write-and-flush the high watermark should be changed to
+  // the buffer limit without the + 1
+  if (options_.per_connection_buffer_limit_bytes_ > 0) {
+    new_connection->setWriteBufferWatermarks(options_.per_connection_buffer_limit_bytes_ / 2,
+                                             options_.per_connection_buffer_limit_bytes_ + 1);
+  }
   cb_.onNewConnection(std::move(new_connection));
 }
 
@@ -111,6 +122,10 @@ void SslListenerImpl::newConnection(int fd, Address::InstanceConstSharedPtr remo
                                                        local_address, ssl_ctx_,
                                                        Ssl::ConnectionImpl::InitialState::Server));
   new_connection->setReadBufferLimit(options_.per_connection_buffer_limit_bytes_);
+  if (options_.per_connection_buffer_limit_bytes_ > 0) {
+    new_connection->setWriteBufferWatermarks(options_.per_connection_buffer_limit_bytes_ / 2,
+                                             options_.per_connection_buffer_limit_bytes_ + 1);
+  }
   cb_.onNewConnection(std::move(new_connection));
 }