fix: Unregister instances when they are dropped

Author: alixinne

src/global.rs

@@ -79,6 +79,10 @@ impl Global {
         self.0.write().await.register_instance(handle);
     }
 
+    pub async fn unregister_instance(&self, id: i32) {
+        self.0.write().await.unregister_instance(id);
+    }
+
     pub async fn get_instance(&self, id: i32) -> Option<InstanceHandle> {
         self.0.read().await.instances.get(&id).cloned()
     }
@@ -156,6 +160,14 @@ impl GlobalData {
     }
 
     fn register_instance(&mut self, handle: InstanceHandle) {
-        self.instances.insert(handle.id(), handle);
+        let id = handle.id();
+        self.instances.insert(id, handle);
+        info!("registered instance {}", id);
+    }
+
+    fn unregister_instance(&mut self, id: i32) {
+        if let Some(_) = self.instances.remove(&id) {
+            info!("unregistered instance {}", id);
+        }
     }
 }

src/main.rs

@@ -32,17 +32,23 @@ async fn run(opts: Opts) -> color_eyre::eyre::Result<()> {
     let global = hyperion::global::GlobalData::new(&config).wrap();
 
     // Initialize and spawn the devices
-    for (_, inst) in &config.instances {
+    for (&id, inst) in &config.instances {
         // Create the instance
         let (inst, handle) = hyperion::instance::Instance::new(global.clone(), inst.clone()).await;
         // Register the instance globally using its handle
         global.register_instance(handle).await;
         // Run the instance futures
-        tokio::spawn(async move {
-            let result = inst.run().await;
+        tokio::spawn({
+            let global = global.clone();
 
-            if let Err(error) = result {
-                error!("Instance error: {:?}", error);
+            async move {
+                let result = inst.run().await;
+
+                if let Err(error) = result {
+                    error!("Instance error: {:?}", error);
+                }
+
+                global.unregister_instance(id).await;
             }
         });
     }