fmt

backend/shard/src/main.rs

@@ -1,19 +1,19 @@
 mod aggregator;
 mod connection;
-mod real_ip;
 mod json_message;
+mod real_ip;
 
 use std::net::IpAddr;
 
-use structopt::StructOpt;
-use http::Uri;
-use simple_logger::SimpleLogger;
-use futures::{StreamExt, SinkExt, channel::mpsc};
-use warp::Filter;
-use warp::filters::ws;
+use aggregator::{Aggregator, FromWebsocket};
 use common::{node_message, LogLevel};
-use aggregator::{ Aggregator, FromWebsocket };
+use futures::{channel::mpsc, SinkExt, StreamExt};
+use http::Uri;
 use real_ip::real_ip;
+use simple_logger::SimpleLogger;
+use structopt::StructOpt;
+use warp::filters::ws;
+use warp::Filter;
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
@@ -27,25 +27,17 @@ struct Opts {
     /// This is the socket address that this shard is listening to. This is restricted to
     /// localhost (127.0.0.1) by default and should be fine for most use cases. If
     /// you are using Telemetry in a container, you likely want to set this to '0.0.0.0:8000'
-    #[structopt(
-        short = "l",
-        long = "listen",
-        default_value = "127.0.0.1:8001",
-    )]
+    #[structopt(short = "l", long = "listen", default_value = "127.0.0.1:8001")]
     socket: std::net::SocketAddr,
     /// The desired log level; one of 'error', 'warn', 'info', 'debug' or 'trace', where
     /// 'error' only logs errors and 'trace' logs everything.
-    #[structopt(
-        required = false,
-        long = "log",
-        default_value = "info",
-    )]
+    #[structopt(required = false, long = "log", default_value = "info")]
     log_level: LogLevel,
     /// Url to the Backend Core endpoint accepting shard connections
     #[structopt(
-    	short = "c",
-    	long = "core",
-    	default_value = "ws://127.0.0.1:8000/shard_submit/",
+        short = "c",
+        long = "core",
+        default_value = "ws://127.0.0.1:8000/shard_submit/"
     )]
     core_url: Uri,
 }
@@ -60,10 +52,7 @@ async fn main() {
         .init()
         .expect("Must be able to start a logger");
 
-    log::info!(
-        "Starting Telemetry Shard version: {}",
-        VERSION
-    );
+    log::info!("Starting Telemetry Shard version: {}", VERSION);
 
     if let Err(e) = start_server(opts).await {
         log::error!("Error starting server: {}", e);
@@ -72,26 +61,21 @@ async fn main() {
 
 /// Declare our routes and start the server.
 async fn start_server(opts: Opts) -> anyhow::Result<()> {
-
     let aggregator = Aggregator::spawn(opts.core_url).await?;
 
     // Handle requests to /health by returning OK.
-    let health_route =
-        warp::path("health")
-        .map(|| "OK");
+    let health_route = warp::path("health").map(|| "OK");
 
     // Handle websocket requests to /submit.
-    let ws_route =
-        warp::path("submit")
-        .and(warp::ws())
-        .and(real_ip())
-        .map(move |ws: ws::Ws, addr: Option<IpAddr>| {
+    let ws_route = warp::path("submit").and(warp::ws()).and(real_ip()).map(
+        move |ws: ws::Ws, addr: Option<IpAddr>| {
             // Send messages from the websocket connection to this sink
             // to have them pass to the aggregator.
             let tx_to_aggregator = aggregator.subscribe_node();
             log::info!("Opening /submit connection from {:?}", addr);
             ws.on_upgrade(move |websocket| async move {
-                let (mut tx_to_aggregator, websocket) = handle_websocket_connection(websocket, tx_to_aggregator, addr).await;
+                let (mut tx_to_aggregator, websocket) =
+                    handle_websocket_connection(websocket, tx_to_aggregator, addr).await;
                 log::info!("Closing /submit connection from {:?}", addr);
                 // Tell the aggregator that this connection has closed, so it can tidy up.
                 let _ = tx_to_aggregator.send(FromWebsocket::Disconnected).await;
@@ -99,7 +83,8 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
                 // a ws::Message using `ws::Message::close_with`, rather than using this method:
                 let _ = websocket.close().await;
             })
-        });
+        },
+    );
 
     // Merge the routes and start our server:
     let routes = ws_route.or(health_route);
@@ -108,8 +93,13 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
 }
 
 /// This takes care of handling messages from an established socket connection.
-async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_aggregator: S, addr: Option<IpAddr>) -> (S, ws::WebSocket)
-    where S: futures::Sink<FromWebsocket, Error = anyhow::Error> + Unpin
+async fn handle_websocket_connection<S>(
+    mut websocket: ws::WebSocket,
+    mut tx_to_aggregator: S,
+    addr: Option<IpAddr>,
+) -> (S, ws::WebSocket)
+where
+    S: futures::Sink<FromWebsocket, Error = anyhow::Error> + Unpin,
 {
     // This could be a oneshot channel, but it's useful to be able to clone
     // messages, and we can't clone oneshot channel senders.
@@ -117,7 +107,7 @@ async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_
 
     // Tell the aggregator about this new connection, and give it a way to close this connection:
     let init_msg = FromWebsocket::Initialize {
-        close_connection: close_connection_tx
+        close_connection: close_connection_tx,
     };
     if let Err(e) = tx_to_aggregator.send(init_msg).await {
         log::error!("Error sending message to aggregator: {}", e);
@@ -179,13 +169,15 @@ async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_
 /// Deserialize an incoming websocket message, returning an error if something
 /// fatal went wrong, [`Some`] message if all went well, and [`None`] if a non-fatal
 /// issue was encountered and the message should simply be ignored.
-fn deserialize_ws_message(msg: Result<ws::Message, warp::Error>) -> anyhow::Result<Option<node_message::NodeMessage>> {
+fn deserialize_ws_message(
+    msg: Result<ws::Message, warp::Error>,
+) -> anyhow::Result<Option<node_message::NodeMessage>> {
     // If we see any errors, log them and end our loop:
     let msg = match msg {
         Err(e) => {
             return Err(anyhow::anyhow!("Error in node websocket connection: {}", e));
-        },
-        Ok(msg) => msg
+        }
+        Ok(msg) => msg,
     };
 
     // If the message isn't something we want to handle, just ignore it.
@@ -202,11 +194,11 @@ fn deserialize_ws_message(msg: Result<ws::Message, warp::Error>) -> anyhow::Resu
             // let bytes: &[u8] = bytes.get(..512).unwrap_or_else(|| &bytes);
             // let msg_start = std::str::from_utf8(bytes).unwrap_or_else(|_| "INVALID UTF8");
             // log::warn!("Failed to parse node message ({}): {}", msg_start, e);
-            return Ok(None)
+            return Ok(None);
         }
     };
 
     // Pull relevant details from the message:
     let node_message: node_message::NodeMessage = node_message.into();
     Ok(Some(node_message))
-}
+}