Select from a list of sockets using futures - async-await

Question

I am trying out the yet-unstable async-await syntax in nightly Rust 1.38 with futures-preview = "0.3.0-alpha.16" and runtime = "0.3.0-alpha.6". It feels really cool, but the docs are (yet) scarce and I got stuck.
To go a bit beyond the basic examples I would like to create an app that:
Accepts TCP connections on a given port;
Broadcasts all the data received from any connection to all active connections.
Existing docs and examples got me this far:
#![feature(async_await)]
#![feature(async_closure)]
use futures::{
prelude::*,
select,
future::select_all,
io::{ReadHalf, WriteHalf, Read},
};
use runtime::net::{TcpListener, TcpStream};
use std::io;
async fn read_stream(mut reader: ReadHalf<TcpStream>) -> (ReadHalf<TcpStream>, io::Result<Box<[u8]>>) {
let mut buffer: Vec<u8> = vec![0; 1024];
match reader.read(&mut buffer).await {
Ok(len) => {
buffer.truncate(len);
(reader, Ok(buffer.into_boxed_slice()))
},
Err(err) => (reader, Err(err)),
}
}
#[runtime::main]
async fn main() -> std::io::Result<()> {
let mut listener = TcpListener::bind("127.0.0.1:8080")?;
println!("Listening on {}", listener.local_addr()?);
let mut incoming = listener.incoming().fuse();
let mut writers: Vec<WriteHalf<TcpStream>> = vec![];
let mut reads = vec![];
loop {
select! {
maybe_stream = incoming.select_next_some() => {
let (mut reader, writer) = maybe_stream?.split();
writers.push(writer);
reads.push(read_stream(reader).fuse());
},
maybe_read = select_all(reads.iter()) => {
match maybe_read {
(reader, Ok(data)) => {
for writer in writers {
writer.write_all(data).await.ok(); // Ignore errors here
}
reads.push(read_stream(reader).fuse());
},
(reader, Err(err)) => {
let reader_addr = reader.peer_addr().unwrap();
writers.retain(|writer| writer.peer_addr().unwrap() != reader_addr);
},
}
}
}
}
}
This fails with:
error: recursion limit reached while expanding the macro `$crate::dispatch`
--> src/main.rs:36:9
|
36 | / select! {
37 | | maybe_stream = incoming.select_next_some() => {
38 | | let (mut reader, writer) = maybe_stream?.split();
39 | | writers.push(writer);
... |
55 | | }
56 | | }
| |_________^
|
= help: consider adding a `#![recursion_limit="128"]` attribute to your crate
= note: this error originates in a macro outside of the current crate (in Nightly builds, run with -Z external-macro-backtrace for more info)
This is very confusing. Maybe I am using select_all() in a wrong way? Any help in making it work is appreciated!
For completeness, my Cargo.toml:
[package]
name = "async-test"
version = "0.1.0"
authors = ["xxx"]
edition = "2018"
[dependencies]
runtime = "0.3.0-alpha.6"
futures-preview = { version = "=0.3.0-alpha.16", features = ["async-await", "nightly"] }

Answer 1

In case someone is following, I hacked it together finally. This code works:
#![feature(async_await)]
#![feature(async_closure)]
#![recursion_limit="128"]
use futures::{
prelude::*,
select,
stream,
io::ReadHalf,
channel::{
oneshot,
mpsc::{unbounded, UnboundedSender},
}
};
use runtime::net::{TcpListener, TcpStream};
use std::{
io,
net::SocketAddr,
collections::HashMap,
};
async fn read_stream(
addr: SocketAddr,
drop: oneshot::Receiver<()>,
mut reader: ReadHalf<TcpStream>,
sender: UnboundedSender<(SocketAddr, io::Result<Box<[u8]>>)>
) {
let mut drop = drop.fuse();
loop {
let mut buffer: Vec<u8> = vec![0; 1024];
select! {
result = reader.read(&mut buffer).fuse() => {
match result {
Ok(len) => {
buffer.truncate(len);
sender.unbounded_send((addr, Ok(buffer.into_boxed_slice())))
.expect("Channel error");
if len == 0 {
return;
}
},
Err(err) => {
sender.unbounded_send((addr, Err(err))).expect("Channel error");
return;
}
}
},
_ = drop => {
return;
},
}
}
}
enum Event {
Connection(io::Result<TcpStream>),
Message(SocketAddr, io::Result<Box<[u8]>>),
}
#[runtime::main]
async fn main() -> std::io::Result<()> {
let mut listener = TcpListener::bind("127.0.0.1:8080")?;
eprintln!("Listening on {}", listener.local_addr()?);
let mut writers = HashMap::new();
let (sender, receiver) = unbounded();
let connections = listener.incoming().map(|maybe_stream| Event::Connection(maybe_stream));
let messages = receiver.map(|(addr, maybe_message)| Event::Message(addr, maybe_message));
let mut events = stream::select(connections, messages);
loop {
match events.next().await {
Some(Event::Connection(Ok(stream))) => {
let addr = stream.peer_addr().unwrap();
eprintln!("New connection from {}", addr);
let (reader, writer) = stream.split();
let (drop_sender, drop_receiver) = oneshot::channel();
writers.insert(addr, (writer, drop_sender));
runtime::spawn(read_stream(addr, drop_receiver, reader, sender.clone()));
},
Some(Event::Message(addr, Ok(message))) => {
if message.len() == 0 {
eprintln!("Connection closed by client: {}", addr);
writers.remove(&addr);
continue;
}
eprintln!("Received {} bytes from {}", message.len(), addr);
if &*message == b"quit\n" {
eprintln!("Dropping client {}", addr);
writers.remove(&addr);
continue;
}
for (&other_addr, (writer, _)) in &mut writers {
if addr != other_addr {
writer.write_all(&message).await.ok(); // Ignore errors
}
}
},
Some(Event::Message(addr, Err(err))) => {
eprintln!("Error reading from {}: {}", addr, err);
writers.remove(&addr);
},
_ => panic!("Event error"),
}
}
}
I use a channel and spawn a reading task for each client. Special care had to be taken to ensure that readers get dropped with writers: this is why oneshot future is used. When oneshot::Sender is dropped, the oneshot::Receiver future resolves to canceled state, which is a notification mechanism for a reading task to know it is time to halt. To demonstrate that it works, we drop a client as soon as we get "quit" message.
Sadly, there is a (seemingly useless) warning regarding an unused JoinHandle from the runtime::spawn call, and I don't really know how to eliminate it.