Determining the End-of-stream - rust

I made a loop for a webserver.
On a windows client I didn't have any problems but on a linux client the server didn't responding to requests.
The problem: I found out that if request_size % buffer_size == 0 then the loop runs once more waiting for more data.
The question: Is there an efficient way of reading data that takes into consideration slow connections, connections that drop packages. (Not just using non_blocking or nodelay.)
let listener = TcpListener::bind("127.0.0.1:80").unwrap();
while let Ok((mut stream, _)) = listener.accept() {
let mut data: Vec<u8> = Vec::new();
let mut buf = [0u8; 32];
while let Ok(size) = stream.read(&mut buf) {
data.extend(buf[..size].iter());
if size != buf.len() { break; }
}
// do something with the data
}
I could increase the buffer size but that wouldn't solve the problem.

First, to detect EOF reliably, you should test the returned size of Read::read against zero and not your buffer size, because if you have a 'slow connections' you might not get enough data to fill the entire buffer at once, causing your loop to quite early with an incomplete message in data.
There are essentially 3 ways to make sure you received the entire message:
Read until EOF
Read a fixed-sized message
Encode some 'content length' and read that many bytes
Notice, that only the last two variants allow your client to eventually send more data over the same stream. Also notice, that these two variants can be implemented comparably easy via Read::read_exact.
Besides notice, if you don't trust your client, it might be helpful to set up TcpStream::set_read_timeout with a reasonably long timeout (e.g. 2 min).
Read until EOF
This is probably the easiest and, according to your title and code, probably the method you are aiming for. However, to generate an EOF, the client must shutdown at least its write channel. So, if your server is stuck in read, I assume you forgot to shutdown your client (tho I have to guess here).
On the server side, if you really want to read until EOF, you don't need a loop yourself, you can simply use the Read::read_to_end utility function. Here is an example for a client & server with the client sending a single message terminated by EOF:
use std::io::Read;
use std::io::Write;
use std::net::TcpListener;
use std::net::TcpStream;
// --- Client code
const SERVER_ADDR: &str = "localhost:1234";
pub fn client() {
let mut socket = TcpStream::connect(SERVER_ADDR).expect("Failed to connect");
// Send a 'single' message, the flushes kinda simulates a very slow connection
for _ in 0..3 {
socket.write(b"Hello").expect("Failed to send");
socket.flush().unwrap();
}
// Instead of shutdow, you can also drop(socket), but than you can't read.
socket.shutdown(std::net::Shutdown::Write).unwrap();
// go reading, or whatever
}
// --- Server code
const SERVER_BIND: &str = "127.0.0.1:1234";
pub fn server() {
let listener = TcpListener::bind(SERVER_BIND).expect("Failed to bind");
while let Ok((stream, _)) = listener.accept() {
let _ = handle_client(stream); // don't care if the client screwed up
}
}
pub fn handle_client(mut socket: TcpStream) -> std::io::Result<()> {
let mut data: Vec<u8> = Vec::new();
// Read all bytes until EOF
socket.read_to_end(&mut data)?;
println!("Data: {:?}", data); // or whatever
Ok(())
}

Related

Rust Tokio mpsc::channel unexpected behavior for multi-task program

In the following program I use Tokio's mpsc channels. The Sender is moved to a task named input_message and the Receiver is moved to another task named printer. Both tasks are tokio::spawn()-ed in the main function. The input_message task is to read the user's input and send it through a Channel. The printer task recv() on the channel to get the user's input and simply prints it to stdout:
use std::error::Error;
use tokio::sync::mpsc;
use std::io::{BufRead, Write};
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let (tx, mut rx) = mpsc::unbounded_channel::<String>();
let printer = tokio::spawn(async move {
loop {
let res = rx.recv().await; // (11) Comment this ..
// let res = rx.try_recv(); // (12) Uncomment this ,,
if let Some(m) = res { // .. and this
// if let Ok(m) = res { // ,, and this
if m.trim() == "q".to_string() {
break;
}
println!("Received: {}", m.trim());
}
}
println!("Printer exited");
});
let input_message = tokio::spawn(async move {
let stdin = std::io::stdin();
let mut bufr = std::io::BufReader::new(stdin);
let mut buf = String::new();
loop {
// Let the printer thread print the string before asking the user's input.
std::thread::sleep(std::time::Duration::from_millis(1));
print!("Enter input: ");
std::io::stdout().flush().unwrap();
bufr.read_line(&mut buf).unwrap();
if buf.trim() == "q".to_string() {
tx.send(buf).unwrap();
break;
}
tx.send(buf).unwrap();
buf = String::new();
}
println!("InputMessage exited");
});
tokio::join!(input_message, printer);
Ok(())
}
The expected behavior of the program is to:
Ask the user a random input (q to quit)
Print that same input to stdout
Using rx.recv().await as in line 11-13 the program seems to buffer the Strings representing the user's input: the various inputs are not received by the printer task that therefore does not print the strings to stdout. Once the quit message (i.e. q) is sent, the input_message task exits and the messages seems to be flushed out of the channel and the receiver processes them all at once, and so the printer task prints all the inputs at once. Here's an example of wrong output:
Enter input: Hello
Enter input: World
Enter input: q
InputMessage exited
Received: Hello
Received: World
Printer exited
My question here is, how is it possible that the channel buffers the messages and processes them in one go only when the sending thread exits, instead of receiving them as they are sent?
What I tried to do is to use the try_recv() function as in line 12-14 and indeed it fixes the problem. The output is correctly printed, here is an example:
Enter input: Hello
Received: Hello
Enter input: World
Received: World
Enter input: q
InputMessage exited
Printer exited
In light of this, I get confused. I get the difference between the recv().await and the try_recv() functions but I think there's something more in this case that I'm ignoring that makes the latter work and the former not work. Is anybody able to shed some light and elaborate on this? Why does try_recv() work and recv().await not, and why should recv().await not work in this scenario? In terms of efficiency is looping on try_recv() bad or "bad practice" at all?
There are a few things to point out here, but first of all, you are waiting for lines on std::io::stdin() which blocks the thread until a line arrives on that stream. While the thread waiting for input, no other future can be executed on this thread, this blog post is a great resource if you want to dive deeper why you shouldn't do that.
Tokio's io module offers an async handle to stdin(), you can work with this as a quick fix, although the documentation explicitly mentions that you should spin up a dedicated (non-async) thread for interactive user input instead of using the async handle.
Swapping std::io::stdin() for tokio::io::stdin() also entails swapping out the standard library BufReader for tokio's implementation that wraps an R: AsyncRead rather than R: Read.
To prevent interleaved writes between the input task and the output task, you can use a responder channel that signals to the input task when the output has been printed. Instead of sending String over the channel, you could send a Message with these fields:
struct Message {
payload: String,
done_tx: oneshot::Sender<()>,
}
After reading an input line, send the Message over the channel to the printer task. The printer task prints the String and signals through the done_tx that the input task can print the input prompt and wait for a new line.
Putting all that together with some other changes like a while loop to wait for messages, you'd end up with something like this:
use std::error::Error;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
use tokio::sync::{mpsc, oneshot};
#[derive(Debug)]
struct Message {
done_tx: oneshot::Sender<()>,
message: String,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
let printer = tokio::spawn(async move {
while let Some(Message {
message: m,
done_tx,
}) = rx.recv().await
{
if m.trim() == "q".to_string() {
break;
}
println!("Received: {}", m.trim());
done_tx.send(()).unwrap();
}
println!("Printer exited");
});
let input_message = tokio::spawn(async move {
let stdin = tokio::io::stdin();
let mut stdout = tokio::io::stdout();
let mut bufr = tokio::io::BufReader::new(stdin);
let mut buf = String::new();
loop {
// Let the printer thread print the string before asking the user's input.
stdout.write(b"Enter input: ").await.unwrap();
stdout.flush().await.unwrap();
bufr.read_line(&mut buf).await.unwrap();
let end = buf.trim() == "q";
let (done_tx, done) = oneshot::channel();
let message = Message {
message: buf,
done_tx,
};
tx.send(message).unwrap();
if end {
break;
}
done.await.unwrap();
buf = String::new();
}
println!("InputMessage exited");
});
tokio::join!(input_message, printer);
Ok(())
}

Non-blocking recv on Tokio mpsc Receiver

I am using Rust and Tokio 1.6 to build an app which can interact with an Elgato StreamDeck via hidapi = "1.2". I want to poll the HID device for events (key down / key up) and send those events on an mpsc channel, while watching a separate mpsc channel for incoming commands to update the device state (reset, change brightness, update image, etc). Since the device handle is not thread safe, I need to do both things from a single thread.
major edits below
This is a rewrite of my original question. I've left my interim answer below, but in the interest of a more self contained example, here is a the basic process using device_query = "0.2":
use device_query::{DeviceState, Keycode};
use std::time::Duration;
use tokio;
use tokio::sync::mpsc::{Receiver, Sender};
use tokio::time::timeout;
#[tokio::main]
async fn main() {
// channel for key press events coming from device loop
let (key_tx, mut key_rx) = tokio::sync::mpsc::channel(32);
// channel for commands sent to device loop
let (dev_tx, mut dev_rx) = tokio::sync::mpsc::channel(32);
start_device_loop(60, key_tx, dev_rx);
println!("Waiting for key presses");
while let Some(k) = key_rx.recv().await {
match k {
Some(ch) => match ch {
Keycode::Q => dev_tx.clone().try_send(String::from("Quit!")).expect("Could not send command"),
ch => println!("{}", ch),
},
_ => (),
}
}
println!("Done.")
}
/// Starts a tokio task, polling the supplied device and sending key events
/// on the supplied mpsc sender
pub fn start_device_loop(hz: u32, tx: Sender<Option<Keycode>>, mut rx: Receiver<String>) {
let poll_wait = 1000 / hz;
let poll_wait = Duration::from_millis(poll_wait as u64);
tokio::task::spawn(async move {
let dev = DeviceState::new();
loop {
let mut keys = dev.query_keymap();
match keys.len() {
0 => (),
1 => tx.clone().try_send(Some(keys.remove(0))).unwrap(),
_ => println!("So many keys..."),
}
match timeout(poll_wait, rx.recv()).await {
Ok(cmd) => println!("Command '{}' received.", cmd.unwrap()),
_ => (),
};
// std::thread::sleep(poll_wait);
}
});
}
Note this does not compile - I get an error future created by async block is not 'Send' and within 'impl Future', the trait 'Send' is not implemented for '*mut x11::xlib::_XDisplay'. My understanding of the error is that because device_query is not thread-safe, and awaiting introduces the possibility of scope moving across threads, nothing may be awaited while a non-thread-safe object is in scope. And indeed, if I comment out the block around match timeout... and uncomment the std::thread::sleep everything compiles and runs.
Which brings me back to the original question; how can I both send and receive messages in a single thread without using await or the apparently forbidden fruit of poll_recv()?
After much hunting I found noop_waker in the futures crate which appears to do what I need in combination with poll_recv:
pub fn start_device_loop(hz: u32, tx: Sender<Option<Keycode>>, mut rx: Receiver<String>) {
let poll_wait = 1000 / hz;
let poll_wait = Duration::from_millis(poll_wait as u64);
tokio::task::spawn_blocking(move || {
let dev = DeviceState::new();
let waker = futures::task::noop_waker();
let mut cx = std::task::Context::from_waker(&waker);
loop {
let mut keys = dev.query_keymap();
match keys.len() {
0 => (),
1 => tx.clone().try_send(Some(keys.remove(0))).unwrap(),
_ => println!("So many keys..."),
}
match rx.poll_recv(&mut cx) {
Poll::Ready(cmd) => println!("Command '{}' received.", cmd.unwrap()),
_ => ()
};
std::thread::sleep(poll_wait);
}
});
}
After digging through docs and tokio source more I can't find anything that suggests poll_recv is supposed to be an internal-only function or that using it here would have any obvious side effects. Letting the process run at 125hz I'm not seeing any excess resource usage either.
I'm leaving the above code for posterity, but since asking this question the try_recv method has been added to Receivers, making this all much cleaner.

Rust: How to fix borrowed value does not live long enough

I have simple client/server application. I am receiving message on the server side from client but I want to send that response to the channel from server to other file and I am receiving error "borrowed value does not live long enough".
I have searched in the stack overflow for similar previous questions but not getting enough understanding of lifetime. Is there a good documentation or if simple example available on this topic?
For now if someone can help me to fix this code (may be edit the portion of code which needs to fix) that would be helpful.
Thanks in advance.
Server side:
use std::os::unix::net::UnixDatagram;
use std::path::Path;
fn unlink_socket (path: impl AsRef<Path>) {
let path = path.as_ref();
if Path::new(path).exists() {
let result = std::fs::remove_file(path);
match result {
Err(e) => {
println!("Couldn't remove the file: {:?}", e);
},
_ => {}
}
}
}
pub fn tcp_datagram_server() {
pub static FILE_PATH: &'static str = "/tmp/datagram.sock";
let (tx, rx) = mpsc::channel();
let mut buf = vec![0; 1024];
unlink_socket(FILE_PATH);
let socket = match UnixDatagram::bind(FILE_PATH) {
Ok(socket) => socket,
Err(e) => {
println!("Couldn't bind: {:?}", e);
return;
}
};
println!("Waiting for client to connect...");
loop {
let received_bytes = socket.recv(buf.as_mut_slice()).expect("recv function failed");
println!("Received {:?}", received_bytes);
let received_message = from_utf8(buf.as_slice()).expect("utf-8 convert failed");
tx.clone().send(received_message);
}
}
fn main() {
tcp_datagram_server();
}
client side:
use std::sync::mpsc;
use std::os::unix::net::UnixDatagram;
use std::path::Path;
use std::io::prelude::*;
pub fn tcp_datagram_client() {
pub static FILE_PATH: &'static str = "/tmp/datagram.sock";
let socket = UnixDatagram::unbound().unwrap();
match socket.connect(FILE_PATH) {
Ok(socket) => socket,
Err(e) => {
println!("Couldn't connect: {:?}", e);
return;
}
};
println!("TCP client Connected to TCP Server {:?}", socket);
loop {
socket.send(b"Hello from client to server").expect("recv function failed");
}
}
fn main() {
tcp_datagram_client();
}
Error I am getting
error[E0597]: `buf` does not live long enough
--> src/unix_datagram_server.rs:38:42
|
38 | let received_message = from_utf8(buf.as_slice()).expect("utf-8 convert failed");
| ^^^ borrowed value does not live long enough
...
41 | }
| -
| |
| `buf` dropped here while still borrowed
| borrow might be used here, when `tx` is dropped and runs the `Drop` code for type `std::sync::mpsc::Sender`
|
= note: values in a scope are dropped in the opposite order they are defined
error: aborting due to previous error; 8 warnings emitted
For now if someone can help me to fix this code (may be edit the portion of code which needs to fix) that would be helpful.
Well the message seems rather clear. send does exactly what it says it does, it sends the parameter through the channel. This means the data must live long enough and remain valid "forever" (it needs to be alive and valid in the channel, as well as when fetched from it by the receiver).
That is not the case here. rustc can't understand that the function never returns, and it can panic anyway which will end up the same: the function will terminate, which will invalidate buf. Since received_message borrows buf, that means received_message can't be valid after the function has terminated. But at that point the message would still be in the channel waiting to be read (or retrieved by the receiver doing who knows what).
Therefore your construction is not allowed.
A second issue is that you're overwriting the buffer data on every loop, which has the same effect of breaking the message you sent during the previous iteration, and thus is not correct either. Though Rust won't let you do that either: if you work around the first error it will tell you that there's an outstanding shared borrow (the message sent through the channel) so you can't modify the backing buffer in the following iteration.
The solution is quite simple: have each iteration create an owned string (copying the current iteration's message) and send that through the channel:
tx.clone().send(received_message.to_string());
Also, these are more style / inefficiency remarks but:
The clone() on tx is completely redundant. The point of having a sender that is Clone is being able to send from multiple threads (hence mp in the channel name, that's for multiple producers). Here you have a single thread, the original sender works fine.
.as_slice() and .as_mut_slice() are rarely used unless necessary, which they aren't here: array references coerce to slices, so you can just use &mut buf and &buf. And why are you calling Path::new on something that's already a path? It doesn't do anything but it's not useful either.
It is rather annoying that your snippet is missing multiple imports and thus doesn't even compile as is.
From more of a unixy perspective, errors are usually printed on stderr. In Rust, eprintln does that for you (otherwise working in the same way println does). And I don't understand the purpose of marking a lexically nested static pub. Since the static is inside the function it's not even visible to the function's siblings, to say nothing of external callers. As a result I'd end up with this:
use std::os::unix::net::UnixDatagram;
use std::path::Path;
use std::sync::mpsc;
use std::str::from_utf8;
fn unlink_socket (path: impl AsRef<Path>) {
let path = path.as_ref();
if path.exists() {
if let Err(e) = std::fs::remove_file(path) {
eprintln!("Couldn't remove the file: {:?}", e);
}
}
}
static FILE_PATH: &'static str = "/tmp/datagram.sock";
pub fn tcp_datagram_server() {
unlink_socket(FILE_PATH);
let socket = match UnixDatagram::bind(FILE_PATH) {
Ok(socket) => socket,
Err(e) => {
eprintln!("Couldn't bind: {:?}", e);
return;
}
};
let (tx, _) = mpsc::channel();
let mut buf = vec![0; 1024];
println!("Waiting for client to connect...");
loop {
let received_bytes = socket.recv(&mut buf).expect("recv function failed");
println!("Received {:?}", received_bytes);
let received_message = from_utf8(&buf).expect("utf-8 convert failed");
tx.send(received_message.to_string());
}
}
There's a hint in the compiler message, that values in a scope are dropped in the opposite order they are defined in, and in the example, buf is defined after tx, which means it will be dropped before tx. Since a reference to buf (in the form of received_message) is passed to tx.send(), then buf should live longer that tx, and therefore switching the definition order will fix this particular error (ie. switch lines 19 and 20).

How can I asynchronously read from both stdout and stderr of a subprocess using Tokio? [duplicate]

I'm making a small ncurses application in Rust that needs to communicate with a child process. I already have a prototype written in Common Lisp. I'm trying to rewrite it because CL uses a huge amount of memory for such a small tool.
I'm having some trouble figuring out how to interact with the sub-process.
What I'm currently doing is roughly this:
Create the process:
let mut program = match Command::new(command)
.args(arguments)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(_) => {
println!("Cannot run program '{}'.", command);
return;
}
};
Pass it to an infinite (until user exits) loop, which reads and handles input and listens for output like this (and writes it to the screen):
fn listen_for_output(program: &mut Child, output_viewer: &TextViewer) {
match program.stdout {
Some(ref mut out) => {
let mut buf_string = String::new();
match out.read_to_string(&mut buf_string) {
Ok(_) => output_viewer.append_string(buf_string),
Err(_) => return,
};
}
None => return,
};
}
The call to read_to_string however blocks the program until the process exits. From what I can see read_to_end and read also seem to block. If I try running something like ls which exits right away, it works, but with something that doesn't exit like python or sbcl it only continues once I kill the subprocess manually.
Based on this answer, I changed the code to use BufReader:
fn listen_for_output(program: &mut Child, output_viewer: &TextViewer) {
match program.stdout.as_mut() {
Some(out) => {
let buf_reader = BufReader::new(out);
for line in buf_reader.lines() {
match line {
Ok(l) => {
output_viewer.append_string(l);
}
Err(_) => return,
};
}
}
None => return,
}
}
However, the problem still remains the same. It will read all lines that are available, and then block. Since the tool is supposed to work with any program, there is no way to guess out when the output will end, before trying to read. There doesn't appear to be a way to set a timeout for BufReader either.
Streams are blocking by default. TCP/IP streams, filesystem streams, pipe streams, they are all blocking. When you tell a stream to give you a chunk of bytes it will stop and wait till it has the given amout of bytes or till something else happens (an interrupt, an end of stream, an error).
The operating systems are eager to return the data to the reading process, so if all you want is to wait for the next line and handle it as soon as it comes in then the method suggested by Shepmaster in Unable to pipe to or from spawned child process more than once (and also in his answer here) works.
Though in theory it doesn't have to work, because an operating system is allowed to make the BufReader wait for more data in read, but in practice the operating systems prefer the early "short reads" to waiting.
This simple BufReader-based approach becomes even more dangerous when you need to handle multiple streams (like the stdout and stderr of a child process) or multiple processes. For example, BufReader-based approach might deadlock when a child process waits for you to drain its stderr pipe while your process is blocked waiting on it's empty stdout.
Similarly, you can't use BufReader when you don't want your program to wait on the child process indefinitely. Maybe you want to display a progress bar or a timer while the child is still working and gives you no output.
You can't use BufReader-based approach if your operating system happens not to be eager in returning the data to the process (prefers "full reads" to "short reads") because in that case a few last lines printed by the child process might end up in a gray zone: the operating system got them, but they're not large enough to fill the BufReader's buffer.
BufReader is limited to what the Read interface allows it to do with the stream, it's no less blocking than the underlying stream is. In order to be efficient it will read the input in chunks, telling the operating system to fill as much of its buffer as it has available.
You might be wondering why reading data in chunks is so important here, why can't the BufReader just read the data byte by byte. The problem is that to read the data from a stream we need the operating system's help. On the other hand, we are not the operating system, we work isolated from it, so as not to mess with it if something goes wrong with our process. So in order to call to the operating system there needs to be a transition to "kernel mode" which might also incur a "context switch". That is why calling the operating system to read every single byte is expensive. We want as few OS calls as possible and so we get the stream data in batches.
To wait on a stream without blocking you'd need a non-blocking stream. MIO promises to have the required non-blocking stream support for pipes, most probably with PipeReader, but I haven't checked it out so far.
The non-blocking nature of a stream should make it possible to read data in chunks regardless of whether the operating system prefers the "short reads" or not. Because non-blocking stream never blocks. If there is no data in the stream it simply tells you so.
In the absense of a non-blocking stream you'll have to resort to spawning threads so that the blocking reads would be performed in a separate thread and thus won't block your primary thread. You might also want to read the stream byte by byte in order to react to the line separator immediately in case the operating system does not prefer the "short reads". Here's a working example: https://gist.github.com/ArtemGr/db40ae04b431a95f2b78.
P.S. Here's an example of a function that allows to monitor the standard output of a program via a shared vector of bytes:
use std::io::Read;
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
/// Pipe streams are blocking, we need separate threads to monitor them without blocking the primary thread.
fn child_stream_to_vec<R>(mut stream: R) -> Arc<Mutex<Vec<u8>>>
where
R: Read + Send + 'static,
{
let out = Arc::new(Mutex::new(Vec::new()));
let vec = out.clone();
thread::Builder::new()
.name("child_stream_to_vec".into())
.spawn(move || loop {
let mut buf = [0];
match stream.read(&mut buf) {
Err(err) => {
println!("{}] Error reading from stream: {}", line!(), err);
break;
}
Ok(got) => {
if got == 0 {
break;
} else if got == 1 {
vec.lock().expect("!lock").push(buf[0])
} else {
println!("{}] Unexpected number of bytes: {}", line!(), got);
break;
}
}
}
})
.expect("!thread");
out
}
fn main() {
let mut cat = Command::new("cat")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.expect("!cat");
let out = child_stream_to_vec(cat.stdout.take().expect("!stdout"));
let err = child_stream_to_vec(cat.stderr.take().expect("!stderr"));
let mut stdin = match cat.stdin.take() {
Some(stdin) => stdin,
None => panic!("!stdin"),
};
}
With a couple of helpers I'm using it to control an SSH session:
try_s! (stdin.write_all (b"echo hello world\n"));
try_s! (wait_forĖ¢ (&out, 0.1, 9., |s| s == "hello world\n"));
P.S. Note that await on a read call in async-std is blocking as well. It's just instead of blocking a system thread it only blocks a chain of futures (a stack-less green thread essentially). The poll_read is the non-blocking interface. In async-std#499 I've asked the developers whether there's a short read guarantee from these APIs.
P.S. There might be a similar concern in Nom: "we would want to tell the IO side to refill according to the parser's result (Incomplete or not)"
P.S. Might be interesting to see how stream reading is implemented in crossterm. For Windows, in poll.rs, they are using the native WaitForMultipleObjects. In unix.rs they are using mio poll.
Tokio's Command
Here is an example of using tokio 0.2:
use std::process::Stdio;
use futures::StreamExt; // 0.3.1
use tokio::{io::BufReader, prelude::*, process::Command}; // 0.2.4, features = ["full"]
#[tokio::main]
async fn main() {
let mut cmd = Command::new("/tmp/slow.bash")
.stdout(Stdio::piped()) // Can do the same for stderr
.spawn()
.expect("cannot spawn");
let stdout = cmd.stdout().take().expect("no stdout");
// Can do the same for stderr
// To print out each line
// BufReader::new(stdout)
// .lines()
// .for_each(|s| async move { println!("> {:?}", s) })
// .await;
// To print out each line *and* collect it all into a Vec
let result: Vec<_> = BufReader::new(stdout)
.lines()
.inspect(|s| println!("> {:?}", s))
.collect()
.await;
println!("All the lines: {:?}", result);
}
Tokio-Threadpool
Here is an example of using tokio 0.1 and tokio-threadpool. We start the process in a thread using the blocking function. We convert that to a stream with stream::poll_fn
use std::process::{Command, Stdio};
use tokio::{prelude::*, runtime::Runtime}; // 0.1.18
use tokio_threadpool; // 0.1.13
fn stream_command_output(
mut command: Command,
) -> impl Stream<Item = Vec<u8>, Error = tokio_threadpool::BlockingError> {
// Ensure that the output is available to read from and start the process
let mut child = command
.stdout(Stdio::piped())
.spawn()
.expect("cannot spawn");
let mut stdout = child.stdout.take().expect("no stdout");
// Create a stream of data
stream::poll_fn(move || {
// Perform blocking IO
tokio_threadpool::blocking(|| {
// Allocate some space to store anything read
let mut data = vec![0; 128];
// Read 1-128 bytes of data
let n_bytes_read = stdout.read(&mut data).expect("cannot read");
if n_bytes_read == 0 {
// Stdout is done
None
} else {
// Only return as many bytes as we read
data.truncate(n_bytes_read);
Some(data)
}
})
})
}
fn main() {
let output_stream = stream_command_output(Command::new("/tmp/slow.bash"));
let mut runtime = Runtime::new().expect("Unable to start the runtime");
let result = runtime.block_on({
output_stream
.map(|d| String::from_utf8(d).expect("Not UTF-8"))
.fold(Vec::new(), |mut v, s| {
print!("> {}", s);
v.push(s);
Ok(v)
})
});
println!("All the lines: {:?}", result);
}
There's numerous possible tradeoffs that can be made here. For example, always allocating 128 bytes isn't ideal, but it's simple to implement.
Support
For reference, here's slow.bash:
#!/usr/bin/env bash
set -eu
val=0
while [[ $val -lt 10 ]]; do
echo $val
val=$(($val + 1))
sleep 1
done
See also:
How do I synchronously return a value calculated in an asynchronous Future in stable Rust?
If Unix support is sufficient, you can also make the two output streams as non-blocking and poll over them as you would do it on TcpStream with the set_nonblocking function.
The ChildStdout and ChildStderr returned by the Command spawn are Stdio (and contain a file descriptor), you can modify directly the read behavior of these handle to make it non-blocking.
Based on the work of jcreekmore/timeout-readwrite-rs and anowell/nonblock-rs, I use this wrapper to modify the stream handles:
extern crate libc;
use std::io::Read;
use std::os::unix::io::AsRawFd;
use libc::{F_GETFL, F_SETFL, fcntl, O_NONBLOCK};
fn set_nonblocking<H>(handle: &H, nonblocking: bool) -> std::io::Result<()>
where
H: Read + AsRawFd,
{
let fd = handle.as_raw_fd();
let flags = unsafe { fcntl(fd, F_GETFL, 0) };
if flags < 0 {
return Err(std::io::Error::last_os_error());
}
let flags = if nonblocking{
flags | O_NONBLOCK
} else {
flags & !O_NONBLOCK
};
let res = unsafe { fcntl(fd, F_SETFL, flags) };
if res != 0 {
return Err(std::io::Error::last_os_error());
}
Ok(())
}
You can manage the two streams as any other non-blocking stream. The following example is based on the polling crate which makes really easy to handle read event and BufReader for line reading:
use std::process::{Command, Stdio};
use std::path::PathBuf;
use std::io::{BufReader, BufRead};
use std::thread;
extern crate polling;
use polling::{Event, Poller};
fn main() -> Result<(), std::io::Error> {
let path = PathBuf::from("./worker.sh").canonicalize()?;
let mut child = Command::new(path)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.expect("Failed to start worker");
let handle = thread::spawn({
let stdout = child.stdout.take().unwrap();
set_nonblocking(&stdout, true)?;
let mut reader_out = BufReader::new(stdout);
let stderr = child.stderr.take().unwrap();
set_nonblocking(&stderr, true)?;
let mut reader_err = BufReader::new(stderr);
move || {
let key_out = 1;
let key_err = 2;
let mut out_closed = false;
let mut err_closed = false;
let poller = Poller::new().unwrap();
poller.add(reader_out.get_ref(), Event::readable(key_out)).unwrap();
poller.add(reader_err.get_ref(), Event::readable(key_err)).unwrap();
let mut line = String::new();
let mut events = Vec::new();
loop {
// Wait for at least one I/O event.
events.clear();
poller.wait(&mut events, None).unwrap();
for ev in &events {
// stdout is ready for reading
if ev.key == key_out {
let len = match reader_out.read_line(&mut line) {
Ok(len) => len,
Err(e) => {
println!("stdout read returned error: {}", e);
0
}
};
if len == 0 {
println!("stdout closed (len is null)");
out_closed = true;
poller.delete(reader_out.get_ref()).unwrap();
} else {
print!("[STDOUT] {}", line);
line.clear();
// reload the poller
poller.modify(reader_out.get_ref(), Event::readable(key_out)).unwrap();
}
}
// stderr is ready for reading
if ev.key == key_err {
let len = match reader_err.read_line(&mut line) {
Ok(len) => len,
Err(e) => {
println!("stderr read returned error: {}", e);
0
}
};
if len == 0 {
println!("stderr closed (len is null)");
err_closed = true;
poller.delete(reader_err.get_ref()).unwrap();
} else {
print!("[STDERR] {}", line);
line.clear();
// reload the poller
poller.modify(reader_err.get_ref(), Event::readable(key_err)).unwrap();
}
}
}
if out_closed && err_closed {
println!("Stream closed, exiting process thread");
break;
}
}
}
});
handle.join().unwrap();
Ok(())
}
Additionally, used with a wrapper over an EventFd, it becomes possible to easily stop the process from another thread without blocking nor active polling and uses and only a single thread.
EDIT: It seems the polling crate sets automatically the polled handles in non-blocking mode following my tests. The set_nonblocking function is still useful in case you want to directly use the nix::poll object.
I have encountered enough use-cases where it was useful to interact with a subprocess over line-delimited text that I wrote a crate for it, interactive_process.
I expect the original problem has long since been solved, but I thought it might be helpful to others.

How to read a little at a time from UnixStream connected to socket?

I am trying to read some data of unknown size from a UnixStream (called socket in this code). The data consists of a header of 6 bytes, and the last two bytes indicate how long the rest of the message should be.
The .into_buf() method I'm calling is from the IntoBuf trait in Tokio's bytes crate.
let mut header = [0u8; 6];
let response1 = self.socket.read(&mut header);
let mut cursor = header.into_buf();
let evt_code = cursor.get_u16_le();
let controller = cursor.get_u16_le();
let param_size = cursor.get_u16_le() as usize;
let mut param = vec![0u8; param_size];
let response2 = self.socket.read(&mut param);
let mut cursor = param.into_buf();
The problem that I'm encountering is that response2 is always an Err E_WOULDBLOCK because my UnixStream is connected to a nonblocking socket. It seems that the first call to read() is reading 6 bytes as intended, but then is just discarding the rest of the content in the stream.
How can I work around this / how can I make read() leave the extra data in the stream?
It seems like the simplest solution would be to just make my initial buffer big and then just read everything at once, but the problem with this is that the largest possible message I could receive is a little over 64KB. Allocating or re-zeroing a 64KB buffer for every read seems quite wasteful, especially because most of the messages are much smaller than this.
I have worked around this by using a raw file descriptor instead of a UnixStream and calling recv() with MSG_PEEK.
let mut header = [0u8; 6];
// need MSG_PEEK otherwise recv() clears the socket and we get EAGAIN when
// we try to read the socket again; this is the main reason for using a raw
// fd instead of a UnixStream
if unsafe { libc::recv(self.fd, header.as_mut_ptr() as *mut c_void, 6, libc::MSG_PEEK) } < 0 {
return Err(io::Error::last_os_error().into());
}
let mut cursor = header.into_buf();
let evt_code = cursor.get_u16_le();
let controller = cursor.get_u16_le();
let param_size = cursor.get_u16_le() as usize;
// since calling recv() with MSG_PEEK doesn't consume the header, we need to make
// this buffer 6 bytes bigger, but that's fine
let mut param = vec![0u8; param_size + 6];
if unsafe { libc::recv(self.fd, param.as_mut_ptr() as *mut c_void, param.len(), 0) } < 0 {
return Err(io::Error::last_os_error().into());
}
let mut cursor = param.into_buf();

Resources