fix unnecessary Arc Mutex on stream

Author: irvingoujAtDevolution

crates/network-scanner/examples/scan.rs

@@ -38,25 +38,21 @@ fn main() -> anyhow::Result<()> {
     let rt = tokio::runtime::Runtime::new()?;
     rt.block_on(async move {
         let scanner = NetworkScanner::new(params).unwrap();
-        let stream = scanner.start()?;
-        let stream_clone = stream.clone();
+        let mut stream = scanner.start()?;
         let now = std::time::Instant::now();
         tokio::task::spawn(async move {
             if tokio::signal::ctrl_c().await.is_ok() {
                 tracing::info!("Ctrl-C received, stopping network scan");
-                stream.stop();
             }
         });
-        while let Ok(Some(res)) = timeout(Duration::from_secs(120), stream_clone.recv())
-            .await
-            .with_context(|| {
-                tracing::error!("Failed to receive from stream");
-                "Failed to receive from stream"
-            })
-        {
+        while let Ok(Some(res)) = timeout(Duration::from_secs(120), stream.recv()).await.with_context(|| {
+            tracing::error!("Failed to receive from stream");
+            "Failed to receive from stream"
+        }) {
             tracing::warn!("Result: {:?}", res);
         }
-        stream_clone.stop();
+
+        stream.stop();
         tracing::warn!("Network Scan finished. elapsed: {:?}", now.elapsed());
         anyhow::Result::<()>::Ok(())
     })?;

crates/network-scanner/src/scanner.rs

@@ -10,7 +10,6 @@ use std::fmt::Display;
 use std::net::IpAddr;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::Mutex;
 use typed_builder::TypedBuilder;
 
 /// Represents a network scanner for discovering devices and their services over a network.
@@ -28,8 +27,8 @@ pub struct NetworkScanner {
 }
 
 impl NetworkScanner {
-    pub fn start(&self) -> anyhow::Result<Arc<NetworkScannerStream>> {
-        let mut task_executor = TaskExecutionRunner::new(self.clone())?;
+    pub fn start(&self) -> anyhow::Result<NetworkScannerStream> {
+        let (mut task_executor, result_receiver) = TaskExecutionRunner::new(self.clone())?;
 
         start_port_scan(&mut task_executor);
 
@@ -46,27 +45,27 @@ impl NetworkScanner {
         }
 
         let TaskExecutionRunner {
-            context:
-                TaskExecutionContext {
-                    result_receiver: port_receiver,
-                    mdns_daemon,
-                    ..
-                },
+            context: TaskExecutionContext { mdns_daemon, .. },
             task_manager,
         } = task_executor;
 
-        let scanner_stream = Arc::new(NetworkScannerStream {
-            result_receiver: port_receiver,
+        let task_manager_clone = task_manager.clone();
+        let mdns_daemon_clone = mdns_daemon.clone();
+
+        let scanner_stream = NetworkScannerStream {
+            result_receiver,
             task_manager,
             mdns_daemon,
-        });
+        };
 
-        let scanner_stream_clone = Arc::clone(&scanner_stream);
         let max_wait_time = Duration::from_millis(self.configs.max_wait_time);
 
         tokio::spawn(async move {
             tokio::time::sleep(max_wait_time).await;
-            scanner_stream_clone.stop();
+            task_manager_clone.stop();
+            if let Some(daemon) = &mdns_daemon_clone {
+                daemon.stop();
+            }
         });
 
         return Ok(scanner_stream);
@@ -394,24 +393,24 @@ pub enum ScanEntry {
 }
 
 pub struct NetworkScannerStream {
-    result_receiver: Arc<Mutex<ScanEntryReceiver>>,
+    result_receiver: ScanEntryReceiver,
     task_manager: TaskManager,
     mdns_daemon: Option<MdnsDaemon>,
 }
 
 impl NetworkScannerStream {
-    pub async fn recv(self: &Arc<Self>) -> Option<ScanEntry> {
+    pub async fn recv(self: &mut Self) -> Option<ScanEntry> {
         // The caller sometimes require Send, hence the Arc is necessary for socket_addr_receiver.
-        self.result_receiver.lock().await.recv().await
+        self.result_receiver.recv().await
     }
 
-    pub async fn recv_timeout(self: &Arc<Self>, duration: Duration) -> anyhow::Result<Option<ScanEntry>> {
-        tokio::time::timeout(duration, self.result_receiver.lock().await.recv())
+    pub async fn recv_timeout(self: &mut Self, duration: Duration) -> anyhow::Result<Option<ScanEntry>> {
+        tokio::time::timeout(duration, self.result_receiver.recv())
             .await
             .context("recv_timeout timed out")
     }
 
-    pub fn stop(self: Arc<Self>) {
+    pub fn stop(self: &Self) {
         self.task_manager.stop();
         if let Some(daemon) = &self.mdns_daemon {
             daemon.stop();

crates/network-scanner/src/task_utils.rs

@@ -59,7 +59,6 @@ pub(crate) struct TaskExecutionContext {
     pub(crate) ip_receiver: Arc<Mutex<IpReceiver>>,
 
     pub(crate) result_sender: ScanEntrySender,
-    pub(crate) result_receiver: Arc<Mutex<ScanEntryReceiver>>,
 
     pub(crate) ip_cache: Arc<parking_lot::RwLock<HashMap<IpAddr, Option<String>>>>,
 
@@ -79,12 +78,11 @@ pub(crate) struct TaskExecutionRunner {
 }
 
 impl TaskExecutionContext {
-    pub(crate) fn new(network_scanner: NetworkScanner) -> anyhow::Result<Self> {
+    pub(crate) fn new(network_scanner: NetworkScanner) -> anyhow::Result<(Self, ScanEntryReceiver)> {
         let (ip_sender, ip_receiver) = tokio::sync::mpsc::channel(5);
         let ip_receiver = Arc::new(Mutex::new(ip_receiver));
 
-        let (port_sender, port_receiver) = tokio::sync::mpsc::channel(100);
-        let port_receiver = Arc::new(Mutex::new(port_receiver));
+        let (result_sender, result_receiver) = tokio::sync::mpsc::channel(100);
 
         let NetworkScanner {
             mdns_daemon,
@@ -98,16 +96,15 @@ impl TaskExecutionContext {
         let res = Self {
             ip_sender,
             ip_receiver,
-            result_sender: port_sender,
-            result_receiver: port_receiver,
+            result_sender,
             ip_cache: Arc::new(parking_lot::RwLock::new(HashMap::new())),
             runtime,
             mdns_daemon,
             configs: ContextConfig::new(configs, &toggles, broadcast_subnet),
             toggles,
         };
 
-        Ok(res)
+        Ok((res, result_receiver))
     }
 }
 
@@ -122,11 +119,15 @@ impl TaskExecutionRunner {
             .spawn_no_sub_task(task(context, self.task_manager.clone()));
     }
 
-    pub(crate) fn new(scanner: NetworkScanner) -> anyhow::Result<Self> {
-        Ok(Self {
-            context: TaskExecutionContext::new(scanner)?,
-            task_manager: TaskManager::new(),
-        })
+    pub(crate) fn new(scanner: NetworkScanner) -> anyhow::Result<(Self, ScanEntryReceiver)> {
+        let (context, receiver) = TaskExecutionContext::new(scanner)?;
+        Ok((
+            Self {
+                context,
+                task_manager: TaskManager::new(),
+            },
+            receiver,
+        ))
     }
 }
 

devolutions-gateway/src/api/net.rs

@@ -49,7 +49,7 @@ pub async fn handle_network_scan(
     })?;
 
     let res = ws.on_upgrade(move |mut websocket| async move {
-        let stream = match scanner.start() {
+        let mut stream = match scanner.start() {
             Ok(stream) => stream,
             Err(e) => {
                 error!(error = format!("{e:#}"), "Failed to start network scan");