Notify a closed non-static stream once it is about to be destroyed via a new interface QuicStream::OnSoonToBeDestroyed() with no-op default implementation.

With this change QuicStream::OnConnectionClosed() will be also called on closed streams still waiting for ACK(zombie streams).

Protected by quic_reloadable_flag_quic_notify_stream_soon_to_destroy.

PiperOrigin-RevId: 704076485
diff --git a/quiche/common/quiche_feature_flags_list.h b/quiche/common/quiche_feature_flags_list.h
index 91dc462..9596fe3 100755
--- a/quiche/common/quiche_feature_flags_list.h
+++ b/quiche/common/quiche_feature_flags_list.h
@@ -46,6 +46,7 @@
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_no_write_control_frame_upon_connection_close, false, true, "If trrue, early return before write control frame in OnCanWrite() if the connection is already closed.")
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_no_write_control_frame_upon_connection_close2, false, false, "If true, QuicSession will block outgoing control frames when the connection is closed.")
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_notify_ack_listener_earlier, false, false, "If true, call QuicAckListenerInterface::OnPacketAcked() before moving the stream to closed stream list.")
+QUICHE_FLAG(bool, quiche_reloadable_flag_quic_notify_stream_soon_to_destroy, false, false, "If true, notify each QUIC stream before it gets destroyed and update ACK listener before that.")
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_optimize_qpack_blocking_manager, false, false, "If true, optimize qpack_blocking_manager for CPU efficiency.")
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_pacing_remove_non_initial_burst, false, false, "If true, remove the non-initial burst in QUIC PacingSender.")
 QUICHE_FLAG(bool, quiche_reloadable_flag_quic_parse_cert_compression_algos_from_chlo, true, true, "If true, parse offered cert compression algorithms from received CHLOs.")
diff --git a/quiche/quic/core/http/quic_spdy_stream_test.cc b/quiche/quic/core/http/quic_spdy_stream_test.cc
index 7512c01..62d7ee1 100644
--- a/quiche/quic/core/http/quic_spdy_stream_test.cc
+++ b/quiche/quic/core/http/quic_spdy_stream_test.cc
@@ -253,6 +253,10 @@
     data_ += std::string(buffer, bytes_read);
   }
 
+  void OnSoonToBeDestroyed() override {
+    on_soon_to_be_destroyed_called_ = true;
+  }
+
   MOCK_METHOD(void, WriteHeadersMock, (bool fin), ());
 
   size_t WriteHeadersImpl(
@@ -282,12 +286,16 @@
   }
 
   size_t headers_payload_length() const { return headers_payload_length_; }
+  bool on_soon_to_be_destroyed_called() const {
+    return on_soon_to_be_destroyed_called_;
+  }
 
  private:
   bool should_process_data_;
   quiche::HttpHeaderBlock saved_headers_;
   std::string data_;
   size_t headers_payload_length_;
+  bool on_soon_to_be_destroyed_called_ = false;
 };
 
 class TestSession : public MockQuicSpdySession {
@@ -1902,7 +1910,7 @@
   EXPECT_EQ(0u, QuicStreamPeer::SendBuffer(stream_).size());
 }
 
-TEST_P(QuicSpdyStreamTest, OnPacketAckedBeforeStreamDestroy) {
+TEST_P(QuicSpdyStreamTest, NotifyOnPacketAckedBeforeStreamDestroy) {
   Initialize(kShouldProcessData);
   quiche::QuicheReferenceCountedPointer<MockAckListener> mock_ack_listener(
       new StrictMock<MockAckListener>);
@@ -1951,8 +1959,10 @@
         if (GetQuicReloadableFlag(quic_notify_ack_listener_earlier)) {
           // Stream is not added to closed stream list yet.
           EXPECT_NE(session_->GetActiveStream(stream_->id()), nullptr);
+          EXPECT_FALSE(stream_->on_soon_to_be_destroyed_called());
         } else {
           EXPECT_EQ(session_->GetActiveStream(stream_->id()), nullptr);
+          EXPECT_TRUE(stream_->on_soon_to_be_destroyed_called());
         }
       }));
   EXPECT_TRUE(stream_->OnStreamFrameAcked(18, 0, true, QuicTime::Delta::Zero(),
@@ -1960,6 +1970,7 @@
                                           &newly_acked_length));
   EXPECT_FALSE(stream_->IsWaitingForAcks());
   EXPECT_EQ(0u, QuicStreamPeer::SendBuffer(stream_).size());
+  EXPECT_TRUE(stream_->on_soon_to_be_destroyed_called());
 }
 
 TEST_P(QuicSpdyStreamTest, StreamDataGetAckedMultipleTimes) {
diff --git a/quiche/quic/core/quic_session.cc b/quiche/quic/core/quic_session.cc
index bde478f..55f8d3a 100644
--- a/quiche/quic/core/quic_session.cc
+++ b/quiche/quic/core/quic_session.cc
@@ -590,17 +590,32 @@
 
   GetMutableCryptoStream()->OnConnectionClosed(frame, source);
 
-  PerformActionOnActiveStreams([this, frame, source](QuicStream* stream) {
-    QuicStreamId id = stream->id();
-    stream->OnConnectionClosed(frame, source);
-    auto it = stream_map_.find(id);
-    if (it != stream_map_.end()) {
-      QUIC_BUG_IF(quic_bug_12435_2, !it->second->IsZombie())
-          << ENDPOINT << "Non-zombie stream " << id
-          << " failed to close under OnConnectionClosed";
-    }
-    return true;
-  });
+  if (!notify_stream_soon_to_destroy_) {
+    PerformActionOnActiveStreams([this, frame, source](QuicStream* stream) {
+      QuicStreamId id = stream->id();
+      stream->OnConnectionClosed(frame, source);
+      auto it = stream_map_.find(id);
+      if (it != stream_map_.end()) {
+        QUIC_BUG_IF(quic_bug_12435_2, !it->second->IsZombie())
+            << ENDPOINT << "Non-zombie stream " << id
+            << " failed to close under OnConnectionClosed";
+      }
+      return true;
+    });
+  } else {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_notify_stream_soon_to_destroy, 1, 2);
+    PerformActionOnNonStaticStreams([this, frame, source](QuicStream* stream) {
+      QuicStreamId id = stream->id();
+      stream->OnConnectionClosed(frame, source);
+      QUIC_BUG_IF(quic_bug_12435_12, stream_map_.contains(id))
+          << ENDPOINT << "Non-static stream " << id
+          << " failed to be moved to closed stream list under "
+             "OnConnectionClosed";
+      return true;
+    });
+    QUIC_BUG_IF(zombie_stream_not_ready_for_destruction,
+                num_zombie_streams_ > 0);
+  }
 
   closed_streams_clean_up_alarm_->Cancel();
   stream_count_reset_alarm_->Cancel();
@@ -612,6 +627,13 @@
   }
 }
 
+void QuicSession::PrepareStreamForDestruction(StreamMap::iterator it) {
+  QUIC_BUG_IF(alive_stream_not_in_stream_map, it == stream_map_.end());
+  it->second->OnSoonToBeDestroyed();
+  closed_streams_.push_back(std::move(it->second));
+  stream_map_.erase(it);
+}
+
 void QuicSession::OnWriteBlocked() {
   if (!connection_->connected()) {
     return;
@@ -1162,8 +1184,7 @@
     // The stream needs to be kept alive because it's waiting for acks.
     ++num_zombie_streams_;
   } else {
-    closed_streams_.push_back(std::move(it->second));
-    stream_map_.erase(it);
+    PrepareStreamForDestruction(it);
     // Do not retransmit data of a closed stream.
     streams_with_pending_retransmission_.erase(stream_id);
     if (!closed_streams_clean_up_alarm_->IsSet()) {
@@ -2407,8 +2428,7 @@
     return;
   }
   --num_zombie_streams_;
-  closed_streams_.push_back(std::move(it->second));
-  stream_map_.erase(it);
+  PrepareStreamForDestruction(it);
 
   if (!closed_streams_clean_up_alarm_->IsSet()) {
     closed_streams_clean_up_alarm_->Set(connection_->clock()->ApproximateNow());
@@ -2848,6 +2868,22 @@
   }
 }
 
+void QuicSession::PerformActionOnNonStaticStreams(
+    quiche::UnretainedCallback<bool(QuicStream*)> action) {
+  std::vector<QuicStream*> streams;
+  for (const auto& it : stream_map_) {
+    if (!it.second->is_static()) {
+      streams.push_back(it.second.get());
+    }
+  }
+
+  for (QuicStream* stream : streams) {
+    if (!action(stream)) {
+      return;
+    }
+  }
+}
+
 void QuicSession::PerformActionOnActiveStreams(
     quiche::UnretainedCallback<bool(QuicStream*)> action) const {
   for (const auto& it : stream_map_) {
diff --git a/quiche/quic/core/quic_session.h b/quiche/quic/core/quic_session.h
index 5ec9838..fee549a 100644
--- a/quiche/quic/core/quic_session.h
+++ b/quiche/quic/core/quic_session.h
@@ -697,6 +697,10 @@
     return enable_stop_sending_for_zombie_streams_;
   }
 
+  bool notify_stream_soon_to_destroy() const {
+    return notify_stream_soon_to_destroy_;
+  }
+
  protected:
   using StreamMap =
       absl::flat_hash_map<QuicStreamId, std::unique_ptr<QuicStream>>;
@@ -977,6 +981,15 @@
 
   bool ExceedsPerLoopStreamLimit() const;
 
+  // Moves the stream pointed by |it| from stream_map_ to closed_streams_.
+  void PrepareStreamForDestruction(StreamMap::iterator it);
+
+  // Called by applications to perform |action| on streams that have received
+  // and sent FIN, but still waiting for ACK. Stream iteration will be stopped
+  // if action returns false.
+  void PerformActionOnNonStaticStreams(
+      quiche::UnretainedCallback<bool(QuicStream*)> action);
+
   // Keep track of highest received byte offset of locally closed streams, while
   // waiting for a definitive final highest offset from the peer.
   absl::flat_hash_map<QuicStreamId, QuicStreamOffset>
@@ -1105,6 +1118,9 @@
 
   const bool enable_stop_sending_for_zombie_streams_ =
       GetQuicReloadableFlag(quic_deliver_stop_sending_to_zombie_streams);
+
+  const bool notify_stream_soon_to_destroy_ =
+      GetQuicReloadableFlag(quic_notify_stream_soon_to_destroy);
 };
 
 }  // namespace quic
diff --git a/quiche/quic/core/quic_session_test.cc b/quiche/quic/core/quic_session_test.cc
index 7f54074..25008ae 100644
--- a/quiche/quic/core/quic_session_test.cc
+++ b/quiche/quic/core/quic_session_test.cc
@@ -238,6 +238,8 @@
               (override));
 
   MOCK_METHOD(bool, HasPendingRetransmission, (), (const, override));
+
+  MOCK_METHOD(void, OnSoonToBeDestroyed, (), (override));
 };
 
 class TestSession : public QuicSession {
@@ -3089,6 +3091,9 @@
   if (GetQuicReloadableFlag(quic_deliver_stop_sending_to_zombie_streams)) {
     EXPECT_CALL(*connection_, SendControlFrame(_)).Times(1);
     EXPECT_CALL(*connection_, OnStreamReset(_, _)).Times(1);
+    if (GetQuicReloadableFlag(quic_notify_stream_soon_to_destroy)) {
+      EXPECT_CALL(*stream, OnSoonToBeDestroyed());
+    }
   } else {
     EXPECT_CALL(*connection_, SendControlFrame(_)).Times(0);
     EXPECT_CALL(*connection_, OnStreamReset(_, _)).Times(0);
@@ -3105,6 +3110,29 @@
   }
 }
 
+TEST_P(QuicSessionTestServer, OnConnectionCloseForZombieStreams) {
+  if (!VersionHasIetfQuicFrames(transport_version()) ||
+      !GetQuicReloadableFlag(quic_notify_stream_soon_to_destroy)) {
+    return;
+  }
+  CompleteHandshake();
+  session_.set_writev_consumes_all_data(true);
+
+  TestStream* stream = session_.CreateOutgoingBidirectionalStream();
+  std::string body(100, '.');
+  QuicStreamPeer::CloseReadSide(stream);
+  stream->WriteOrBufferData(body, true, nullptr);
+  EXPECT_TRUE(stream->IsWaitingForAcks());
+  // Verify that the stream is a zombie.
+  EXPECT_TRUE(stream->IsZombie());
+  ASSERT_EQ(0u, session_.closed_streams()->size());
+
+  EXPECT_CALL(*stream, OnSoonToBeDestroyed());
+  connection_->ReallyCloseConnection(QUIC_NO_ERROR, "Testing",
+                                     ConnectionCloseBehavior::SILENT_CLOSE);
+  EXPECT_EQ(0, session_.GetNumActiveStreams());
+}
+
 // If stream is closed, return true and do not close the connection.
 TEST_P(QuicSessionTestServer, OnStopSendingClosedStream) {
   if (!VersionHasIetfQuicFrames(transport_version())) {
diff --git a/quiche/quic/core/quic_stream.cc b/quiche/quic/core/quic_stream.cc
index dbce695..7f10791 100644
--- a/quiche/quic/core/quic_stream.cc
+++ b/quiche/quic/core/quic_stream.cc
@@ -541,6 +541,10 @@
   return true;
 }
 
+void QuicStream::OnSoonToBeDestroyed() {
+  QUICHE_DCHECK(write_side_closed() && read_side_closed());
+}
+
 int QuicStream::num_frames_received() const {
   return sequencer_.num_frames_received();
 }
@@ -616,18 +620,24 @@
 
 void QuicStream::OnConnectionClosed(const QuicConnectionCloseFrame& frame,
                                     ConnectionCloseSource /*source*/) {
-  if (read_side_closed_ && write_side_closed_) {
-    return;
-  }
-  auto error_code = frame.quic_error_code;
-  if (error_code != QUIC_NO_ERROR) {
-    stream_error_ =
-        QuicResetStreamError::FromInternal(QUIC_STREAM_CONNECTION_ERROR);
-    connection_error_ = error_code;
+  if (!read_side_closed_ || !write_side_closed_) {
+    auto error_code = frame.quic_error_code;
+    if (error_code != QUIC_NO_ERROR) {
+      stream_error_ =
+          QuicResetStreamError::FromInternal(QUIC_STREAM_CONNECTION_ERROR);
+      connection_error_ = error_code;
+    }
+
+    CloseWriteSide();
+    CloseReadSide();
   }
 
-  CloseWriteSide();
-  CloseReadSide();
+  if (session_->notify_stream_soon_to_destroy()) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_notify_stream_soon_to_destroy, 2, 2);
+    // The stream may still be waiting for ACKs, but no ACK will be received any
+    // more. So move it into the closed stream list.
+    session_->MaybeCloseZombieStream(id_);
+  }
 }
 
 void QuicStream::OnFinRead() {
diff --git a/quiche/quic/core/quic_stream.h b/quiche/quic/core/quic_stream.h
index 4da8b7f..f4a9f58 100644
--- a/quiche/quic/core/quic_stream.h
+++ b/quiche/quic/core/quic_stream.h
@@ -412,6 +412,9 @@
   // indicating it can start processing data.
   void OnStreamCreatedFromPendingStream();
 
+  // Called by the sessionwhen a closed stream is about to be destroyed.
+  virtual void OnSoonToBeDestroyed();
+
   void DisableConnectionFlowControlForThisStream() {
     stream_contributes_to_connection_flow_control_ = false;
   }