diff --git a/quic/core/quic_dispatcher.cc b/quic/core/quic_dispatcher.cc
index 3e0846c..8439f1e 100644
--- a/quic/core/quic_dispatcher.cc
+++ b/quic/core/quic_dispatcher.cc
@@ -738,7 +738,7 @@
   return kFateProcess;
 }
 
-void QuicDispatcher::CleanUpSession(SessionMap::iterator it,
+void QuicDispatcher::CleanUpSession(QuicConnectionId server_connection_id,
                                     QuicConnection* connection,
                                     ConnectionCloseSource /*source*/) {
   write_blocked_list_.erase(connection);
@@ -770,18 +770,16 @@
           // |action| argument is not used by this call to
           // StatelesslyTerminateConnection().
           action);
-      session_map_.erase(it);
       return;
     }
     QUIC_CODE_COUNT(quic_v44_add_to_time_wait_list_with_stateless_reset);
   }
   time_wait_list_manager_->AddConnectionIdToTimeWait(
-      it->first, action,
+      server_connection_id, action,
       TimeWaitConnectionInfo(
           connection->version().HasIetfInvariantHeader(),
           connection->termination_packets(),
           connection->sent_packet_manager().GetRttStats()->smoothed_rtt()));
-  session_map_.erase(it);
 }
 
 void QuicDispatcher::StartAcceptingNewConnections() {
@@ -885,7 +883,8 @@
     }
     closed_session_list_.push_back(std::move(it->second));
   }
-  CleanUpSession(it, connection, source);
+  CleanUpSession(it->first, connection, source);
+  session_map_.erase(it);
 }
 
 void QuicDispatcher::OnWriteBlocked(
diff --git a/quic/core/quic_dispatcher.h b/quic/core/quic_dispatcher.h
index ca83a79..d2711e0 100644
--- a/quic/core/quic_dispatcher.h
+++ b/quic/core/quic_dispatcher.h
@@ -261,11 +261,12 @@
       QuicBufferedPacketStore::EnqueuePacketResult result,
       QuicConnectionId server_connection_id);
 
-  // Removes the session from the session map and write blocked list, and adds
-  // the ConnectionId to the time-wait list.
-  virtual void CleanUpSession(SessionMap::iterator it,
-                              QuicConnection* connection,
-                              ConnectionCloseSource source);
+  // Removes the session from the write blocked list, and adds the ConnectionId
+  // to the time-wait list.  The caller needs to manually remove the session
+  // from the map after that.
+  void CleanUpSession(QuicConnectionId server_connection_id,
+                      QuicConnection* connection,
+                      ConnectionCloseSource source);
 
   // Called to terminate a connection statelessly. Depending on |format|, either
   // 1) send connection close with |error_code| and |error_details| and add