From e2dc6df8a290f13c224f5f09c46cb397e8ade3de Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Wed, 28 Jul 2021 14:04:34 -0700
Subject: [PATCH 1/2] Use unordered_map to store sessions for performance
 reasons

---
 src/session/Key.cc     | 10 ++++++++++
 src/session/Key.h      | 14 ++++++++++++++
 src/session/Manager.cc | 12 ++++++++++--
 src/session/Manager.h  |  4 ++--
 4 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/session/Key.cc b/src/session/Key.cc
index 5c7317505c..f87acc76b4 100644
--- a/src/session/Key.cc
+++ b/src/session/Key.cc
@@ -70,4 +70,14 @@ bool Key::operator<(const Key& rhs) const
 	return memcmp(data, rhs.data, size) < 0;
 	}
 
+bool Key::operator==(const Key& rhs) const
+	{
+	if ( size != rhs.size )
+		return false;
+	else if ( type != rhs.type )
+		return false;
+
+	return memcmp(data, rhs.data, size) == 0;
+	}
+
 } // namespace zeek::session::detail
diff --git a/src/session/Key.h b/src/session/Key.h
index 7201fbcaa5..ca37d31088 100644
--- a/src/session/Key.h
+++ b/src/session/Key.h
@@ -4,9 +4,12 @@
 
 #include <cstddef>
 #include <cstdint>
+#include "zeek/Hash.h"
 
 namespace zeek::session::detail {
 
+struct KeyHash;
+
 /**
  * This type is used as the key for the map in SessionManager. It represents a
  * raw block of memory that points to a key of some type for a session, such as
@@ -57,12 +60,23 @@ public:
 	void CopyData();
 
 	bool operator<(const Key& rhs) const;
+	bool operator==(const Key& rhs) const;
+
+	std::size_t Hash() const {
+		return zeek::detail::HashKey::HashBytes(data, size);
+	}
 
 private:
+	friend struct KeyHash;
+
 	const uint8_t* data = nullptr;
 	size_t size = 0;
 	size_t type = CONNECTION_KEY_TYPE;
 	bool copied = false;
 };
 
+struct KeyHash {
+	std::size_t operator()(const Key& k) const { return k.Hash(); }
+};
+
 } // namespace zeek::session::detail
diff --git a/src/session/Manager.cc b/src/session/Manager.cc
index 527adfed74..1792a212ed 100644
--- a/src/session/Manager.cc
+++ b/src/session/Manager.cc
@@ -218,9 +218,17 @@ void Manager::Insert(Session* s, bool remove_existing)
 
 void Manager::Drain()
 	{
-	for ( const auto& entry : session_map )
+	std::vector<const detail::Key*> keys;
+	keys.reserve(session_map.size());
+
+	for ( auto& entry : session_map )
+		keys.push_back(&(entry.first));
+	std::sort(keys.begin(), keys.end(), [](const detail::Key* a, const detail::Key* b) {
+		                                    return *a < *b; });
+
+	for ( const auto* k : keys )
 		{
-		Session* tc = entry.second;
+		Session* tc = session_map.at(*k);
 		tc->Done();
 		tc->RemovalEvent();
 		}
diff --git a/src/session/Manager.h b/src/session/Manager.h
index 976cc7f10d..e824e0b14b 100644
--- a/src/session/Manager.h
+++ b/src/session/Manager.h
@@ -3,7 +3,7 @@
 #pragma once
 
 #include <sys/types.h> // for u_char
-#include <map>
+#include <unordered_map>
 #include <utility>
 
 #include "zeek/Frag.h"
@@ -119,7 +119,7 @@ public:
 
 private:
 
-	using SessionMap = std::map<detail::Key, Session*>;
+	using SessionMap = std::unordered_map<detail::Key, Session*, detail::KeyHash>;
 
 	// Inserts a new connection into the sessions map. If a connection with
 	// the same key already exists in the map, it will be overwritten by

From 201617540ddd14a8d0d93499336dacfb402b11c7 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Tue, 3 Aug 2021 10:07:47 -0700
Subject: [PATCH 2/2] Only sort sessions during Drain() if a random seed is set

---
 src/session/Manager.cc | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/session/Manager.cc b/src/session/Manager.cc
index 1792a212ed..6fdf275004 100644
--- a/src/session/Manager.cc
+++ b/src/session/Manager.cc
@@ -218,19 +218,34 @@ void Manager::Insert(Session* s, bool remove_existing)
 
 void Manager::Drain()
 	{
-	std::vector<const detail::Key*> keys;
-	keys.reserve(session_map.size());
-
-	for ( auto& entry : session_map )
-		keys.push_back(&(entry.first));
-	std::sort(keys.begin(), keys.end(), [](const detail::Key* a, const detail::Key* b) {
-		                                    return *a < *b; });
-
-	for ( const auto* k : keys )
+	// If a random seed was passed in, we're most likely in testing mode and need the
+	// order of the sessions to be consistent. Sort the keys to force that order
+	// every run.
+	if ( zeek::util::detail::have_random_seed() )
 		{
-		Session* tc = session_map.at(*k);
-		tc->Done();
-		tc->RemovalEvent();
+		std::vector<const detail::Key*> keys;
+		keys.reserve(session_map.size());
+
+		for ( auto& entry : session_map )
+			keys.push_back(&(entry.first));
+		std::sort(keys.begin(), keys.end(), [](const detail::Key* a, const detail::Key* b) {
+			                                    return *a < *b; });
+
+		for ( const auto* k : keys )
+			{
+			Session* tc = session_map.at(*k);
+			tc->Done();
+			tc->RemovalEvent();
+			}
+		}
+	else
+		{
+		for ( const auto& entry : session_map )
+			{
+			Session* tc = entry.second;
+			tc->Done();
+			tc->RemovalEvent();
+			}
 		}
 	}