From 2988d0c5cf38e90e2d9cf4bc5f86c50d4cc4f14c Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@mycelium.engineering>
Date: Tue, 13 Feb 2024 17:01:59 +0000
Subject: [PATCH] network: implement proper backend destruction

---
 network/src/backend.rs | 31 +++++++++++++++++-------
 network/src/lib.rs     | 53 +++++++++++++++++++++++++++++-------------
 2 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/network/src/backend.rs b/network/src/backend.rs
index 8fa5384..9fc6589 100644
--- a/network/src/backend.rs
+++ b/network/src/backend.rs
@@ -16,6 +16,7 @@ use smoltcp::time::Instant;
 use smoltcp::wire::{HardwareAddress, IpCidr};
 use tokio::select;
 use tokio::sync::mpsc::{channel, Receiver};
+use tokio::task::JoinHandle;
 
 const TX_CHANNEL_BUFFER_LEN: usize = 300;
 
@@ -43,7 +44,7 @@ struct NetworkStack<'a> {
 }
 
 impl NetworkStack<'_> {
-    async fn poll(&mut self) -> Result<()> {
+    async fn poll(&mut self) -> Result<bool> {
         let what = select! {
             x = self.kdev.receiver.recv() => NetworkStackSelect::Receive(x),
             x = self.bridge.from_bridge_receiver.recv() => NetworkStackSelect::Send(x),
@@ -76,13 +77,14 @@ impl NetworkStack<'_> {
                 }
             }
 
-            NetworkStackSelect::Receive(None) => {}
-            NetworkStackSelect::Send(None) => {}
+            NetworkStackSelect::Receive(None) | NetworkStackSelect::Send(None) => {
+                return Ok(false);
+            }
 
             NetworkStackSelect::Reclaim => {}
         }
 
-        Ok(())
+        Ok(true)
     }
 }
 
@@ -112,8 +114,11 @@ impl NetworkBackend {
     pub async fn run(&self) -> Result<()> {
         let mut stack = self.create_network_stack().await?;
         loop {
-            stack.poll().await?;
+            if !stack.poll().await? {
+                break;
+            }
         }
+        Ok(())
     }
 
     async fn create_network_stack(&self) -> Result<NetworkStack> {
@@ -152,8 +157,8 @@ impl NetworkBackend {
         })
     }
 
-    pub async fn launch(self) -> Result<()> {
-        tokio::task::spawn(async move {
+    pub async fn launch(self) -> Result<JoinHandle<()>> {
+        Ok(tokio::task::spawn(async move {
             info!(
                 "lauched network backend for hypha guest {}",
                 self.metadata.uuid
@@ -164,7 +169,15 @@ impl NetworkBackend {
                     self.metadata.uuid, error
                 );
             }
-        });
-        Ok(())
+        }))
+    }
+}
+
+impl Drop for NetworkBackend {
+    fn drop(&mut self) {
+        info!(
+            "destroyed network backend for hypha guest {}",
+            self.metadata.uuid
+        );
     }
 }
diff --git a/network/src/lib.rs b/network/src/lib.rs
index 032d02a..37f85fd 100644
--- a/network/src/lib.rs
+++ b/network/src/lib.rs
@@ -1,10 +1,10 @@
-use std::time::Duration;
+use std::{collections::HashMap, time::Duration};
 
 use anyhow::Result;
 use autonet::{AutoNetworkChangeset, AutoNetworkCollector, NetworkMetadata};
 use futures::{future::join_all, TryFutureExt};
 use log::warn;
-use tokio::time::sleep;
+use tokio::{task::JoinHandle, time::sleep};
 use uuid::Uuid;
 use vbridge::VirtualBridge;
 
@@ -21,12 +21,14 @@ pub mod raw_socket;
 pub mod vbridge;
 
 pub struct NetworkService {
+    pub backends: HashMap<Uuid, JoinHandle<()>>,
     pub bridge: VirtualBridge,
 }
 
 impl NetworkService {
     pub fn new() -> Result<NetworkService> {
         Ok(NetworkService {
+            backends: HashMap::new(),
             bridge: VirtualBridge::new()?,
         })
     }
@@ -47,30 +49,47 @@ impl NetworkService {
         collector: &mut AutoNetworkCollector,
         changeset: AutoNetworkChangeset,
     ) -> Result<()> {
+        for removal in &changeset.removed {
+            if let Some(handle) = self.backends.remove(&removal.uuid) {
+                handle.abort();
+            }
+        }
+
         let futures = changeset
             .added
             .iter()
             .map(|metadata| {
-                self.add_network_backend(metadata.clone())
+                self.add_network_backend(metadata)
                     .map_err(|x| (metadata.clone(), x))
             })
             .collect::<Vec<_>>();
 
-        let failed = futures::executor::block_on(async move {
+        let (launched, failed) = futures::executor::block_on(async move {
             let mut failed: Vec<Uuid> = Vec::new();
+            let mut launched: Vec<(Uuid, JoinHandle<()>)> = Vec::new();
             let results = join_all(futures).await;
             for result in results {
-                if let Err((metadata, error)) = result {
-                    warn!(
-                        "failed to launch network backend for hypha guest {}: {}",
-                        metadata.uuid, error
-                    );
-                    failed.push(metadata.uuid);
-                }
+                match result {
+                    Ok(launch) => {
+                        launched.push(launch);
+                    }
+
+                    Err((metadata, error)) => {
+                        warn!(
+                            "failed to launch network backend for hypha guest {}: {}",
+                            metadata.uuid, error
+                        );
+                        failed.push(metadata.uuid);
+                    }
+                };
             }
-            failed
+            (launched, failed)
         });
 
+        for (uuid, handle) in launched {
+            self.backends.insert(uuid, handle);
+        }
+
         for uuid in failed {
             collector.mark_unknown(uuid)?;
         }
@@ -78,10 +97,12 @@ impl NetworkService {
         Ok(())
     }
 
-    async fn add_network_backend(&self, metadata: NetworkMetadata) -> Result<()> {
-        let mut network = NetworkBackend::new(metadata, self.bridge.clone())?;
+    async fn add_network_backend(
+        &self,
+        metadata: &NetworkMetadata,
+    ) -> Result<(Uuid, JoinHandle<()>)> {
+        let mut network = NetworkBackend::new(metadata.clone(), self.bridge.clone())?;
         network.init().await?;
-        network.launch().await?;
-        Ok(())
+        Ok((metadata.uuid, network.launch().await?))
     }
 }