Remove connection_id added to connection_dispatcher_map_ via either OnConnectionIdRetired or OnConnectionAddedToTimeWaitList in gfe_quic_dispatcher when multiple CIDs are supported.

Protected by FLAGS_quic_restart_flag_quic_time_wait_list_support_multiple_cid_v2.

PiperOrigin-RevId: 351362306
Change-Id: I6c49f460cbe4c1ca219d6529caef596803f1c02d
diff --git a/quic/core/quic_dispatcher.cc b/quic/core/quic_dispatcher.cc
index 28f3923..a4b4099 100644
--- a/quic/core/quic_dispatcher.cc
+++ b/quic/core/quic_dispatcher.cc
@@ -328,7 +328,7 @@
   }
   if (support_multiple_cid_per_connection_) {
     QUIC_RESTART_FLAG_COUNT(
-        quic_dispatcher_support_multiple_cid_per_connection);
+        quic_dispatcher_support_multiple_cid_per_connection_v2);
   }
   QUIC_BUG_IF(GetSupportedVersions().empty())
       << "Trying to create dispatcher without any supported versions";
diff --git a/quic/core/quic_dispatcher.h b/quic/core/quic_dispatcher.h
index 75cdacf..14ca07c 100644
--- a/quic/core/quic_dispatcher.h
+++ b/quic/core/quic_dispatcher.h
@@ -455,8 +455,9 @@
       GetQuicRestartFlag(quic_use_reference_counted_sesssion_map);
   const bool support_multiple_cid_per_connection_ =
       use_reference_counted_session_map_ &&
-      GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid) &&
-      GetQuicRestartFlag(quic_dispatcher_support_multiple_cid_per_connection);
+      GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2) &&
+      GetQuicRestartFlag(
+          quic_dispatcher_support_multiple_cid_per_connection_v2);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_dispatcher_test.cc b/quic/core/quic_dispatcher_test.cc
index a614199..b801228 100644
--- a/quic/core/quic_dispatcher_test.cc
+++ b/quic/core/quic_dispatcher_test.cc
@@ -1974,8 +1974,8 @@
   QuicDispatcherSupportMultipleConnectionIdPerConnectionTest()
       : QuicDispatcherTestBase(crypto_test_utils::ProofSourceForTesting()) {
     SetQuicRestartFlag(quic_use_reference_counted_sesssion_map, true);
-    SetQuicRestartFlag(quic_time_wait_list_support_multiple_cid, true);
-    SetQuicRestartFlag(quic_dispatcher_support_multiple_cid_per_connection,
+    SetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2, true);
+    SetQuicRestartFlag(quic_dispatcher_support_multiple_cid_per_connection_v2,
                        true);
     dispatcher_ = std::make_unique<NiceMock<TestDispatcher>>(
         &config_, &crypto_config_, &version_manager_,
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 8a1a48a..b99ca34 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -70,7 +70,7 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_send_quic_fallback_server_config_on_leto_error, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_dont_fetch_quic_private_keys_from_leto, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_http2_parse_priority_update_frame, false)
-QUIC_FLAG(FLAGS_quic_restart_flag_quic_dispatcher_support_multiple_cid_per_connection, false)
+QUIC_FLAG(FLAGS_quic_restart_flag_quic_dispatcher_support_multiple_cid_per_connection_v2, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_enable_zero_rtt_for_tls_v2, true)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_offload_pacing_to_usps2, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_server_temporarily_retain_tls_zero_rtt_keys, true)
@@ -79,5 +79,5 @@
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_support_release_time_for_gso, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_testonly_default_false, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_testonly_default_true, true)
-QUIC_FLAG(FLAGS_quic_restart_flag_quic_time_wait_list_support_multiple_cid, false)
+QUIC_FLAG(FLAGS_quic_restart_flag_quic_time_wait_list_support_multiple_cid_v2, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_use_reference_counted_sesssion_map, false)
diff --git a/quic/core/quic_time_wait_list_manager.cc b/quic/core/quic_time_wait_list_manager.cc
index 11fca70..b59c5bc 100644
--- a/quic/core/quic_time_wait_list_manager.cc
+++ b/quic/core/quic_time_wait_list_manager.cc
@@ -82,7 +82,7 @@
       writer_(writer),
       visitor_(visitor) {
   if (use_indirect_connection_id_map_) {
-    QUIC_RESTART_FLAG_COUNT(quic_time_wait_list_support_multiple_cid);
+    QUIC_RESTART_FLAG_COUNT(quic_time_wait_list_support_multiple_cid_v2);
   }
   SetConnectionIdCleanUpAlarm();
 }
@@ -153,9 +153,14 @@
       GetQuicFlag(FLAGS_quic_time_wait_list_max_connections);
   DCHECK(connection_id_map_.empty() ||
          num_connections() < static_cast<size_t>(max_connections));
+  if (use_indirect_connection_id_map_ && new_connection_id) {
+    for (const auto& cid : info.active_connection_ids) {
+      visitor_->OnConnectionAddedToTimeWaitList(cid);
+    }
+  }
   AddConnectionIdDataToMap(canonical_connection_id, num_packets, action,
                            std::move(info));
-  if (new_connection_id) {
+  if (!use_indirect_connection_id_map_ && new_connection_id) {
     visitor_->OnConnectionAddedToTimeWaitList(canonical_connection_id);
   }
 }
diff --git a/quic/core/quic_time_wait_list_manager.h b/quic/core/quic_time_wait_list_manager.h
index ad6d45c..6dbd737 100644
--- a/quic/core/quic_time_wait_list_manager.h
+++ b/quic/core/quic_time_wait_list_manager.h
@@ -332,7 +332,7 @@
   // When this is default true, remove the connection_id argument of
   // AddConnectionIdToTimeWait.
   bool use_indirect_connection_id_map_ =
-      GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid);
+      GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_time_wait_list_manager_test.cc b/quic/core/quic_time_wait_list_manager_test.cc
index 5ece763..189d3f0 100644
--- a/quic/core/quic_time_wait_list_manager_test.cc
+++ b/quic/core/quic_time_wait_list_manager_test.cc
@@ -447,13 +447,14 @@
 
 TEST_F(QuicTimeWaitListManagerTest,
        CleanUpOldConnectionIdsForMultipleConnectionIdsPerConnection) {
-  if (!GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid)) {
+  if (!GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2)) {
     return;
   }
 
   connection_id_ = TestConnectionId(7);
   const size_t kConnectionCloseLength = 100;
   EXPECT_CALL(visitor_, OnConnectionAddedToTimeWaitList(connection_id_));
+  EXPECT_CALL(visitor_, OnConnectionAddedToTimeWaitList(TestConnectionId(8)));
   std::vector<std::unique_ptr<QuicEncryptedPacket>> termination_packets;
   termination_packets.push_back(
       std::unique_ptr<QuicEncryptedPacket>(new QuicEncryptedPacket(
@@ -711,13 +712,14 @@
 
 TEST_F(QuicTimeWaitListManagerTest,
        SendConnectionClosePacketsForMultipleConnectionIds) {
-  if (!GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid)) {
+  if (!GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2)) {
     return;
   }
 
   connection_id_ = TestConnectionId(7);
   const size_t kConnectionCloseLength = 100;
   EXPECT_CALL(visitor_, OnConnectionAddedToTimeWaitList(connection_id_));
+  EXPECT_CALL(visitor_, OnConnectionAddedToTimeWaitList(TestConnectionId(8)));
   std::vector<std::unique_ptr<QuicEncryptedPacket>> termination_packets;
   termination_packets.push_back(
       std::unique_ptr<QuicEncryptedPacket>(new QuicEncryptedPacket(