ouisync/network/

connection.rs

use super::{
    peer_addr::PeerAddr,
    peer_info::PeerInfo,
    peer_source::{ConnectionDirection, PeerSource},
    peer_state::PeerState,
    runtime_id::PublicRuntimeId,
    stats::{ByteCounters, StatsTracker},
};
use crate::sync::{AwaitDrop, DropAwaitable, WatchSenderExt};
use portable_atomic::{AtomicU64, Ordering};
use std::{
    collections::{HashMap, hash_map::Entry},
    fmt,
    sync::Arc,
    time::SystemTime,
};
use tokio::sync::watch;

/// Container for known connections.
pub(super) struct ConnectionSet {
    connections: watch::Sender<HashMap<ConnectionKey, ConnectionData>>,
}

impl ConnectionSet {
    pub fn new() -> Self {
        Self {
            connections: watch::Sender::new(HashMap::default()),
        }
    }

    /// Attempt to reserve an connection to the given peer. If the connection hasn't been reserved
    /// yet, it returns a `ConnectionPermit` which keeps the connection reserved as long as it
    /// lives. Otherwise it returns `None`. To release a connection the permit needs to be dropped.
    /// Also returns a notification object that can be used to wait until the permit gets released.
    pub fn reserve(&self, addr: PeerAddr, source: PeerSource) -> ReserveResult {
        let key = ConnectionKey {
            addr,
            dir: source.direction(),
        };

        self.connections
            .send_if_modified_return(|connections| match connections.entry(key) {
                Entry::Vacant(entry) => {
                    let id = ConnectionId::next();

                    entry.insert(ConnectionData {
                        id,
                        state: InnerPeerState::Known,
                        source,
                        stats_tracker: StatsTracker::default(),
                        on_release: DropAwaitable::new(),
                    });

                    (
                        true,
                        ReserveResult::Permit(ConnectionPermit {
                            connections: self.connections.clone(),
                            key,
                            id,
                        }),
                    )
                }
                Entry::Occupied(entry) => {
                    let data = entry.get();

                    (
                        false,
                        ReserveResult::Occupied(data.on_release.subscribe(), data.source, data.id),
                    )
                }
            })
    }

    pub fn peer_info_collector(&self) -> PeerInfoCollector {
        PeerInfoCollector(self.connections.clone())
    }

    pub fn get_peer_info(&self, addr: PeerAddr) -> Option<PeerInfo> {
        self.with(addr, |data| data.peer_info(addr))
    }

    pub fn get_peer_key(&self, addr: PeerAddr) -> Option<usize> {
        self.with(addr, |data| match data.state {
            InnerPeerState::Active { peer_key, .. } => Some(peer_key),
            InnerPeerState::Known | InnerPeerState::Connecting | InnerPeerState::Handshaking => {
                None
            }
        })
        .flatten()
    }

    fn with<F, R>(&self, addr: PeerAddr, f: F) -> Option<R>
    where
        F: FnOnce(&ConnectionData) -> R,
    {
        let connections = self.connections.borrow();

        connections
            .get(&ConnectionKey {
                addr,
                dir: ConnectionDirection::Incoming,
            })
            .or_else(|| {
                connections.get(&ConnectionKey {
                    addr,
                    dir: ConnectionDirection::Outgoing,
                })
            })
            .map(f)
    }

    pub fn subscribe(&self) -> watch::Receiver<HashMap<ConnectionKey, ConnectionData>> {
        self.connections.subscribe()
    }
}

/// Unique identifier of a connection. Connections are mostly already identified by the peer address
/// and direction (incoming / outgoing), but this type allows to distinguish even connections with
/// the same address/direction but that were established in two separate occasions.
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
#[repr(transparent)]
pub(super) struct ConnectionId(u64);

impl ConnectionId {
    pub fn next() -> Self {
        static NEXT: AtomicU64 = AtomicU64::new(0);
        Self(NEXT.fetch_add(1, Ordering::Relaxed))
    }
}

impl fmt::Display for ConnectionId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

pub(super) enum ReserveResult {
    Permit(ConnectionPermit),
    // Use the receiver to get notified when the existing permit is destroyed.
    Occupied(AwaitDrop, PeerSource, ConnectionId),
}

#[derive(Clone)]
pub struct PeerInfoCollector(watch::Sender<HashMap<ConnectionKey, ConnectionData>>);

impl PeerInfoCollector {
    pub fn collect(&self) -> Vec<PeerInfo> {
        self.0
            .borrow()
            .iter()
            .map(|(key, data)| data.peer_info(key.addr))
            .collect()
    }
}

/// Connection permit that prevents another connection to the same peer (socket address) to be
/// established as long as it remains in scope.
pub(super) struct ConnectionPermit {
    connections: watch::Sender<HashMap<ConnectionKey, ConnectionData>>,
    key: ConnectionKey,
    id: ConnectionId,
}

impl ConnectionPermit {
    pub fn mark_as_connecting(&self) {
        self.set_state(InnerPeerState::Connecting);
    }

    pub fn mark_as_handshaking(&self) {
        self.set_state(InnerPeerState::Handshaking);
    }

    pub fn mark_as_active(&self, runtime_id: PublicRuntimeId, peer_key: usize) {
        self.set_state(InnerPeerState::Active {
            id: runtime_id,
            since: SystemTime::now(),
            peer_key,
        });
    }

    fn set_state(&self, new_state: InnerPeerState) {
        self.connections.send_if_modified(|connections| {
            // unwrap is ok because if `self` exists then the entry should exists as well.
            let peer = connections.get_mut(&self.key).unwrap();

            if peer.state != new_state {
                peer.state = new_state;
                true
            } else {
                false
            }
        });
    }

    /// Returns a `AwaitDrop` that gets notified when this permit gets released.
    pub fn released(&self) -> AwaitDrop {
        self.with(|data| data.on_release.subscribe())
    }

    pub fn addr(&self) -> PeerAddr {
        self.key.addr
    }

    pub fn id(&self) -> ConnectionId {
        self.id
    }

    pub fn source(&self) -> PeerSource {
        self.with(|data| data.source)
    }

    pub fn byte_counters(&self) -> Arc<ByteCounters> {
        self.with(|data| data.stats_tracker.bytes.clone())
    }

    fn with<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&ConnectionData) -> R,
    {
        // unwrap is ok because if `self` exists then the entry exists as well.
        f(self.connections.borrow().get(&self.key).unwrap())
    }
}

impl fmt::Debug for ConnectionPermit {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("ConnectionPermit")
            .field("key", &self.key)
            .field("id", &self.id)
            .finish_non_exhaustive()
    }
}

impl Drop for ConnectionPermit {
    fn drop(&mut self) {
        self.connections.send_if_modified(|connections| {
            let Entry::Occupied(entry) = connections.entry(self.key) else {
                return false;
            };

            if entry.get().id != self.id {
                return false;
            }

            entry.remove();
            true
        });
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub(super) struct ConnectionKey {
    addr: PeerAddr,
    dir: ConnectionDirection,
}

pub(super) struct ConnectionData {
    id: ConnectionId,
    state: InnerPeerState,
    source: PeerSource,
    stats_tracker: StatsTracker,
    on_release: DropAwaitable,
}

impl ConnectionData {
    fn peer_info(&self, addr: PeerAddr) -> PeerInfo {
        let stats = self.stats_tracker.read();

        PeerInfo {
            addr,
            source: self.source,
            state: self.state.to_peer_state(),
            stats,
        }
    }
}

#[derive(Eq, PartialEq)]
enum InnerPeerState {
    Known,
    Connecting,
    Handshaking,
    Active {
        id: PublicRuntimeId,
        since: SystemTime,
        peer_key: usize,
    },
}

impl InnerPeerState {
    fn to_peer_state(&self) -> PeerState {
        match self {
            Self::Known => PeerState::Known,
            Self::Connecting => PeerState::Connecting,
            Self::Handshaking => PeerState::Handshaking,
            Self::Active { id, since, .. } => PeerState::Active {
                id: *id,
                since: *since,
            },
        }
    }
}