teanet-rust/src/lib.rs

use std::{
    collections::VecDeque,
    net::{SocketAddr, UdpSocket},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc::{Receiver, channel},
    },
    thread::spawn,
    time::Duration,
};

use serde::{Serialize, de::DeserializeOwned};
use thiserror::*;

use crate::{
    connections::{Connection, ConnectionError, ConnectionManager},
    listener::{Listener, ListenerMessage},
    package::CloseReason,
    stats::NetStats,
};

pub mod connections;
pub(crate) mod listener;
pub mod package;
pub mod stats;

/// Error for the Peer
#[derive(Error, Debug)]
pub enum PeerError {
    #[error("Error binding to port: {0}")]
    BindError(std::io::Error),
    #[error("Error listening to messages: {0}")]
    ListenerError(std::io::Error),
}

/// Possible messages from occurring events from the peer
pub enum PeerMessage<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> {
    /// A new connection has been connected
    NewConnection(Connection<T>),
    /// An existing connection has disconnected, with an optional error
    Disconnected(Connection<T>, Option<ConnectionError>, CloseReason),
    /// The Peer has been closed
    Closed,
    /// A single message of type T
    Message(Connection<T>, T),
}

/// Optional configuration available for peers, mainly configuration of
/// different intervals
#[derive(Debug, Clone)]
pub struct PeerConfig {
    ping_interval: Duration,
    timeout: Duration,
    disconnect_timeout: Duration,
    message_retry: Duration,
    identifier: String,
}

impl Default for PeerConfig {
    fn default() -> Self {
        Self {
            ping_interval: Duration::from_millis(100),
            timeout: Duration::from_millis(2000),
            disconnect_timeout: Duration::from_millis(500),
            message_retry: Duration::from_millis(100),
            identifier: "ExamplePeer".to_string(),
        }
    }
}

impl PeerConfig {
    /// Set the desired interval between "ping"s
    pub fn with_ping_interval(self, interval: Duration) -> PeerConfig {
        PeerConfig {
            ping_interval: interval,
            ..self
        }
    }

    /// Set the length of timeout; that is to say the duration of time which is
    /// acceptable to occur between pings before a Timeout-error occurs
    pub fn with_timeout(self, timeout: Duration) -> PeerConfig {
        PeerConfig { timeout, ..self }
    }

    /// Set the length of time that is awaited since the last "Closed" message
    /// during disconnection before disconnection actually occurs. Should be
    /// significantly higher than ping interval.
    pub fn with_disconnect(self, timeout: Duration) -> PeerConfig {
        PeerConfig {
            disconnect_timeout: timeout,
            ..self
        }
    }

    /// Sets the duration which is awaited before any queued reliable messages
    /// are sent again. This timer is also reset every time any new messages are
    /// sent manually.
    pub fn with_retry(self, retry: Duration) -> PeerConfig {
        PeerConfig {
            message_retry: retry,
            ..self
        }
    }

    /// Sets the identifier string which must be same for both connecting peers
    /// in order for connection to succeed.
    pub fn with_identifier(self, ident: String) -> PeerConfig {
        PeerConfig {
            identifier: ident,
            ..self
        }
    }
}

pub struct Peer<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static> {
    connection_mgr: ConnectionManager<T>,
    closed: Arc<AtomicBool>,
    receiver: Receiver<ListenerMessage<T>>,
    messages: VecDeque<PeerMessage<T>>,
}
impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Drop
    for Peer<T>
{
    fn drop(&mut self) {
        self.closed.store(true, Ordering::Relaxed);
    }
}

impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Peer<T> {
    /// Bind given port to listen to incoming messages. Creates a new Peer that
    /// can establish new connections.
    pub fn listen(port: Option<u16>, config: PeerConfig) -> Result<Peer<T>, PeerError> {
        let socket = Arc::new(
            UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], port.unwrap_or(0))))
                .map_err(|e| PeerError::BindError(e))?,
        );
        let closed = Arc::new(AtomicBool::new(false));

        let (sender, receiver) = channel();

        spawn({
            let socket = socket.clone();
            let closed = closed.clone();

            if let Err(e) = socket.set_nonblocking(true) {
                sender.send(ListenerMessage::Error(e)).ok();
            }

            move || {
                let mut listener = Listener::new(socket, sender);

                while !closed.load(Ordering::Relaxed) {
                    listener.poll();
                }
            }
        });

        Ok(Peer {
            connection_mgr: ConnectionManager::new(socket, config.clone()),
            closed,
            receiver,
            messages: VecDeque::new(),
        })
    }

    /// Connect the peer to the given address.
    pub fn connect_to(&mut self, addr: SocketAddr) {
        self.connection_mgr.connect_to(addr);
    }

    /// Send pings and poll for events from the listener thread. Returns
    /// possible events, such as connections and messages. Should be called as
    /// often as possible.
    #[must_use]
    pub fn poll(&mut self) -> Result<Option<PeerMessage<T>>, PeerError> {
        if self.connection_mgr.udp.is_closed() {
            if self.closed.load(Ordering::Acquire) {
                self.messages.push_back(PeerMessage::Closed);
            }
            self.closed.store(true, Ordering::Release);
        }

        self.connection_mgr.send_pings();
        self.connection_mgr.send_queued_messages();

        let mut msg;
        while {
            msg = match self.receiver.try_recv() {
                Ok(msg) => Some(msg),
                Err(_) => None,
            };

            msg.is_some()
        } {
            match msg.unwrap() {
                ListenerMessage::Package(package, socket_addr, bytes) => {
                    self.messages.extend(self.connection_mgr.handle_package(
                        package,
                        &socket_addr,
                        bytes,
                    ));
                }
                ListenerMessage::PackageError(listener_error, socket_addr) => {
                    println!("Error: {}", listener_error);
                    self.connection_mgr.error_connection(
                        &socket_addr,
                        ConnectionError::ListenerError(listener_error),
                    );
                }
                ListenerMessage::Error(error) => {
                    println!("Error {}", error);
                    self.close();
                    return Err(PeerError::ListenerError(error));
                }
            }
        }

        self.connection_mgr.close_silent_connections();

        self.messages
            .extend(self.connection_mgr.purge_old_connections());

        Ok(self.messages.pop_front())
    }

    /// Send a reliable message, which is retried until the remote peer
    /// acknowledges it.
    pub fn send_reliable(&mut self, to: &SocketAddr, message: T) {
        self.connection_mgr.send(to, message, true);
    }

    /// Sends an unreliable message, which is never retried again.
    pub fn send_unreliable(&mut self, to: &SocketAddr, message: T) {
        self.connection_mgr.send(to, message, false);
    }

    /// Send a reliable message to every connection currently connected
    pub fn broadcast_reliable(&mut self, message: T) {
        for connection in self.connections() {
            self.connection_mgr
                .send(&connection.address, message.clone(), true);
        }
    }

    /// Send an unreliable message to every connection currently connected
    pub fn broadcast_unreliable(&mut self, message: T) {
        for connection in self.connections() {
            self.connection_mgr
                .send(&connection.address, message.clone(), false);
        }
    }

    /// Close the Peer, disconnecting all connections.
    pub fn close(&mut self) {
        self.connection_mgr.close();
    }

    pub fn close_connection(&mut self, addr: &SocketAddr, reason: CloseReason) {
        self.connection_mgr.close_connection(addr, reason);
    }

    /// List of all current connections
    pub fn connections(&self) -> Vec<Connection<T>> {
        self.connection_mgr.connections()
    }

    pub fn set_accepting_connections(&mut self, accepting: bool) {
        self.connection_mgr.set_accepting_connections(accepting);
    }

    pub fn statistics(&self) -> NetStats {
        let conns = self.connection_mgr.connections();
        let all_stats = conns.iter().map(|c| c.statistics());
        let mut total_stats = NetStats::default();
        for stats in all_stats {
            total_stats.bytes_rx += stats.bytes_rx;
            total_stats.bytes_tx += stats.bytes_tx;
            total_stats.messages_expected += stats.messages_expected;
            total_stats.messages_received += stats.messages_received;
        }
        total_stats
    }
}