Implemented H2 stream level buffer accounting.

include/envoy/buffer/buffer.h

@@ -131,7 +131,8 @@ class Instance {
 
   /**
    * Binds the account to be charged for resources used by the buffer. This
-   * should only be called once.
+   * should only be called when the buffer is empty as existing slices
+   * within the buffer won't retroactively get tagged.
    *
    * @param account a shared_ptr to the account to charge.
    */
@@ -449,6 +450,7 @@ class Instance {
    * @param watermark supplies the buffer high watermark size threshold, in bytes.
    */
   virtual void setWatermarks(uint32_t watermark) PURE;
+
   /**
    * Returns the configured high watermark. A return value of 0 indicates that watermark
    * functionality is disabled.

include/envoy/http/codec.h

@@ -380,6 +380,14 @@ class Stream {
    * small window updates as satisfying the idle timeout as this is a potential DoS vector.
    */
   virtual void setFlushTimeout(std::chrono::milliseconds timeout) PURE;
+
+  /**
+   * Sets the account for this stream, propagating it to all of its buffers.
+   * This should only be called on client streams since server streams create
+   * their own account.
+   * @param the account to assign this stream.
+   */
+  virtual void setAccount(Buffer::BufferMemoryAccountSharedPtr account) PURE;
 };
 
 /**

include/envoy/http/filter.h

@@ -6,6 +6,7 @@
 #include <string>
 
 #include "envoy/access_log/access_log.h"
+#include "envoy/buffer/buffer.h"
 #include "envoy/common/scope_tracker.h"
 #include "envoy/event/dispatcher.h"
 #include "envoy/grpc/status.h"
@@ -544,6 +545,11 @@ class StreamDecoderFilterCallbacks : public virtual StreamFilterCallbacks {
    */
   virtual uint32_t decoderBufferLimit() PURE;
 
+  /**
+   * @return the account, if any, used by this stream.
+   */
+  virtual Buffer::BufferMemoryAccountSharedPtr account() const PURE;
+
   /**
    * Takes a stream, and acts as if the headers are newly arrived.
    * On success, this will result in a creating a new filter chain and likely

include/envoy/router/router.h

@@ -1327,6 +1327,12 @@ class GenericUpstream {
    * @param reason supplies the reset reason.
    */
   virtual void resetStream() PURE;
+
+  /**
+   * Sets the upstream to use the following account.
+   * @param the account to assign the generic upstream.
+   */
+  virtual void setAccount(Buffer::BufferMemoryAccountSharedPtr account) PURE;
 };
 
 using GenericConnPoolPtr = std::unique_ptr<GenericConnPool>;

source/common/buffer/buffer_impl.cc

@@ -43,11 +43,11 @@ void OwnedImpl::addDrainTracker(std::function<void()> drain_tracker) {
 
 void OwnedImpl::bindAccount(BufferMemoryAccountSharedPtr account) {
   ASSERT(slices_.empty());
-  // We don't yet have an account bound.
-  ASSERT(!account_);
   account_ = std::move(account);
 }
 
+BufferMemoryAccountSharedPtr OwnedImpl::getAccountForTest() { return account_; }
+
 void OwnedImpl::add(const void* data, uint64_t size) { addImpl(data, size); }
 
 void OwnedImpl::addBufferFragment(BufferFragment& fragment) {

source/common/buffer/buffer_impl.h

@@ -710,6 +710,11 @@ class OwnedImpl : public LibEventInstance {
    */
   virtual void appendSliceForTest(absl::string_view data);
 
+  /**
+   * @return the BufferMemoryAccount bound to this buffer, if any.
+   */
+  BufferMemoryAccountSharedPtr getAccountForTest();
+
   // Does not implement watermarking.
   // TODO(antoniovicente) Implement watermarks by merging the OwnedImpl and WatermarkBuffer
   // implementations. Also, make high-watermark config a constructor argument.

source/common/buffer/watermark_buffer.h

@@ -47,7 +47,7 @@ class WatermarkBuffer : public OwnedImpl {
 
 protected:
   virtual void checkHighAndOverflowWatermarks();
-  void checkLowWatermark();
+  virtual void checkLowWatermark();
 
 private:
   void commit(uint64_t length, absl::Span<RawSlice> slices,

source/common/http/async_client_impl.h

@@ -365,6 +365,8 @@ class AsyncStreamImpl : public AsyncClient::Stream,
   Upstream::ClusterInfoConstSharedPtr clusterInfo() override { return parent_.cluster_; }
   void clearRouteCache() override {}
   uint64_t streamId() const override { return stream_id_; }
+  // TODO(kbaichoo): Implement this?
+  Buffer::BufferMemoryAccountSharedPtr account() const override { return nullptr; }
   Tracing::Span& activeSpan() override { return active_span_; }
   const Tracing::Config& tracingConfig() override { return tracing_config_; }
   void continueDecoding() override { NOT_IMPLEMENTED_GCOVR_EXCL_LINE; }

source/common/http/conn_manager_impl.cc

@@ -259,7 +259,14 @@ RequestDecoder& ConnectionManagerImpl::newStream(ResponseEncoder& response_encod
   }
 
   ENVOY_CONN_LOG(debug, "new stream", read_callbacks_->connection());
-  ActiveStreamPtr new_stream(new ActiveStream(*this, response_encoder.getStream().bufferLimit()));
+
+  // Set the account to start accounting.
+  Buffer::BufferMemoryAccountSharedPtr downstream_request_account =
+      std::make_shared<Buffer::BufferMemoryAccountImpl>();
+  response_encoder.getStream().setAccount(downstream_request_account);
+
+  ActiveStreamPtr new_stream(new ActiveStream(*this, response_encoder.getStream().bufferLimit(),
+                                              std::move(downstream_request_account)));
   new_stream->state_.is_internally_created_ = is_internally_created;
   new_stream->response_encoder_ = &response_encoder;
   new_stream->response_encoder_->getStream().addCallbacks(*new_stream);
@@ -571,13 +578,14 @@ void ConnectionManagerImpl::RdsRouteConfigUpdateRequester::requestSrdsUpdate(
 }
 
 ConnectionManagerImpl::ActiveStream::ActiveStream(ConnectionManagerImpl& connection_manager,
-                                                  uint32_t buffer_limit)
+                                                  uint32_t buffer_limit,
+                                                  Buffer::BufferMemoryAccountSharedPtr account)
     : connection_manager_(connection_manager),
       stream_id_(connection_manager.random_generator_.random()),
       filter_manager_(*this, connection_manager_.read_callbacks_->connection().dispatcher(),
                       connection_manager_.read_callbacks_->connection(), stream_id_,
-                      connection_manager_.config_.proxy100Continue(), buffer_limit,
-                      connection_manager_.config_.filterFactory(),
+                      std::move(account), connection_manager_.config_.proxy100Continue(),
+                      buffer_limit, connection_manager_.config_.filterFactory(),
                       connection_manager_.config_.localReply(),
                       connection_manager_.codec_->protocol(), connection_manager_.timeSource(),
                       connection_manager_.read_callbacks_->connection().streamInfo().filterState(),

source/common/http/conn_manager_impl.h

@@ -156,7 +156,8 @@ class ConnectionManagerImpl : Logger::Loggable<Logger::Id::http>,
                               public Tracing::Config,
                               public ScopeTrackedObject,
                               public FilterManagerCallbacks {
-    ActiveStream(ConnectionManagerImpl& connection_manager, uint32_t buffer_limit);
+    ActiveStream(ConnectionManagerImpl& connection_manager, uint32_t buffer_limit,
+                 Buffer::BufferMemoryAccountSharedPtr account);
     void completeRequest();
 
     const Network::Connection* connection();

source/common/http/filter_manager.cc

@@ -1556,5 +1556,9 @@ void ActiveStreamFilterBase::resetStream() { parent_.filter_manager_callbacks_.r
 
 uint64_t ActiveStreamFilterBase::streamId() const { return parent_.streamId(); }
 
+Buffer::BufferMemoryAccountSharedPtr ActiveStreamDecoderFilter::account() const {
+  return parent_.account();
+}
+
 } // namespace Http
 } // namespace Envoy

source/common/http/filter_manager.h

@@ -2,6 +2,7 @@
 
 #include <memory>
 
+#include "envoy/buffer/buffer.h"
 #include "envoy/common/optref.h"
 #include "envoy/extensions/filters/common/matcher/action/v3/skip_action.pb.h"
 #include "envoy/extensions/filters/network/http_connection_manager/v3/http_connection_manager.pb.h"
@@ -275,6 +276,7 @@ struct ActiveStreamDecoderFilter : public ActiveStreamFilterBase,
   void addUpstreamSocketOptions(const Network::Socket::OptionsSharedPtr& options) override;
 
   Network::Socket::OptionsSharedPtr getUpstreamSocketOptions() const override;
+  Buffer::BufferMemoryAccountSharedPtr account() const override;
 
   // Each decoder filter instance checks if the request passed to the filter is gRPC
   // so that we can issue gRPC local responses to gRPC requests. Filter's decodeHeaders()
@@ -638,15 +640,16 @@ class FilterManager : public ScopeTrackedObject,
                       Logger::Loggable<Logger::Id::http> {
 public:
   FilterManager(FilterManagerCallbacks& filter_manager_callbacks, Event::Dispatcher& dispatcher,
-                const Network::Connection& connection, uint64_t stream_id, bool proxy_100_continue,
+                const Network::Connection& connection, uint64_t stream_id,
+                Buffer::BufferMemoryAccountSharedPtr account, bool proxy_100_continue,
                 uint32_t buffer_limit, FilterChainFactory& filter_chain_factory,
                 const LocalReply::LocalReply& local_reply, Http::Protocol protocol,
                 TimeSource& time_source, StreamInfo::FilterStateSharedPtr parent_filter_state,
                 StreamInfo::FilterState::LifeSpan filter_state_life_span)
       : filter_manager_callbacks_(filter_manager_callbacks), dispatcher_(dispatcher),
-        connection_(connection), stream_id_(stream_id), proxy_100_continue_(proxy_100_continue),
-        buffer_limit_(buffer_limit), filter_chain_factory_(filter_chain_factory),
-        local_reply_(local_reply),
+        connection_(connection), stream_id_(stream_id), account_(std::move(account)),
+        proxy_100_continue_(proxy_100_continue), buffer_limit_(buffer_limit),
+        filter_chain_factory_(filter_chain_factory), local_reply_(local_reply),
         stream_info_(protocol, time_source, connection.addressProviderSharedPtr(),
                      parent_filter_state, filter_state_life_span) {}
   ~FilterManager() override {
@@ -911,6 +914,7 @@ class FilterManager : public ScopeTrackedObject,
   const Network::Connection* connection() const { return &connection_; }
 
   uint64_t streamId() const { return stream_id_; }
+  Buffer::BufferMemoryAccountSharedPtr account() const { return account_; }
 
   Buffer::InstancePtr& bufferedRequestData() { return buffered_request_data_; }
 
@@ -982,6 +986,7 @@ class FilterManager : public ScopeTrackedObject,
   Event::Dispatcher& dispatcher_;
   const Network::Connection& connection_;
   const uint64_t stream_id_;
+  Buffer::BufferMemoryAccountSharedPtr account_;
   const bool proxy_100_continue_;
 
   std::list<ActiveStreamDecoderFilterPtr> decoder_filters_;

source/common/http/http1/codec_impl.h

@@ -69,6 +69,10 @@ class StreamEncoderImpl : public virtual StreamEncoder,
     // require a flush timeout not already covered by other timeouts.
   }
 
+  void setAccount(Buffer::BufferMemoryAccountSharedPtr) override {
+    // TODO(kbaichoo): implement account tracking for H1.
+  }
+
   void setIsResponseToHeadRequest(bool value) { is_response_to_head_request_ = value; }
   void setIsResponseToConnectRequest(bool value) { is_response_to_connect_request_ = value; }
   void setDetails(absl::string_view details) { details_ = details; }

source/common/http/http2/codec_impl.cc

@@ -130,8 +130,17 @@ template <typename T> static T* removeConst(const void* object) {
 }
 
 ConnectionImpl::StreamImpl::StreamImpl(ConnectionImpl& parent, uint32_t buffer_limit)
-    : parent_(parent), local_end_stream_sent_(false), remote_end_stream_(false),
-      data_deferred_(false), received_noninformational_headers_(false),
+    : parent_(parent),
+      pending_recv_data_(parent_.connection_.dispatcher().getWatermarkFactory().create(
+          [this]() -> void { this->pendingRecvBufferLowWatermark(); },
+          [this]() -> void { this->pendingRecvBufferHighWatermark(); },
+          []() -> void { /* TODO(adisuissa): Handle overflow watermark */ })),
+      pending_send_data_(parent_.connection_.dispatcher().getWatermarkFactory().create(
+          [this]() -> void { this->pendingSendBufferLowWatermark(); },
+          [this]() -> void { this->pendingSendBufferHighWatermark(); },
+          []() -> void { /* TODO(adisuissa): Handle overflow watermark */ })),
+      local_end_stream_sent_(false), remote_end_stream_(false), data_deferred_(false),
+      received_noninformational_headers_(false),
       pending_receive_buffer_high_watermark_called_(false),
       pending_send_buffer_high_watermark_called_(false), reset_due_to_messaging_error_(false) {
   parent_.stats_.streams_active_.inc();
@@ -145,7 +154,7 @@ ConnectionImpl::StreamImpl::~StreamImpl() { ASSERT(stream_idle_timer_ == nullptr
 void ConnectionImpl::StreamImpl::destroy() {
   disarmStreamIdleTimer();
   parent_.stats_.streams_active_.dec();
-  parent_.stats_.pending_send_bytes_.sub(pending_send_data_.length());
+  parent_.stats_.pending_send_bytes_.sub(pending_send_data_->length());
 }
 
 static void insertHeader(std::vector<nghttp2_nv>& headers, const HeaderEntry& header) {
@@ -253,7 +262,7 @@ void ConnectionImpl::ServerStreamImpl::encodeHeaders(const ResponseHeaderMap& he
 void ConnectionImpl::StreamImpl::encodeTrailersBase(const HeaderMap& trailers) {
   ASSERT(!local_end_stream_);
   local_end_stream_ = true;
-  if (pending_send_data_.length() > 0) {
+  if (pending_send_data_->length() > 0) {
     // In this case we want trailers to come after we release all pending body data that is
     // waiting on window updates. We need to save the trailers so that we can emit them later.
     // However, for empty trailers, we don't need to to save the trailers.
@@ -409,13 +418,13 @@ void ConnectionImpl::StreamImpl::submitMetadata(uint8_t flags) {
 }
 
 ssize_t ConnectionImpl::StreamImpl::onDataSourceRead(uint64_t length, uint32_t* data_flags) {
-  if (pending_send_data_.length() == 0 && !local_end_stream_) {
+  if (pending_send_data_->length() == 0 && !local_end_stream_) {
     ASSERT(!data_deferred_);
     data_deferred_ = true;
     return NGHTTP2_ERR_DEFERRED;
   } else {
     *data_flags |= NGHTTP2_DATA_FLAG_NO_COPY;
-    if (local_end_stream_ && pending_send_data_.length() <= length) {
+    if (local_end_stream_ && pending_send_data_->length() <= length) {
       *data_flags |= NGHTTP2_DATA_FLAG_EOF;
       if (pending_trailers_to_encode_) {
         // We need to tell the library to not set end stream so that we can emit the trailers.
@@ -425,7 +434,7 @@ ssize_t ConnectionImpl::StreamImpl::onDataSourceRead(uint64_t length, uint32_t*
       }
     }
 
-    return std::min(length, pending_send_data_.length());
+    return std::min(length, pending_send_data_->length());
   }
 }
 
@@ -446,7 +455,7 @@ void ConnectionImpl::StreamImpl::onDataSourceSend(const uint8_t* framehd, size_t
   }
 
   parent_.stats_.pending_send_bytes_.sub(length);
-  output.move(pending_send_data_, length);
+  output.move(*pending_send_data_, length);
   parent_.connection_.write(output, false);
 }
 
@@ -491,7 +500,9 @@ void ConnectionImpl::StreamImpl::onPendingFlushTimer() {
 
 void ConnectionImpl::StreamImpl::encodeData(Buffer::Instance& data, bool end_stream) {
   ASSERT(!local_end_stream_);
-  encodeDataHelper(data, end_stream, /*skip_encoding_empty_trailers=*/false);
+  encodeDataHelper(data, end_stream,
+                   /*skip_encoding_empty_trailers=*/
+                   false);
 }
 
 void ConnectionImpl::StreamImpl::encodeDataHelper(Buffer::Instance& data, bool end_stream,
@@ -502,7 +513,7 @@ void ConnectionImpl::StreamImpl::encodeDataHelper(Buffer::Instance& data, bool e
 
   local_end_stream_ = end_stream;
   parent_.stats_.pending_send_bytes_.add(data.length());
-  pending_send_data_.move(data);
+  pending_send_data_->move(data);
   if (data_deferred_) {
     int rc = nghttp2_session_resume_data(parent_.session_, stream_id_);
     ASSERT(rc == 0);
@@ -514,7 +525,7 @@ void ConnectionImpl::StreamImpl::encodeDataHelper(Buffer::Instance& data, bool e
     // Intended to check through coverage that this error case is tested
     return;
   }
-  if (local_end_stream_ && pending_send_data_.length() > 0) {
+  if (local_end_stream_ && pending_send_data_->length() > 0) {
     createPendingFlushTimer();
   }
 }
@@ -576,6 +587,12 @@ void ConnectionImpl::StreamImpl::onMetadataDecoded(MetadataMapPtr&& metadata_map
   }
 }
 
+void ConnectionImpl::StreamImpl::setAccount(Buffer::BufferMemoryAccountSharedPtr account) {
+  buffer_memory_account_ = account;
+  pending_recv_data_->bindAccount(buffer_memory_account_);
+  pending_send_data_->bindAccount(buffer_memory_account_);
+}
+
 ConnectionImpl::ConnectionImpl(Network::Connection& connection, CodecStats& stats,
                                Random::RandomGenerator& random_generator,
                                const envoy::config::core::v3::Http2ProtocolOptions& http2_options,
@@ -707,7 +724,7 @@ int ConnectionImpl::onData(int32_t stream_id, const uint8_t* data, size_t len) {
   StreamImpl* stream = getStream(stream_id);
   // If this results in buffering too much data, the watermark buffer will call
   // pendingRecvBufferHighWatermark, resulting in ++read_disable_count_
-  stream->pending_recv_data_.add(data, len);
+  stream->pending_recv_data_->add(data, len);
   // Update the window to the peer unless some consumer of this stream's data has hit a flow control
   // limit and disabled reads on this stream
   if (!stream->buffersOverrun()) {
@@ -862,10 +879,10 @@ Status ConnectionImpl::onFrameReceived(const nghttp2_frame* frame) {
     // It's possible that we are waiting to send a deferred reset, so only raise data if local
     // is not complete.
     if (!stream->deferred_reset_) {
-      stream->decoder().decodeData(stream->pending_recv_data_, stream->remote_end_stream_);
+      stream->decoder().decodeData(*stream->pending_recv_data_, stream->remote_end_stream_);
     }
 
-    stream->pending_recv_data_.drain(stream->pending_recv_data_.length());
+    stream->pending_recv_data_->drain(stream->pending_recv_data_->length());
     break;
   }
   case NGHTTP2_RST_STREAM: {