Simply SetRetransmissionAlarm by consolidating more of the changes from cl/381566145 and cl/314387174 into SetRetransmissionAlarm to ensure the logic is applied consistently.

(This change is based on Ian's cl/381657413)

Protected by FLAGS_quic_reloadable_flag_quic_simplify_set_retransmission_alarm.

PiperOrigin-RevId: 439976832
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 0716c52..510e096 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -1942,6 +1942,7 @@
 
 bool QuicConnection::ShouldSetRetransmissionAlarmOnPacketSent(
     bool in_flight, EncryptionLevel level) const {
+  QUICHE_DCHECK(!sent_packet_manager_.simplify_set_retransmission_alarm());
   if (!retransmission_alarm_->IsSet()) {
     return true;
   }
@@ -3660,8 +3661,13 @@
       return true;
     }
   }
-  if (ShouldSetRetransmissionAlarmOnPacketSent(in_flight,
-                                               packet->encryption_level)) {
+  if (sent_packet_manager_.simplify_set_retransmission_alarm()) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_simplify_set_retransmission_alarm, 1, 2);
+    if (in_flight || !retransmission_alarm_->IsSet()) {
+      SetRetransmissionAlarm();
+    }
+  } else if (ShouldSetRetransmissionAlarmOnPacketSent(
+                 in_flight, packet->encryption_level)) {
     SetRetransmissionAlarm();
   }
   SetPingAlarm();
@@ -4334,8 +4340,11 @@
   undecryptable_packets_.emplace_back(packet, decryption_level,
                                       last_received_packet_info_);
   if (perspective_ == Perspective::IS_CLIENT) {
-    if (!retransmission_alarm_->IsSet() ||
-        GetRetransmissionDeadline() < retransmission_alarm_->deadline()) {
+    if (sent_packet_manager_.simplify_set_retransmission_alarm()) {
+      SetRetransmissionAlarm();
+    } else if (!retransmission_alarm_->IsSet() ||
+               GetRetransmissionDeadline() <
+                   retransmission_alarm_->deadline()) {
       // Re-arm PTO only if we can make it sooner to speed up recovery.
       SetRetransmissionAlarm();
     }
@@ -4403,8 +4412,12 @@
     undecryptable_packets_.clear();
   }
   if (perspective_ == Perspective::IS_CLIENT) {
-    if (!retransmission_alarm_->IsSet() || undecryptable_packets_.empty() ||
-        GetRetransmissionDeadline() < retransmission_alarm_->deadline()) {
+    if (sent_packet_manager_.simplify_set_retransmission_alarm()) {
+      SetRetransmissionAlarm();
+    } else if (!retransmission_alarm_->IsSet() ||
+               undecryptable_packets_.empty() ||
+               GetRetransmissionDeadline() <
+                   retransmission_alarm_->deadline()) {
       // 1) If there is still undecryptable packet, only re-arm PTO to make it
       // sooner to speed up recovery.
       // 2) If all undecryptable packets get processed, re-arm (which may
@@ -4760,6 +4773,27 @@
     retransmission_alarm_->Cancel();
     return;
   }
+  PacketNumberSpace packet_number_space;
+  if (sent_packet_manager_.simplify_set_retransmission_alarm() &&
+      SupportsMultiplePacketNumberSpaces() && !IsHandshakeConfirmed() &&
+      !sent_packet_manager_
+           .GetEarliestPacketSentTimeForPto(&packet_number_space)
+           .IsInitialized()) {
+    // Before handshake gets confirmed, GetEarliestPacketSentTimeForPto
+    // returning 0 indicates no packets are in flight or only application data
+    // is in flight.
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_simplify_set_retransmission_alarm, 2, 2);
+    if (perspective_ == Perspective::IS_SERVER) {
+      // No need to arm PTO on server side.
+      retransmission_alarm_->Cancel();
+      return;
+    }
+    if (retransmission_alarm_->IsSet() &&
+        GetRetransmissionDeadline() > retransmission_alarm_->deadline()) {
+      // Do not postpone armed PTO on the client side.
+      return;
+    }
+  }
 
   retransmission_alarm_->Update(GetRetransmissionDeadline(), kAlarmGranularity);
 }
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 6cab381..6f7b4e2 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -13473,7 +13473,6 @@
 
 TEST_P(QuicConnectionTest, MigratePath) {
   EXPECT_CALL(visitor_, GetHandshakeState())
-      .Times(testing::AtMost(2))
       .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
   EXPECT_CALL(visitor_, OnPathDegrading());
   connection_.OnPathDegradingDetected();
@@ -15778,6 +15777,58 @@
                   "Already sent connection close");
 }
 
+// Regression test for b/157895910.
+TEST_P(QuicConnectionTest, EarliestSentTimeNotInitializedWhenPtoFires) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  set_perspective(Perspective::IS_SERVER);
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
+  EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(AnyNumber());
+  use_tagging_decrypter();
+
+  // Received INITIAL 1.
+  ProcessCryptoPacketAtLevel(1, ENCRYPTION_INITIAL);
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x03));
+  SetDecrypter(ENCRYPTION_HANDSHAKE,
+               std::make_unique<StrictTaggingDecrypter>(0x03));
+  connection_.SetEncrypter(ENCRYPTION_FORWARD_SECURE,
+                           std::make_unique<TaggingEncrypter>(0x04));
+  {
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    // Send INITIAL 1.
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+    connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_INITIAL);
+    // Send HANDSHAKE 2.
+    EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+    connection_.SendCryptoDataWithString(std::string(200, 'a'), 0,
+                                         ENCRYPTION_HANDSHAKE);
+    // Send half RTT data.
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+    connection_.SendStreamDataWithString(0, std::string(2000, 'b'), 0, FIN);
+  }
+
+  // Received ACKs for both INITIAL and HANDSHAKE packets.
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _))
+      .Times(AnyNumber());
+  QuicFrames frames1;
+  QuicAckFrame ack_frame1 = InitAckFrame(1);
+  frames1.push_back(QuicFrame(&ack_frame1));
+
+  QuicFrames frames2;
+  QuicAckFrame ack_frame2 =
+      InitAckFrame({{QuicPacketNumber(2), QuicPacketNumber(3)}});
+  frames2.push_back(QuicFrame(&ack_frame2));
+  ProcessCoalescedPacket(
+      {{2, frames1, ENCRYPTION_INITIAL}, {3, frames2, ENCRYPTION_HANDSHAKE}});
+  // Verify PTO is not armed given the only outstanding data is half RTT data.
+  EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index ac0e7d9..66d5dc9 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -37,6 +37,8 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_allow_client_enabled_bbr_v2, false)
 // If true, close read side but not write side in QuicSpdyStream::OnStreamReset().
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_fix_on_stream_reset, true)
+// If true, consolidate more logic into SetRetransmissionAlarm to ensure the logic is applied consistently.
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_simplify_set_retransmission_alarm, true)
 // If true, default on PTO which unifies TLP + RTO loss recovery.
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_default_on_pto2, true)
 // If true, default-enable 5RTO blachole detection.
diff --git a/quic/core/quic_sent_packet_manager.cc b/quic/core/quic_sent_packet_manager.cc
index f7e91b9..6fabf5b 100644
--- a/quic/core/quic_sent_packet_manager.cc
+++ b/quic/core/quic_sent_packet_manager.cc
@@ -1292,7 +1292,8 @@
     return QuicTime::Zero();
   }
   PacketNumberSpace packet_number_space;
-  if (supports_multiple_packet_number_spaces() &&
+  if (!simplify_set_retransmission_alarm_ &&
+      supports_multiple_packet_number_spaces() &&
       unacked_packets_.perspective() == Perspective::IS_SERVER &&
       !GetEarliestPacketSentTimeForPto(&packet_number_space).IsInitialized()) {
     // Do not set the timer on the server side if the only in flight packets are
diff --git a/quic/core/quic_sent_packet_manager.h b/quic/core/quic_sent_packet_manager.h
index 22067ff..f31cc1c 100644
--- a/quic/core/quic_sent_packet_manager.h
+++ b/quic/core/quic_sent_packet_manager.h
@@ -495,6 +495,10 @@
 
   bool use_lower_min_irtt() const { return use_lower_min_irtt_; }
 
+  bool simplify_set_retransmission_alarm() const {
+    return simplify_set_retransmission_alarm_;
+  }
+
  private:
   friend class test::QuicConnectionPeer;
   friend class test::QuicSentPacketManagerPeer;
@@ -770,6 +774,9 @@
   // Latched value of --quic_use_lower_min_for_trusted_irtt.
   bool use_lower_min_irtt_ =
       GetQuicReloadableFlag(quic_use_lower_min_for_trusted_irtt);
+
+  const bool simplify_set_retransmission_alarm_ =
+      GetQuicReloadableFlag(quic_simplify_set_retransmission_alarm);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index bf141e4..6a33fda 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -4167,6 +4167,9 @@
 
 // Regression test for b/157895910.
 TEST_F(QuicSentPacketManagerTest, EarliestSentTimeNotInitializedWhenPtoFires) {
+  if (GetQuicReloadableFlag(quic_simplify_set_retransmission_alarm)) {
+    return;
+  }
   manager_.EnableMultiplePacketNumberSpacesSupport();
   EXPECT_CALL(*send_algorithm_, PacingRate(_))
       .WillRepeatedly(Return(QuicBandwidth::Zero()));