Replace RwLock<HashMap> with DashMap in handshake

2020-06-19 23:45:56 +02:00
parent c1dfc848c4
commit 6e307fc70e
5 changed files with 84 additions and 29 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16,6 +16,14 @@ dependencies = [
 "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

+[[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"
@@ -157,6 +165,24 @@ dependencies = [
 "bitflags 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

+[[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"
@@ -207,6 +233,16 @@ dependencies = [
 "zeroize 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

+[[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"
@@ -593,6 +629,11 @@ name = "ppv-lite86"
 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"
@@ -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"
--- 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]
--- 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());
+            // write lock the shard and insert
+            match self.id_map.entry(id) {
+                Entry::Vacant(entry) => {
+                    entry.insert(*pk.as_bytes());
                    return id;
                }
+                _ => (),
+            };
        }
    }
 }
--- 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

--- 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};