Replace RwLock<HashMap> with DashMap in handshake

5 files changed, 84 insertions, 29 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 7bf08ab..662120d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -17,6 +17,14 @@ dependencies = [
 ]
 
 [[package]]
+name = "ahash"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "const-random 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
 name = "aho-corasick"
 version = "0.7.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -158,6 +166,24 @@ dependencies = [
 ]
 
 [[package]]
+name = "const-random"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "const-random-macro 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)",
+ "proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
+name = "const-random-macro"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "getrandom 0.1.14 (registry+https://github.com/rust-lang/crates.io-index)",
+ "proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
 name = "cpuprofiler"
 version = "0.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -208,6 +234,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "dashmap"
+version = "3.11.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "ahash 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
+ "cfg-if 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)",
+ "num_cpus 1.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
 name = "digest"
 version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -594,6 +630,11 @@ version = "0.2.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
+name = "proc-macro-hack"
+version = "0.5.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
 name = "proc-macro2"
 version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1134,6 +1175,7 @@ dependencies = [
  "clear_on_drop 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "cpuprofiler 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "crossbeam-channel 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "dashmap 3.11.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "digest 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1216,6 +1258,7 @@ dependencies = [
 [metadata]
 "checksum addr2line 0.12.1 (registry+https://github.com/rust-lang/crates.io-index)" = "a49806b9dadc843c61e7c97e72490ad7f7220ae249012fbda9ad0609457c0543"
 "checksum aead 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "4cf01b9b56e767bb57b94ebf91a58b338002963785cdd7013e21c0d4679471e4"
+"checksum ahash 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "e8fd72866655d1904d6b0997d0b07ba561047d070fbe29de039031c641b61217"
 "checksum aho-corasick 0.7.10 (registry+https://github.com/rust-lang/crates.io-index)" = "8716408b8bc624ed7f65d223ddb9ac2d044c0547b6fa4b0d554f3a9540496ada"
 "checksum arraydeque 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "f0ffd3d69bd89910509a5d31d1f1353f38ccffdd116dd0099bbd6627f7bd8ad8"
 "checksum atty 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8"
@@ -1236,11 +1279,14 @@ dependencies = [
 "checksum chacha20poly1305 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "48901293601228db2131606f741db33561f7576b5d19c99cd66222380a7dc863"
 "checksum clear_on_drop 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "97276801e127ffb46b66ce23f35cc96bd454fa311294bced4bbace7baa8b1d17"
 "checksum cloudabi 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f"
+"checksum const-random 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "2f1af9ac737b2dd2d577701e59fd09ba34822f6f2ebdb30a7647405d9e55e16a"
+"checksum const-random-macro 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "25e4c606eb459dd29f7c57b2e0879f2b6f14ee130918c2b78ccb58a9624e6c7a"
 "checksum cpuprofiler 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "43f8479dbcfd2bbaa0c0c26779b913052b375981cdf533091f2127ea3d42e52b"
 "checksum crossbeam-channel 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "cced8691919c02aac3cb0a1bc2e9b73d89e832bf9a06fc579d4e71b68a2da061"
 "checksum crossbeam-utils 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c3c7c73a2d1e9fc0886a08b93e98eb643461230d5f1925e4036204d5f2e261a8"
 "checksum crypto-mac 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "4434400df11d95d556bac068ddfedd482915eb18fe8bea89bc80b6e4b1c179e5"
 "checksum curve25519-dalek 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "26778518a7f6cffa1d25a44b602b62b979bd88adb9e99ffec546998cf3404839"
+"checksum dashmap 3.11.4 (registry+https://github.com/rust-lang/crates.io-index)" = "8cfcd41ae02d60edded204341d2798ba519c336c51a37330aa4b98a1128def32"
 "checksum digest 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "f3d0c8c8752312f9713efd397ff63acb9f85585afbf179282e720e7704954dd5"
 "checksum env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)" = "44533bbbb3bb3c1fa17d9f2e4e38bbbaf8396ba82193c4cb1b6445d711445d36"
 "checksum error-chain 0.12.2 (registry+https://github.com/rust-lang/crates.io-index)" = "d371106cc88ffdfb1eabd7111e432da544f16f3e2d7bf1dfe8bf575f1df045cd"
@@ -1289,6 +1335,7 @@ dependencies = [
 "checksum pnet_transport 0.25.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1b75ccaee7b5daba9f9a7d47bceeb73cc32edde9952dc5409460d6621ec667b6"
 "checksum poly1305 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b5829f50f48e9ddb79f3f7c3097029d0caee30f8286accb241416df603b080b8"
 "checksum ppv-lite86 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)" = "237a5ed80e274dbc66f86bd59c1e25edc039660be53194b5fe0a482e0f2612ea"
+"checksum proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)" = "7e0456befd48169b9f13ef0f0ad46d492cf9d2dbb918bcf38e01eed4ce3ec5e4"
 "checksum proc-macro2 1.0.14 (registry+https://github.com/rust-lang/crates.io-index)" = "de40dd4ff82d9c9bab6dae29dbab1167e515f8df9ed17d2987cb6012db206933"
 "checksum proptest 0.9.6 (registry+https://github.com/rust-lang/crates.io-index)" = "01c477819b845fe023d33583ebf10c9f62518c8d79a0960ba5c36d6ac8a55a5b"
 "checksum quick-error 1.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
diff --git a/Cargo.toml b/Cargo.toml
index 625e830..8da2f84 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ clear_on_drop = "0.2.3"
 env_logger = "0.7"
 num_cpus = "^1.10"
 crossbeam-channel = "0.4"
+dashmap = "3.11"
 cpuprofiler = { version = "*", optional = true }
 
 [dependencies.treebitmap]
diff --git a/src/wireguard/handshake/device.rs b/src/wireguard/handshake/device.rs
index 3a3d023..5e69921 100644
--- a/src/wireguard/handshake/device.rs
+++ b/src/wireguard/handshake/device.rs
@@ -1,11 +1,12 @@
-use spin::RwLock;
 use std::collections::hash_map;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::Mutex;
-use zerocopy::AsBytes;
 
 use byteorder::{ByteOrder, LittleEndian};
+use dashmap::mapref::entry::Entry;
+use dashmap::DashMap;
+use zerocopy::AsBytes;
 
 use rand::prelude::{CryptoRng, RngCore};
 use rand::Rng;
@@ -36,7 +37,7 @@ pub struct KeyState {
 /// (the instance is a Peer object in the parent module)
 pub struct Device<O> {
     keyst: Option<KeyState>,
-    id_map: RwLock<HashMap<u32, [u8; 32]>>,
+    id_map: DashMap<u32, [u8; 32]>, // concurrent map
     pk_map: HashMap<[u8; 32], Peer<O>>,
     limiter: Mutex<RateLimiter>,
 }
@@ -62,7 +63,7 @@ impl<'a, O> Iterator for Iter<'a, O> {
  */
 impl<O> Device<O> {
     pub fn clear(&mut self) {
-        self.id_map.write().clear();
+        self.id_map.clear();
         self.pk_map.clear();
     }
 
@@ -96,7 +97,7 @@ impl<O> Device<O> {
     pub fn new() -> Device<O> {
         Device {
             keyst: None,
-            id_map: RwLock::new(HashMap::new()),
+            id_map: DashMap::new(),
             pk_map: HashMap::new(),
             limiter: Mutex::new(RateLimiter::new()),
         }
@@ -208,16 +209,14 @@ impl<O> Device<O> {
     ///
     /// The call might fail if the public key is not found
     pub fn remove(&mut self, pk: &PublicKey) -> Result<(), ConfigError> {
-        // take write-lock on receive id table
-        let mut id_map = self.id_map.write();
-
         // remove the peer
         self.pk_map
             .remove(pk.as_bytes())
             .ok_or(ConfigError::new("Public key not in device"))?;
 
-        // purge the id map (linear scan)
-        id_map.retain(|_, v| v != pk.as_bytes());
+        // remove every id entry for the peer in the public key map
+        // O(n) operations, however it is rare: only when removing peers.
+        self.id_map.retain(|_, v| v != pk.as_bytes());
         Ok(())
     }
 
@@ -265,9 +264,8 @@ impl<O> Device<O> {
     ///
     /// * `id` - The (sender) id to release
     pub fn release(&self, id: u32) {
-        let mut m = self.id_map.write();
-        debug_assert!(m.contains_key(&id), "Releasing id not allocated");
-        m.remove(&id);
+        let old = self.id_map.remove(&id);
+        assert!(old.is_some(), "released id not allocated");
     }
 
     /// Begin a new handshake
@@ -446,32 +444,40 @@ impl<O> Device<O> {
     //
     // Return the peer currently associated with the receiver identifier
     pub(super) fn lookup_id(&self, id: u32) -> Result<(&Peer<O>, PublicKey), HandshakeError> {
-        let im = self.id_map.read();
-        let pk = im.get(&id).ok_or(HandshakeError::UnknownReceiverId)?;
-        match self.pk_map.get(pk) {
+        // obtain a read reference to entry in the id_map
+        let pk = self
+            .id_map
+            .get(&id)
+            .ok_or(HandshakeError::UnknownReceiverId)?;
+
+        // lookup the public key from the pk map
+        match self.pk_map.get(&*pk) {
             Some(peer) => Ok((peer, PublicKey::from(*pk))),
-            _ => unreachable!(), // if the id-lookup succeeded, the peer should exist
+            _ => unreachable!(),
         }
     }
 
     // Internal function
     //
-    // Allocated a new receiver identifier for the peer
+    // Allocated a new receiver identifier for the peer.
+    // Implemented via rejection sampling.
     fn allocate<R: RngCore + CryptoRng>(&self, rng: &mut R, pk: &PublicKey) -> u32 {
         loop {
             let id = rng.gen();
 
-            // check membership with read lock
-            if self.id_map.read().contains_key(&id) {
+            // read lock the shard and do quick check
+            if self.id_map.contains_key(&id) {
                 continue;
             }
 
-            // take write lock and add index
-            let mut m = self.id_map.write();
-            if !m.contains_key(&id) {
-                m.insert(id, *pk.as_bytes());
-                return id;
-            }
+            // write lock the shard and insert
+            match self.id_map.entry(id) {
+                Entry::Vacant(entry) => {
+                    entry.insert(*pk.as_bytes());
+                    return id;
+                }
+                _ => (),
+            };
         }
     }
 }
diff --git a/src/wireguard/peer.rs b/src/wireguard/peer.rs
index 27d39bd..37b4684 100644
--- a/src/wireguard/peer.rs
+++ b/src/wireguard/peer.rs
@@ -22,13 +22,15 @@ pub struct PeerInner<T: Tun, B: UDP> {
     // wireguard device state
     pub wg: WireGuard<T, B>,
 
+    // TODO: eliminate
+    pub pk: PublicKey,
+
     // handshake state
     pub walltime_last_handshake: Mutex<Option<SystemTime>>, // walltime for last handshake (for UAPI status)
     pub last_handshake_sent: Mutex<Instant>,                // instant for last handshake
-    pub handshake_queued: AtomicBool, // is a handshake job currently queued for the peer?
+    pub handshake_queued: AtomicBool,                       // is a handshake job currently queued?
 
     // stats and configuration
-    pub pk: PublicKey, // public key (TODO: there has to be a way to remove this)
     pub rx_bytes: AtomicU64, // received bytes
     pub tx_bytes: AtomicU64, // transmitted bytes
 
diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 1a12abb..62ef932 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -3,7 +3,6 @@ use std::ops::Deref;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use std::thread;
-use std::time::Instant;
 
 use log;
 use spin::{Mutex, RwLock};