EasyTier/easytier-core/src/peers/conn/peer_map.rs

use std::{
    collections::{BTreeSet, HashMap, HashSet},
    net::{Ipv4Addr, Ipv6Addr},
    sync::Arc,
};

use anyhow::Context;
use dashmap::{DashMap, DashSet};
use parking_lot::Mutex;
use tokio::sync::RwLock;

use crate::{
    config::PeerId,
    foundation::small_map_shard_amount,
    packet::ZCPacket,
    peers::{
        context::{ArcPeerContext, NetworkIdentity, PeerEvent},
        error::Error,
        util::shrink_dashmap,
    },
    proto::{
        core_peer::peer::{PeerConnInfo, Route as CoreRoute},
        peer_rpc::{
            DirectConnectedPeerInfo, PeerIdentityType, PeerInfoForGlobalMap, RoutePeerInfo,
        },
    },
    tunnel::TunnelError,
};

use super::{
    peer::Peer,
    peer_conn::{PeerConn, PeerConnId},
};
use crate::peers::{
    PacketRecvChan,
    route::{ArcRoute, NextHopPolicy},
};

pub struct PeerMap {
    context: ArcPeerContext,
    my_peer_id: PeerId,
    peer_map: DashMap<PeerId, Arc<Peer>>,
    packet_send: PacketRecvChan,
    routes: RwLock<Vec<ArcRoute>>,
    alive_client_urls: Arc<Mutex<HashMap<url::Url, HashSet<PeerConnId>>>>,
}

impl PeerMap {
    pub(crate) fn new(
        packet_send: PacketRecvChan,
        context: ArcPeerContext,
        my_peer_id: PeerId,
    ) -> Self {
        PeerMap {
            context,
            my_peer_id,
            peer_map: DashMap::with_shard_amount(small_map_shard_amount()),
            packet_send,
            routes: RwLock::new(Vec::new()),
            alive_client_urls: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    async fn add_new_peer(&self, peer: Peer) {
        let peer_id = peer.peer_node_id;
        self.peer_map.insert(peer_id, Arc::new(peer));
        self.context.issue_event(PeerEvent::PeerAdded(peer_id));
    }

    pub async fn add_new_peer_conn(&self, peer_conn: PeerConn) -> Result<(), Error> {
        let _ = self.maintain_alive_client_urls(&peer_conn);
        let peer_id = peer_conn.get_peer_id();
        let no_entry = self.peer_map.get(&peer_id).is_none();
        if no_entry {
            let new_peer = Peer::new(peer_id, self.packet_send.clone(), self.context.clone());
            new_peer.add_peer_conn(peer_conn).await?;
            self.add_new_peer(new_peer).await;
        } else {
            let peer = self.peer_map.get(&peer_id).unwrap().clone();
            peer.add_peer_conn(peer_conn).await?;
        }
        Ok(())
    }

    fn maintain_alive_client_urls(&self, peer_conn: &PeerConn) -> Option<()> {
        let conn_info = peer_conn.get_conn_info();
        if !conn_info.is_client {
            return None;
        }

        let close_notifier = peer_conn.get_close_notifier();
        let alive_conns_weak = Arc::downgrade(&self.alive_client_urls);
        let conn_id = close_notifier.get_conn_id();
        let alive_client_url: url::Url = conn_info.tunnel?.remote_addr?.into();
        self.alive_client_urls
            .lock()
            .entry(alive_client_url.clone())
            .or_default()
            .insert(conn_id);

        tokio::spawn(async move {
            if let Some(mut waiter) = close_notifier.get_waiter().await {
                let _ = waiter.recv().await;
            }
            let Some(alive_conns) = alive_conns_weak.upgrade() else {
                return;
            };
            let mut guard = alive_conns.lock();
            if let Some(conn_ids) = guard.get_mut(&alive_client_url) {
                conn_ids.retain(|id| id != &conn_id);
                if conn_ids.is_empty() {
                    guard.remove(&alive_client_url);
                }
            }
            let alive_conn_count = guard.len();
            drop(guard);
            tracing::debug!(
                ?conn_id,
                "peer conn is closed, current alive conns: {}",
                alive_conn_count
            );
        });

        Some(())
    }

    pub fn is_client_url_alive(&self, url: &url::Url) -> bool {
        self.alive_client_urls.lock().contains_key(url)
    }

    pub fn get_peer_by_id(&self, peer_id: PeerId) -> Option<Arc<Peer>> {
        self.peer_map.get(&peer_id).map(|v| v.clone())
    }

    pub fn get_directly_connections_by_peer_id(&self, peer_id: PeerId) -> DashSet<uuid::Uuid> {
        if let Some(peer) = self.get_peer_by_id(peer_id) {
            return peer.get_directly_connections();
        }

        DashSet::new()
    }

    pub fn has_peer(&self, peer_id: PeerId) -> bool {
        peer_id == self.my_peer_id || self.peer_map.contains_key(&peer_id)
    }

    pub(crate) fn has_direct_attached_peer(&self, peer_id: PeerId) -> bool {
        self.get_peer_by_id(peer_id)
            .is_some_and(|peer| peer.has_direct_attached_conn())
    }

    pub(crate) fn is_self(&self, peer_id: PeerId) -> bool {
        peer_id == self.my_peer_id
    }

    pub async fn send_msg_directly(&self, msg: ZCPacket, dst_peer_id: PeerId) -> Result<(), Error> {
        if dst_peer_id == self.my_peer_id {
            let packet_send = self.packet_send.clone();
            tokio::spawn(async move {
                let ret = packet_send
                    .send(msg)
                    .await
                    .with_context(|| "send msg to self failed");
                if ret.is_err() {
                    tracing::error!("send msg to self failed: {:?}", ret);
                }
            });
            return Ok(());
        }

        match self.get_peer_by_id(dst_peer_id) {
            Some(peer) => {
                peer.send_msg(msg).await?;
            }
            None => {
                tracing::error!("no peer for dst_peer_id: {}", dst_peer_id);
                return Err(Error::RouteError(Some(format!(
                    "peer map sengmsg directly no connected dst_peer_id: {}",
                    dst_peer_id
                ))));
            }
        }

        Ok(())
    }

    pub async fn get_gateway_peer_id(
        &self,
        dst_peer_id: PeerId,
        policy: NextHopPolicy,
    ) -> Option<PeerId> {
        if dst_peer_id == self.my_peer_id {
            return Some(dst_peer_id);
        }

        if self.has_peer(dst_peer_id) && matches!(policy, NextHopPolicy::LeastHop) {
            return Some(dst_peer_id);
        }

        // get route info
        for route in self.routes.read().await.iter() {
            if let Some(gateway_peer_id) = route
                .get_next_hop_with_policy(dst_peer_id, policy.clone())
                .await
            {
                // NOTIC: for foreign network, gateway_peer_id may not connect to me
                return Some(gateway_peer_id);
            }
        }

        None
    }

    pub async fn list_peers_own_foreign_network(
        &self,
        network_identity: &NetworkIdentity,
    ) -> Vec<PeerId> {
        let mut ret = Vec::new();
        for route in self.routes.read().await.iter() {
            let peers = route.list_peers_own_foreign_network(network_identity).await;
            ret.extend(peers);
        }
        ret
    }

    pub async fn send_msg(
        &self,
        msg: ZCPacket,
        dst_peer_id: PeerId,
        policy: NextHopPolicy,
    ) -> Result<(), Error> {
        let Some(gateway_peer_id) = self.get_gateway_peer_id(dst_peer_id, policy).await else {
            return Err(Error::RouteError(Some(format!(
                "peer map sengmsg no gateway for dst_peer_id: {}",
                dst_peer_id
            ))));
        };

        self.send_msg_directly(msg, gateway_peer_id).await?;
        Ok(())
    }

    pub async fn get_peer_id_by_ipv4(&self, ipv4: &Ipv4Addr) -> Option<PeerId> {
        for route in self.routes.read().await.iter() {
            let peer_id = route.get_peer_id_by_ipv4(ipv4).await;
            if peer_id.is_some() {
                return peer_id;
            }
        }
        None
    }

    pub async fn get_peer_id_by_ipv6(&self, ipv6: &Ipv6Addr) -> Option<PeerId> {
        for route in self.routes.read().await.iter() {
            let peer_id = route.get_peer_id_by_ipv6(ipv6).await;
            if peer_id.is_some() {
                return peer_id;
            }
        }
        None
    }

    pub async fn get_route_peer_info(&self, peer_id: PeerId) -> Option<RoutePeerInfo> {
        for route in self.routes.read().await.iter() {
            if let Some(info) = route.get_peer_info(peer_id).await {
                return Some(info);
            }
        }
        None
    }

    pub async fn get_origin_my_peer_id(
        &self,
        network_name: &str,
        foreign_my_peer_id: PeerId,
    ) -> Option<PeerId> {
        for route in self.routes.read().await.iter() {
            let origin_peer_id = route
                .get_origin_my_peer_id(network_name, foreign_my_peer_id)
                .await;
            if origin_peer_id.is_some() {
                return origin_peer_id;
            }
        }
        None
    }

    pub fn is_empty(&self) -> bool {
        self.peer_map.is_empty()
    }

    pub fn list_peers(&self) -> Vec<PeerId> {
        let mut ret = Vec::new();
        for item in self.peer_map.iter() {
            let peer_id = item.key();
            ret.push(*peer_id);
        }
        ret
    }

    pub async fn list_peers_with_conn(&self) -> Vec<PeerId> {
        let mut ret = Vec::new();
        for item in self.peer_map.iter() {
            if item.value().has_live_conns() {
                ret.push(*item.key());
            }
        }
        ret
    }

    pub async fn list_peer_conns(&self, peer_id: PeerId) -> Option<Vec<PeerConnInfo>> {
        if let Some(p) = self.get_peer_by_id(peer_id) {
            Some(p.list_peer_conns().await)
        } else {
            None
        }
    }

    pub async fn get_peer_default_conn_id(&self, peer_id: PeerId) -> Option<PeerConnId> {
        self.get_peer_by_id(peer_id)
            .map(|p| p.get_default_conn_id())
    }

    pub fn get_peer_identity_type(&self, peer_id: PeerId) -> Option<PeerIdentityType> {
        self.get_peer_by_id(peer_id)
            .and_then(|p| p.get_peer_identity_type())
    }

    pub fn get_peer_public_key(&self, peer_id: PeerId) -> Option<Vec<u8>> {
        self.get_peer_by_id(peer_id)
            .and_then(|p| p.get_peer_public_key())
    }

    pub async fn close_peer_conn(
        &self,
        peer_id: PeerId,
        conn_id: &PeerConnId,
    ) -> Result<(), Error> {
        if let Some(p) = self.get_peer_by_id(peer_id) {
            p.close_peer_conn(conn_id).await
        } else {
            Err(Error::NotFound)
        }
    }

    pub async fn close_peer(&self, peer_id: PeerId) -> Result<(), TunnelError> {
        let remove_ret = self.peer_map.remove(&peer_id);
        shrink_dashmap(&self.peer_map, None);

        self.context.issue_event(PeerEvent::PeerRemoved(peer_id));
        tracing::info!(
            ?peer_id,
            has_old_value = ?remove_ret.is_some(),
            peer_ref_counter = ?remove_ret.map(|v| Arc::strong_count(&v.1)),
            "peer is closed"
        );
        Ok(())
    }

    pub async fn add_route(&self, route: ArcRoute) {
        let mut routes = self.routes.write().await;
        routes.insert(0, route);
    }

    pub(crate) async fn clear_resources(&self) {
        for peer_id in self.list_peers() {
            let _ = self.close_peer(peer_id).await;
        }
        let routes = {
            let mut routes = self.routes.write().await;
            std::mem::take(&mut *routes)
        };
        for route in routes {
            route.close().await;
        }
        self.alive_client_urls.lock().clear();
    }

    pub async fn clean_peer_without_conn(&self) {
        let mut to_remove = vec![];

        for peer_id in self.list_peers() {
            let conns = self.list_peer_conns(peer_id).await;
            if conns.is_none() || conns.as_ref().unwrap().is_empty() {
                to_remove.push(peer_id);
            }
        }

        for peer_id in to_remove {
            self.close_peer(peer_id).await.unwrap();
        }
    }

    pub async fn list_routes(&self) -> DashMap<PeerId, PeerId> {
        let route_map = DashMap::new();
        for route in self.routes.read().await.iter() {
            for item in route.list_routes().await.iter() {
                route_map.insert(item.peer_id, item.next_hop_peer_id);
            }
        }
        route_map
    }

    pub async fn list_route_infos(&self) -> Vec<CoreRoute> {
        if let Some(route) = self.routes.read().await.iter().next() {
            return route.list_routes().await;
        }
        vec![]
    }

    pub async fn need_relay_by_foreign_network(&self, dst_peer_id: PeerId) -> Result<bool, Error> {
        // if gateway_peer_id is not connected to me, means need relay by foreign network
        let gateway_id = self
            .get_gateway_peer_id(dst_peer_id, NextHopPolicy::LeastHop)
            .await
            .ok_or(Error::RouteError(Some(format!(
                "peer map need_relay_by_foreign_network no gateway for dst_peer_id: {}",
                dst_peer_id
            ))))?;

        Ok(!self.has_peer(gateway_id))
    }

    pub fn my_peer_id(&self) -> PeerId {
        self.my_peer_id
    }
}

impl Drop for PeerMap {
    fn drop(&mut self) {
        tracing::debug!(
            self.my_peer_id,
            network = ?self.context.network_identity(),
            "PeerMap is dropped"
        );
    }
}

/// Aggregates the directly-connected peers of several peer maps into the
/// peer-center reporting format, keeping the lowest observed latency per peer.
pub(crate) async fn direct_peer_info(peer_maps: &[Arc<PeerMap>]) -> PeerInfoForGlobalMap {
    let mut peers = BTreeSet::new();
    for peer_map in peer_maps {
        peers.extend(peer_map.list_peers());
    }

    let mut ret = PeerInfoForGlobalMap::default();
    for peer in peers {
        let mut conns = None;
        for peer_map in peer_maps {
            if let Some(found) = peer_map.list_peer_conns(peer).await {
                conns = Some(found);
                break;
            }
        }
        let Some(min_lat) = min_measured_latency_us(
            conns
                .into_iter()
                .flatten()
                .map(|conn| conn.stats.as_ref().unwrap().latency_us),
        ) else {
            continue;
        };

        ret.direct_peers.insert(
            peer,
            DirectConnectedPeerInfo {
                latency_ms: std::cmp::max(1, (min_lat as u32 / 1000) as i32),
            },
        );
    }
    ret
}

/// Lowest latency among the connections that already have a latency sample.
///
/// A freshly established connection reports `latency_us == 0` until its first
/// ping completes. Treating that as a real value would advertise the peer as
/// 1 ms away to the peer center, so latency-first routing could detour through
/// a slow or lossy link until the next report. Unmeasured connections are
/// skipped instead; a peer with none measured is left out of the report and is
/// reported once its first sample arrives, because the peer list changes then.
fn min_measured_latency_us(latencies_us: impl IntoIterator<Item = u64>) -> Option<u64> {
    latencies_us.into_iter().filter(|lat| *lat > 0).min()
}

#[cfg(test)]
mod tests {
    use super::min_measured_latency_us;

    #[test]
    fn unmeasured_connections_are_not_reported_as_fastest() {
        assert_eq!(min_measured_latency_us([0, 306_000]), Some(306_000));
        assert_eq!(min_measured_latency_us([306_000, 0, 93_000]), Some(93_000));
    }

    #[test]
    fn peer_without_measured_connection_is_skipped() {
        assert_eq!(min_measured_latency_us([0, 0]), None);
        assert_eq!(min_measured_latency_us([]), None);
    }

    #[test]
    fn sub_millisecond_lan_latency_is_still_reported() {
        assert_eq!(min_measured_latency_us([400]), Some(400));
    }
}