RTR server以及client的开发、test的补充

src/rtr/server/service.rs

@@ -0,0 +1,154 @@
+use std::net::SocketAddr;
+use std::path::Path;
+use std::sync::{
+    Arc,
+    atomic::{AtomicUsize, Ordering},
+};
+
+use tokio::sync::{broadcast, watch, Semaphore};
+use tokio::task::JoinHandle;
+use tracing::error;
+
+use crate::rtr::cache::SharedRtrCache;
+use crate::rtr::server::config::RtrServiceConfig;
+use crate::rtr::server::listener::RtrServer;
+use crate::rtr::server::notifier::RtrNotifier;
+
+pub struct RtrService {
+    cache: SharedRtrCache,
+    notify_tx: broadcast::Sender<()>,
+    shutdown_tx: watch::Sender<bool>,
+    connection_limiter: Arc<Semaphore>,
+    active_connections: Arc<AtomicUsize>,
+    config: RtrServiceConfig,
+}
+
+impl RtrService {
+    pub fn new(cache: SharedRtrCache) -> Self {
+        Self::with_config(cache, RtrServiceConfig::default())
+    }
+
+    pub fn with_config(cache: SharedRtrCache, config: RtrServiceConfig) -> Self {
+        let (notify_tx, _) = broadcast::channel(config.notify_queue_size);
+        let (shutdown_tx, _) = watch::channel(false);
+
+        Self {
+            cache,
+            notify_tx,
+            shutdown_tx,
+            connection_limiter: Arc::new(Semaphore::new(config.max_connections)),
+            active_connections: Arc::new(AtomicUsize::new(0)),
+            config,
+        }
+    }
+
+    pub fn cache(&self) -> SharedRtrCache {
+        self.cache.clone()
+    }
+
+    pub fn notifier(&self) -> RtrNotifier {
+        RtrNotifier::new(self.notify_tx.clone())
+    }
+
+    pub fn notify_cache_updated(&self) {
+        let _ = self.notify_tx.send(());
+    }
+
+    pub fn active_connections(&self) -> usize {
+        self.active_connections.load(Ordering::Relaxed)
+    }
+
+    pub fn max_connections(&self) -> usize {
+        self.config.max_connections
+    }
+
+    pub fn tcp_server(&self, bind_addr: SocketAddr) -> RtrServer {
+        RtrServer::new(
+            bind_addr,
+            self.cache.clone(),
+            self.notify_tx.clone(),
+            self.shutdown_tx.clone(),
+            self.connection_limiter.clone(),
+            self.active_connections.clone(),
+        )
+    }
+
+    pub fn tls_server(&self, bind_addr: SocketAddr) -> RtrServer {
+        RtrServer::new(
+            bind_addr,
+            self.cache.clone(),
+            self.notify_tx.clone(),
+            self.shutdown_tx.clone(),
+            self.connection_limiter.clone(),
+            self.active_connections.clone(),
+        )
+    }
+
+    pub fn spawn_tcp(&self, bind_addr: SocketAddr) -> JoinHandle<()> {
+        let server = self.tcp_server(bind_addr);
+        tokio::spawn(async move {
+            if let Err(err) = server.run_tcp().await {
+                error!("RTR TCP server {} exited with error: {:?}", bind_addr, err);
+            }
+        })
+    }
+
+    pub fn spawn_tls_from_pem(
+        &self,
+        bind_addr: SocketAddr,
+        cert_path: impl AsRef<Path>,
+        key_path: impl AsRef<Path>,
+    ) -> JoinHandle<()> {
+        let cert_path = cert_path.as_ref().to_path_buf();
+        let key_path = key_path.as_ref().to_path_buf();
+        let server = self.tls_server(bind_addr);
+
+        tokio::spawn(async move {
+            if let Err(err) = server.run_tls_from_pem(cert_path, key_path).await {
+                error!("RTR TLS server {} exited with error: {:?}", bind_addr, err);
+            }
+        })
+    }
+
+    pub fn spawn_tcp_and_tls_from_pem(
+        &self,
+        tcp_bind_addr: SocketAddr,
+        tls_bind_addr: SocketAddr,
+        cert_path: impl AsRef<Path>,
+        key_path: impl AsRef<Path>,
+    ) -> RunningRtrService {
+        let tcp_handle = self.spawn_tcp(tcp_bind_addr);
+        let tls_handle = self.spawn_tls_from_pem(tls_bind_addr, cert_path, key_path);
+
+        RunningRtrService {
+            shutdown_tx: self.shutdown_tx.clone(),
+            handles: vec![tcp_handle, tls_handle],
+        }
+    }
+
+    pub fn spawn_tcp_only(&self, tcp_bind_addr: SocketAddr) -> RunningRtrService {
+        let tcp_handle = self.spawn_tcp(tcp_bind_addr);
+
+        RunningRtrService {
+            shutdown_tx: self.shutdown_tx.clone(),
+            handles: vec![tcp_handle],
+        }
+    }
+}
+
+pub struct RunningRtrService {
+    shutdown_tx: watch::Sender<bool>,
+    handles: Vec<JoinHandle<()>>,
+}
+
+impl RunningRtrService {
+    pub fn shutdown(&self) {
+        let _ = self.shutdown_tx.send(true);
+    }
+
+    pub async fn wait(self) {
+        for handle in self.handles {
+            let _ = handle.await;
+        }
+    }
+}