diff --git a/quic/core/quic_dispatcher.cc b/quic/core/quic_dispatcher.cc
index 0b0707a..631d8d9 100644
--- a/quic/core/quic_dispatcher.cc
+++ b/quic/core/quic_dispatcher.cc
@@ -8,6 +8,7 @@
 #include <string>
 #include <utility>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/strings/string_view.h"
 #include "net/third_party/quiche/src/quic/core/chlo_extractor.h"
 #include "net/third_party/quiche/src/quic/core/crypto/crypto_protocol.h"
@@ -318,6 +319,9 @@
       expected_server_connection_id_length_(
           expected_server_connection_id_length),
       should_update_expected_server_connection_id_length_(false) {
+  if (use_reference_counted_session_map_) {
+    QUIC_RESTART_FLAG_COUNT(quic_use_reference_counted_sesssion_map);
+  }
   QUIC_BUG_IF(GetSupportedVersions().empty())
       << "Trying to create dispatcher without any supported versions";
   QUIC_DLOG(INFO) << "Created QuicDispatcher with versions: "
@@ -325,8 +329,13 @@
 }
 
 QuicDispatcher::~QuicDispatcher() {
-  session_map_.clear();
-  closed_session_list_.clear();
+  if (use_reference_counted_session_map_) {
+    reference_counted_session_map_.clear();
+    closed_ref_counted_session_list_.clear();
+  } else {
+    session_map_.clear();
+    closed_session_list_.clear();
+  }
 }
 
 void QuicDispatcher::InitializeWithWriter(QuicPacketWriter* writer) {
@@ -491,29 +500,58 @@
 
   // Packets with connection IDs for active connections are processed
   // immediately.
-  auto it = session_map_.find(server_connection_id);
-  if (it != session_map_.end()) {
-    DCHECK(!buffered_packets_.HasBufferedPackets(server_connection_id));
-    if (packet_info.version_flag &&
-        packet_info.version != it->second->version() &&
-        packet_info.version == LegacyVersionForEncapsulation()) {
-      // This packet is using the Legacy Version Encapsulation version but the
-      // corresponding session isn't, attempt extraction of inner packet.
-      ChloAlpnExtractor alpn_extractor;
-      if (ChloExtractor::Extract(packet_info.packet, packet_info.version,
-                                 config_->create_session_tag_indicators(),
-                                 &alpn_extractor,
-                                 server_connection_id.length())) {
-        if (MaybeHandleLegacyVersionEncapsulation(this, &alpn_extractor,
-                                                  packet_info)) {
-          return true;
+  if (use_reference_counted_session_map_) {
+    auto it = reference_counted_session_map_.find(server_connection_id);
+    if (it != reference_counted_session_map_.end()) {
+      DCHECK(!buffered_packets_.HasBufferedPackets(server_connection_id));
+      if (packet_info.version_flag &&
+          packet_info.version != it->second->version() &&
+          packet_info.version == LegacyVersionForEncapsulation()) {
+        // This packet is using the Legacy Version Encapsulation version but the
+        // corresponding session isn't, attempt extraction of inner packet.
+        ChloAlpnExtractor alpn_extractor;
+        if (ChloExtractor::Extract(packet_info.packet, packet_info.version,
+                                   config_->create_session_tag_indicators(),
+                                   &alpn_extractor,
+                                   server_connection_id.length())) {
+          if (MaybeHandleLegacyVersionEncapsulation(this, &alpn_extractor,
+                                                    packet_info)) {
+            return true;
+          }
         }
       }
+      it->second->ProcessUdpPacket(packet_info.self_address,
+                                   packet_info.peer_address,
+                                   packet_info.packet);
+      return true;
     }
-    it->second->ProcessUdpPacket(packet_info.self_address,
-                                 packet_info.peer_address, packet_info.packet);
-    return true;
-  } else if (packet_info.version.IsKnown()) {
+  } else {
+    auto it = session_map_.find(server_connection_id);
+    if (it != session_map_.end()) {
+      DCHECK(!buffered_packets_.HasBufferedPackets(server_connection_id));
+      if (packet_info.version_flag &&
+          packet_info.version != it->second->version() &&
+          packet_info.version == LegacyVersionForEncapsulation()) {
+        // This packet is using the Legacy Version Encapsulation version but the
+        // corresponding session isn't, attempt extraction of inner packet.
+        ChloAlpnExtractor alpn_extractor;
+        if (ChloExtractor::Extract(packet_info.packet, packet_info.version,
+                                   config_->create_session_tag_indicators(),
+                                   &alpn_extractor,
+                                   server_connection_id.length())) {
+          if (MaybeHandleLegacyVersionEncapsulation(this, &alpn_extractor,
+                                                    packet_info)) {
+            return true;
+          }
+        }
+      }
+      it->second->ProcessUdpPacket(packet_info.self_address,
+                                   packet_info.peer_address,
+                                   packet_info.packet);
+      return true;
+    }
+  }
+  if (packet_info.version.IsKnown()) {
     // We did not find the connection ID, check if we've replaced it.
     // This is only performed for supported versions because packets with
     // unsupported versions can flow through this function in order to send
@@ -523,14 +561,26 @@
     QuicConnectionId replaced_connection_id = MaybeReplaceServerConnectionId(
         server_connection_id, packet_info.version);
     if (replaced_connection_id != server_connection_id) {
-      // Search for the replacement.
-      auto it2 = session_map_.find(replaced_connection_id);
-      if (it2 != session_map_.end()) {
-        DCHECK(!buffered_packets_.HasBufferedPackets(replaced_connection_id));
-        it2->second->ProcessUdpPacket(packet_info.self_address,
-                                      packet_info.peer_address,
-                                      packet_info.packet);
-        return true;
+      if (use_reference_counted_session_map_) {
+        // Search for the replacement.
+        auto it2 = reference_counted_session_map_.find(replaced_connection_id);
+        if (it2 != reference_counted_session_map_.end()) {
+          DCHECK(!buffered_packets_.HasBufferedPackets(replaced_connection_id));
+          it2->second->ProcessUdpPacket(packet_info.self_address,
+                                        packet_info.peer_address,
+                                        packet_info.packet);
+          return true;
+        }
+      } else {
+        // Search for the replacement.
+        auto it2 = session_map_.find(replaced_connection_id);
+        if (it2 != session_map_.end()) {
+          DCHECK(!buffered_packets_.HasBufferedPackets(replaced_connection_id));
+          it2->second->ProcessUdpPacket(packet_info.self_address,
+                                        packet_info.peer_address,
+                                        packet_info.packet);
+          return true;
+        }
       }
     }
   }
@@ -797,26 +847,68 @@
 
 void QuicDispatcher::PerformActionOnActiveSessions(
     std::function<void(QuicSession*)> operation) const {
-  for (auto const& kv : session_map_) {
-    operation(kv.second.get());
+  if (use_reference_counted_session_map_) {
+    absl::flat_hash_set<QuicSession*> visited_session;
+    visited_session.reserve(reference_counted_session_map_.size());
+    for (auto const& kv : reference_counted_session_map_) {
+      QuicSession* session = kv.second.get();
+      if (visited_session.insert(session).second) {
+        operation(session);
+      }
+    }
+  } else {
+    for (auto const& kv : session_map_) {
+      operation(kv.second.get());
+    }
   }
 }
 
+// Get a snapshot of all sessions.
+std::vector<std::shared_ptr<QuicSession>> QuicDispatcher::GetSessionsSnapshot()
+    const {
+  DCHECK(use_reference_counted_session_map_);
+  std::vector<std::shared_ptr<QuicSession>> snapshot;
+  snapshot.reserve(reference_counted_session_map_.size());
+  absl::flat_hash_set<QuicSession*> visited_session;
+  visited_session.reserve(reference_counted_session_map_.size());
+  for (auto const& kv : reference_counted_session_map_) {
+    QuicSession* session = kv.second.get();
+    if (visited_session.insert(session).second) {
+      snapshot.push_back(kv.second);
+    }
+  }
+  return snapshot;
+}
+
 std::unique_ptr<QuicPerPacketContext> QuicDispatcher::GetPerPacketContext()
     const {
   return nullptr;
 }
 
 void QuicDispatcher::DeleteSessions() {
-  if (!write_blocked_list_.empty()) {
-    for (const std::unique_ptr<QuicSession>& session : closed_session_list_) {
-      if (write_blocked_list_.erase(session->connection()) != 0) {
-        QUIC_BUG << "QuicConnection was in WriteBlockedList before destruction "
-                 << session->connection()->connection_id();
+  if (use_reference_counted_session_map_) {
+    if (!write_blocked_list_.empty()) {
+      for (const auto& session : closed_ref_counted_session_list_) {
+        if (write_blocked_list_.erase(session->connection()) != 0) {
+          QUIC_BUG
+              << "QuicConnection was in WriteBlockedList before destruction "
+              << session->connection()->connection_id();
+        }
       }
     }
+    closed_ref_counted_session_list_.clear();
+  } else {
+    if (!write_blocked_list_.empty()) {
+      for (const std::unique_ptr<QuicSession>& session : closed_session_list_) {
+        if (write_blocked_list_.erase(session->connection()) != 0) {
+          QUIC_BUG
+              << "QuicConnection was in WriteBlockedList before destruction "
+              << session->connection()->connection_id();
+        }
+      }
+    }
+    closed_session_list_.clear();
   }
-  closed_session_list_.clear();
 }
 
 void QuicDispatcher::OnCanWrite() {
@@ -852,14 +944,27 @@
 }
 
 void QuicDispatcher::Shutdown() {
-  while (!session_map_.empty()) {
-    QuicSession* session = session_map_.begin()->second.get();
-    session->connection()->CloseConnection(
-        QUIC_PEER_GOING_AWAY, "Server shutdown imminent",
-        ConnectionCloseBehavior::SEND_CONNECTION_CLOSE_PACKET);
-    // Validate that the session removes itself from the session map on close.
-    DCHECK(session_map_.empty() ||
-           session_map_.begin()->second.get() != session);
+  if (use_reference_counted_session_map_) {
+    while (!reference_counted_session_map_.empty()) {
+      QuicSession* session =
+          reference_counted_session_map_.begin()->second.get();
+      session->connection()->CloseConnection(
+          QUIC_PEER_GOING_AWAY, "Server shutdown imminent",
+          ConnectionCloseBehavior::SEND_CONNECTION_CLOSE_PACKET);
+      // Validate that the session removes itself from the session map on close.
+      DCHECK(reference_counted_session_map_.empty() ||
+             reference_counted_session_map_.begin()->second.get() != session);
+    }
+  } else {
+    while (!session_map_.empty()) {
+      QuicSession* session = session_map_.begin()->second.get();
+      session->connection()->CloseConnection(
+          QUIC_PEER_GOING_AWAY, "Server shutdown imminent",
+          ConnectionCloseBehavior::SEND_CONNECTION_CLOSE_PACKET);
+      // Validate that the session removes itself from the session map on close.
+      DCHECK(session_map_.empty() ||
+             session_map_.begin()->second.get() != session);
+    }
   }
   DeleteSessions();
 }
@@ -868,32 +973,61 @@
                                         QuicErrorCode error,
                                         const std::string& error_details,
                                         ConnectionCloseSource source) {
-  auto it = session_map_.find(server_connection_id);
-  if (it == session_map_.end()) {
-    QUIC_BUG << "ConnectionId " << server_connection_id
-             << " does not exist in the session map.  Error: "
-             << QuicErrorCodeToString(error);
-    QUIC_BUG << QuicStackTrace();
-    return;
-  }
-
-  QUIC_DLOG_IF(INFO, error != QUIC_NO_ERROR)
-      << "Closing connection (" << server_connection_id
-      << ") due to error: " << QuicErrorCodeToString(error)
-      << ", with details: " << error_details;
-
-  QuicConnection* connection = it->second->connection();
-  if (ShouldDestroySessionAsynchronously()) {
-    // Set up alarm to fire immediately to bring destruction of this session
-    // out of current call stack.
-    if (closed_session_list_.empty()) {
-      delete_sessions_alarm_->Update(helper()->GetClock()->ApproximateNow(),
-                                     QuicTime::Delta::Zero());
+  if (use_reference_counted_session_map_) {
+    auto it = reference_counted_session_map_.find(server_connection_id);
+    if (it == reference_counted_session_map_.end()) {
+      QUIC_BUG << "ConnectionId " << server_connection_id
+               << " does not exist in the session map.  Error: "
+               << QuicErrorCodeToString(error);
+      QUIC_BUG << QuicStackTrace();
+      return;
     }
-    closed_session_list_.push_back(std::move(it->second));
+
+    QUIC_DLOG_IF(INFO, error != QUIC_NO_ERROR)
+        << "Closing connection (" << server_connection_id
+        << ") due to error: " << QuicErrorCodeToString(error)
+        << ", with details: " << error_details;
+
+    QuicConnection* connection = it->second->connection();
+    if (ShouldDestroySessionAsynchronously()) {
+      // Set up alarm to fire immediately to bring destruction of this session
+      // out of current call stack.
+      if (closed_ref_counted_session_list_.empty()) {
+        delete_sessions_alarm_->Update(helper()->GetClock()->ApproximateNow(),
+                                       QuicTime::Delta::Zero());
+      }
+      closed_ref_counted_session_list_.push_back(std::move(it->second));
+    }
+    CleanUpSession(it->first, connection, source);
+    reference_counted_session_map_.erase(it);
+  } else {
+    auto it = session_map_.find(server_connection_id);
+    if (it == session_map_.end()) {
+      QUIC_BUG << "ConnectionId " << server_connection_id
+               << " does not exist in the session map.  Error: "
+               << QuicErrorCodeToString(error);
+      QUIC_BUG << QuicStackTrace();
+      return;
+    }
+
+    QUIC_DLOG_IF(INFO, error != QUIC_NO_ERROR)
+        << "Closing connection (" << server_connection_id
+        << ") due to error: " << QuicErrorCodeToString(error)
+        << ", with details: " << error_details;
+
+    QuicConnection* connection = it->second->connection();
+    if (ShouldDestroySessionAsynchronously()) {
+      // Set up alarm to fire immediately to bring destruction of this session
+      // out of current call stack.
+      if (closed_session_list_.empty()) {
+        delete_sessions_alarm_->Update(helper()->GetClock()->ApproximateNow(),
+                                       QuicTime::Delta::Zero());
+      }
+      closed_session_list_.push_back(std::move(it->second));
+    }
+    CleanUpSession(it->first, connection, source);
+    session_map_.erase(it);
   }
-  CleanUpSession(it->first, connection, source);
-  session_map_.erase(it);
 }
 
 void QuicDispatcher::OnWriteBlocked(
@@ -1023,12 +1157,24 @@
     }
     QUIC_DLOG(INFO) << "Created new session for " << server_connection_id;
 
-    auto insertion_result = session_map_.insert(
-        std::make_pair(server_connection_id, std::move(session)));
-    QUIC_BUG_IF(!insertion_result.second)
-        << "Tried to add a session to session_map with existing connection id: "
-        << server_connection_id;
-    DeliverPacketsToSession(packets, insertion_result.first->second.get());
+    if (use_reference_counted_session_map_) {
+      auto insertion_result = reference_counted_session_map_.insert(
+          std::make_pair(server_connection_id,
+                         std::shared_ptr<QuicSession>(std::move(session))));
+      QUIC_BUG_IF(!insertion_result.second)
+          << "Tried to add a session to session_map with existing connection "
+             "id: "
+          << server_connection_id;
+      DeliverPacketsToSession(packets, insertion_result.first->second.get());
+    } else {
+      auto insertion_result = session_map_.insert(
+          std::make_pair(server_connection_id, std::move(session)));
+      QUIC_BUG_IF(!insertion_result.second)
+          << "Tried to add a session to session_map with existing connection "
+             "id: "
+          << server_connection_id;
+      DeliverPacketsToSession(packets, insertion_result.first->second.get());
+    }
   }
 }
 
@@ -1124,12 +1270,24 @@
   QUIC_DLOG(INFO) << "Created new session for "
                   << packet_info->destination_connection_id;
 
-  auto insertion_result = session_map_.insert(std::make_pair(
-      packet_info->destination_connection_id, std::move(session)));
-  QUIC_BUG_IF(!insertion_result.second)
-      << "Tried to add a session to session_map with existing connection id: "
-      << packet_info->destination_connection_id;
-  QuicSession* session_ptr = insertion_result.first->second.get();
+  QuicSession* session_ptr;
+  if (use_reference_counted_session_map_) {
+    auto insertion_result =
+        reference_counted_session_map_.insert(std::make_pair(
+            packet_info->destination_connection_id,
+            std::shared_ptr<QuicSession>(std::move(session.release()))));
+    QUIC_BUG_IF(!insertion_result.second)
+        << "Tried to add a session to session_map with existing connection id: "
+        << packet_info->destination_connection_id;
+    session_ptr = insertion_result.first->second.get();
+  } else {
+    auto insertion_result = session_map_.insert(std::make_pair(
+        packet_info->destination_connection_id, std::move(session)));
+    QUIC_BUG_IF(!insertion_result.second)
+        << "Tried to add a session to session_map with existing connection id: "
+        << packet_info->destination_connection_id;
+    session_ptr = insertion_result.first->second.get();
+  }
   std::list<BufferedPacket> packets =
       buffered_packets_.DeliverPackets(packet_info->destination_connection_id)
           .buffered_packets;
diff --git a/quic/core/quic_dispatcher.h b/quic/core/quic_dispatcher.h
index 58d7932..0a920ff 100644
--- a/quic/core/quic_dispatcher.h
+++ b/quic/core/quic_dispatcher.h
@@ -27,6 +27,7 @@
 #include "net/third_party/quiche/src/quic/core/quic_time_wait_list_manager.h"
 #include "net/third_party/quiche/src/quic/core/quic_version_manager.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_containers.h"
+#include "net/third_party/quiche/src/quic/platform/api/quic_reference_counted.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_socket_address.h"
 
 namespace quic {
@@ -109,7 +110,17 @@
                                  std::unique_ptr<QuicSession>,
                                  QuicConnectionIdHash>;
 
-  size_t NumSessions() const { return session_map_.size(); }
+  using ReferenceCountedSessionMap = QuicHashMap<QuicConnectionId,
+                                                 std::shared_ptr<QuicSession>,
+                                                 QuicConnectionIdHash>;
+
+  // TODO(haoyuewang) Update this function when multiple CIDs per connection are
+  // supported.
+  size_t NumSessions() const {
+    return use_reference_counted_session_map_
+               ? reference_counted_session_map_.size()
+               : session_map_.size();
+  }
 
   const SessionMap& session_map() const { return session_map_; }
 
@@ -141,8 +152,15 @@
   void PerformActionOnActiveSessions(
       std::function<void(QuicSession*)> operation) const;
 
+  // Get a snapshot of all sessions.
+  std::vector<std::shared_ptr<QuicSession>> GetSessionsSnapshot() const;
+
   bool accept_new_connections() const { return accept_new_connections_; }
 
+  bool use_reference_counted_session_map() const {
+    return use_reference_counted_session_map_;
+  }
+
  protected:
   virtual std::unique_ptr<QuicSession> CreateQuicSession(
       QuicConnectionId server_connection_id,
@@ -359,12 +377,14 @@
   WriteBlockedList write_blocked_list_;
 
   SessionMap session_map_;
+  ReferenceCountedSessionMap reference_counted_session_map_;
 
   // Entity that manages connection_ids in time wait state.
   std::unique_ptr<QuicTimeWaitListManager> time_wait_list_manager_;
 
   // The list of closed but not-yet-deleted sessions.
   std::vector<std::unique_ptr<QuicSession>> closed_session_list_;
+  std::vector<std::shared_ptr<QuicSession>> closed_ref_counted_session_list_;
 
   // The helper used for all connections.
   std::unique_ptr<QuicConnectionHelperInterface> helper_;
@@ -416,6 +436,9 @@
   // If true, change expected_server_connection_id_length_ to be the received
   // destination connection ID length of all IETF long headers.
   bool should_update_expected_server_connection_id_length_;
+
+  const bool use_reference_counted_session_map_ =
+      GetQuicRestartFlag(quic_use_reference_counted_sesssion_map);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 912455f..8e99f37 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -73,3 +73,4 @@
 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_use_reference_counted_sesssion_map, false)
diff --git a/quic/test_tools/quic_dispatcher_peer.cc b/quic/test_tools/quic_dispatcher_peer.cc
index 072b801..be428ed 100644
--- a/quic/test_tools/quic_dispatcher_peer.cc
+++ b/quic/test_tools/quic_dispatcher_peer.cc
@@ -116,10 +116,17 @@
 // static
 QuicSession* QuicDispatcherPeer::GetFirstSessionIfAny(
     QuicDispatcher* dispatcher) {
-  if (dispatcher->session_map_.empty()) {
-    return nullptr;
+  if (dispatcher->use_reference_counted_session_map()) {
+    if (dispatcher->reference_counted_session_map_.empty()) {
+      return nullptr;
+    }
+    return dispatcher->reference_counted_session_map_.begin()->second.get();
+  } else {
+    if (dispatcher->session_map_.empty()) {
+      return nullptr;
+    }
+    return dispatcher->session_map_.begin()->second.get();
   }
-  return dispatcher->session_map_.begin()->second.get();
 }
 
 }  // namespace test