Only set packet state if all data gets retransmitted.

Make "Failed to ConsumeCryptoData" fire only if not on soft max packet length. Also, remove the added member variable for debugging only.

Protected by FLAGS_quic_restart_flag_quic_set_packet_state_if_all_data_retransmitted.

PiperOrigin-RevId: 430955887
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index f563fde..d487eb0 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -15269,6 +15269,94 @@
   }
 }
 
+// Regression test for b/216133388.
+TEST_P(QuicConnectionTest, FailedToConsumeCryptoData) {
+  if (!version().HasIetfQuicFrames()) {
+    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);
+  EXPECT_TRUE(connection_.HasPendingAcks());
+
+  peer_framer_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                            std::make_unique<TaggingEncrypter>(0x02));
+
+  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));
+  SetDecrypter(ENCRYPTION_ZERO_RTT,
+               std::make_unique<StrictTaggingDecrypter>(0x02));
+  connection_.SetEncrypter(ENCRYPTION_FORWARD_SECURE,
+                           std::make_unique<TaggingEncrypter>(0x04));
+  // Received ENCRYPTION_ZERO_RTT 1.
+  ProcessDataPacketAtLevel(1, !kHasStopWaiting, ENCRYPTION_ZERO_RTT);
+  {
+    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 1-RTT 3.
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+    connection_.SendStreamDataWithString(0, std::string(40, 'a'), 0, NO_FIN);
+  }
+  // Received HANDSHAKE Ping, hence discard INITIAL keys.
+  peer_framer_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                            std::make_unique<TaggingEncrypter>(0x03));
+  connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
+  connection_.NeuterUnencryptedPackets();
+  ProcessCryptoPacketAtLevel(1, ENCRYPTION_HANDSHAKE);
+  clock_.AdvanceTime(kAlarmGranularity);
+  {
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    // Sending this 1-RTT data would leave the coalescer only have space to
+    // accommodate the HANDSHAKE ACK. The crypto data cannot be bundled with the
+    // ACK.
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+    connection_.SendStreamDataWithString(0, std::string(1395, 'a'), 40, NO_FIN);
+  }
+  // Verify retransmission alarm is armed.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  const QuicTime retransmission_time =
+      connection_.GetRetransmissionAlarm()->deadline();
+  clock_.AdvanceTime(retransmission_time - clock_.Now());
+  connection_.GetRetransmissionAlarm()->Fire();
+
+  if (GetQuicRestartFlag(quic_set_packet_state_if_all_data_retransmitted)) {
+    // Verify the retransmission is a coalesced packet with HANDSHAKE 2 and
+    // 1-RTT 3.
+    EXPECT_EQ(0x04040404u, writer_->final_bytes_of_last_packet());
+    // Only the first packet in the coalesced packet has been processed.
+    EXPECT_EQ(1u, writer_->crypto_frames().size());
+    // Process the coalesced 1-RTT packet.
+    ASSERT_TRUE(writer_->coalesced_packet() != nullptr);
+    auto packet = writer_->coalesced_packet()->Clone();
+    writer_->framer()->ProcessPacket(*packet);
+    EXPECT_EQ(1u, writer_->stream_frames().size());
+    ASSERT_TRUE(writer_->coalesced_packet() == nullptr);
+  } else {
+    // Although packet 2 has not been retransmitted, it has been marked PTOed
+    // and a HANDHSAKE PING gets retransmitted.
+    EXPECT_EQ(0x03030303u, writer_->final_bytes_of_last_packet());
+    EXPECT_EQ(1u, writer_->ping_frames().size());
+    EXPECT_TRUE(writer_->stream_frames().empty());
+    ASSERT_TRUE(writer_->coalesced_packet() == nullptr);
+  }
+  // Verify retransmission alarm is still armed.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_crypto_stream.cc b/quic/core/quic_crypto_stream.cc
index 99fb2cd..b579ff5 100644
--- a/quic/core/quic_crypto_stream.cc
+++ b/quic/core/quic_crypto_stream.cc
@@ -378,7 +378,7 @@
       crypto_frame->offset, crypto_frame->data_length);
 }
 
-void QuicCryptoStream::RetransmitData(QuicCryptoFrame* crypto_frame,
+bool QuicCryptoStream::RetransmitData(QuicCryptoFrame* crypto_frame,
                                       TransmissionType type) {
   QUIC_BUG_IF(quic_bug_12573_6,
               !QuicVersionUsesCryptoFrames(session()->transport_version()))
@@ -389,7 +389,7 @@
       &substreams_[crypto_frame->level].send_buffer;
   retransmission.Difference(send_buffer->bytes_acked());
   if (retransmission.Empty()) {
-    return;
+    return true;
   }
   for (const auto& interval : retransmission) {
     size_t retransmission_offset = interval.min();
@@ -400,9 +400,10 @@
     send_buffer->OnStreamDataRetransmitted(retransmission_offset,
                                            bytes_consumed);
     if (bytes_consumed < retransmission_length) {
-      break;
+      return false;
     }
   }
+  return true;
 }
 
 void QuicCryptoStream::WriteBufferedCryptoFrames() {
diff --git a/quic/core/quic_crypto_stream.h b/quic/core/quic_crypto_stream.h
index 8197a5f..cada7db 100644
--- a/quic/core/quic_crypto_stream.h
+++ b/quic/core/quic_crypto_stream.h
@@ -217,8 +217,9 @@
   void OnCryptoFrameLost(QuicCryptoFrame* crypto_frame);
 
   // Called to retransmit any outstanding data in the range indicated by the
-  // encryption level, offset, and length in |crypto_frame|.
-  void RetransmitData(QuicCryptoFrame* crypto_frame, TransmissionType type);
+  // encryption level, offset, and length in |crypto_frame|. Returns true if all
+  // data gets retransmitted.
+  bool RetransmitData(QuicCryptoFrame* crypto_frame, TransmissionType type);
 
   // Called to write buffered crypto frames.
   void WriteBufferedCryptoFrames();
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index c5c31c5..bfe20e2 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -107,6 +107,8 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_conservative_cwnd_and_pacing_gains, false)
 // If true, validate that peer owns the new address once the server detects peer migration or is probed from that address, and also apply anti-amplification limit while sending to that address.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_server_reverse_validate_new_path3, true)
+// If true, when a packet is forced retransmitted, only set packet state if all data gets retransmitted.
+QUIC_FLAG(FLAGS_quic_restart_flag_quic_set_packet_state_if_all_data_retransmitted, false)
 // If true, when client attempts TLS resumption, use token in session_cache_ instead of cached_states_ in QuicCryptoClientConfig.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_tls_use_token_in_session_cache, true)
 // When the flag is true, exit STARTUP after the same number of loss events as PROBE_UP.
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index d34d1e6..885dc48 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -1451,17 +1451,16 @@
             level, write_length - total_bytes_consumed,
             offset + total_bytes_consumed, fully_pad_crypto_handshake_packets_,
             next_transmission_type_, &frame)) {
-      // The only pending data in the packet is non-retransmittable frames. I'm
-      // assuming here that they won't occupy so much of the packet that a
+      // The only pending data in the packet is non-retransmittable frames.
+      // I'm assuming here that they won't occupy so much of the packet that a
       // CRYPTO frame won't fit.
-      const std::string error_message = absl::StrCat(
+      QUIC_BUG_IF(quic_bug_10752_26, !HasSoftMaxPacketLength()) << absl::StrCat(
           ENDPOINT, "Failed to ConsumeCryptoData at level ", level,
           ", pending_frames: ", GetPendingFramesInfo(),
           ", has_soft_max_packet_length: ", HasSoftMaxPacketLength(),
           ", max_packet_length: ", max_packet_length_, ", transmission_type: ",
           TransmissionTypeToString(next_transmission_type_),
           ", packet_number: ", packet_number().ToString());
-      QUIC_BUG(quic_bug_10752_26) << error_message;
       return 0;
     }
     total_bytes_consumed += frame.crypto_frame->data_length;
diff --git a/quic/core/quic_sent_packet_manager.cc b/quic/core/quic_sent_packet_manager.cc
index af86d2f..ff40ba4 100644
--- a/quic/core/quic_sent_packet_manager.cc
+++ b/quic/core/quic_sent_packet_manager.cc
@@ -670,8 +670,42 @@
   // Handshake packets should never be sent as probing retransmissions.
   QUICHE_DCHECK(!transmission_info->has_crypto_handshake ||
                 transmission_type != PROBING_RETRANSMISSION);
+  if (ShouldForceRetransmission(transmission_type)) {
+    const bool retransmitted = unacked_packets_.RetransmitFrames(
+        QuicFrames(transmission_info->retransmittable_frames),
+        transmission_type);
+    if (GetQuicRestartFlag(quic_set_packet_state_if_all_data_retransmitted)) {
+      QUIC_RESTART_FLAG_COUNT(quic_set_packet_state_if_all_data_retransmitted);
+      if (!retransmitted) {
+        // Do not set packet state if the data is not fully retransmitted.
+        // This should only happen if packet payload size decreases which can be
+        // caused by:
+        // 1) connection tries to opportunistically retransmit data
+        // when sending a packet of a different packet number space, or
+        // 2) path MTU decreases, or
+        // 3) packet header size increases (e.g., packet number length
+        // increases).
+        QUIC_CODE_COUNT(quic_retransmit_frames_failed);
+        return;
+      }
+      QUIC_CODE_COUNT(quic_retransmit_frames_succeeded);
+    }
+  } else {
+    unacked_packets_.NotifyFramesLost(*transmission_info, transmission_type);
 
-  HandleRetransmission(transmission_type, transmission_info);
+    if (!transmission_info->retransmittable_frames.empty()) {
+      if (transmission_type == LOSS_RETRANSMISSION) {
+        // Record the first packet sent after loss, which allows to wait 1
+        // more RTT before giving up on this lost packet.
+        transmission_info->first_sent_after_loss =
+            unacked_packets_.largest_sent_packet() + 1;
+      } else {
+        // Clear the recorded first packet sent after loss when version or
+        // encryption changes.
+        transmission_info->first_sent_after_loss.Clear();
+      }
+    }
+  }
 
   // Get the latest transmission_info here as it can be invalidated after
   // HandleRetransmission adding new sent packets into unacked_packets_.
@@ -683,44 +717,6 @@
       QuicUtils::RetransmissionTypeToPacketState(transmission_type);
 }
 
-void QuicSentPacketManager::HandleRetransmission(
-    TransmissionType transmission_type,
-    QuicTransmissionInfo* transmission_info) {
-  if (ShouldForceRetransmission(transmission_type)) {
-    // TODO(fayang): Consider to make RTO and PROBING retransmission
-    // strategies be configurable by applications. Today, TLP, RTO and PROBING
-    // retransmissions are handled similarly, i.e., always retranmist the
-    // oldest outstanding data. This is not ideal in general because different
-    // applications may want different strategies. For example, some
-    // applications may want to use higher priority stream data for bandwidth
-    // probing, and some applications want to consider RTO is an indication of
-    // loss, etc.
-    // transmission_info owning these frames may be deallocated after each
-    // retransimission. Make a copy of retransmissible frames to prevent the
-    // invalidation.
-    unacked_packets_.RetransmitFrames(
-        QuicFrames(transmission_info->retransmittable_frames),
-        transmission_type);
-    return;
-  }
-
-  unacked_packets_.NotifyFramesLost(*transmission_info, transmission_type);
-  if (transmission_info->retransmittable_frames.empty()) {
-    return;
-  }
-
-  if (transmission_type == LOSS_RETRANSMISSION) {
-    // Record the first packet sent after loss, which allows to wait 1
-    // more RTT before giving up on this lost packet.
-    transmission_info->first_sent_after_loss =
-        unacked_packets_.largest_sent_packet() + 1;
-  } else {
-    // Clear the recorded first packet sent after loss when version or
-    // encryption changes.
-    transmission_info->first_sent_after_loss.Clear();
-  }
-}
-
 void QuicSentPacketManager::RecordOneSpuriousRetransmission(
     const QuicTransmissionInfo& info) {
   stats_->bytes_spuriously_retransmitted += info.bytes_sent;
diff --git a/quic/core/quic_sent_packet_manager.h b/quic/core/quic_sent_packet_manager.h
index dac8ccf..cb4fdef 100644
--- a/quic/core/quic_sent_packet_manager.h
+++ b/quic/core/quic_sent_packet_manager.h
@@ -548,12 +548,6 @@
   void MarkForRetransmission(QuicPacketNumber packet_number,
                              TransmissionType transmission_type);
 
-  // Performs whatever work is need to retransmit the data correctly, either
-  // by retransmitting the frames directly or by notifying that the frames
-  // are lost.
-  void HandleRetransmission(TransmissionType transmission_type,
-                            QuicTransmissionInfo* transmission_info);
-
   // Called after packets have been marked handled with last received ack frame.
   void PostProcessNewlyAckedPackets(QuicPacketNumber ack_packet_number,
                                     EncryptionLevel ack_decrypted_level,
diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index cceceb7..8335103 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -47,22 +47,19 @@
 
 class MockDebugDelegate : public QuicSentPacketManager::DebugDelegate {
  public:
-  MOCK_METHOD(void,
-              OnSpuriousPacketRetransmission,
+  MOCK_METHOD(void, OnSpuriousPacketRetransmission,
               (TransmissionType transmission_type, QuicByteCount byte_size),
               (override));
-  MOCK_METHOD(void,
-              OnPacketLoss,
+  MOCK_METHOD(void, OnPacketLoss,
               (QuicPacketNumber lost_packet_number,
                EncryptionLevel encryption_level,
-               TransmissionType transmission_type,
-               QuicTime detection_time),
+               TransmissionType transmission_type, QuicTime detection_time),
               (override));
 };
 
 class QuicSentPacketManagerTest : public QuicTest {
  public:
-  void RetransmitCryptoPacket(uint64_t packet_number) {
+  bool RetransmitCryptoPacket(uint64_t packet_number) {
     EXPECT_CALL(
         *send_algorithm_,
         OnPacketSent(_, BytesInFlight(), QuicPacketNumber(packet_number),
@@ -73,10 +70,10 @@
     packet.has_crypto_handshake = IS_HANDSHAKE;
     manager_.OnPacketSent(&packet, clock_.Now(), HANDSHAKE_RETRANSMISSION,
                           HAS_RETRANSMITTABLE_DATA, true);
+    return true;
   }
 
-  void RetransmitDataPacket(uint64_t packet_number,
-                            TransmissionType type,
+  bool RetransmitDataPacket(uint64_t packet_number, TransmissionType type,
                             EncryptionLevel level) {
     EXPECT_CALL(
         *send_algorithm_,
@@ -86,20 +83,18 @@
     packet.encryption_level = level;
     manager_.OnPacketSent(&packet, clock_.Now(), type, HAS_RETRANSMITTABLE_DATA,
                           true);
+    return true;
   }
 
-  void RetransmitDataPacket(uint64_t packet_number, TransmissionType type) {
-    RetransmitDataPacket(packet_number, type, ENCRYPTION_INITIAL);
+  bool RetransmitDataPacket(uint64_t packet_number, TransmissionType type) {
+    return RetransmitDataPacket(packet_number, type, ENCRYPTION_INITIAL);
   }
 
  protected:
   const CongestionControlType kInitialCongestionControlType = kCubicBytes;
   QuicSentPacketManagerTest()
-      : manager_(Perspective::IS_SERVER,
-                 &clock_,
-                 QuicRandom::GetInstance(),
-                 &stats_,
-                 kInitialCongestionControlType),
+      : manager_(Perspective::IS_SERVER, &clock_, QuicRandom::GetInstance(),
+                 &stats_, kInitialCongestionControlType),
         send_algorithm_(new StrictMock<MockSendAlgorithm>),
         network_change_visitor_(new StrictMock<MockNetworkChangeVisitor>) {
     QuicSentPacketManagerPeer::SetSendAlgorithm(&manager_, send_algorithm_);
@@ -174,8 +169,7 @@
     EXPECT_CALL(*network_change_visitor_, OnCongestionChange());
   }
 
-  void ExpectAckAndLoss(bool rtt_updated,
-                        uint64_t largest_observed,
+  void ExpectAckAndLoss(bool rtt_updated, uint64_t largest_observed,
                         uint64_t lost_packet) {
     EXPECT_CALL(
         *send_algorithm_,
@@ -186,10 +180,8 @@
   }
 
   // |packets_acked| and |packets_lost| should be in packet number order.
-  void ExpectAcksAndLosses(bool rtt_updated,
-                           uint64_t* packets_acked,
-                           size_t num_packets_acked,
-                           uint64_t* packets_lost,
+  void ExpectAcksAndLosses(bool rtt_updated, uint64_t* packets_acked,
+                           size_t num_packets_acked, uint64_t* packets_lost,
                            size_t num_packets_lost) {
     std::vector<QuicPacketNumber> ack_vector;
     for (size_t i = 0; i < num_packets_acked; ++i) {
@@ -224,7 +216,7 @@
       EXPECT_CALL(notifier_, RetransmitFrames(_, _))
           .WillOnce(WithArgs<1>(
               Invoke([this, new_packet_number](TransmissionType type) {
-                RetransmitDataPacket(new_packet_number, type);
+                return RetransmitDataPacket(new_packet_number, type);
               })));
     } else {
       EXPECT_CALL(notifier_, OnFrameLost(_)).Times(1);
@@ -311,8 +303,7 @@
     SendAckPacket(packet_number, largest_acked, ENCRYPTION_INITIAL);
   }
 
-  void SendAckPacket(uint64_t packet_number,
-                     uint64_t largest_acked,
+  void SendAckPacket(uint64_t packet_number, uint64_t largest_acked,
                      EncryptionLevel level) {
     EXPECT_CALL(
         *send_algorithm_,
@@ -398,8 +389,9 @@
 TEST_F(QuicSentPacketManagerTest, RetransmitThenAckBeforeSend) {
   SendDataPacket(1);
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(2, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(2, type);
+      })));
   QuicSentPacketManagerPeer::MarkForRetransmission(&manager_, 1,
                                                    TLP_RETRANSMISSION);
   // Ack 1.
@@ -421,7 +413,7 @@
 
 TEST_F(QuicSentPacketManagerTest, RetransmitThenStopRetransmittingBeforeSend) {
   SendDataPacket(1);
-  EXPECT_CALL(notifier_, RetransmitFrames(_, _));
+  EXPECT_CALL(notifier_, RetransmitFrames(_, _)).WillRepeatedly(Return(true));
   QuicSentPacketManagerPeer::MarkForRetransmission(&manager_, 1,
                                                    TLP_RETRANSMISSION);
 
@@ -852,16 +844,18 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(2, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(2, type);
+      })));
   manager_.MaybeRetransmitTailLossProbe();
 
   // The second tail loss probe retransmits 1 packet.
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(3, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(3, type);
+      })));
   manager_.MaybeRetransmitTailLossProbe();
   EXPECT_CALL(*send_algorithm_, CanSend(_)).WillOnce(Return(false));
   EXPECT_EQ(QuicTime::Delta::Infinite(), manager_.TimeUntilSend(clock_.Now()));
@@ -921,8 +915,9 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(101, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(101, type);
+      })));
   manager_.MaybeRetransmitTailLossProbe();
   EXPECT_CALL(*send_algorithm_, CanSend(_)).WillOnce(Return(false));
   EXPECT_EQ(QuicTime::Delta::Infinite(), manager_.TimeUntilSend(clock_.Now()));
@@ -932,8 +927,9 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(102, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(102, type);
+      })));
   EXPECT_TRUE(manager_.MaybeRetransmitTailLossProbe());
   EXPECT_CALL(*send_algorithm_, CanSend(_)).WillOnce(Return(false));
   EXPECT_EQ(QuicTime::Delta::Infinite(), manager_.TimeUntilSend(clock_.Now()));
@@ -948,10 +944,12 @@
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(103, type); })))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(104, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(103, type);
+      })))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(104, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(2u, stats_.tlp_count);
   EXPECT_EQ(1u, stats_.rto_count);
@@ -998,8 +996,10 @@
   // The first retransmits 2 packets.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(6); }))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(7); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(6); }))
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(7); }));
   manager_.OnRetransmissionTimeout();
   // Expect all 4 handshake packets to be in flight and 3 data packets.
   EXPECT_EQ(7 * kDefaultLength, manager_.GetBytesInFlight());
@@ -1008,8 +1008,10 @@
   // The second retransmits 2 packets.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(8); }))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(9); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(8); }))
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(9); }));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(9 * kDefaultLength, manager_.GetBytesInFlight());
   EXPECT_TRUE(manager_.HasUnackedCryptoPackets());
@@ -1041,12 +1043,14 @@
 
   // Retransmit the crypto packet as 2.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(2); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(2); }));
   manager_.OnRetransmissionTimeout();
 
   // Retransmit the crypto packet as 3.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(3); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(3); }));
   manager_.OnRetransmissionTimeout();
 
   // Now ack the second crypto packet, and ensure the first gets removed, but
@@ -1079,8 +1083,10 @@
   // Retransmit 2 crypto packets, but not the serialized packet.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(4); }))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(5); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(4); }))
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(5); }));
   manager_.OnRetransmissionTimeout();
   EXPECT_TRUE(manager_.HasUnackedCryptoPackets());
 }
@@ -1094,13 +1100,15 @@
 
   // Retransmit the crypto packet as 2.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(2); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(2); }));
   manager_.OnRetransmissionTimeout();
   EXPECT_TRUE(manager_.HasUnackedCryptoPackets());
 
   // Retransmit the crypto packet as 3.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(3); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(3); }));
   manager_.OnRetransmissionTimeout();
   EXPECT_TRUE(manager_.HasUnackedCryptoPackets());
 
@@ -1145,10 +1153,12 @@
   EXPECT_FALSE(manager_.MaybeRetransmitTailLossProbe());
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(101, type); })))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(102, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(101, type);
+      })))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(102, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(102 * kDefaultLength, manager_.GetBytesInFlight());
 
@@ -1210,8 +1220,9 @@
   EXPECT_FALSE(manager_.MaybeRetransmitTailLossProbe());
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(1)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(101, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(101, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(101 * kDefaultLength, manager_.GetBytesInFlight());
 }
@@ -1244,10 +1255,12 @@
   EXPECT_FALSE(manager_.MaybeRetransmitTailLossProbe());
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(101, type); })))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(102, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(101, type);
+      })))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(102, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(102 * kDefaultLength, manager_.GetBytesInFlight());
 
@@ -1279,15 +1292,17 @@
   SendDataPacket(1);
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(2, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(2, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(2 * kDefaultLength, manager_.GetBytesInFlight());
 
   // Rto a second time.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(3, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(3, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(3 * kDefaultLength, manager_.GetBytesInFlight());
 
@@ -1313,15 +1328,17 @@
   SendDataPacket(1);
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(2, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(2, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(2 * kDefaultLength, manager_.GetBytesInFlight());
 
   // Rto a second time.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(3, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(3, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(3 * kDefaultLength, manager_.GetBytesInFlight());
 
@@ -1363,7 +1380,8 @@
   // Retransmit the packet by invoking the retransmission timeout.
   clock_.AdvanceTime(1.5 * srtt);
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(2); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(2); }));
   // When session decides what to write, crypto_packet_send_time gets updated.
   crypto_packet_send_time = clock_.Now();
   manager_.OnRetransmissionTimeout();
@@ -1375,7 +1393,8 @@
   // Retransmit the packet for the 2nd time.
   clock_.AdvanceTime(2 * 1.5 * srtt);
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(3); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(3); }));
   // When session decides what to write, crypto_packet_send_time gets updated.
   crypto_packet_send_time = clock_.Now();
   manager_.OnRetransmissionTimeout();
@@ -1419,7 +1438,8 @@
   // Retransmit the packet by invoking the retransmission timeout.
   clock_.AdvanceTime(2 * srtt);
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(2); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(2); }));
   crypto_packet_send_time = clock_.Now();
   manager_.OnRetransmissionTimeout();
 
@@ -1454,8 +1474,9 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(3, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(3, type);
+      })));
   EXPECT_TRUE(manager_.MaybeRetransmitTailLossProbe());
   EXPECT_CALL(*send_algorithm_, CanSend(_)).WillOnce(Return(false));
   EXPECT_EQ(QuicTime::Delta::Infinite(), manager_.TimeUntilSend(clock_.Now()));
@@ -1486,10 +1507,12 @@
   clock_.AdvanceTime(expected_rto_delay);
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(5, type); })))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(6, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(5, type);
+      })))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(6, type);
+      })));
   manager_.OnRetransmissionTimeout();
   // All previous packets are inflight, plus two rto retransmissions.
   EXPECT_EQ(6 * kDefaultLength, manager_.GetBytesInFlight());
@@ -1541,7 +1564,7 @@
     delay = delay + delay;
     EXPECT_CALL(notifier_, RetransmitFrames(_, _))
         .WillOnce(WithArgs<1>(Invoke([this, i](TransmissionType type) {
-          RetransmitDataPacket(i + 2, type);
+          return RetransmitDataPacket(i + 2, type);
         })));
     manager_.OnRetransmissionTimeout();
   }
@@ -1575,7 +1598,7 @@
     delay = delay + delay;
     EXPECT_CALL(notifier_, RetransmitFrames(_, _))
         .WillOnce(WithArgs<1>(Invoke([this, i](TransmissionType type) {
-          RetransmitDataPacket(i + 2, type);
+          return RetransmitDataPacket(i + 2, type);
         })));
     manager_.OnRetransmissionTimeout();
   }
@@ -2612,7 +2635,8 @@
   SendAckPacket(3, 1, ENCRYPTION_FORWARD_SECURE);
   // Retransmit SHLO.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(InvokeWithoutArgs([this]() { RetransmitCryptoPacket(4); }));
+      .WillOnce(
+          InvokeWithoutArgs([this]() { return RetransmitCryptoPacket(4); }));
   manager_.OnRetransmissionTimeout();
 
   // Successfully decrypt a forward secure packet.
@@ -2625,7 +2649,7 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(6, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(6, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeRetransmitTailLossProbe();
 
@@ -2664,10 +2688,12 @@
   }
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(11, type); })))
-      .WillOnce(WithArgs<1>(Invoke(
-          [this](TransmissionType type) { RetransmitDataPacket(12, type); })));
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(11, type);
+      })))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        return RetransmitDataPacket(12, type);
+      })));
   manager_.OnRetransmissionTimeout();
   EXPECT_EQ(1u, stats_.rto_count);
   EXPECT_EQ(0u, manager_.pending_timer_transmission_count());
@@ -2723,10 +2749,10 @@
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify PTO period gets set to twice the current value.
@@ -2789,7 +2815,7 @@
   // Verify one probe packet gets sent.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
 }
@@ -2903,7 +2929,7 @@
   // Verify 1 probe packets get sent and packet number gets skipped.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify PTO period gets set to twice the current value. Also, ack delay is
@@ -3024,10 +3050,10 @@
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify no exponential backoff.
@@ -3044,10 +3070,10 @@
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(5, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(5, type, ENCRYPTION_FORWARD_SECURE);
       })))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(6, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(6, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify still no exponential backoff.
@@ -3064,10 +3090,10 @@
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(7, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(7, type, ENCRYPTION_FORWARD_SECURE);
       })))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(8, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(8, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify exponential backoff starts.
@@ -3084,10 +3110,10 @@
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .Times(2)
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(9, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(9, type, ENCRYPTION_FORWARD_SECURE);
       })))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(10, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(10, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify exponential backoff continues.
@@ -3144,7 +3170,7 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeRetransmitTailLossProbe();
 
@@ -3152,7 +3178,7 @@
   manager_.OnRetransmissionTimeout();
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(4, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeRetransmitTailLossProbe();
 
@@ -3161,6 +3187,7 @@
       .WillOnce(WithArgs<0>(Invoke([&crypto_frame](const QuicFrames& frames) {
         EXPECT_EQ(1u, frames.size());
         EXPECT_NE(crypto_frame, frames[0].stream_frame);
+        return true;
       })));
   manager_.OnRetransmissionTimeout();
 }
@@ -3200,7 +3227,7 @@
   // Verify 1 probe packets get sent and packet number gets skipped.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
 
@@ -3248,7 +3275,7 @@
   // Verify 1 probe packets get sent and packet number gets skipped.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
 
@@ -3265,7 +3292,7 @@
   // Verify 1 probe packets get sent and packet number gets skipped.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(5, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(5, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   expected_pto_delay =
@@ -3332,7 +3359,7 @@
   // Verify probe packet gets sent.
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_HANDSHAKE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_HANDSHAKE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify PTO period gets set to twice the current value.
@@ -3492,7 +3519,7 @@
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify PTO period gets set to twice the current value and based on packet3.
@@ -3569,7 +3596,7 @@
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
+        return RetransmitDataPacket(3, type, ENCRYPTION_FORWARD_SECURE);
       })));
   manager_.MaybeSendProbePackets();
   // Verify PTO period gets set to twice the expected value and based on
@@ -4004,7 +4031,8 @@
   manager_.OnRetransmissionTimeout();
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke([this]() { RetransmitCryptoPacket(3); })));
+      .WillOnce(
+          WithArgs<1>(Invoke([this]() { return RetransmitCryptoPacket(3); })));
   manager_.MaybeSendProbePackets();
   if (GetQuicReloadableFlag(quic_default_on_pto)) {
     manager_.AdjustPendingTimerTransmissions();
@@ -4033,7 +4061,8 @@
   manager_.OnRetransmissionTimeout();
 
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
-      .WillOnce(WithArgs<1>(Invoke([this]() { RetransmitCryptoPacket(3); })));
+      .WillOnce(
+          WithArgs<1>(Invoke([this]() { return RetransmitCryptoPacket(3); })));
   manager_.MaybeSendProbePackets();
   if (GetQuicReloadableFlag(quic_default_on_pto)) {
     manager_.AdjustPendingTimerTransmissions();
@@ -4188,7 +4217,7 @@
   clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(4, type, ENCRYPTION_INITIAL);
+        return RetransmitDataPacket(4, type, ENCRYPTION_INITIAL);
       })));
   manager_.RetransmitDataOfSpaceIfAny(INITIAL_DATA);
   // Verify PTO is re-armed based on packet 2.
@@ -4199,7 +4228,7 @@
   clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
   EXPECT_CALL(notifier_, RetransmitFrames(_, _))
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
-        RetransmitDataPacket(5, type, ENCRYPTION_INITIAL);
+        return RetransmitDataPacket(5, type, ENCRYPTION_INITIAL);
       })));
   manager_.RetransmitDataOfSpaceIfAny(INITIAL_DATA);
   // Verify PTO does not change.
@@ -4259,8 +4288,7 @@
 }
 
 SerializedPacket MakePacketWithAckFrequencyFrame(
-    int packet_number,
-    int ack_frequency_sequence_number,
+    int packet_number, int ack_frequency_sequence_number,
     QuicTime::Delta max_ack_delay) {
   auto* ack_frequency_frame = new QuicAckFrequencyFrame();
   ack_frequency_frame->max_ack_delay = max_ack_delay;
diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index 89fa564..90c84de 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -2280,7 +2280,7 @@
   }
 }
 
-void QuicSession::RetransmitFrames(const QuicFrames& frames,
+bool QuicSession::RetransmitFrames(const QuicFrames& frames,
                                    TransmissionType type) {
   QuicConnection::ScopedPacketFlusher retransmission_flusher(connection_);
   for (const QuicFrame& frame : frames) {
@@ -2289,12 +2289,17 @@
       continue;
     }
     if (frame.type == CRYPTO_FRAME) {
-      GetMutableCryptoStream()->RetransmitData(frame.crypto_frame, type);
+      const bool data_retransmitted =
+          GetMutableCryptoStream()->RetransmitData(frame.crypto_frame, type);
+      if (GetQuicRestartFlag(quic_set_packet_state_if_all_data_retransmitted) &&
+          !data_retransmitted) {
+        return false;
+      }
       continue;
     }
     if (frame.type != STREAM_FRAME) {
       if (!control_frame_manager_.RetransmitControlFrame(frame, type)) {
-        break;
+        return false;
       }
       continue;
     }
@@ -2303,9 +2308,10 @@
         !stream->RetransmitStreamData(frame.stream_frame.offset,
                                       frame.stream_frame.data_length,
                                       frame.stream_frame.fin, type)) {
-      break;
+      return false;
     }
   }
+  return true;
 }
 
 bool QuicSession::IsFrameOutstanding(const QuicFrame& frame) const {
diff --git a/quic/core/quic_session.h b/quic/core/quic_session.h
index 2734c6e..b8f175e 100644
--- a/quic/core/quic_session.h
+++ b/quic/core/quic_session.h
@@ -193,7 +193,7 @@
                     QuicTime receive_timestamp) override;
   void OnStreamFrameRetransmitted(const QuicStreamFrame& frame) override;
   void OnFrameLost(const QuicFrame& frame) override;
-  void RetransmitFrames(const QuicFrames& frames,
+  bool RetransmitFrames(const QuicFrames& frames,
                         TransmissionType type) override;
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;
diff --git a/quic/core/quic_unacked_packet_map.cc b/quic/core/quic_unacked_packet_map.cc
index dada106..1c6f303 100644
--- a/quic/core/quic_unacked_packet_map.cc
+++ b/quic/core/quic_unacked_packet_map.cc
@@ -477,9 +477,9 @@
   }
 }
 
-void QuicUnackedPacketMap::RetransmitFrames(const QuicFrames& frames,
+bool QuicUnackedPacketMap::RetransmitFrames(const QuicFrames& frames,
                                             TransmissionType type) {
-  session_notifier_->RetransmitFrames(frames, type);
+  return session_notifier_->RetransmitFrames(frames, type);
 }
 
 void QuicUnackedPacketMap::MaybeAggregateAckedStreamFrame(
diff --git a/quic/core/quic_unacked_packet_map.h b/quic/core/quic_unacked_packet_map.h
index b8f8ffe..885dfb9 100644
--- a/quic/core/quic_unacked_packet_map.h
+++ b/quic/core/quic_unacked_packet_map.h
@@ -60,7 +60,8 @@
                         TransmissionType type);
 
   // Notifies session_notifier to retransmit frames with |transmission_type|.
-  void RetransmitFrames(const QuicFrames& frames, TransmissionType type);
+  // Returns true if all data gets retransmitted.
+  bool RetransmitFrames(const QuicFrames& frames, TransmissionType type);
 
   // Marks |info| as no longer in flight.
   void RemoveFromInFlight(QuicTransmissionInfo* info);
diff --git a/quic/core/session_notifier_interface.h b/quic/core/session_notifier_interface.h
index fc16b63..a3c29ed 100644
--- a/quic/core/session_notifier_interface.h
+++ b/quic/core/session_notifier_interface.h
@@ -28,8 +28,9 @@
   // Called when |frame| is considered as lost.
   virtual void OnFrameLost(const QuicFrame& frame) = 0;
 
-  // Called to retransmit |frames| with transmission |type|.
-  virtual void RetransmitFrames(const QuicFrames& frames,
+  // Called to retransmit |frames| with transmission |type|. Returns true if all
+  // data gets retransmitted.
+  virtual bool RetransmitFrames(const QuicFrames& frames,
                                 TransmissionType type) = 0;
 
   // Returns true if |frame| is outstanding and waiting to be acked.
diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index 705f951..92aa72f 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h
@@ -1388,7 +1388,7 @@
   MOCK_METHOD(void, OnStreamFrameRetransmitted, (const QuicStreamFrame&),
               (override));
   MOCK_METHOD(void, OnFrameLost, (const QuicFrame&), (override));
-  MOCK_METHOD(void, RetransmitFrames,
+  MOCK_METHOD(bool, RetransmitFrames,
               (const QuicFrames&, TransmissionType type), (override));
   MOCK_METHOD(bool, IsFrameOutstanding, (const QuicFrame&), (const, override));
   MOCK_METHOD(bool, HasUnackedCryptoData, (), (const, override));
diff --git a/quic/test_tools/simple_session_notifier.cc b/quic/test_tools/simple_session_notifier.cc
index ce498ca..4eca1f7 100644
--- a/quic/test_tools/simple_session_notifier.cc
+++ b/quic/test_tools/simple_session_notifier.cc
@@ -204,6 +204,9 @@
     const bool can_bundle_fin =
         state.fin_buffered && (state.bytes_sent + length == state.bytes_total);
     connection_->SetTransmissionType(NOT_RETRANSMISSION);
+    QuicConnection::ScopedEncryptionLevelContext context(
+        connection_,
+        connection_->framer().GetEncryptionLevelToSendApplicationData());
     QuicConsumedData consumed = connection_->SendStreamData(
         pair.first, length, state.bytes_sent, can_bundle_fin ? FIN : NO_FIN);
     QUIC_DVLOG(1) << "Tries to write stream_id: " << pair.first << " ["
@@ -333,7 +336,7 @@
   state->fin_lost = fin_lost;
 }
 
-void SimpleSessionNotifier::RetransmitFrames(const QuicFrames& frames,
+bool SimpleSessionNotifier::RetransmitFrames(const QuicFrames& frames,
                                              TransmissionType type) {
   QuicConnection::ScopedPacketFlusher retransmission_flusher(connection_);
   connection_->SetTransmissionType(type);
@@ -353,7 +356,7 @@
         size_t consumed = connection_->SendCryptoData(frame.crypto_frame->level,
                                                       length, offset);
         if (consumed < length) {
-          break;
+          return false;
         }
       }
       connection_->SetDefaultEncryptionLevel(current_encryption_level);
@@ -366,7 +369,7 @@
       if (!connection_->SendControlFrame(copy)) {
         // Connection is write blocked.
         DeleteFrame(&copy);
-        return;
+        return false;
       }
       continue;
     }
@@ -379,7 +382,6 @@
         frame.stream_frame.offset + frame.stream_frame.data_length);
     EncryptionLevel retransmission_encryption_level =
         connection_->encryption_level();
-    EncryptionLevel current_encryption_level = connection_->encryption_level();
     if (QuicUtils::IsCryptoStreamId(connection_->transport_version(),
                                     frame.stream_frame.stream_id)) {
       for (size_t i = 0; i < NUM_ENCRYPTION_LEVELS; ++i) {
@@ -399,11 +401,13 @@
       const bool can_bundle_fin =
           retransmit_fin &&
           (retransmission_offset + retransmission_length == state.bytes_sent);
-      if (QuicUtils::IsCryptoStreamId(connection_->transport_version(),
-                                      frame.stream_frame.stream_id)) {
-        // Set appropriate encryption level for crypto stream.
-        connection_->SetDefaultEncryptionLevel(retransmission_encryption_level);
-      }
+      QuicConnection::ScopedEncryptionLevelContext context(
+          connection_,
+          QuicUtils::IsCryptoStreamId(connection_->transport_version(),
+                                      frame.stream_frame.stream_id)
+              ? retransmission_encryption_level
+              : connection_->framer()
+                    .GetEncryptionLevelToSendApplicationData());
       consumed = connection_->SendStreamData(
           frame.stream_frame.stream_id, retransmission_length,
           retransmission_offset, can_bundle_fin ? FIN : NO_FIN);
@@ -416,15 +420,10 @@
       if (can_bundle_fin) {
         retransmit_fin = !consumed.fin_consumed;
       }
-      if (QuicUtils::IsCryptoStreamId(connection_->transport_version(),
-                                      frame.stream_frame.stream_id)) {
-        // Restore encryption level.
-        connection_->SetDefaultEncryptionLevel(current_encryption_level);
-      }
       if (consumed.bytes_consumed < retransmission_length ||
           (can_bundle_fin && !consumed.fin_consumed)) {
         // Connection is write blocked.
-        return;
+        return false;
       }
     }
     if (retransmit_fin) {
@@ -432,8 +431,12 @@
                     << " retransmits fin only frame.";
       consumed = connection_->SendStreamData(frame.stream_frame.stream_id, 0,
                                              state.bytes_sent, FIN);
+      if (!consumed.fin_consumed) {
+        return false;
+      }
     }
   }
+  return true;
 }
 
 bool SimpleSessionNotifier::IsFrameOutstanding(const QuicFrame& frame) const {
diff --git a/quic/test_tools/simple_session_notifier.h b/quic/test_tools/simple_session_notifier.h
index 5712884..0d82c70 100644
--- a/quic/test_tools/simple_session_notifier.h
+++ b/quic/test_tools/simple_session_notifier.h
@@ -83,7 +83,7 @@
                     QuicTime receive_timestamp) override;
   void OnStreamFrameRetransmitted(const QuicStreamFrame& /*frame*/) override {}
   void OnFrameLost(const QuicFrame& frame) override;
-  void RetransmitFrames(const QuicFrames& frames,
+  bool RetransmitFrames(const QuicFrames& frames,
                         TransmissionType type) override;
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;
diff --git a/quic/test_tools/simple_session_notifier_test.cc b/quic/test_tools/simple_session_notifier_test.cc
index a88cb80..552ff5b 100644
--- a/quic/test_tools/simple_session_notifier_test.cc
+++ b/quic/test_tools/simple_session_notifier_test.cc
@@ -317,6 +317,9 @@
 
 TEST_F(SimpleSessionNotifierTest, RetransmitFrames) {
   InSequence s;
+  connection_.SetEncrypter(
+      ENCRYPTION_FORWARD_SECURE,
+      std::make_unique<NullEncrypter>(Perspective::IS_CLIENT));
   // Send stream 3 data [0, 10) and fin.
   EXPECT_CALL(connection_, SendStreamData(3, 10, 0, FIN))
       .WillOnce(Return(QuicConsumedData(10, true)));
diff --git a/quic/test_tools/simulator/quic_endpoint.cc b/quic/test_tools/simulator/quic_endpoint.cc
index fd5c5b7..ee3c089 100644
--- a/quic/test_tools/simulator/quic_endpoint.cc
+++ b/quic/test_tools/simulator/quic_endpoint.cc
@@ -189,10 +189,10 @@
   notifier_->OnFrameLost(frame);
 }
 
-void QuicEndpoint::RetransmitFrames(const QuicFrames& frames,
+bool QuicEndpoint::RetransmitFrames(const QuicFrames& frames,
                                     TransmissionType type) {
   QUICHE_DCHECK(notifier_);
-  notifier_->RetransmitFrames(frames, type);
+  return notifier_->RetransmitFrames(frames, type);
 }
 
 bool QuicEndpoint::IsFrameOutstanding(const QuicFrame& frame) const {
diff --git a/quic/test_tools/simulator/quic_endpoint.h b/quic/test_tools/simulator/quic_endpoint.h
index 3e445dc..7fc4d87 100644
--- a/quic/test_tools/simulator/quic_endpoint.h
+++ b/quic/test_tools/simulator/quic_endpoint.h
@@ -125,7 +125,7 @@
                     QuicTime receive_timestamp) override;
   void OnStreamFrameRetransmitted(const QuicStreamFrame& /*frame*/) override {}
   void OnFrameLost(const QuicFrame& frame) override;
-  void RetransmitFrames(const QuicFrames& frames,
+  bool RetransmitFrames(const QuicFrames& frames,
                         TransmissionType type) override;
   bool IsFrameOutstanding(const QuicFrame& frame) const override;
   bool HasUnackedCryptoData() const override;