tidy up main.rs

src/main.rs

@@ -1,10 +1,11 @@
 #[macro_use]
 extern crate log;
 
-use anyhow::{Context, Error, Result};
-use future::ok;
+use anyhow::{Context, Result};
+use future::ready;
 use futures::prelude::*;
-use net::{ClientMessage, ServerMessage, server::greet};
+use futures::try_join;
+use net::{server::greet, ClientMessage, ServerMessage};
 use serde_json::{from_str, to_string};
 use std::net::ToSocketAddrs;
 use stream::FuturesUnordered;
@@ -29,11 +30,15 @@ async fn main() -> Result<()> {
     let args = Args::from_args();
 
     // dispatch websockets
-    let socket_handler = ws().map(|upgrade: Ws| upgrade.on_upgrade(|ws| async {
-        if let Err(error) = handle_socket(ws).await {
-            warn!("Websocket connection lost: {:#}", error);
-        }
-    }));
+    let socket_handler = ws().map(|upgrade: Ws| {
+        upgrade.on_upgrade(|ws| {
+            async {
+                if let Err(error) = handle_socket(ws).await {
+                    warn!("Websocket connection lost: {:#}", error);
+                }
+            }
+        })
+    });
 
     let addrs = args
         .listen
@@ -54,24 +59,29 @@ async fn handle_socket(websocket: WebSocket) -> Result<()> {
     let (sink, mut source) = websocket.split();
 
     let mut sink = sink.with(|msg: ServerMessage| {
-        let json = to_string(&msg).expect("JSON encoding shouldn't fail");
-        ok::<Message, Error>(Message::text(json))
+        ready(
+            to_string(&msg)
+                .context("JSON encoding shouldn't fail")
+                .map(|json| Message::text(json)),
+        )
     });
 
     greet(&mut sink).await?;
 
-    loop {
-        match source.next().await {
-            Some(Ok(msg)) => match msg.to_str() {
-                Ok(json) => {
-                    let msg = from_str::<ClientMessage>(json).context("Parsing JSON")?;
-                    debug!("Client message: {:?}", &msg);
-                }
-                Err(()) => debug!("Non-text message"),
-            },
-            Some(Err(error)) => return Err(error.into()),
-            None => break,
+    let input_task = async {
+        loop {
+            match source.next().await {
+                Some(Ok(msg)) => match msg.to_str() {
+                    Ok(json) => {
+                        let msg = from_str::<ClientMessage>(json).context("Parsing JSON")?;
+                        debug!("Client message: {:?}", &msg);
+                    }
+                    Err(()) => debug!("Non-text message"),
+                },
+                Some(Err(error)) => return Err(error.into()),
+                None => break Ok(()),
+            }
         }
     };
-    Ok(())
+    try_join!(input_task).map(|((),)| ())
 }