Add type alias for common MessageIdMap

shotover · Feb 19, 2024 · 370669e · 370669e
1 parent 8629f70
commit 370669e
Show file tree

Hide file tree

Showing 7 changed files with 16 additions and 13 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/shotover/Cargo.toml b/shotover/Cargo.toml
@@ -119,6 +119,7 @@ string = { version = "0.3.0", optional = true }
 xxhash-rust = { version = "0.8.6", features = ["xxh3"], optional = true }
 dashmap = { version = "5.4.0", optional = true }
 atoi = { version = "2.0.0", optional = true }
+fnv = "1.0.7"
 
 [dev-dependencies]
 criterion = { version = "0.5.0", features = ["async_tokio"] }

diff --git a/shotover/src/message/mod.rs b/shotover/src/message/mod.rs
@@ -13,11 +13,15 @@ use bytes::Bytes;
 #[cfg(feature = "cassandra")]
 use cassandra_protocol::compression::Compression;
 use derivative::Derivative;
+use fnv::FnvBuildHasher;
 use nonzero_ext::nonzero;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use std::num::NonZeroU32;
 use std::time::Instant;
 
+pub type MessageIdMap<T> = HashMap<MessageId, T, FnvBuildHasher>;
+
 pub enum Metadata {
     #[cfg(feature = "cassandra")]
     Cassandra(CassandraMetadata),

diff --git a/shotover/src/transforms/cassandra/sink_cluster/mod.rs b/shotover/src/transforms/cassandra/sink_cluster/mod.rs
@@ -2,7 +2,7 @@ use self::node_pool::{get_accessible_owned_connection, NodePoolBuilder, Prepared
 use self::rewrite::MessageRewriter;
 use crate::frame::cassandra::{CassandraMetadata, Tracing};
 use crate::frame::{CassandraFrame, CassandraOperation, CassandraResult, Frame};
-use crate::message::{Message, Messages, Metadata};
+use crate::message::{Message, MessageIdMap, Messages, Metadata};
 use crate::tls::{TlsConnector, TlsConnectorConfig};
 use crate::transforms::cassandra::connection::{CassandraConnection, Response, ResponseError};
 use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
@@ -137,7 +137,7 @@ impl CassandraSinkClusterBuilder {
             shotover_peers,
             local_shotover_node,
             to_rewrite: vec![],
-            prepare_requests_to_destination_nodes: HashMap::new(),
+            prepare_requests_to_destination_nodes: MessageIdMap::default(),
         };
 
         Self {

diff --git a/shotover/src/transforms/cassandra/sink_cluster/rewrite.rs b/shotover/src/transforms/cassandra/sink_cluster/rewrite.rs
@@ -2,7 +2,7 @@ use super::node::ConnectionFactory;
 use super::node_pool::NodePool;
 use super::ShotoverNode;
 use crate::frame::{CassandraFrame, CassandraOperation, CassandraResult, Frame};
-use crate::message::{Message, Messages};
+use crate::message::{Message, MessageIdMap, Messages};
 use crate::{
     frame::{
         cassandra::{parse_statement_single, Tracing},
@@ -20,7 +20,6 @@ use cql3_parser::common::{
 use cql3_parser::select::{Select, SelectElement};
 use futures::future::try_join_all;
 use itertools::Itertools;
-use std::collections::HashMap;
 use std::fmt::Write;
 use std::net::{IpAddr, Ipv4Addr};
 use uuid::Uuid;
@@ -57,7 +56,7 @@ pub struct MessageRewriter {
     pub shotover_peers: Vec<ShotoverNode>,
     pub local_shotover_node: ShotoverNode,
     pub to_rewrite: Vec<TableToRewrite>,
-    pub prepare_requests_to_destination_nodes: HashMap<MessageId, Uuid>,
+    pub prepare_requests_to_destination_nodes: MessageIdMap<Uuid>,
 }
 
 impl MessageRewriter {

diff --git a/shotover/src/transforms/mod.rs b/shotover/src/transforms/mod.rs
@@ -1,11 +1,11 @@
 //! Various types required for defining a transform
 
-use crate::message::{Message, MessageId, Messages};
+use self::chain::TransformAndMetrics;
+use crate::message::{Message, MessageIdMap, Messages};
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use core::fmt;
 use futures::Future;
-use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::iter::Rev;
 use std::net::SocketAddr;
@@ -14,8 +14,6 @@ use std::slice::IterMut;
 use tokio::sync::mpsc;
 use tokio::time::Instant;
 
-use self::chain::TransformAndMetrics;
-
 #[cfg(feature = "cassandra")]
 pub mod cassandra;
 pub mod chain;
@@ -188,7 +186,7 @@ impl<'a> Wrapper<'a> {
         result
     }
 
-    pub fn clone_requests_into_hashmap(&self, destination: &mut HashMap<MessageId, Message>) {
+    pub fn clone_requests_into_hashmap(&self, destination: &mut MessageIdMap<Message>) {
         for request in &self.requests {
             destination.insert(request.id(), request.clone());
         }

diff --git a/shotover/src/transforms/protect/mod.rs b/shotover/src/transforms/protect/mod.rs
@@ -1,7 +1,7 @@
 use crate::frame::{
     value::GenericValue, CassandraFrame, CassandraOperation, CassandraResult, Frame,
 };
-use crate::message::{Message, MessageId, Messages};
+use crate::message::{Message, MessageIdMap, Messages};
 use crate::transforms::protect::key_management::KeyManager;
 pub use crate::transforms::protect::key_management::KeyManagerConfig;
 use crate::transforms::{Transform, TransformBuilder, Wrapper};
@@ -56,7 +56,7 @@ impl TransformConfig for ProtectConfig {
                 .collect(),
             key_source: self.key_manager.build().await?,
             key_id: "XXXXXXX".to_string(),
-            requests: HashMap::new(),
+            requests: MessageIdMap::default(),
         }))
     }
 }
@@ -69,7 +69,7 @@ pub struct Protect {
     // TODO this should be a function to create key_ids based on "something", e.g. primary key
     // for the moment this is just a string
     key_id: String,
-    requests: HashMap<MessageId, Message>,
+    requests: MessageIdMap<Message>,
 }
 
 impl TransformBuilder for Protect {