fix: Introduce ServerHandle type, stop listening on abort

Author: alixinne

src/instance.rs

@@ -4,14 +4,13 @@ use thiserror::Error;
 use tokio::select;
 use tokio::sync::broadcast;
 
-use crate::api::types::PriorityInfo;
-use crate::color::color_to16;
-use crate::models::Color;
 use crate::{
-    color::{ChannelAdjustments, ChannelAdjustmentsBuilder},
+    api::types::PriorityInfo,
+    color::{color_to16, ChannelAdjustments, ChannelAdjustmentsBuilder},
     global::Global,
     image::RawImage,
-    models::{self, DeviceConfig, InstanceConfig},
+    models::{Color, Color16, DeviceConfig, InstanceConfig},
+    servers::ServerHandle,
 };
 
 mod black_border_detector;
@@ -40,10 +39,11 @@ pub struct Instance {
     config: Arc<InstanceConfig>,
     device: Device,
     muxer: PriorityMuxer,
-    color_data: Vec<models::Color16>,
+    color_data: Vec<Color16>,
     black_border_detector: BlackBorderDetector,
     channel_adjustments: ChannelAdjustments,
     smoothing: Smoothing,
+    _boblight_server: ServerHandle,
 }
 
 impl Instance {
@@ -61,40 +61,36 @@ impl Instance {
 
         let config = Arc::new(config);
 
-        // TODO: Terminate BoblightServer on Instance drop?
-        tokio::spawn({
-            let instance = config.clone();
-            let config = instance.boblight_server.clone();
-
-            async move {
-                let result = crate::servers::bind(config, global, move |tcp, global| {
+        let _boblight_server = ServerHandle::spawn(
+            "Boblight",
+            config.boblight_server.clone(),
+            global.clone(),
+            {
+                let instance = config.clone();
+                move |tcp, global| {
                     crate::servers::boblight::handle_client(
                         tcp,
                         tx.clone(),
                         instance.clone(),
                         global,
                     )
-                })
-                .await;
-
-                if let Err(error) = result {
-                    error!("Boblight server terminated: {:?}", error);
                 }
-            }
-        });
+            },
+        );
 
         Ok(Self {
             config,
             device,
             muxer,
-            color_data: vec![models::Color16::default(); led_count],
+            color_data: vec![Color16::default(); led_count],
             black_border_detector,
             channel_adjustments,
             smoothing,
+            _boblight_server,
         })
     }
 
-    fn handle_color(&mut self, color: models::Color) {
+    fn handle_color(&mut self, color: Color) {
         let color = crate::color::color_to16(color);
         self.color_data.iter_mut().map(|x| *x = color).count();
     }
@@ -150,7 +146,7 @@ impl Instance {
                 }
             }
 
-            *value = models::Color16::from_components((
+            *value = Color16::from_components((
                 (r_acc / cnt.max(1)).max(0).min(u16::MAX as _) as u16,
                 (g_acc / cnt.max(1)).max(0).min(u16::MAX as _) as u16,
                 (b_acc / cnt.max(1)).max(0).min(u16::MAX as _) as u16,

src/main.rs

@@ -1,6 +1,7 @@
 #[macro_use]
 extern crate log;
 
+use hyperion::servers::ServerHandle;
 use structopt::StructOpt;
 use tokio::runtime::Builder;
 use tokio::signal;
@@ -79,60 +80,37 @@ async fn run(opts: Opts) -> color_eyre::eyre::Result<()> {
         });
     }
 
-    // Start the flatbuffers servers
-    if config.global.flatbuffers_server.enable {
-        tokio::spawn({
-            let global = global.clone();
-            let config = config.global.flatbuffers_server.clone();
-
-            async move {
-                let result =
-                    hyperion::servers::bind(config, global, hyperion::servers::flat::handle_client)
-                        .await;
-
-                if let Err(error) = result {
-                    error!("Flatbuffers server terminated: {:?}", error);
-                }
-            }
-        });
-    }
+    // Start the Flatbuffers servers
+    let _flatbuffers_server = if config.global.flatbuffers_server.enable {
+        Some(ServerHandle::spawn(
+            "Flatbuffers",
+            config.global.flatbuffers_server.clone(),
+            global.clone(),
+            hyperion::servers::flat::handle_client,
+        ))
+    } else {
+        None
+    };
 
     // Start the JSON server
-    tokio::spawn({
-        let global = global.clone();
-        let config = config.global.json_server.clone();
-
-        async move {
-            let result =
-                hyperion::servers::bind(config, global, hyperion::servers::json::handle_client)
-                    .await;
-
-            if let Err(error) = result {
-                error!("JSON server terminated: {:?}", error);
-            }
-        }
-    });
+    let _json_server = ServerHandle::spawn(
+        "JSON",
+        config.global.json_server.clone(),
+        global.clone(),
+        hyperion::servers::json::handle_client,
+    );
 
     // Start the Protobuf server
-    if config.global.proto_server.enable {
-        tokio::spawn({
-            let global = global.clone();
-            let config = config.global.proto_server.clone();
-
-            async move {
-                let result = hyperion::servers::bind(
-                    config,
-                    global,
-                    hyperion::servers::proto::handle_client,
-                )
-                .await;
-
-                if let Err(error) = result {
-                    error!("Protobuf server terminated: {:?}", error);
-                }
-            }
-        });
-    }
+    let _proto_server = if config.global.proto_server.enable {
+        Some(ServerHandle::spawn(
+            "Protobuf",
+            config.global.proto_server.clone(),
+            global.clone(),
+            hyperion::servers::proto::handle_client,
+        ))
+    } else {
+        None
+    };
 
     // Should we continue running?
     let mut abort = false;

src/servers.rs

@@ -1,6 +1,9 @@
 use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 
-use tokio::net::{TcpListener, TcpStream};
+use tokio::{
+    net::{TcpListener, TcpStream},
+    task::JoinHandle,
+};
 
 use crate::{global::Global, models::ServerConfig};
 
@@ -9,7 +12,7 @@ pub mod flat;
 pub mod json;
 pub mod proto;
 
-pub async fn bind<T, E, F>(
+async fn bind<T, E, F>(
     options: T,
     global: Global,
     handle_client: impl Fn((TcpStream, SocketAddr), Global) -> F,
@@ -49,3 +52,38 @@ where
         });
     }
 }
+
+pub struct ServerHandle {
+    join_handle: JoinHandle<()>,
+}
+
+impl ServerHandle {
+    pub fn spawn<T, E, F, H>(
+        name: &'static str,
+        options: T,
+        global: Global,
+        handle_client: H,
+    ) -> Self
+    where
+        T: ServerConfig + Send + 'static,
+        F: futures::Future<Output = Result<(), E>> + Send + 'static,
+        E: From<std::io::Error> + std::fmt::Display + Send + 'static,
+        H: Fn((TcpStream, SocketAddr), Global) -> F + Send + 'static,
+    {
+        Self {
+            join_handle: tokio::spawn(async move {
+                let result = bind(options, global, handle_client).await;
+
+                if let Err(error) = result {
+                    error!("{} server terminated: {}", name, error);
+                }
+            }),
+        }
+    }
+}
+
+impl Drop for ServerHandle {
+    fn drop(&mut self) {
+        self.join_handle.abort();
+    }
+}