use std::{future::Future, sync::Arc}; use log::{debug, error, info}; use karyons_core::{ async_util::{TaskGroup, TaskResult}, crypto::KeyPair, GlobalExecutor, }; use karyons_net::{listen, tls, Conn, Endpoint, Listener as NetListener}; use crate::{ monitor::{ConnEvent, Monitor}, slots::ConnectionSlots, tls_config::tls_server_config, Error, Result, }; /// Responsible for creating inbound connections with other peers. pub struct Listener { /// Identity Key pair key_pair: KeyPair, /// Managing spawned tasks. task_group: TaskGroup<'static>, /// Manages available inbound slots. connection_slots: Arc, /// Enables secure connection. enable_tls: bool, /// Responsible for network and system monitoring. monitor: Arc, } impl Listener { /// Creates a new Listener pub fn new( key_pair: &KeyPair, connection_slots: Arc, enable_tls: bool, monitor: Arc, ex: GlobalExecutor, ) -> Arc { Arc::new(Self { key_pair: key_pair.clone(), connection_slots, task_group: TaskGroup::new(ex), enable_tls, monitor, }) } /// Starts a listener on the given `endpoint`. For each incoming connection /// that is accepted, it invokes the provided `callback`, and pass the /// connection to the callback. /// /// Returns the resloved listening endpoint. pub async fn start( self: &Arc, endpoint: Endpoint, // https://github.com/rust-lang/rfcs/pull/2132 callback: impl FnOnce(Conn) -> Fut + Clone + Send + 'static, ) -> Result where Fut: Future> + Send + 'static, { let listener = match self.listend(&endpoint).await { Ok(listener) => { self.monitor .notify(&ConnEvent::Listening(endpoint.clone()).into()) .await; listener } Err(err) => { error!("Failed to listen on {endpoint}: {err}"); self.monitor .notify(&ConnEvent::ListenFailed(endpoint).into()) .await; return Err(err); } }; let resolved_endpoint = listener.local_endpoint()?; info!("Start listening on {resolved_endpoint}"); let selfc = self.clone(); self.task_group .spawn(selfc.listen_loop(listener, callback), |_| async {}); Ok(resolved_endpoint) } /// Shuts down the listener pub async fn shutdown(&self) { self.task_group.cancel().await; } async fn listen_loop( self: Arc, listener: Box, callback: impl FnOnce(Conn) -> Fut + Clone + Send + 'static, ) where Fut: Future> + Send + 'static, { loop { // Wait for an available inbound slot. self.connection_slots.wait_for_slot().await; let result = listener.accept().await; let (conn, endpoint) = match result { Ok(c) => { let endpoint = match c.peer_endpoint() { Ok(e) => e, Err(err) => { self.monitor.notify(&ConnEvent::AcceptFailed.into()).await; error!("Failed to accept a new connection: {err}"); continue; } }; self.monitor .notify(&ConnEvent::Accepted(endpoint.clone()).into()) .await; (c, endpoint) } Err(err) => { error!("Failed to accept a new connection: {err}"); self.monitor.notify(&ConnEvent::AcceptFailed.into()).await; continue; } }; self.connection_slots.add(); let selfc = self.clone(); let on_disconnect = |res| async move { if let TaskResult::Completed(Err(err)) = res { debug!("Inbound connection dropped: {err}"); } selfc .monitor .notify(&ConnEvent::Disconnected(endpoint).into()) .await; selfc.connection_slots.remove().await; }; let callback = callback.clone(); self.task_group.spawn(callback(conn), on_disconnect); } } async fn listend(&self, endpoint: &Endpoint) -> Result> { if self.enable_tls { let tls_config = tls_server_config(&self.key_pair)?; tls::listen(endpoint, tls_config).await } else { listen(endpoint).await } .map_err(Error::KaryonsNet) } }