update server

Author: matklad

src/bin/server.rs

@@ -1,16 +1,19 @@
 #![feature(async_await)]
 
-use std::{collections::HashMap, net::ToSocketAddrs};
+use std::{
+    net::ToSocketAddrs,
+    sync::Arc,
+    collections::hash_map::{HashMap, Entry},
+};
 
 use futures::{
-    io::{BufReader, WriteHalf, AsyncRead},
-    stream::{FuturesUnordered, Stream},
-    channel::mpsc::{self, unbounded},
+    channel::mpsc,
     SinkExt,
     select,
 };
 
 use async_std::{
+    io::BufReader,
     prelude::*,
     task,
     net::{TcpListener, TcpStream},
@@ -20,119 +23,156 @@ type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>
 type Sender<T> = mpsc::UnboundedSender<T>;
 type Receiver<T> = mpsc::UnboundedReceiver<T>;
 
+#[derive(Debug)]
+enum Void {}
+
 fn main() -> Result<()> {
     task::block_on(server("127.0.0.1:8080"))
 }
 
-#[derive(Debug)]
-enum Event {
-    NewPeer {
-        name: String,
-        stream: WriteHalf<TcpStream>,
-    },
-    Message {
-        from: String,
-        to: Vec<String>,
-        msg: String,
-    },
-}
-
 async fn server(addr: impl ToSocketAddrs) -> Result<()> {
     let listener = TcpListener::bind(addr).await?;
 
-    let (broker_sender, broker_receiver) = unbounded::<Event>();
+    let (broker_sender, broker_receiver) = mpsc::unbounded();
     let broker = task::spawn(broker(broker_receiver));
-    let mut incoming = listener.incoming().fuse();
-    let mut tasks = FuturesUnordered::new();
-    loop {
-        let stream = select! {
-            stream = incoming.next() => stream,
-            res = tasks.next() => {
-                if let Some(res) = res {
-                    res?
-                }
-                continue;
-            },
-            complete => break,
-        };
-
-        let stream = stream.unwrap()?;
+    let mut incoming = listener.incoming();
+    while let Some(stream) = incoming.next().await {
+        let stream = stream?;
         println!("Accepting from: {}", stream.peer_addr()?);
-
-        let handle = task::spawn(client(broker_sender.clone(), stream));
-        tasks.push(handle);
+        spawn_and_log_error(client(broker_sender.clone(), stream));
     }
-    broker.await?;
+    drop(broker_sender);
+    broker.await;
     Ok(())
 }
 
 async fn client(mut broker: Sender<Event>, stream: TcpStream) -> Result<()> {
-    let (reader, writer) = stream.split();
-    let reader = BufReader::new(reader);
+    let stream = Arc::new(stream);
+    let reader = BufReader::new(&*stream);
     let mut lines = reader.lines();
+
     let name = match lines.next().await {
         None => Err("peer disconnected immediately")?,
         Some(line) => line?,
     };
-    broker.send(Event::NewPeer { name: name.clone(), stream: writer }).await.unwrap();
+    let (_shutdown_sender, shutdown_receiver) = mpsc::unbounded::<Void>();
+    broker.send(Event::NewPeer {
+        name: name.clone(),
+        stream: Arc::clone(&stream),
+        shutdown: shutdown_receiver,
+    }).await.unwrap();
 
     while let Some(line) = lines.next().await {
         let line = line?;
         let (dest, msg) = match line.find(':') {
             None => continue,
             Some(idx) => (&line[..idx], line[idx + 1 ..].trim()),
         };
+        let dest: Vec<String> = dest.split(',').map(|name| name.trim().to_string()).collect();
+        let msg: String = msg.trim().to_string();
+
         broker.send(Event::Message {
             from: name.clone(),
-            to: dest.split(',').map(|name| name.trim().to_string()).collect(),
-            msg: msg.trim().to_string(),
+            to: dest,
+            msg,
         }).await.unwrap();
     }
+
     Ok(())
 }
 
 async fn client_writer(
-    mut receiver: Receiver<String>,
-    mut writer: WriteHalf<TcpStream>,
+    messages: &mut Receiver<String>,
+    stream: Arc<TcpStream>,
+    mut shutdown: Receiver<Void>,
 ) -> Result<()> {
-    while let Some(msg) = receiver.next().await {
-        writer.write_all(msg.as_bytes()).await?;
+    let mut stream = &*stream;
+    loop {
+        select! {
+            msg = messages.next() => match msg {
+                Some(msg) => stream.write_all(msg.as_bytes()).await?,
+                None => break,
+            },
+            void = shutdown.next() => match void {
+                Some(void) => match void {},
+                None => break,
+            }
+        }
     }
     Ok(())
 }
 
-async fn broker(mut events: Receiver<Event>) -> Result<()> {
+#[derive(Debug)]
+enum Event {
+    NewPeer {
+        name: String,
+        stream: Arc<TcpStream>,
+        shutdown: Receiver<Void>,
+    },
+    Message {
+        from: String,
+        to: Vec<String>,
+        msg: String,
+    },
+}
+
+async fn broker(mut events: Receiver<Event>) {
+    let (disconnect_sender, mut disconnect_receiver) =
+        mpsc::unbounded::<(String, Receiver<String>)>();
     let mut peers: HashMap<String, Sender<String>> = HashMap::new();
-    let mut tasks = FuturesUnordered::new();
 
     loop {
         let event = select! {
-            event = events.next() => event,
-            res = tasks.next() => {
-                if let Some(res) = res {
-                    res?;
-                }
+            event = events.next() => match event {
+                None => break,
+                Some(event) => event,
+            },
+            disconnect = disconnect_receiver.next() => {
+                let (name, _pending_messages) = disconnect.unwrap();
+                assert!(peers.remove(&name).is_some());
                 continue;
-            }
-            complete => break,
+            },
         };
-
         match event {
-            Some(Event::Message { from, to, msg } )=> {
+            Event::Message { from, to, msg } => {
                 for addr in to {
                     if let Some(peer) = peers.get_mut(&addr) {
-                        peer.send(format!("from {}: {}\n", from, msg)).await?
+                        peer.send(format!("from {}: {}\n", from, msg)).await
+                            .unwrap()
                     }
                 }
             }
-            Some(Event::NewPeer { name, stream}) => {
-                let (client_sender, client_receiver) = unbounded();
-                peers.insert(name.clone(), client_sender);
-                let handle = task::spawn(client_writer(client_receiver, stream));
-                tasks.push(handle)
+            Event::NewPeer { name, stream, shutdown } => {
+                match peers.entry(name.clone()) {
+                    Entry::Occupied(..) => (),
+                    Entry::Vacant(entry) => {
+                        let (client_sender, mut client_receiver) = mpsc::unbounded();
+                        entry.insert(client_sender);
+                        let mut disconnect_sender = disconnect_sender.clone();
+                        spawn_and_log_error(async move {
+                            let res = client_writer(&mut client_receiver, stream, shutdown).await;
+                            disconnect_sender.send((name, client_receiver)).await
+                                .unwrap();
+                            res
+                        });
+                    }
+                }
             }
-            None => continue,
         }
     }
-    Ok(())
+    drop(peers);
+    drop(disconnect_sender);
+    while let Some((_name, _pending_messages)) = disconnect_receiver.next().await {
+    }
+}
+
+fn spawn_and_log_error<F>(fut: F) -> task::JoinHandle<()>
+where
+    F: Future<Output = Result<()>> + Send + 'static,
+{
+    task::spawn(async move {
+        if let Err(e) = fut.await {
+            eprintln!("{}", e)
+        }
+    })
 }