Replace HashMap with DashMap (#3995)

This commit is contained in:
Drazen Urch
2023-10-16 15:40:26 +02:00
committed by GitHub
parent c598082335
commit e4dda5e541
4 changed files with 36 additions and 39 deletions
@@ -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<SocketAddr, Client>;
pub(crate) type ClientRegistry = DashMap<SocketAddr, Client>;
+15 -16
View File
@@ -14,8 +14,7 @@ use crate::node::http::ApiState;
async fn process_final_message(client: Client, state: Arc<ApiState>) -> 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<ApiState>) -> 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<ApiState>) -> StatusCo
async fn process_init_message(init_message: InitMessage, state: Arc<ApiState>) -> 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<Arc<ApiState>>,
) -> (StatusCode, Json<Vec<ClientPublicKey>>) {
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::<Vec<ClientPublicKey>>(),
),
@@ -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
}
+15 -16
View File
@@ -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<RwLock<ClientRegistry>>,
client_registry: Arc<ClientRegistry>,
sphinx_key_pair: Arc<encryption::KeyPair>,
registration_in_progress: Arc<RwLock<HashMap<ClientPublicKey, u64>>>,
registration_in_progress: Arc<DashMap<ClientPublicKey, u64>>,
}
fn router_with_state(state: Arc<ApiState>) -> Router {
@@ -33,7 +33,7 @@ fn router_with_state(state: Arc<ApiState>) -> Router {
}
pub(crate) async fn start_http_api(
client_registry: Arc<RwLock<ClientRegistry>>,
client_registry: Arc<ClientRegistry>,
sphinx_key_pair: Arc<encryption::KeyPair>,
) {
// 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<u64> =
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::<Vec<ClientPublicKey>>(),
clients
)
+4 -5
View File
@@ -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<St = PersistentStorage> {
sphinx_keypair: Arc<encryption::KeyPair>,
storage: St,
client_registry: Arc<RwLock<ClientRegistry>>,
client_registry: Arc<ClientRegistry>,
}
impl<St> Gateway<St> {
@@ -107,7 +106,7 @@ impl<St> Gateway<St> {
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<St> Gateway<St> {
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()),
}
}