I'm learning Rust and Tokio and I suspect I may be going in the wrong direction.
I'm trying to open a connection to a remote server and perform a handshake. I want to use non-blocking IO so I'm using Tokio's thread pool. The handshake needs to be performed quickly or the remote will close the socket so I'm trying to chain the message exchange in a single block_on section:
let result: Result<(), Box<dyn std::error::Error>> = session
.runtime()
.borrow_mut()
.block_on(async {
let startup = startup(session.configuration());
stream.write_all(startup.as_ref()).await?;
let mut buffer:Vec<u8> = Vec::new();
let mut tmp = [0u8; 1];
loop {
let total = stream.read(&mut tmp).await;
/*
if total == 0 {
break;
}
*/
if total.is_err() {
break;
}
buffer.extend(&tmp);
}
Ok(())
});
My problem is what to do when there are no more bytes in the socket to read. My current implementation reads the response and after the last byte hangs, I believe because the socket is not closed. I thought checking for 0 bytes read would be enough but the call to read() never returns.
What's the best way to handle this?
From your comment:
Nope, the connection is meant to remain open.
If you read from an open connection, the read will block until there are enough bytes to satisfy it or the other end closes the connection, similar to how blocking reads work in C. Tokio is working as-intended.
If closing the stream does not signal the end of a message, then you will have to do your own work to figure out when to stop reading and start processing. A simple way would to just prefix the request with a length, and only read that many bytes.
Note that you'd have to do the above no matter what API you'd use. The fact that you use tokio or not doesn't really answer the fundamental question of "when is the message over".
Related
I am working on a TCP proxy server where a client connects to it and it starts a new connection to a backend and forward all packets (bidirectional). The function that handles this part looks like this.
pub fn route(source: TcpStream, worker: Worker) -> Result<()> {
// `source` and `destination` are both `TcpStream`
let mut source = source;
let mut destination = TcpStream::connect(worker.address)?;
let mut source_copy = source.try_clone()?;
let mut destination_copy = destination.try_clone()?;
let src2dst = std::thread::spawn(move || {
std::io::copy(&mut source_copy, &mut destination_copy).unwrap();
});
let dst2src = std::thread::spawn(move || {
std::io::copy(&mut destination, &mut source).unwrap();
// source.shutdown(Shutdown::Both).unwrap(); <--- if this line is commented, it will stuck
});
src2dst.join().unwrap();
dst2src.join().unwrap();
Ok(())
}
However, in its current form, this function will stuck. In particular, if I shutdown the source when destination stops writing, it won't block. But I am still not sure why this works (or why it does not). I am currently using it only to proxy HTTP traffic, and it does not appear to have problems. But I am not sure if it works for generic TCP. What is the proper way to do this?
std::io::copy reads from the reader until it gets an end-of-file condition. The way to signal an end-of-file condition on a TCP connection is to send a FIN packet, which is realized in software by calling shutdown to shut down the write half of the connection.
HTTP allows reusing a single connection to send multiple requests, one after the other. There is enough information in a request for the recipient to determine where the request ends, and likewise for responses. In your case, it appears the backend shuts down its write half of the connection, which causes your io::copy to return. If the HTTP response has Connection: close in its headers, then you must shut down the write half, or else the client will hang waiting for the response body because it never receives a FIN packet.
TL;DR: When the copy from A to B returns normally, you should shut down the write half of B to forward the end-of-file condition (FIN packet). This goes for the copies in both directions, and is valid for TCP in general. Shutting down the read half of sockets doesn't appear to be necessary (it doesn't generate any TCP packets).
Based on tokio's example at https://github.com/tokio-rs/tokio/blob/master/examples/proxy.rs
let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split();
let client_to_server = io::copy(&mut ri, &mut wo);
let server_to_client = io::copy(&mut ro, &mut wi);
try_join(client_to_server, server_to_client).await?;
Ok(())
I have a modified version so that I can handle the termination of each connection as in:
// Server will disconnect their side normally 8s later, from what I've observed
let server_to_client = io::copy(&mut ro, &mut wi).map(|f| {
server_session_time = server_start_time.elapsed().unwrap();
f
});
// Normally, this will stop first, as the client disconnects as soon as he has the results...
let client_to_server = io::copy(&mut ri, &mut wo).map(|f| {
client_session_time = client_start_time.elapsed().unwrap();
f
});
// Join on both
match try_join(client_to_server, server_to_client).await {...}
This has allowed me to time correctly the connected time for the client side, since the clients immediately close connection upon receiving the answer, while the proxied server seems to take (in my case 8s) to close.
Given this structure of code, is there any possibility to terminate the downstream connection from server_to_client, once I exit the future of the client_to_server (i.e. not wait the 8s that I observe that it takes to be shutdown)?
Ok with a few more examples, was able to understand what I had to do.
For any people coming back to this question in the future, what is needed is that you implement the bidirectional copy yourself based on the 4 futures of each of the reads and writes with tokio::select!.
That will allow to access to all the streams and when one of them terminates, it is your option if you want to complete processing the others or just stop.
As it is above there is no way to "cancel" the "other" copy...
You can look both at the implementation of io::copy https://github.com/tokio-rs/tokio-io/blob/master/src/copy.rs and tokio::select https://docs.rs/tokio/0.2.20/tokio/macro.select.html, to build your 4-way select.
I want to terminate reading from a tokio::io::lines stream. I merged it with a oneshot future and terminated it, but tokio::run was still working.
use futures::{sync::oneshot, *}; // 0.1.27
use std::{io::BufReader, time::Duration};
use tokio::prelude::*; // 0.1.21
fn main() {
let (tx, rx) = oneshot::channel::<()>();
let lines = tokio::io::lines(BufReader::new(tokio::io::stdin()));
let lines = lines.for_each(|item| {
println!("> {:?}", item);
Ok(())
});
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(5000));
println!("system shutting down");
let _ = tx.send(());
});
let lines = lines.select2(rx);
tokio::run(lines.map(|_| ()).map_err(|_| ()));
}
How can I stop reading from this?
There's nothing wrong with your strategy, but it will only work with futures that don't execute a blocking operation via Tokio's blocking (the traditional kind of blocking should never be done inside a future).
You can test this by replacing the tokio::io::lines(..) future with a simple interval future:
let lines = Interval::new(Instant::now(), Duration::from_secs(1));
The problem is that tokio::io::Stdin internally uses tokio_threadpool::blocking .
When you use Tokio thread pool blocking (emphasis mine):
NB: The entire task that called blocking is blocked whenever the
supplied closure blocks, even if you have used future combinators such
as select - the other futures in this task will not make progress
until the closure returns. If this is not desired, ensure that
blocking runs in its own task (e.g. using
futures::sync::oneshot::spawn).
Since this will block every other future in the combinator, your Receiver will not be able to get a signal from the Senderuntil the blocking ends.
Please see How can I read non-blocking from stdin? or you can use tokio-stdin-stdout, which creates a channel to consume data from stdin thread. It also has a line-by-line example.
Thank you for your comment and correcting my sentences.
I tried to stop this non-blocking Future and succeeded.
let lines = Interval::new(Instant::now(), Duration::from_secs(1));
My understating is that it would work for this case to wrap the blocking Future with tokio threadpool::blocking.
I'll try it later.
Thank you very much.
What I want to do:
... write a (1) server/ (N) clients (network-game-)architecture that uses UDP sockets as underlying base for communication.
Messages are sent as Vec<u8>, encoded via bincode (crate)
I also want to be able to occasionally send datagrams that can exceed the typical max MTU of ~1500 bytes and be correctly assembled on receiver end, including sending of ack-messages etc. (I assume I'll have to implement that myself, right?)
For the UdpSocket I thought about using tokio's implementation and maybe framed. I am not sure whether this is a good choice though, as it seems that this would introduce an unnecessary step of mapping Vec<u8> (serialized by bincode) to Vec<u8> (needed by UdpCodec of tokio) (?)
Consider this minimal code-example:
Cargo.toml (server)
bincode = "1.0"
futures = "0.1"
tokio-core = "^0.1"
(Serde and serde-derive are used in shared crate where the protocol is defined!)
(I want to replace tokio-core with tokio asap)
fn main() -> () {
let addr = format!("127.0.0.1:{port}", port = 8080);
let addr = addr.parse::<SocketAddr>().expect(&format!("Couldn't create valid SocketAddress out of {}", addr));
let mut core = Core::new().unwrap();
let handle = core.handle();
let socket = UdpSocket::bind(&addr, &handle).expect(&format!("Couldn't bind socket to address {}", addr));
let udp_future = socket.framed(MyCodec {}).for_each(|(addr, data)| {
socket.send_to(&data, &addr); // Just echo back the data
Ok(())
});
core.run(udp_future).unwrap();
}
struct MyCodec;
impl UdpCodec for MyCodec {
type In = (SocketAddr, Vec<u8>);
type Out = (SocketAddr, Vec<u8>);
fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result<Self::In> {
Ok((*src, buf.to_vec()))
}
fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
let (addr, mut data) = msg;
buf.append(&mut data);
addr
}
}
The problem here is:
let udp_future = socket.framed(MyCodec {}).for_each(|(addr, data)| {
| ------ value moved here ^^^^^^^^^^^^^^ value captured here after move
|
= note: move occurs because socket has type tokio_core::net::UdpSocket, which does not implement the Copy trait
The error makes total sense, yet I am not sure how I would create such a simple echo-service. In reality, the handling of a message involves a bit more logic ofc, but for the sake of a minimal example, this should be enough to give a rough idea.
My workaround is an ugly hack: creating a second socket.
Here's the signature of UdpSocket::framed from Tokio's documentation:
pub fn framed<C: UdpCodec>(self, codec: C) -> UdpFramed<C>
Note that it takes self, not &self; that is, calling this function consumes the socket. The UdpFramed wrapper owns the underlying socket when you call this. Your compilation error is telling you that you're moving socket when you call this method, but you're also trying to borrow socket inside your closure (to call send_to).
This probably isn't what you want for real code. The whole point of using framed() is to turn your socket into something higher-level, so you can send your codec's items directly instead of having to assemble datagrams. Using send or send_to directly on the socket will probably break the framing of your message protocol. In this code, where you're trying to implement a simple echo server, you don't need to use framed at all. But if you do want to have your cake and eat it and use both framed and send_to, luckily UdpFramed still allows you to borrow the underlying UdpSocket, using get_ref. You can fix your problem this way:
let framed = {
let socket = UdpSocket::bind(&addr, &handle).expect(&format!("Couldn't bind socket to address {}", addr));
socket.framed(MyCodec {})
}
let udp_future = framed.for_each(|(addr, data)| {
info!(self.logger, "Udp packet received from {}: length: {}", addr, data.len());
framed.get_ref().send_to(&data, &addr); // Just echo back the data
Ok(())
});
I haven't checked this code, since (as Shepmaster rightly pointed out) your code snippet has other problems, but it should give you the idea anyway. I'll repeat my warning from earlier: if you do this in real code, it will break the network protocol you're using. get_ref's documentation puts it like this:
Note that care should be taken to not tamper with the underlying stream of data coming in as it may corrupt the stream of frames otherwise being worked with.
To answer the new part of your question: yes, you need to handle reassembly yourself, which means your codec does actually need to do some framing on the bytes you're sending. Typically this might involve a start sequence which cannot occur in the Vec<u8>. The start sequence lets you recognise the start of the next message after a packet was lost (which happens a lot with UDP). If there's no byte sequence that can't occur in the Vec<u8>, you need to escape it when it does occur. You might then send the length of the message, followed by the data itself; or just the data, followed by an end sequence and a checksum so you know none was lost. There are pros and cons to these designs, and it's a big topic in itself.
You also need your UdpCodec to contain data: a map from SocketAddr to the partially-reassembled message that's currently in progress. In decode, if you are given the start of a message, copy it into the map and return Ok. If you are given the middle of a message, and you already have the start of a message in the map (for that SocketAddr), append the buffer to the existing buffer and return Ok. When you get to the end of the message, return the whole thing and empty the buffer. The methods on UdpCodec take &mut self in order to enable this use case. (NB In theory, you should also deal with packets arriving out of order, but that's actually quite rare in the real world.)
encode is a lot simpler: you just need to add the same framing and copy the message into the buffer.
Let me reiterate here that you don't need to and shouldn't use the underlying socket after calling framed() on it. UdpFramed is both a source and a sink, so you use that one object to send the replies as well. You can even use split() to get separate Stream and Sink implementations out of it, if that makes the ownership easier in your application.
Overall, now I've seen how much of the problem you're struggling with, I'd recommend just using several TCP sockets instead of UDP. If you want a connection-oriented, reliable protocol, TCP already exists and does that for you. It's very easy to spend a lot of time making a "reliable" layer on top of UDP that is both slower and less reliable than TCP.
Is there a way to check whether data is available on stdin in Rust, or to do a read that returns immediately with the currently available data?
My goal is to be able to read the input produced for instance by cursor keys in a shell that is setup to return all read data immediately. For instance with an equivalent to: stty -echo -echok -icanon min 1 time 0.
I suppose one solution would be to use ncurses or similar libraries, but I would like to avoid any kind of large dependencies.
So far, I got only blocking input, which is not what I want:
let mut reader = stdin();
let mut s = String::new();
match reader.read_to_string(&mut s) {...} // this blocks :(
Converting OP's comment into an answer:
You can spawn a thread and send data over a channel. You can then poll that channel in the main thread using try_recv.
use std::io;
use std::sync::mpsc;
use std::sync::mpsc::Receiver;
use std::sync::mpsc::TryRecvError;
use std::{thread, time};
fn main() {
let stdin_channel = spawn_stdin_channel();
loop {
match stdin_channel.try_recv() {
Ok(key) => println!("Received: {}", key),
Err(TryRecvError::Empty) => println!("Channel empty"),
Err(TryRecvError::Disconnected) => panic!("Channel disconnected"),
}
sleep(1000);
}
}
fn spawn_stdin_channel() -> Receiver<String> {
let (tx, rx) = mpsc::channel::<String>();
thread::spawn(move || loop {
let mut buffer = String::new();
io::stdin().read_line(&mut buffer).unwrap();
tx.send(buffer).unwrap();
});
rx
}
fn sleep(millis: u64) {
let duration = time::Duration::from_millis(millis);
thread::sleep(duration);
}
Most operating systems default to work with the standard input and output in a blocking way. No wonder then that the Rust library follows in stead.
To read from a blocking stream in a non-blocking way you might create a separate thread, so that the extra thread blocks instead of the main one. Checking whether a blocking file descriptor produced some input is similar: spawn a thread, make it read the data, check whether it produced any data so far.
Here's a piece of code that I use with a similar goal of processing a pipe output interactively and that can hopefully serve as an example. It sends the data over a channel, which supports the try_recv method - allowing you to check whether the data is available or not.
Someone has told me that mio might be used to read from a pipe in a non-blocking way, so you might want to check it out too. I suspect that passing the stdin file descriptor (0) to Receiver::from_raw_fd should just work.
You could also potentially look at using ncurses (also on crates.io) which would allow you read in raw mode. There are a few examples in the Github repository which show how to do this.