Handle sigint and sigterm

Author: Virv12

pixie-server/Cargo.lock

@@ -28,6 +28,7 @@ http-body-util = "0.1.3"
 futures = "0.3.30"
 tokio-stream = { version = "0.1.17", features = ["sync"] }
 chrono = "0.4.41"
+tokio-util = "0.7.15"
 
 [dependencies.pixie-shared]
 path = "../pixie-shared"

pixie-server/Cargo.toml

@@ -139,6 +139,7 @@ pub async fn main(state: Arc<State>) -> Result<()> {
         tokio::select! {
             ret = units_rx.changed() => ret.unwrap(),
             ret = hostmap_rx.changed() => ret.unwrap(),
+            _ = state.cancel_token.cancelled() => break,
         }
 
         let hosts2 = get_hosts(
@@ -151,4 +152,5 @@ pub async fn main(state: Arc<State>) -> Result<()> {
             dnsmasq.reload()?;
         }
     }
+    Ok(())
 }

pixie-server/src/dnsmasq.rs

@@ -162,13 +162,15 @@ async fn status(extract::State(state): extract::State<Arc<State>>) -> impl IntoR
     let image_rx = WatchStream::new(state.subscribe_images());
     let hostmap_rx = WatchStream::new(state.subscribe_hostmap());
 
-    let messages = futures::stream::iter(initial_messages).chain(futures::stream::select(
-        futures::stream::select(
-            image_rx.map(StatusUpdate::ImagesStats),
-            units_rx.map(StatusUpdate::Units),
-        ),
-        hostmap_rx.map(StatusUpdate::HostMap),
-    ));
+    let messages = futures::stream::iter(initial_messages)
+        .chain(futures::stream::select(
+            futures::stream::select(
+                image_rx.map(StatusUpdate::ImagesStats),
+                units_rx.map(StatusUpdate::Units),
+            ),
+            hostmap_rx.map(StatusUpdate::HostMap),
+        ))
+        .take_until(state.cancel_token.clone().cancelled_owned());
     let lines = messages.map(|msg| serde_json::to_string(&msg).map(|x| x + "\n"));
 
     Response::builder()
@@ -208,8 +210,11 @@ pub async fn main(state: Arc<State>) -> Result<()> {
     }
     router = router.layer(TraceLayer::new_for_http());
 
+    let shutdown_token = state.cancel_token.clone().cancelled_owned();
     let listener = TcpListener::bind(listen_on).await?;
-    axum::serve(listener, router.with_state(state)).await?;
+    axum::serve(listener, router.with_state(state))
+        .with_graceful_shutdown(shutdown_token)
+        .await?;
 
     Ok(())
 }

pixie-server/src/http.rs

@@ -19,7 +19,7 @@ use std::{
     process::{Child, Command, Stdio},
     sync::Arc,
 };
-use tokio::task::JoinHandle;
+use tokio::{signal::unix::SignalKind, task::JoinHandle};
 
 /// Finds the mac address for the given ip.
 ///
@@ -130,11 +130,26 @@ async fn main() -> Result<()> {
 
     let state = Arc::new(State::load(options.storage_dir)?);
 
+    tokio::spawn({
+        let cancel_token = state.cancel_token.clone();
+        let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
+            .expect("failed to register SIGNINT handler");
+        let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
+            .expect("failed to register SIGNTERM handler");
+        async move {
+            let _token_guard = cancel_token.drop_guard();
+            tokio::select! {
+                _ = sigint.recv() => log::info!("Received SIGINT."),
+                _ = sigterm.recv() => log::info!("Received SIGTERM."),
+            };
+        }
+    });
+
     let state2 = state.clone();
     tokio::spawn(async move {
-        let mut signal = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup())
-            .expect("failed to register signal handler");
-        while let Some(()) = signal.recv().await {
+        let mut sighup = tokio::signal::unix::signal(SignalKind::hangup())
+            .expect("failed to register SIGHUP handler");
+        while let Some(()) = sighup.recv().await {
             state2.reload().unwrap();
         }
     });

pixie-server/src/main.rs

@@ -19,7 +19,10 @@ pub async fn main(state: Arc<State>) -> Result<()> {
 
     let mut buf = [0; PACKET_LEN];
     loop {
-        let (len, peer_addr) = socket.recv_from(&mut buf).await?;
+        let (len, peer_addr) = tokio::select! {
+            x = socket.recv_from(&mut buf) => x?,
+            _ = state.cancel_token.cancelled() => break,
+        };
         let IpAddr::V4(peer_ip) = peer_addr.ip() else {
             bail!("IPv6 is not supported")
         };
@@ -37,4 +40,5 @@ pub async fn main(state: Arc<State>) -> Result<()> {
             .as_secs();
         state.set_unit_ping(UnitSelector::MacAddr(peer_mac), time, &buf[..len]);
     }
+    Ok(())
 }

pixie-server/src/ping.rs

@@ -29,6 +29,7 @@ use std::{
     },
 };
 use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
 
 pub use units::UnitSelector;
 
@@ -99,6 +100,9 @@ pub struct State {
     registration_hint: Mutex<Option<RegistrationInfo>>,
     images_stats: watch::Sender<ImagesStats>,
     chunks_stats: Mutex<ChunksStats>,
+
+    /// Token to shutdown the entire server.
+    pub cancel_token: CancellationToken,
 }
 
 impl State {
@@ -206,6 +210,8 @@ impl State {
         let run_dir = PathBuf::from(format!("/run/pixie-{}", std::process::id()));
         std::fs::create_dir(&run_dir)?;
 
+        let cancel_token = CancellationToken::new();
+
         Ok(Self {
             storage_dir,
             run_dir,
@@ -215,6 +221,7 @@ impl State {
             registration_hint: Mutex::new(None),
             images_stats: watch::Sender::new(images_stats),
             chunks_stats: Mutex::new(chunks_stats),
+            cancel_token,
         })
     }
 

pixie-server/src/state/mod.rs

@@ -87,12 +87,16 @@ pub async fn main(state: Arc<State>) -> Result<()> {
     let listener = TcpListener::bind((Ipv4Addr::UNSPECIFIED, ACTION_PORT)).await?;
     log::info!("Listening on {}", listener.local_addr()?);
     loop {
-        let (stream, addr) = listener.accept().await?;
+        let (stream, addr) = tokio::select! {
+            x = listener.accept() => x?,
+            _ = state.cancel_token.cancelled() => break,
+        };
         let state = state.clone();
         tokio::spawn(async move {
             if let Err(e) = handle_connection(state, stream, addr).await {
                 log::error!("Error handling tcp connection: {}", e);
             }
         });
     }
+    Ok(())
 }

pixie-server/src/tcp.rs

@@ -33,25 +33,34 @@ async fn broadcast_chunks(
     let mut index = [0; 32];
 
     loop {
-        while let Ok(hash) = rx.try_recv() {
-            queue.insert(hash);
-        }
+        let get_index = async {
+            while let Ok(hash) = rx.try_recv() {
+                queue.insert(hash);
+            }
+
+            let hash = queue
+                .range((Bound::Excluded(index), Bound::Unbounded))
+                .next()
+                .or_else(|| queue.iter().next())
+                .copied();
 
-        let hash = queue
-            .range((Bound::Excluded(index), Bound::Unbounded))
-            .next()
-            .or_else(|| queue.iter().next())
-            .copied();
+            match hash {
+                Some(hash) => {
+                    queue.remove(&hash);
+                    Some(hash)
+                }
+                None => rx.recv().await,
+            }
+        };
 
-        index = match hash {
-            Some(hash) => {
-                queue.remove(&hash);
-                hash
+        tokio::select! {
+            hash = get_index => {
+                let Some(hash) = hash else {
+                    break;
+                };
+                index = hash;
             }
-            None => match rx.recv().await {
-                Some(hash) => hash,
-                None => break,
-            },
+            _ = state.cancel_token.cancelled() => break,
         };
 
         let Some(data) = state
@@ -162,6 +171,10 @@ fn compute_hint(state: &State) -> Result<RegistrationInfo> {
 
 async fn broadcast_hint(state: &State, socket: &UdpSocket, ip: Ipv4Addr) -> Result<()> {
     loop {
+        tokio::select! {
+            _ = time::sleep(Duration::from_secs(1)) => {}
+            _ = state.cancel_token.cancelled() => break,
+        }
         let hint = HintPacket {
             station: compute_hint(state)?,
             images: state.config.images.clone(),
@@ -170,14 +183,17 @@ async fn broadcast_hint(state: &State, socket: &UdpSocket, ip: Ipv4Addr) -> Resu
         let data = postcard::to_allocvec(&hint)?;
         let hint_addr = SocketAddrV4::new(ip, HINT_PORT);
         socket.send_to(&data, hint_addr).await?;
-        time::sleep(Duration::from_secs(1)).await;
     }
+    Ok(())
 }
 
 async fn handle_requests(state: &State, socket: &UdpSocket, tx: Sender<[u8; 32]>) -> Result<()> {
     let mut buf = [0; PACKET_LEN];
-    'recv_packet: loop {
-        let (len, addr) = socket.recv_from(&mut buf).await?;
+    loop {
+        let (len, addr) = tokio::select! {
+            x = socket.recv_from(&mut buf) => x?,
+            _ = state.cancel_token.cancelled() => break,
+        };
         let req: postcard::Result<UdpRequest> = postcard::from_bytes(&buf[..len]);
         match req {
             Ok(UdpRequest::Discover) => {
@@ -187,14 +203,14 @@ async fn handle_requests(state: &State, socket: &UdpSocket, tx: Sender<[u8; 32]>
                 let IpAddr::V4(peer_ip) = addr.ip() else {
                     bail!("IPv6 is not supported")
                 };
-                let peer_mac = match find_mac(peer_ip) {
-                    Ok(peer_mac) => peer_mac,
+                match find_mac(peer_ip) {
+                    Ok(peer_mac) => {
+                        state.set_unit_progress(UnitSelector::MacAddr(peer_mac), Some((frac, tot)));
+                    }
                     Err(err) => {
                         log::error!("Error handling udp packet: {}", err);
-                        continue 'recv_packet;
                     }
                 };
-                state.set_unit_progress(UnitSelector::MacAddr(peer_mac), Some((frac, tot)));
             }
             Ok(UdpRequest::RequestChunks(chunks)) => {
                 for hash in chunks {
@@ -206,6 +222,7 @@ async fn handle_requests(state: &State, socket: &UdpSocket, tx: Sender<[u8; 32]>
             }
         }
     }
+    Ok(())
 }
 
 pub async fn main(state: Arc<State>) -> Result<()> {