improve cancellation

src/server/capture_task.rs

@@ -123,15 +123,22 @@ async fn do_capture(
         .await;
 
     // FIXME DUPLICATES
-    // let clients = server
-    //     .client_manager
-    //     .borrow()
-    //     .get_client_states()
-    //     .map(|(h, (c, _))| (h, c.pos))
-    //     .collect::<Vec<_>>();
-    // for (handle, pos) in clients {
-    //     capture.create(handle, pos.into()).await?;
-    // }
+    let clients = server
+        .client_manager
+        .borrow()
+        .get_client_states()
+        .map(|(h, s)| (h, s.clone()))
+        .collect::<Vec<_>>();
+    log::info!("{clients:?}");
+    let clients = server
+        .client_manager
+        .borrow()
+        .get_client_states()
+        .map(|(h, (c, _))| (h, c.pos))
+        .collect::<Vec<_>>();
+    for (handle, pos) in clients {
+        capture.create(handle, pos.into()).await?;
+    }
 
     let mut pressed_keys = HashSet::new();
     loop {

src/server/network_task.rs

@@ -6,6 +6,7 @@ use tokio::{
     sync::mpsc::{Receiver, Sender},
     task::JoinHandle,
 };
+use tokio_util::sync::CancellationToken;
 
 use crate::frontend::FrontendEvent;
 use input_event::{Event, ProtocolError};
@@ -15,6 +16,7 @@ use super::Server;
 pub async fn new(
     server: Server,
     frontend_notify_tx: Sender<FrontendEvent>,
+    cancellation_token: CancellationToken,
 ) -> io::Result<(
     JoinHandle<()>,
     Sender<(Event, SocketAddr)>,
@@ -38,8 +40,9 @@ pub async fn new(
                 _ = udp_sender => break, /* channel closed */
                 port = port_rx.recv() => match port {
                     Some(port) => update_port(&server, &frontend_notify_tx, &mut socket, port).await,
-                    _ => break,
-                }
+                    _ => continue,
+                },
+                _ = cancellation_token.cancelled() => break, /* cancellation requested */
             }
         }
     });
@@ -80,18 +83,13 @@ async fn udp_receiver(
 ) {
     loop {
         let event = receive_event(socket).await;
-        if receiver_tx.send(event).await.is_err() {
-            break;
-        }
+        receiver_tx.send(event).await.expect("channel closed");
     }
 }
 
 async fn udp_sender(socket: &UdpSocket, rx: &mut Receiver<(Event, SocketAddr)>) {
     loop {
-        let (event, addr) = match rx.recv().await {
-            Some(e) => e,
-            None => return,
-        };
+        let (event, addr) = rx.recv().await.expect("channel closed");
         if let Err(e) = send_event(socket, event, addr) {
             log::warn!("udp send failed: {e}");
         };

src/server/resolver_task.rs

@@ -1,6 +1,10 @@
 use std::collections::HashSet;
 
-use tokio::{sync::mpsc::Sender, task::JoinHandle};
+use tokio::{
+    sync::mpsc::{Receiver, Sender},
+    task::JoinHandle,
+};
+use tokio_util::sync::CancellationToken;
 
 use crate::{client::ClientHandle, dns::DnsResolver, frontend::FrontendEvent};
 
@@ -14,52 +18,65 @@ pub struct DnsRequest {
 
 pub fn new(
     resolver: DnsResolver,
-    mut server: Server,
-    mut frontend: Sender<FrontendEvent>,
+    server: Server,
+    frontend: Sender<FrontendEvent>,
+    cancellation_token: CancellationToken,
 ) -> (JoinHandle<()>, Sender<DnsRequest>) {
-    let (dns_tx, mut dns_rx) = tokio::sync::mpsc::channel::<DnsRequest>(32);
+    let (dns_tx, dns_rx) = tokio::sync::mpsc::channel::<DnsRequest>(32);
     let resolver_task = tokio::task::spawn_local(async move {
-        loop {
-            let (host, handle) = match dns_rx.recv().await {
-                Some(r) => (r.hostname, r.handle),
-                None => break,
-            };
-
-            /* update resolving status */
-            if let Some((_, s)) = server.client_manager.borrow_mut().get_mut(handle) {
-                s.resolving = true;
-            }
-            notify_state_change(&mut frontend, &mut server, handle).await;
-
-            let ips = match resolver.resolve(&host).await {
-                Ok(ips) => ips,
-                Err(e) => {
-                    log::warn!("could not resolve host '{host}': {e}");
-                    vec![]
-                }
-            };
-
-            /* update ips and resolving state */
-            if let Some((c, s)) = server.client_manager.borrow_mut().get_mut(handle) {
-                let mut addrs = HashSet::from_iter(c.fix_ips.iter().cloned());
-                for ip in ips {
-                    addrs.insert(ip);
-                }
-                s.ips = addrs;
-                s.resolving = false;
-            }
-            notify_state_change(&mut frontend, &mut server, handle).await;
+        tokio::select! {
+            _ = cancellation_token.cancelled() => {},
+            _ = do_dns(resolver, server, frontend, dns_rx) => {},
         }
     });
     (resolver_task, dns_tx)
 }
 
+async fn do_dns(
+    resolver: DnsResolver,
+    server: Server,
+    frontend: Sender<FrontendEvent>,
+    mut dns_rx: Receiver<DnsRequest>,
+) {
+    loop {
+        let (host, handle) = match dns_rx.recv().await {
+            Some(r) => (r.hostname, r.handle),
+            None => break,
+        };
+
+        /* update resolving status */
+        if let Some((_, s)) = server.client_manager.borrow_mut().get_mut(handle) {
+            s.resolving = true;
+        }
+        notify_state_change(&frontend, &server, handle).await;
+
+        let ips = match resolver.resolve(&host).await {
+            Ok(ips) => ips,
+            Err(e) => {
+                log::warn!("could not resolve host '{host}': {e}");
+                vec![]
+            }
+        };
+
+        /* update ips and resolving state */
+        if let Some((c, s)) = server.client_manager.borrow_mut().get_mut(handle) {
+            let mut addrs = HashSet::from_iter(c.fix_ips.iter().cloned());
+            for ip in ips {
+                addrs.insert(ip);
+            }
+            s.ips = addrs;
+            s.resolving = false;
+        }
+        notify_state_change(&frontend, &server, handle).await;
+    }
+}
+
 async fn notify_state_change(
-    frontend: &mut Sender<FrontendEvent>,
-    server: &mut Server,
+    frontend: &Sender<FrontendEvent>,
+    server: &Server,
     handle: ClientHandle,
 ) {
-    let state = server.client_manager.borrow_mut().get_mut(handle).cloned();
+    let state = server.client_manager.borrow().get(handle).cloned();
     if let Some((config, state)) = state {
         let _ = frontend
             .send(FrontendEvent::State(handle, config, state))