Use unordered_map to store sessions for performance reasons

This commit is contained in:
Tim Wojtulewicz 2021-07-28 14:04:34 -07:00
parent 2a717e05cc
commit e2dc6df8a2
4 changed files with 36 additions and 4 deletions

View file

@ -70,4 +70,14 @@ bool Key::operator<(const Key& rhs) const
return memcmp(data, rhs.data, size) < 0; 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 } // namespace zeek::session::detail

View file

@ -4,9 +4,12 @@
#include <cstddef> #include <cstddef>
#include <cstdint> #include <cstdint>
#include "zeek/Hash.h"
namespace zeek::session::detail { namespace zeek::session::detail {
struct KeyHash;
/** /**
* This type is used as the key for the map in SessionManager. It represents a * 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 * raw block of memory that points to a key of some type for a session, such as
@ -57,12 +60,23 @@ public:
void CopyData(); void CopyData();
bool operator<(const Key& rhs) const; 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: private:
friend struct KeyHash;
const uint8_t* data = nullptr; const uint8_t* data = nullptr;
size_t size = 0; size_t size = 0;
size_t type = CONNECTION_KEY_TYPE; size_t type = CONNECTION_KEY_TYPE;
bool copied = false; bool copied = false;
}; };
struct KeyHash {
std::size_t operator()(const Key& k) const { return k.Hash(); }
};
} // namespace zeek::session::detail } // namespace zeek::session::detail

View file

@ -218,9 +218,17 @@ void Manager::Insert(Session* s, bool remove_existing)
void Manager::Drain() 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->Done();
tc->RemovalEvent(); tc->RemovalEvent();
} }

View file

@ -3,7 +3,7 @@
#pragma once #pragma once
#include <sys/types.h> // for u_char #include <sys/types.h> // for u_char
#include <map> #include <unordered_map>
#include <utility> #include <utility>
#include "zeek/Frag.h" #include "zeek/Frag.h"
@ -119,7 +119,7 @@ public:
private: 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 // 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 // the same key already exists in the map, it will be overwritten by