refactor: Split server binding and accepting loop

Author: alixinne

src/instance.rs

@@ -10,7 +10,7 @@ use crate::{
     global::Global,
     image::RawImage,
     models::{Color, Color16, DeviceConfig, InstanceConfig},
-    servers::ServerHandle,
+    servers::{self, ServerHandle},
 };
 
 mod black_border_detector;
@@ -27,6 +27,8 @@ use smoothing::*;
 
 #[derive(Debug, Error)]
 pub enum InstanceError {
+    #[error("i/o error: {0}")]
+    Io(#[from] std::io::Error),
     #[error("device error: {0}")]
     Device(#[from] DeviceError),
     #[error("recv error: {0}")]
@@ -59,22 +61,18 @@ impl Instance {
 
         let config = Arc::new(config);
 
-        let _boblight_server = ServerHandle::spawn(
+        let _boblight_server = servers::bind(
             "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,
-                    )
+                    servers::boblight::handle_client(tcp, tx.clone(), instance.clone(), global)
                 }
             },
-        );
+        )
+        .await?;
 
         Ok(Self {
             config,

src/main.rs

@@ -1,7 +1,6 @@
 #[macro_use]
 extern crate log;
 
-use hyperion::servers::ServerHandle;
 use structopt::StructOpt;
 use tokio::runtime::Builder;
 use tokio::signal;
@@ -82,32 +81,39 @@ async fn run(opts: Opts) -> color_eyre::eyre::Result<()> {
 
     // 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,
-        ))
+        Some(
+            hyperion::servers::bind(
+                "Flatbuffers",
+                config.global.flatbuffers_server.clone(),
+                global.clone(),
+                hyperion::servers::flat::handle_client,
+            )
+            .await?,
+        )
     } else {
         None
     };
 
     // Start the JSON server
-    let _json_server = ServerHandle::spawn(
+    let _json_server = hyperion::servers::bind(
         "JSON",
         config.global.json_server.clone(),
         global.clone(),
         hyperion::servers::json::handle_client,
-    );
+    )
+    .await?;
 
     // Start the Protobuf server
     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,
-        ))
+        Some(
+            hyperion::servers::bind(
+                "Protobuf",
+                config.global.proto_server.clone(),
+                global.clone(),
+                hyperion::servers::proto::handle_client,
+            )
+            .await?,
+        )
     } else {
         None
     };

src/servers.rs

@@ -12,15 +12,21 @@ pub mod flat;
 pub mod json;
 pub mod proto;
 
-async fn bind<T, E, F>(
+pub struct ServerHandle {
+    join_handle: JoinHandle<()>,
+}
+
+pub async fn bind<T, E, F, H>(
+    name: &'static str,
     options: T,
     global: Global,
-    handle_client: impl Fn((TcpStream, SocketAddr), Global) -> F,
-) -> Result<(), E>
+    handle_client: H,
+) -> std::io::Result<ServerHandle>
 where
-    T: ServerConfig,
+    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,
 {
     // Compute binding address
     let address = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), options.port());
@@ -31,55 +37,39 @@ where
     // Notify we are listening
     info!("server listening on {}", address);
 
-    loop {
-        let incoming = listener.accept().await?;
-        tokio::spawn({
-            let peer_addr = incoming.1.clone();
-            let ft = handle_client(incoming, global.clone());
+    // Spawn accepting loop
+    let join_handle = tokio::spawn(async move {
+        let result: Result<(), _> = loop {
+            match listener.accept().await {
+                Ok(incoming) => {
+                    tokio::spawn({
+                        let peer_addr = incoming.1.clone();
+                        let ft = handle_client(incoming, global.clone());
 
-            async move {
-                let result = ft.await;
+                        async move {
+                            let result = ft.await;
 
-                match result {
-                    Ok(_) => {
-                        info!("({}) client disconnected", peer_addr);
-                    }
-                    Err(error) => {
-                        error!("({}) client error:{}", peer_addr, error);
-                    }
+                            match result {
+                                Ok(_) => {
+                                    info!("({}) client disconnected", peer_addr);
+                                }
+                                Err(error) => {
+                                    error!("({}) client error:{}", peer_addr, error);
+                                }
+                            }
+                        }
+                    });
                 }
+                Err(error) => break Err(error),
             }
-        });
-    }
-}
-
-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);
-                }
-            }),
+        if let Err(error) = result {
+            error!("{} server terminated: {}", name, error);
         }
-    }
+    });
+
+    Ok(ServerHandle { join_handle })
 }
 
 impl Drop for ServerHandle {