From e4dda5e541d4d55f4383a949f45131e37bd772cf Mon Sep 17 00:00:00 2001 From: Drazen Urch Date: Mon, 16 Oct 2023 15:40:26 +0200 Subject: [PATCH] Replace HashMap with DashMap (#3995) --- .../client_handling/client_registration.rs | 4 +-- .../src/node/http/api/v1/client_registry.rs | 31 +++++++++---------- gateway/src/node/http/mod.rs | 31 +++++++++---------- gateway/src/node/mod.rs | 9 +++--- 4 files changed, 36 insertions(+), 39 deletions(-) diff --git a/gateway/src/node/client_handling/client_registration.rs b/gateway/src/node/client_handling/client_registration.rs index b8fa5a1cbe..3dd97b7b45 100644 --- a/gateway/src/node/client_handling/client_registration.rs +++ b/gateway/src/node/client_handling/client_registration.rs @@ -1,5 +1,4 @@ use std::{ - collections::HashMap, fmt, hash::{Hash, Hasher}, net::SocketAddr, @@ -8,6 +7,7 @@ use std::{ }; use base64::{engine::general_purpose, Engine}; +use dashmap::DashMap; use hmac::{Hmac, Mac}; use nym_crypto::asymmetric::encryption::PrivateKey; use serde::{Deserialize, Serialize}; @@ -178,4 +178,4 @@ impl<'de> Deserialize<'de> for ClientPublicKey { } } -pub(crate) type ClientRegistry = HashMap; +pub(crate) type ClientRegistry = DashMap; diff --git a/gateway/src/node/http/api/v1/client_registry.rs b/gateway/src/node/http/api/v1/client_registry.rs index 8e0a63b03a..07750d57f1 100644 --- a/gateway/src/node/http/api/v1/client_registry.rs +++ b/gateway/src/node/http/api/v1/client_registry.rs @@ -14,8 +14,7 @@ use crate::node::http::ApiState; async fn process_final_message(client: Client, state: Arc) -> StatusCode { let preshared_nonce = { - let in_progress_ro = state.registration_in_progress.read().await; - if let Some(nonce) = in_progress_ro.get(client.pub_key()) { + if let Some(nonce) = state.registration_in_progress.get(client.pub_key()) { *nonce } else { return StatusCode::BAD_REQUEST; @@ -27,12 +26,10 @@ async fn process_final_message(client: Client, state: Arc) -> StatusCo .is_ok() { { - let mut in_progress_rw = state.registration_in_progress.write().await; - in_progress_rw.remove(client.pub_key()); + state.registration_in_progress.remove(client.pub_key()); } { - let mut registry_rw = state.client_registry.write().await; - registry_rw.insert(client.socket(), client); + state.client_registry.insert(client.socket(), client); } return StatusCode::OK; } @@ -42,8 +39,9 @@ async fn process_final_message(client: Client, state: Arc) -> StatusCo async fn process_init_message(init_message: InitMessage, state: Arc) -> u64 { let nonce: u64 = fastrand::u64(..); - let mut registry_rw = state.registration_in_progress.write().await; - registry_rw.insert(init_message.pub_key().clone(), nonce); + state + .registration_in_progress + .insert(init_message.pub_key().clone(), nonce); nonce } @@ -67,12 +65,12 @@ pub(crate) async fn register_client( pub(crate) async fn get_all_clients( State(state): State>, ) -> (StatusCode, Json>) { - let registry_ro = state.client_registry.read().await; ( StatusCode::OK, Json( - registry_ro - .values() + state + .client_registry + .iter() .map(|c| c.pub_key().clone()) .collect::>(), ), @@ -87,12 +85,13 @@ pub(crate) async fn get_client( Ok(pub_key) => pub_key, Err(_) => return (StatusCode::BAD_REQUEST, Json(vec![])), }; - let registry_ro = state.client_registry.read().await; - let clients = registry_ro + let clients = state + .client_registry .iter() - .filter_map(|(_, c)| { - if c.pub_key() == &pub_key { - Some(c.clone()) + .filter_map(|r| { + let client = r.value(); + if client.pub_key() == &pub_key { + Some(client.clone()) } else { None } diff --git a/gateway/src/node/http/mod.rs b/gateway/src/node/http/mod.rs index 9d091f232f..d9fcd8a331 100644 --- a/gateway/src/node/http/mod.rs +++ b/gateway/src/node/http/mod.rs @@ -1,12 +1,12 @@ -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; use axum::{ routing::{get, post}, Router, }; +use dashmap::DashMap; use log::info; use nym_crypto::asymmetric::encryption; -use tokio::sync::RwLock; mod api; use api::v1::client_registry::*; @@ -16,9 +16,9 @@ use super::client_handling::client_registration::{ClientPublicKey, ClientRegistr const ROUTE_PREFIX: &str = "/api/v1/gateway/client-interfaces/wireguard"; pub struct ApiState { - client_registry: Arc>, + client_registry: Arc, sphinx_key_pair: Arc, - registration_in_progress: Arc>>, + registration_in_progress: Arc>, } fn router_with_state(state: Arc) -> Router { @@ -33,7 +33,7 @@ fn router_with_state(state: Arc) -> Router { } pub(crate) async fn start_http_api( - client_registry: Arc>, + client_registry: Arc, sphinx_key_pair: Arc, ) { // Port should be 80 post smoosh @@ -46,7 +46,7 @@ pub(crate) async fn start_http_api( let state = Arc::new(ApiState { client_registry, sphinx_key_pair, - registration_in_progress: Arc::new(RwLock::new(HashMap::new())), + registration_in_progress: Arc::new(DashMap::new()), }); let routes = router_with_state(state); @@ -62,17 +62,17 @@ pub(crate) async fn start_http_api( mod test { use std::net::SocketAddr; use std::str::FromStr; - use std::{collections::HashMap, sync::Arc}; + use std::sync::Arc; use axum::body::Body; use axum::http::Request; use axum::http::StatusCode; + use dashmap::DashMap; use hmac::Mac; use tower::Service; use tower::ServiceExt; use nym_crypto::asymmetric::encryption; - use tokio::sync::RwLock; use x25519_dalek::{PublicKey, StaticSecret}; use crate::node::client_handling::client_registration::{ @@ -105,8 +105,8 @@ mod test { let client_dh = client_static_private.diffie_hellman(&gateway_static_public); - let registration_in_progress = Arc::new(RwLock::new(HashMap::new())); - let client_registry = Arc::new(RwLock::new(HashMap::new())); + let registration_in_progress = Arc::new(DashMap::new()); + let client_registry = Arc::new(DashMap::new()); let state = Arc::new(ApiState { client_registry: Arc::clone(&client_registry), @@ -136,7 +136,7 @@ mod test { .unwrap(); assert_eq!(response.status(), StatusCode::OK); - assert!(!registration_in_progress.read().await.is_empty()); + assert!(!registration_in_progress.is_empty()); let nonce: Option = serde_json::from_slice(&hyper::body::to_bytes(response.into_body()).await.unwrap()) @@ -171,7 +171,7 @@ mod test { .unwrap(); assert_eq!(response.status(), StatusCode::OK); - assert!(!client_registry.read().await.is_empty()); + assert!(!client_registry.is_empty()); let clients_request = Request::builder() .method("GET") @@ -194,11 +194,10 @@ mod test { assert!(!clients.is_empty()); - let ro_clients = client_registry.read().await.clone(); assert_eq!( - ro_clients - .values() - .map(|c| c.pub_key().clone()) + client_registry + .iter() + .map(|c| c.value().pub_key().clone()) .collect::>(), clients ) diff --git a/gateway/src/node/mod.rs b/gateway/src/node/mod.rs index 132671c8a1..4cc1f4042c 100644 --- a/gateway/src/node/mod.rs +++ b/gateway/src/node/mod.rs @@ -18,6 +18,7 @@ use crate::node::mixnet_handling::receiver::connection_handler::ConnectionHandle use crate::node::statistics::collector::GatewayStatisticsCollector; use crate::node::storage::Storage; use anyhow::bail; +use dashmap::DashMap; use futures::channel::{mpsc, oneshot}; use log::*; use nym_crypto::asymmetric::{encryption, identity}; @@ -29,12 +30,10 @@ use nym_task::{TaskClient, TaskManager}; use nym_validator_client::{nyxd, DirectSigningHttpRpcNyxdClient}; use rand::seq::SliceRandom; use rand::thread_rng; -use std::collections::HashMap; use std::error::Error; use std::net::SocketAddr; use std::path::PathBuf; use std::sync::Arc; -use tokio::sync::RwLock; pub(crate) mod client_handling; pub(crate) mod helpers; @@ -91,7 +90,7 @@ pub(crate) struct Gateway { sphinx_keypair: Arc, storage: St, - client_registry: Arc>, + client_registry: Arc, } impl Gateway { @@ -107,7 +106,7 @@ impl Gateway { sphinx_keypair: Arc::new(helpers::load_sphinx_keys(&config)?), config, network_requester_opts, - client_registry: Arc::new(RwLock::new(HashMap::new())), + client_registry: Arc::new(DashMap::new()), }) } @@ -125,7 +124,7 @@ impl Gateway { identity_keypair: Arc::new(identity_keypair), sphinx_keypair: Arc::new(sphinx_keypair), storage, - client_registry: Arc::new(RwLock::new(HashMap::new())), + client_registry: Arc::new(DashMap::new()), } }