wip

crates/lavina-core/src/clustering/mod.rs

@@ -6,7 +6,7 @@ use prost::Message;
 use serde::Deserialize;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
-use tokio::spawn;
+use tokio::{select, pin, spawn};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::RwLock;
 use tokio::task::JoinHandle;
@@ -75,7 +75,7 @@ impl ClusterRegistry {
 
             let handshake_response = read_message::<proto::HandshakeResponse>(&mut stream).await?;
             info!("Connection to the bootstrap server succeeded");
-            peer_launch(stream , inner.clone()).await;
+            peer_launch(stream, inner.clone()).await;
         }
 
         Ok(ClusterRegistry {
@@ -93,16 +93,19 @@ impl ClusterRegistry {
             if let Err(e) = k.send(PeerMsg::Close).await {
                 warn!("Failed to stop peer actor");
             }
-
         }
         Ok(())
     }
 }
 
 async fn start_connection_listener(listener: TcpListener, cluster: Arc<RwLock<ClusterRegistryInner>>) -> Terminator {
-    Terminator::spawn(|_| async move {
+    Terminator::spawn(|terminator| async move {
+        pin!(terminator);
         loop {
-            let (mut stream, addr) = match listener.accept().await {
+            select! {
+                _ = &mut terminator => break,
+                res = listener.accept() => {
+                    let (mut stream, addr) = match res {
                         Ok(stream) => stream,
                         Err(e) => {
                             warn!("Error accepting connection: {}", e);
@@ -136,6 +139,9 @@ async fn start_connection_listener(listener: TcpListener, cluster: Arc<RwLock<Cl
                         }
                     });
                 }
+            }
+        }
+        Ok(())
     })
 }