Internal QUICHE change

PiperOrigin-RevId: 302100818
Change-Id: Ic15954c478f38fa61d0059a737dff6560de74454

quic/core/quic_connection.cc

diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index fd5e539..5f2b031 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2322,13 +2322,7 @@
         << ", previous_validated_mtu_:" << previous_validated_mtu_
         << ", max_packet_length():" << max_packet_length()
         << ", is_mtu_discovery:" << is_mtu_discovery;
-    if (previous_validated_mtu_ != 0 &&
-        GetQuicReloadableFlag(quic_ignore_one_write_error_after_mtu_probe)) {
-      QUIC_CODE_COUNT(quic_ignore_one_write_error_after_mtu_probe);
-      SetMaxPacketLength(previous_validated_mtu_);
-      mtu_discoverer_.Disable();
-      mtu_discovery_alarm_->Cancel();
-      previous_validated_mtu_ = 0;
+    if (ShouldIgnoreWriteError()) {
       return true;
     }
 
@@ -2420,7 +2414,7 @@
     return;
   }
 
-  if (IsWriteError(result.status)) {
+  if (IsWriteError(result.status) && !ShouldIgnoreWriteError()) {
     OnWriteError(result.error_code);
   }
 }
@@ -2451,6 +2445,20 @@
   return false;
 }
 
+bool QuicConnection::ShouldIgnoreWriteError() {
+  if (!GetQuicReloadableFlag(quic_ignore_one_write_error_after_mtu_probe) ||
+      previous_validated_mtu_ == 0) {
+    return false;
+  }
+
+  QUIC_CODE_COUNT(quic_ignore_one_write_error_after_mtu_probe);
+  SetMaxPacketLength(previous_validated_mtu_);
+  mtu_discoverer_.Disable();
+  mtu_discovery_alarm_->Cancel();
+  previous_validated_mtu_ = 0;
+  return true;
+}
+
 void QuicConnection::OnWriteError(int error_code) {
   if (write_error_occurred_) {
     // A write error already occurred. The connection is being closed.

quic/core/quic_connection.h

diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index f86c52c..fcd177a 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1216,6 +1216,11 @@
   // Whether connection is limited by amplification factor.
   bool LimitedByAmplificationFactor() const;
 
+  // We've got a packet write error, should we ignore it?
+  // NOTE: This is not a const function - if return true, the max packet size is
+  // reverted to a previous(smaller) value to avoid write errors in the future.
+  bool ShouldIgnoreWriteError();
+
   QuicFramer framer_;
 
   // Contents received in the current packet, especially used to identify
@@ -1439,7 +1444,7 @@
   // - MTU discovery has never been enabled, or
   // - MTU discovery has been enabled, but the connection got a packet write
   //   error with a new (successfully probed) MTU, so it reverted
-  //   |long_term_mtu_| to the value before the increase.
+  //   |long_term_mtu_| to the value before the last increase.
   QuicPacketLength previous_validated_mtu_;
   // The value of the MTU regularly used by the connection. This is different
   // from the value returned by max_packet_size(), as max_packet_size() returns

quic/core/quic_connection_test.cc

diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index f73b736..cd83e94 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -452,6 +452,9 @@
       SetWriteBlocked();
       return WriteResult(WRITE_STATUS_BLOCKED, /*errno*/ -1);
     }
+    if (write_should_fail_) {
+      return WriteResult(WRITE_STATUS_ERROR, /*errno*/ -1);
+    }
     int bytes_flushed = bytes_buffered_;
     bytes_buffered_ = 0;
     return WriteResult(WRITE_STATUS_OK, bytes_flushed);
@@ -5142,6 +5145,79 @@
   }
 }
 
+// After a successful MTU probe, one and only one write error should be ignored
+// if it happened in QuicConnection::FlushPacket.
+TEST_P(QuicConnectionTest,
+       MtuDiscoveryIgnoreOneWriteErrorInFlushAfterSuccessfulProbes) {
+  MtuDiscoveryTestInit();
+  writer_->SetBatchMode(true);
+
+  const QuicPacketCount packets_between_probes_base = 5;
+  set_packets_between_probes_base(packets_between_probes_base);
+
+  connection_.EnablePathMtuDiscovery(send_algorithm_);
+
+  const QuicByteCount original_max_packet_length =
+      connection_.max_packet_length();
+  // Send enough packets so that the next one triggers path MTU discovery.
+  for (QuicPacketCount i = 0; i < packets_between_probes_base - 1; i++) {
+    SendStreamDataToPeer(3, ".", i, NO_FIN, nullptr);
+    ASSERT_FALSE(connection_.GetMtuDiscoveryAlarm()->IsSet());
+  }
+
+  // Trigger the probe.
+  SendStreamDataToPeer(3, "!", packets_between_probes_base - 1, NO_FIN,
+                       nullptr);
+  ASSERT_TRUE(connection_.GetMtuDiscoveryAlarm()->IsSet());
+  QuicByteCount probe_size;
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _))
+      .WillOnce(SaveArg<3>(&probe_size));
+  connection_.GetMtuDiscoveryAlarm()->Fire();
+
+  EXPECT_THAT(probe_size, InRange(connection_.max_packet_length(),
+                                  kMtuDiscoveryTargetPacketSizeHigh));
+
+  const QuicPacketNumber probe_packet_number =
+      FirstSendingPacketNumber() + packets_between_probes_base;
+  ASSERT_EQ(probe_packet_number, creator_->packet_number());
+
+  // Acknowledge all packets sent so far.
+  QuicAckFrame probe_ack = InitAckFrame(probe_packet_number);
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(true, _, _, _, _))
+      .Times(AnyNumber());
+  ProcessAckPacket(&probe_ack);
+  EXPECT_EQ(probe_size, connection_.max_packet_length());
+  EXPECT_EQ(0u, connection_.GetBytesInFlight());
+
+  EXPECT_EQ(1u, connection_.mtu_probe_count());
+
+  if (GetQuicReloadableFlag(quic_ignore_one_write_error_after_mtu_probe)) {
+    writer_->SetShouldWriteFail();
+
+    // Ignore PACKET_WRITE_ERROR once.
+    {
+      QuicConnection::ScopedPacketFlusher flusher(&connection_);
+      // flusher's destructor will call connection_.FlushPackets, which should
+      // get a WRITE_STATUS_ERROR from the writer and ignore it.
+    }
+    EXPECT_EQ(original_max_packet_length, connection_.max_packet_length());
+    EXPECT_TRUE(connection_.connected());
+
+    // Close connection on another PACKET_WRITE_ERROR.
+    EXPECT_CALL(visitor_, OnConnectionClosed(_, _))
+        .WillOnce(Invoke(this, &QuicConnectionTest::SaveConnectionCloseFrame));
+    {
+      QuicConnection::ScopedPacketFlusher flusher(&connection_);
+      // flusher's destructor will call connection_.FlushPackets, which should
+      // get a WRITE_STATUS_ERROR from the writer and ignore it.
+    }
+    EXPECT_EQ(original_max_packet_length, connection_.max_packet_length());
+    EXPECT_FALSE(connection_.connected());
+    EXPECT_THAT(saved_connection_close_frame_.quic_error_code,
+                IsError(QUIC_PACKET_WRITE_ERROR));
+  }
+}
+
 // Simulate the case where the first attempt to send a probe is write blocked,
 // and after unblock, the second attempt returns a MSG_TOO_BIG error.
 TEST_P(QuicConnectionTest, MtuDiscoveryWriteBlocked) {