Overload: Reset H2 server stream only use codec level reset mechanism

envoy/buffer/buffer.h

@@ -123,6 +123,16 @@ class BufferMemoryAccount {
    * should trigger a reset of the corresponding upstream stream if it exists.
    */
   virtual void resetDownstream() PURE;
+
+  /**
+   * Should be invoked when the Envoy level stream is complete.
+   */
+  virtual void onEnvoyStreamComplete() PURE;
+
+  /**
+   * Whether the Envoy level stream was complete.
+   */
+  virtual bool sawEnvoyStreamComplete() PURE;
 };
 
 using BufferMemoryAccountSharedPtr = std::shared_ptr<BufferMemoryAccount>;

source/common/buffer/watermark_buffer.h

@@ -123,6 +123,9 @@ class BufferMemoryAccountImpl : public BufferMemoryAccount {
     }
   }
 
+  void onEnvoyStreamComplete() override { saw_envoy_stream_complete_ = true; }
+  bool sawEnvoyStreamComplete() override { return saw_envoy_stream_complete_; }
+
   // The number of memory classes the Account expects to exists. See
   // *WatermarkBufferFactory* for details on the memory classes.
   static constexpr uint32_t NUM_MEMORY_CLASSES_ = 8;
@@ -143,7 +146,7 @@ class BufferMemoryAccountImpl : public BufferMemoryAccount {
   absl::optional<uint32_t> current_bucket_idx_{};
 
   WatermarkBufferFactory* factory_ = nullptr;
-
+  bool saw_envoy_stream_complete_ = false;
   OptRef<Http::StreamResetHandler> reset_handler_;
   // Keep a copy of the shared_ptr pointing to this account. We opted to go this
   // route rather than enable_shared_from_this to avoid wasteful atomic

source/common/http/filter_manager.h

@@ -777,6 +777,10 @@ class FilterManager : public ScopeTrackedObject,
         filter->handle_->onStreamComplete();
       }
     }
+
+    if (account_) {
+      account_->onEnvoyStreamComplete();
+    }
   }
 
   void destroyFilters() {

source/common/http/http2/codec_impl.cc

@@ -537,7 +537,16 @@ void ConnectionImpl::ServerStreamImpl::resetStream(StreamResetReason reason) {
     buffer_memory_account_->clearDownstream();
   }
 
-  StreamImpl::resetStream(reason);
+  if (buffer_memory_account_ && buffer_memory_account_->sawEnvoyStreamComplete()) {
+    // We should use a lower level reset mechanism to reset the stream.
+    resetStreamWorker(reason);
+    if (parent_.sendPendingFramesAndHandleError()) {
+      // Intended to check through coverage that this error case is tested
+      return;
+    }
+  } else {
+    StreamImpl::resetStream(reason);
+  }
 }
 
 void ConnectionImpl::StreamImpl::resetStream(StreamResetReason reason) {

test/integration/buffer_accounting_integration_test.cc

@@ -611,4 +611,89 @@ TEST_P(Http2OverloadManagerIntegrationTest,
   EXPECT_EQ(smallest_response->headers().getStatusValue(), "200");
 }
 
+TEST_P(Http2OverloadManagerIntegrationTest, CanResetStreamIfEnvoyLevelStreamEnded) {
+  initializeOverloadManagerInBootstrap(
+      TestUtility::parseYaml<envoy::config::overload::v3::OverloadAction>(R"EOF(
+      name: "envoy.overload_actions.reset_high_memory_stream"
+      triggers:
+        - name: "envoy.resource_monitors.testonly.fake_resource_monitor"
+          scaled:
+            scaling_threshold: 0.90
+            saturation_threshold: 0.98
+    )EOF"));
+  initialize();
+
+  // Set 10MiB receive window for the client.
+  const int downstream_window_size = 10 * 1024 * 1024;
+  envoy::config::core::v3::Http2ProtocolOptions http2_options =
+      ::Envoy::Http2::Utility::initializeAndValidateOptions(
+          envoy::config::core::v3::Http2ProtocolOptions());
+  http2_options.mutable_initial_stream_window_size()->set_value(downstream_window_size);
+  http2_options.mutable_initial_connection_window_size()->set_value(downstream_window_size);
+  codec_client_ = makeRawHttpConnection(makeClientConnection(lookupPort("http")), http2_options);
+
+  // Makes us have Envoy's writes to downstream return EAGAIN
+  writev_matcher_->setSourcePort(lookupPort("http"));
+  writev_matcher_->setWritevReturnsEgain();
+
+  // Send a request
+  auto encoder_decoder = codec_client_->startRequest(Http::TestRequestHeaderMapImpl{
+      {":method", "POST"},
+      {":path", "/"},
+      {":scheme", "http"},
+      {":authority", "host"},
+      {"content-length", "10"},
+  });
+  auto& encoder = encoder_decoder.first;
+  const std::string data(10, 'a');
+  codec_client_->sendData(encoder, data, true);
+  auto response = std::move(encoder_decoder.second);
+
+  waitForNextUpstreamRequest();
+  FakeStreamPtr upstream_request_for_response = std::move(upstream_request_);
+
+  // Send the responses back. It is larger than the downstream's receive window
+  // size. Thus, the codec will not end the stream, but the Envoy level stream
+  // should.
+  upstream_request_for_response->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "200"}},
+                                               false);
+  const int response_size = downstream_window_size + 1024; // Slightly over the window size.
+  upstream_request_for_response->encodeData(response_size, true);
+
+  // Wait for the stream to have seen complete.
+  if (streamBufferAccounting()) {
+    EXPECT_TRUE(
+        buffer_factory_->waitUntilExpectedNumberOfActiveAccountsThatSawEnvoyStreamComplete(1));
+  }
+
+  // Set the pressure so the overload action kills the response if doing stream
+  // accounting
+  updateResource(0.95);
+  test_server_->waitForGaugeEq(
+      "overload.envoy.overload_actions.reset_high_memory_stream.scale_percent", 62);
+
+  if (streamBufferAccounting()) {
+    test_server_->waitForCounterGe("envoy.overload_actions.reset_high_memory_stream.count", 1);
+  }
+
+  // Reduce resource pressure
+  updateResource(0.80);
+  test_server_->waitForGaugeEq(
+      "overload.envoy.overload_actions.reset_high_memory_stream.scale_percent", 0);
+
+  // Resume writes to downstream.
+  writev_matcher_->setResumeWrites();
+
+  if (streamBufferAccounting()) {
+    EXPECT_TRUE(response->waitForReset());
+    EXPECT_TRUE(response->reset());
+  } else {
+    // If we're not doing the accounting, we didn't end up resetting the
+    // streams.
+    ASSERT_TRUE(response->waitForEndStream());
+    ASSERT_TRUE(response->complete());
+    EXPECT_EQ(response->headers().getStatusValue(), "200");
+  }
+}
+
 } // namespace Envoy

test/integration/tracked_watermark_buffer.cc

@@ -274,6 +274,25 @@ bool TrackedWatermarkBufferFactory::waitUntilExpectedNumberOfAccountsAndBoundBuf
   return mutex_.AwaitWithTimeout(absl::Condition(&predicate), absl::FromChrono(timeout));
 }
 
+bool TrackedWatermarkBufferFactory::
+    waitUntilExpectedNumberOfActiveAccountsThatSawEnvoyStreamComplete(
+        uint32_t expected_num_accounts, std::chrono::milliseconds timeout) {
+  absl::MutexLock lock(&mutex_);
+  auto predicate = [this, expected_num_accounts]() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    mutex_.AssertHeld();
+    // Remove non-active accounts
+    removeDanglingAccounts();
+    uint32_t num_accounts_seen_envoy_stream_complete = 0;
+    for (auto& entry : account_infos_) {
+      if (entry.first->sawEnvoyStreamComplete()) {
+        ++num_accounts_seen_envoy_stream_complete;
+      }
+    }
+    return num_accounts_seen_envoy_stream_complete == expected_num_accounts;
+  };
+  return mutex_.AwaitWithTimeout(absl::Condition(&predicate), absl::FromChrono(timeout));
+}
+
 void TrackedWatermarkBufferFactory::checkIfExpectedBalancesMet() {
   if (!expected_balances_ || expected_balances_met_.HasBeenNotified()) {
     return;

test/integration/tracked_watermark_buffer.h

@@ -129,6 +129,9 @@ class TrackedWatermarkBufferFactory : public WatermarkBufferFactory {
       uint32_t num_accounts, uint32_t num_bound_buffers,
       std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
 
+  bool waitUntilExpectedNumberOfActiveAccountsThatSawEnvoyStreamComplete(
+      uint32_t num_accounts, std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
+
   using AccountToBoundBuffersMap =
       absl::flat_hash_map<BufferMemoryAccountSharedPtr,
                           absl::flat_hash_set<TrackedWatermarkBuffer*>>;