fix peer center for latency report

easytier/src/peer_center/server.rs

@@ -1,45 +1,48 @@
 use std::{
+    collections::BinaryHeap,
     hash::{Hash, Hasher},
     sync::Arc,
 };
 
+use crossbeam::atomic::AtomicCell;
 use dashmap::DashMap;
 use once_cell::sync::Lazy;
-use tokio::{sync::RwLock, task::JoinSet};
+use tokio::{task::JoinSet};
 
-use crate::common::PeerId;
+use crate::{common::PeerId, rpc::DirectConnectedPeerInfo};
 
 use super::{
     service::{GetGlobalPeerMapResponse, GlobalPeerMap, PeerCenterService, PeerInfoForGlobalMap},
     Digest, Error,
 };
 
-pub(crate) struct PeerCenterServerGlobalData {
-    pub global_peer_map: GlobalPeerMap,
-    pub digest: Digest,
-    pub update_time: std::time::Instant,
-    pub peer_update_time: DashMap<PeerId, std::time::Instant>,
+#[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, Hash)]
+pub(crate) struct SrcDstPeerPair {
+    src: PeerId,
+    dst: PeerId,
 }
 
-impl PeerCenterServerGlobalData {
-    fn new() -> Self {
-        PeerCenterServerGlobalData {
-            global_peer_map: GlobalPeerMap::new(),
-            digest: Digest::default(),
-            update_time: std::time::Instant::now(),
-            peer_update_time: DashMap::new(),
-        }
-    }
+#[derive(Debug, Clone)]
+pub(crate) struct PeerCenterInfoEntry {
+    info: DirectConnectedPeerInfo,
+    update_time: std::time::Instant,
+}
+
+#[derive(Default)]
+pub(crate) struct PeerCenterServerGlobalData {
+    pub(crate) global_peer_map: DashMap<SrcDstPeerPair, PeerCenterInfoEntry>,
+    pub(crate) peer_report_time: DashMap<PeerId, std::time::Instant>,
+    pub(crate) digest: AtomicCell<Digest>,
 }
 
 // a global unique instance for PeerCenterServer
-pub(crate) static GLOBAL_DATA: Lazy<DashMap<PeerId, Arc<RwLock<PeerCenterServerGlobalData>>>> =
+pub(crate) static GLOBAL_DATA: Lazy<DashMap<PeerId, Arc<PeerCenterServerGlobalData>>> =
     Lazy::new(DashMap::new);
 
-pub(crate) fn get_global_data(node_id: PeerId) -> Arc<RwLock<PeerCenterServerGlobalData>> {
+pub(crate) fn get_global_data(node_id: PeerId) -> Arc<PeerCenterServerGlobalData> {
     GLOBAL_DATA
         .entry(node_id)
-        .or_insert_with(|| Arc::new(RwLock::new(PeerCenterServerGlobalData::new())))
+        .or_insert_with(|| Arc::new(PeerCenterServerGlobalData::default()))
         .value()
         .clone()
 }
@@ -48,8 +51,6 @@ pub(crate) fn get_global_data(node_id: PeerId) -> Arc<RwLock<PeerCenterServerGlo
 pub struct PeerCenterServer {
     // every peer has its own server, so use per-struct dash map is ok.
     my_node_id: PeerId,
-    digest_map: DashMap<PeerId, Digest>,
-
     tasks: Arc<JoinSet<()>>,
 }
 
@@ -65,26 +66,32 @@ impl PeerCenterServer {
 
         PeerCenterServer {
             my_node_id,
-            digest_map: DashMap::new(),
-
             tasks: Arc::new(tasks),
         }
     }
 
     async fn clean_outdated_peer(my_node_id: PeerId) {
         let data = get_global_data(my_node_id);
-        let mut locked_data = data.write().await;
-        let now = std::time::Instant::now();
-        let mut to_remove = Vec::new();
-        for kv in locked_data.peer_update_time.iter() {
-            if now.duration_since(*kv.value()).as_secs() > 20 {
-                to_remove.push(*kv.key());
-            }
-        }
-        for peer_id in to_remove {
-            locked_data.global_peer_map.map.remove(&peer_id);
-            locked_data.peer_update_time.remove(&peer_id);
-        }
+        data.peer_report_time.retain(|_, v| {
+            std::time::Instant::now().duration_since(*v) < std::time::Duration::from_secs(180)
+        });
+        data.global_peer_map.retain(|_, v| {
+            std::time::Instant::now().duration_since(v.update_time)
+                < std::time::Duration::from_secs(180)
+        });
+    }
+
+    fn calc_global_digest(my_node_id: PeerId) -> Digest {
+        let data = get_global_data(my_node_id);
+        let mut hasher = std::collections::hash_map::DefaultHasher::new();
+        data.global_peer_map
+            .iter()
+            .map(|v| v.key().clone())
+            .collect::<BinaryHeap<_>>()
+            .into_sorted_vec()
+            .into_iter()
+            .for_each(|v| v.hash(&mut hasher));
+        hasher.finish()
     }
 }
 
@@ -95,39 +102,28 @@ impl PeerCenterService for PeerCenterServer {
         self,
         _: tarpc::context::Context,
         my_peer_id: PeerId,
-        peers: Option<PeerInfoForGlobalMap>,
-        digest: Digest,
+        peers: PeerInfoForGlobalMap,
     ) -> Result<(), Error> {
-        tracing::trace!("receive report_peers");
+        tracing::debug!("receive report_peers");
 
         let data = get_global_data(self.my_node_id);
-        let mut locked_data = data.write().await;
-        locked_data
-            .peer_update_time
+        data.peer_report_time
             .insert(my_peer_id, std::time::Instant::now());
 
-        let old_digest = self.digest_map.get(&my_peer_id);
-        // if digest match, no need to update
-        if let Some(old_digest) = old_digest {
-            if *old_digest == digest {
-                return Ok(());
-            }
+        for (peer_id, peer_info) in peers.direct_peers {
+            let pair = SrcDstPeerPair {
+                src: my_peer_id,
+                dst: peer_id,
+            };
+            let entry = PeerCenterInfoEntry {
+                info: peer_info,
+                update_time: std::time::Instant::now(),
+            };
+            data.global_peer_map.insert(pair, entry);
         }
 
-        if peers.is_none() {
-            return Err(Error::DigestMismatch);
-        }
-
-        self.digest_map.insert(my_peer_id, digest);
-        locked_data
-            .global_peer_map
-            .map
-            .insert(my_peer_id, peers.unwrap());
-
-        let mut hasher = std::collections::hash_map::DefaultHasher::new();
-        locked_data.global_peer_map.map.hash(&mut hasher);
-        locked_data.digest = hasher.finish() as Digest;
-        locked_data.update_time = std::time::Instant::now();
+        data.digest
+            .store(PeerCenterServer::calc_global_digest(self.my_node_id));
 
         Ok(())
     }
@@ -138,15 +134,26 @@ impl PeerCenterService for PeerCenterServer {
         digest: Digest,
     ) -> Result<Option<GetGlobalPeerMapResponse>, Error> {
         let data = get_global_data(self.my_node_id);
-        if digest == data.read().await.digest {
+        if digest == data.digest.load() && digest != 0 {
             return Ok(None);
         }
 
-        let data = get_global_data(self.my_node_id);
-        let locked_data = data.read().await;
+        let mut global_peer_map = GlobalPeerMap::new();
+        for item in data.global_peer_map.iter() {
+            let (pair, entry) = item.pair();
+            global_peer_map
+                .map
+                .entry(pair.src)
+                .or_insert_with(|| PeerInfoForGlobalMap {
+                    direct_peers: Default::default(),
+                })
+                .direct_peers
+                .insert(pair.dst, entry.info.clone());
+        }
+
         Ok(Some(GetGlobalPeerMapResponse {
-            global_peer_map: locked_data.global_peer_map.clone(),
-            digest: locked_data.digest,
+            global_peer_map,
+            digest: data.digest.load(),
         }))
     }
 }