simple broker

Author: matklad

src/main.rs

@@ -1,6 +1,9 @@
 #![feature(async_await)]
 
-use std::{net::ToSocketAddrs, sync::Arc};
+use std::{net::ToSocketAddrs, sync::Arc, collections::HashMap};
+
+use futures::channel::mpsc;
+use futures::SinkExt;
 
 use async_std::{
     io::BufReader,
@@ -10,31 +13,39 @@ use async_std::{
 };
 
 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>;
+
 
 fn main() -> Result<()> {
     task::block_on(server("127.0.0.1:8080"))
 }
 
 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 incoming = listener.incoming();
     while let Some(stream) = incoming.next().await {
         let stream = stream?;
         println!("Accepting from: {}", stream.peer_addr()?);
-        let _handle = task::spawn(client(stream));
+        let _handle = task::spawn(client(broker_sender.clone(), stream));
     }
+    broker.await?;
     Ok(())
 }
 
-async fn client(stream: TcpStream) -> Result<()> {
-    let reader = BufReader::new(&stream);
+async fn client(mut broker: Sender<Event>, stream: TcpStream) -> Result<()> {
+    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?,
     };
-    println!("name = {}", name);
+    broker.send(Event::NewPeer { name: name.clone(), stream: Arc::clone(&stream) }).await.unwrap();
 
     while let Some(line) = lines.next().await {
         let line = line?;
@@ -44,6 +55,58 @@ async fn client(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 {
+            from: name.clone(),
+            to: dest,
+            msg,
+        }).await.unwrap();
+    }
+    Ok(())
+}
+
+async fn client_writer(
+    mut messages: Receiver<String>,
+    stream: Arc<TcpStream>,
+) -> Result<()> {
+    let mut stream = &*stream;
+    while let Some(msg) = messages.next().await {
+        stream.write_all(msg.as_bytes()).await?;
+    }
+    Ok(())
+}
+
+#[derive(Debug)]
+enum Event {
+    NewPeer {
+        name: String,
+        stream: Arc<TcpStream>,
+    },
+    Message {
+        from: String,
+        to: Vec<String>,
+        msg: String,
+    },
+}
+
+async fn broker(mut events: Receiver<Event>) -> Result<()> {
+    let mut peers: HashMap<String, Sender<String>> = HashMap::new();
+
+    while let Some(event) = events.next().await {
+        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));
+            }
+        }
     }
     Ok(())
 }

tutorial.adoc

@@ -223,4 +223,218 @@ That way, a client can directly `.write_all` a message to recipients.
 However, this would be wrong: if Alice sends `bob: foo`, and Charley sends `bob: bar`, Bob might actually receive `fobaor`.
 Sending a message over a socket might require several syscalls, so two concurrent ``.write_all``'s might interfere with each other!
 
+As a rule of thumb, only a single task should write to each `TcpStream`.
+So let's create a `client_writer` task which receives messages over a channel and writes them to the socket.
+This task would be the point of serialization of messages.
+if Alice and Charley send two messages to Bob at the same time, Bob will see the messages in the same order as they arrive in the channel.
 
+[source,rust]
+----
+use futures::channel::mpsc; <1>
+use futures::SinkExt;
+
+type Sender<T> = mpsc::UnboundedSender<T>; <2>
+type Receiver<T> = mpsc::UnboundedReceiver<T>;
+
+async fn client_writer(
+    mut messages: Receiver<String>,
+    stream: Arc<TcpStream>, <3>
+) -> Result<()> {
+    let mut stream = &*stream;
+    while let Some(msg) = messages.next().await {
+        stream.write_all(msg.as_bytes()).await?;
+    }
+    Ok(())
+}
+----
+
+<1> We will use channels from the `futures` crate.
+<2> For simplicity, we will use `unbounded` channels, and won't be discussing backpressure in this tutorial.
+<3> As `client` and `client_writer` share the same `TcpStream`, we need to put it into an `Arc`.
+    Note that because `client` only reads from and `client_writer` only writes to the stream, so we don't get a race here.
+
+
+== Connecting Readers and Writers
+
+So how we make sure that messages read in `client` flow into the relevant `client_writer`?
+We should somehow maintain an `peers: HashMap<String, Sender<String>>` map which allows a client to find destination channels.
+However, this map would be a bit of shared mutable state, so we'll have to wrap an `RwLock` over it and answer tough questions of what should happen if the client joins at the same moment as it receives a message.
+
+One trick to make reasoning about state simpler comes from the actor model.
+We can create a dedicated broker tasks which owns the `peers` map and communicates with other tasks by channels.
+By hiding `peers` inside such "actor" task, we remove the need for mutxes and also make serialization point explicit.
+The order of events "Bob sends message to Alice" and "Alice joins" is determined by the order of the corresponding events in the broker's event queue.
+
+[source,rust]
+----
+#[derive(Debug)]
+enum Event { <1>
+    NewPeer {
+        name: String,
+        stream: Arc<TcpStream>,
+    },
+    Message {
+        from: String,
+        to: Vec<String>,
+        msg: String,
+    },
+}
+
+async fn broker(mut events: Receiver<Event>) -> Result<()> {
+    let mut peers: HashMap<String, Sender<String>> = HashMap::new(); <2>
+
+    while let Some(event) = events.next().await {
+        match event {
+            Event::Message { from, to, msg } => {  <3>
+                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); <4>
+                let _handle = task::spawn(client_writer(client_receiver, stream)); <5>
+            }
+        }
+    }
+    Ok(())
+}
+----
+
+<1> Broker should handle two types of events: a message or an arrival of a new peer.
+<2> Internal state of the broker is a `HashMap`.
+    Note how we don't need a `Mutex` here and can confidently say, at each iteration of the broker's loop, what is the current set of peers
+<3> To handle a message we send it over a channel to each destination
+<4> To handle new peer, we first register it in the peer's map ...
+<4> ... and then spawn a dedicated task to actually write the messages to the socket.
+
+== All Together
+
+At this point, we only need to start broker to get a fully-functioning (in the happy case!) chat:
+
+[source,rust]
+----
+#![feature(async_await)]
+
+use std::{net::ToSocketAddrs, sync::Arc, collections::HashMap};
+
+use futures::channel::mpsc;
+use futures::SinkExt;
+
+use async_std::{
+    io::BufReader,
+    prelude::*,
+    task,
+    net::{TcpListener, TcpStream},
+};
+
+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>;
+
+
+fn main() -> Result<()> {
+    task::block_on(server("127.0.0.1:8080"))
+}
+
+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 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 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>
+        .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 { <5>
+            from: name.clone(),
+            to: dest,
+            msg,
+        }).await.unwrap();
+    }
+    Ok(())
+}
+
+async fn client_writer(
+    mut messages: Receiver<String>,
+    stream: Arc<TcpStream>,
+) -> Result<()> {
+    let mut stream = &*stream;
+    while let Some(msg) = messages.next().await {
+        stream.write_all(msg.as_bytes()).await?;
+    }
+    Ok(())
+}
+
+#[derive(Debug)]
+enum Event {
+    NewPeer {
+        name: String,
+        stream: Arc<TcpStream>,
+    },
+    Message {
+        from: String,
+        to: Vec<String>,
+        msg: String,
+    },
+}
+
+async fn broker(mut events: Receiver<Event>) -> Result<()> {
+    let mut peers: HashMap<String, Sender<String>> = HashMap::new();
+
+    while let Some(event) = events.next().await {
+        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));
+            }
+        }
+    }
+    Ok(())
+}
+----
+
+<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.
+    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.