Tcp flow control

docs/configuration/cluster_manager/cluster_stats.rst

@@ -53,6 +53,8 @@ Every cluster has a statistics tree rooted at *cluster.<name>.* with the followi
   upstream_rq_retry, Counter, Total request retries
   upstream_rq_retry_success, Counter, Total request retry successes
   upstream_rq_retry_overflow, Counter, Total requests not retried due to circuit breaking
+  upstream_flow_control_paused_reading_total, Counter, Total number of times flow control paused reading from upstream.
+  upstream_flow_control_resumed_reading_total, Counter, Total number of times flow control resumed reading from upstream.
   membership_change, Counter, Total cluster membership changes
   membership_healthy, Gauge, Current cluster healthy total (inclusive of both health checking and outlier detection)
   membership_total, Gauge, Current cluster membership total

docs/configuration/network_filters/tcp_proxy_filter.rst

@@ -138,4 +138,6 @@ statistics are rooted at *tcp.<stat_prefix>.* with the following statistics:
   downstream_cx_no_route, Counter, Number of connections for which no matching route was found.
   downstream_cx_tx_bytes_total, Counter, Total bytes written to the downstream connection.
   downstream_cx_tx_bytes_buffered, Gauge, Total bytes currently buffered to the downstream connection.
+  downstream_flow_control_paused_reading_total, Counter, Total number of times flow control paused reading from downstream.
+  downstream_flow_control_resumed_reading_total, Counter, Total number of times flow control resumed reading from downstream.
 

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;
 };
 
 /**
@@ -153,15 +164,21 @@ class Connection : public Event::DeferredDeletable, public FilterManager {
   virtual void write(Buffer::Instance& data) PURE;
 
   /**
-   * Set a soft limit on the size of the read buffer prior to flushing to further stages in the
+   * Set a soft limit on the size of buffers for the connection.
+   * For the read buffer, this limits the bytes read prior to flushing to further stages in the
    * processing pipeline.
+   * For the write buffer, it sets watermarks.  When enough data is buffered it triggers a call to
+   * onAboveWriteBufferHighWatermark, which allows subscribers to enforce flow control by disabling
+   * reads on the socket funneling data to the write buffer.  When enough data is drained from the
+   * write buffer, onBelowWriteBufferHighWatermark is called which similarly allows subscribers
+   * resuming reading.
    */
-  virtual void setReadBufferLimit(uint32_t limit) PURE;
+  virtual void setBufferLimits(uint32_t limit) PURE;
 
   /**
-   * Get the value set with setReadBufferLimit.
+   * Get the value set with setBufferLimits.
    */
-  virtual uint32_t readBufferLimit() const PURE;
+  virtual uint32_t bufferLimit() const PURE;
 };
 
 typedef std::unique_ptr<Connection> ConnectionPtr;

include/envoy/upstream/upstream.h

@@ -201,6 +201,8 @@ class HostSet {
   COUNTER(upstream_rq_retry)                                                                       \
   COUNTER(upstream_rq_retry_success)                                                               \
   COUNTER(upstream_rq_retry_overflow)                                                              \
+  COUNTER(upstream_flow_control_paused_reading_total)                                              \
+  COUNTER(upstream_flow_control_resumed_reading_total)                                             \
   GAUGE  (max_host_weight)                                                                         \
   COUNTER(membership_change)                                                                       \
   GAUGE  (membership_healthy)                                                                      \

source/common/buffer/BUILD

@@ -8,6 +8,16 @@ load(
 
 envoy_package()
 
+envoy_cc_library(
+    name = "watermark_buffer_lib",
+    srcs = ["watermark_buffer.cc"],
+    hdrs = ["watermark_buffer.h"],
+    deps = [
+        "//include/envoy/buffer:buffer_interface",
+        "//source/common/common:assert_lib",
+    ],
+)
+
 envoy_cc_library(
     name = "buffer_lib",
     srcs = ["buffer_impl.cc"],

source/common/buffer/watermark_buffer.cc

@@ -0,0 +1,89 @@
+#include "common/buffer/watermark_buffer.h"
+
+#include "common/common/assert.h"
+
+namespace Envoy {
+namespace Buffer {
+
+void WatermarkBuffer::add(const void* data, uint64_t size) {
+  wrapped_buffer_->add(data, size);
+  checkHighWatermark();
+}
+
+void WatermarkBuffer::add(const std::string& data) {
+  wrapped_buffer_->add(data);
+  checkHighWatermark();
+}
+
+void WatermarkBuffer::add(const Instance& data) {
+  wrapped_buffer_->add(data);
+  checkHighWatermark();
+}
+
+void WatermarkBuffer::commit(RawSlice* iovecs, uint64_t num_iovecs) {
+  wrapped_buffer_->commit(iovecs, num_iovecs);
+  checkHighWatermark();
+}
+
+void WatermarkBuffer::drain(uint64_t size) {
+  wrapped_buffer_->drain(size);
+  checkLowWatermark();
+}
+
+void WatermarkBuffer::move(Instance& rhs) {
+  wrapped_buffer_->move(rhs);
+  checkHighWatermark();
+}
+
+void WatermarkBuffer::move(Instance& rhs, uint64_t length) {
+  wrapped_buffer_->move(rhs, length);
+  checkHighWatermark();
+}
+
+int WatermarkBuffer::read(int fd, uint64_t max_length) {
+  int bytes_read = wrapped_buffer_->read(fd, max_length);
+  checkHighWatermark();
+  return bytes_read;
+}
+
+uint64_t WatermarkBuffer::reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iovecs) {
+  uint64_t bytes_reserved = wrapped_buffer_->reserve(length, iovecs, num_iovecs);
+  checkHighWatermark();
+  return bytes_reserved;
+}
+
+int WatermarkBuffer::write(int fd) {
+  int bytes_written = wrapped_buffer_->write(fd);
+  checkLowWatermark();
+  return bytes_written;
+}
+
+void WatermarkBuffer::setWatermarks(uint32_t low_watermark, uint32_t high_watermark) {
+  ASSERT(low_watermark < high_watermark);
+  low_watermark_ = low_watermark;
+  high_watermark_ = high_watermark;
+  checkHighWatermark();
+  checkLowWatermark();
+}
+
+void WatermarkBuffer::checkLowWatermark() {
+  if (!above_high_watermark_called_ || wrapped_buffer_->length() >= low_watermark_) {
+    return;
+  }
+
+  above_high_watermark_called_ = false;
+  below_low_watermark_();
+}
+
+void WatermarkBuffer::checkHighWatermark() {
+  if (above_high_watermark_called_ || high_watermark_ == 0 ||
+      wrapped_buffer_->length() <= high_watermark_) {
+    return;
+  }
+
+  above_high_watermark_called_ = true;
+  above_high_watermark_();
+}
+
+} // namespace Buffer
+} // namespace Envoy

source/common/buffer/watermark_buffer.h

@@ -0,0 +1,64 @@
+#pragma once
+
+#include <string>
+
+#include "envoy/buffer/buffer.h"
+
+namespace Envoy {
+namespace Buffer {
+
+// A wrapper for an underlying buffer which does watermark validation.
+// The underlying buffer's ownership is transfered to the Watermark buffer.   Each time the inner
+// buffer is resized (written to or drained), the watermarks are checked.  As the buffer size
+// transitions from under the low watermark to above the high watermark, the above_high_watermark
+// function is called one time. It will not be called again until the buffer is drained below the
+// low watermark, at which point the below_low_watermark function is called.
+class WatermarkBuffer : public Instance {
+public:
+  WatermarkBuffer(InstancePtr&& buffer, std::function<void()> below_low_watermark,
+                  std::function<void()> above_high_watermark)
+      : wrapped_buffer_(std::move(buffer)), below_low_watermark_(below_low_watermark),
+        above_high_watermark_(above_high_watermark) {}
+
+  // Instance
+  void add(const void* data, uint64_t size) override;
+  void add(const std::string& data) override;
+  void add(const Instance& data) override;
+  void commit(RawSlice* iovecs, uint64_t num_iovecs) override;
+  void drain(uint64_t size) override;
+  uint64_t getRawSlices(RawSlice* out, uint64_t out_size) const override {
+    return wrapped_buffer_->getRawSlices(out, out_size);
+  }
+  uint64_t length() const override { return wrapped_buffer_->length(); }
+  void* linearize(uint32_t size) override { return wrapped_buffer_->linearize(size); }
+  void move(Instance& rhs) override;
+  void move(Instance& rhs, uint64_t length) override;
+  int read(int fd, uint64_t max_length) override;
+  uint64_t reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iovecs) override;
+  ssize_t search(const void* data, uint64_t size, size_t start) const override {
+    return wrapped_buffer_->search(data, size, start);
+  }
+  int write(int fd) override;
+
+  void setWatermarks(uint32_t low_watermark, uint32_t high_watermark);
+
+private:
+  void checkHighWatermark();
+  void checkLowWatermark();
+
+  InstancePtr wrapped_buffer_;
+  std::function<void()> below_low_watermark_;
+  std::function<void()> above_high_watermark_;
+
+  // Used for enforcing buffer limits (off by default).  If these are set to non-zero by a call to
+  // setWatermarks() the watermark callbacks will be called as described above.
+  uint32_t high_watermark_{0};
+  uint32_t low_watermark_{0};
+  // Tracks the latest state of watermark callbacks.
+  // True between the time above_high_watermark_ has been called until above_high_watermark_ has
+  // been called.
+  bool above_high_watermark_called_{false};
+};
+
+} // namespace Buffer
+} // namespace Envoy

source/common/event/dispatcher_impl.cc

@@ -10,6 +10,7 @@
 #include "envoy/network/listen_socket.h"
 #include "envoy/network/listener.h"
 
+#include "common/buffer/buffer_impl.h"
 #include "common/event/file_event_impl.h"
 #include "common/event/signal_impl.h"
 #include "common/event/timer_impl.h"

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,58 @@ void TcpProxy::initializeReadFilterCallbacks(Network::ReadFilterCallbacks& callb
                                                 config_->stats().downstream_cx_tx_bytes_buffered_});
 }
 
+void TcpProxy::readDisableUpstream(bool disable) {
+  upstream_connection_->readDisable(disable);
+  if (disable) {
+    read_callbacks_->upstreamHost()
+        ->cluster()
+        .stats()
+        .upstream_flow_control_paused_reading_total_.inc();
+  } else {
+    read_callbacks_->upstreamHost()
+        ->cluster()
+        .stats()
+        .upstream_flow_control_resumed_reading_total_.inc();
+  }
+}
+
+void TcpProxy::readDisableDownstream(bool disable) {
+  read_callbacks_->connection().readDisable(disable);
+  if (disable) {
+    config_->stats().downstream_flow_control_paused_reading_total_.inc();
+  } else {
+    config_->stats().downstream_flow_control_resumed_reading_total_.inc();
+  }
+}
+
+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

@@ -30,7 +30,9 @@ namespace Filter {
   COUNTER(downstream_cx_tx_bytes_total)                                                            \
   GAUGE  (downstream_cx_tx_bytes_buffered)                                                         \
   COUNTER(downstream_cx_total)                                                                     \
-  COUNTER(downstream_cx_no_route)
+  COUNTER(downstream_cx_no_route)                                                                  \
+  COUNTER(downstream_flow_control_paused_reading_total)                                            \
+  COUNTER(downstream_flow_control_resumed_reading_total)
 // clang-format on
 
 /**
@@ -94,14 +96,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 +120,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 +130,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;