graceful shutdown

Author: matklad

src/main.rs

@@ -26,11 +26,16 @@ async fn server(addr: impl ToSocketAddrs) -> Result<()> {
 
     let (broker_sender, broker_receiver) = mpsc::unbounded();
     let broker = task::spawn(broker(broker_receiver));
+    let mut readers = Vec::new();
     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));
+        let handle = task::spawn(client(broker_sender.clone(), stream));
+        readers.push(handle);
+    }
+    for reader in readers {
+        reader.await?;
     }
     broker.await?;
     Ok(())
@@ -91,6 +96,7 @@ enum Event {
 
 async fn broker(mut events: Receiver<Event>) -> Result<()> {
     let mut peers: HashMap<String, Sender<String>> = HashMap::new();
+    let mut writers = Vec::new();
 
     while let Some(event) = events.next().await {
         match event {
@@ -104,9 +110,14 @@ async fn broker(mut events: Receiver<Event>) -> Result<()> {
             Event::NewPeer { name, stream} => {
                 let (client_sender, client_receiver) = mpsc::unbounded();
                 peers.insert(name.clone(), client_sender);
-                let _handle = task::spawn(client_writer(client_receiver, stream));
+                let handle = task::spawn(client_writer(client_receiver, stream));
+                writers.push(handle);
             }
         }
     }
+    drop(peers);
+    for writer in writers {
+        writer.await?;
+    }
     Ok(())
 }

tutorial.adoc

@@ -343,27 +343,26 @@ async fn server(addr: impl ToSocketAddrs) -> Result<()> {
     let listener = TcpListener::bind(addr).await?;
 
     let (broker_sender, broker_receiver) = mpsc::unbounded(); <1>
-    let broker = task::spawn(broker(broker_receiver));
+    let _broker_handle = task::spawn(broker(broker_receiver));
     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));
     }
-    broker.await?; <2>
     Ok(())
 }
 
 async fn client(mut broker: Sender<Event>, stream: TcpStream) -> Result<()> {
-    let stream = Arc::new(stream); <3>
+    let stream = Arc::new(stream); <2>
     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: Arc::clone(&stream) }).await <4>
+    broker.send(Event::NewPeer { name: name.clone(), stream: Arc::clone(&stream) }).await <3>
         .unwrap();
 
     while let Some(line) = lines.next().await {
@@ -375,7 +374,7 @@ async fn client(mut broker: Sender<Event>, stream: TcpStream) -> Result<()> {
         let dest: Vec<String> = dest.split(',').map(|name| name.trim().to_string()).collect();
         let msg: String = msg.trim().to_string();
 
-        broker.send(Event::Message { <5>
+        broker.send(Event::Message { <4>
             from: name.clone(),
             to: dest,
             msg,
@@ -432,9 +431,97 @@ async fn broker(mut events: Receiver<Event>) -> Result<()> {
 ----
 
 <1> Inside the `server`, we create broker's channel and `task`.
-<2> After ``server``'s end, we join broker to make sure all pending messages are delivered.
-    Note that this doesn't quite do the trick yet, as we are not joining readers and writers themselves.
-<3> Inside `client`, we need to wrap `TcpStream` into an `Arc`, to be able to share it with the `client_writer`.
-<4> On login, we notify the broker.
+<2> Inside `client`, we need to wrap `TcpStream` into an `Arc`, to be able to share it with the `client_writer`.
+<3> On login, we notify the broker.
     Note that we `.unwrap` on send: broker should outlive all the clients and if that's not the case the broker probably panicked, so we can escalate the panic as well.
-<5> Similarly, we forward parsed messages to the broker, assuming that it is alive.
+<4> Similarly, we forward parsed messages to the broker, assuming that it is alive.
+
+== Clean Shutdown
+
+On of the problems of the current implementation is that it doesn't handle graceful shutdown.
+If we break from the accept loop for some reason, all in-flight tasks are just dropped on the floor.
+A more correct shutdown sequence would be:
+
+. Stop accepting new clients
+. Deliver all pending messages
+. Exit the process
+
+A clean shutdown in a channel based architecture is easy, although it can appear a magic trick at first.
+In Rust, receiver side of a channel is closed as soon as all senders are dropped.
+That is, as soon as producers exit and drop their senders, the rest of the system shutdowns naturally.
+In `async_std` this translates to two rules:
+
+. Make sure that channels form an acyclic graph.
+. Take care to wait, in the correct order, until intermediate layers of the system process pending messages.
+
+In `a-chat`, we already have an unidirectional flow of messages: `reader -> broker -> writer`.
+However, we never wait for broker and writers, which might cause some messages to get dropped.
+Let's add waiting to the server:
+
+
+[source,rust]
+----
+async fn server(addr: impl ToSocketAddrs) -> Result<()> {
+    let listener = TcpListener::bind(addr).await?;
+
+    let (broker_sender, broker_receiver) = mpsc::unbounded();
+    let broker = task::spawn(broker(broker_receiver));
+    let mut readers = Vec::new();
+    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));
+        readers.push(handle);
+    }
+    drop(broker_sender);
+    for reader in readers {
+        reader.await?; <1>
+    }
+    broker.await?; <5>
+    Ok(())
+}
+----
+
+And to the broker:
+
+[source,rust]
+----
+async fn broker(mut events: Receiver<Event>) -> Result<()> {
+    let mut writers = Vec::new();
+    let mut peers: HashMap<String, Sender<String>> = HashMap::new();
+
+    while let Some(event) = events.next().await { <2>
+        match event {
+            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?
+                    }
+                }
+            }
+            Event::NewPeer { name, stream} => {
+                let (client_sender, client_receiver) = mpsc::unbounded();
+                peers.insert(name.clone(), client_sender);
+                let handle = task::spawn(client_writer(client_receiver, stream));
+                writers.push(handle);
+            }
+        }
+    }
+    drop(peers); <3>
+    for writer in writers {
+        writer.await?;
+    }
+    Ok(())
+}
+----
+
+Notice what happens with all of the channels once we exit the accept loop:
+
+<1> First, we drop the main broker's sender and join all of the readers.
+    When the readers are done, there's no sender for the broker's channel.
+<2> Next, the broker exits `while let Some(event) = events.next().await` loop.
+<3> It's crucial that, at this stage, we drop the `peers` map.
+    This drops writer's senders.
+<4> Now we can join all of the writers.
+<5> Finally, we join the broker, which also gurantess that all the writes have terminated.