What is `take()` doing here, and why do I need it? - rust

This code spawns a child process, consuming its stderr and stdout line by line, and logging each appropriately. It compiles and works.
use std::error::Error;
use std::process::{Stdio};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Command, Child};
use tracing::{info, warn};
macro_rules! relay_pipe_lines {
($pipe:expr, $handler:expr) => {
tokio::spawn(async move {
let mut reader = BufReader::new($pipe).lines();
loop {
let line = reader
.next_line()
.await
.unwrap_or_else(|_| Some(String::new()));
match line {
None => break,
Some(line) => $handler(line)
}
}
});
};
}
pub fn start_and_log_command(mut command: Command) -> Result<Child, Box<dyn Error>> {
command.stdout(Stdio::piped()).stderr(Stdio::piped());
let mut child = command.spawn()?;
let child_stdout = child.stdout.take().unwrap(); // remove `take` from here
let child_stderr = child.stderr.take().unwrap(); // .. or from here and it fails
let child_pid = child.id().unwrap();
relay_pipe_lines!(child_stdout, |line|info!("[pid {}:stdout]: {}", child_pid, line));
relay_pipe_lines!(child_stderr, |line|warn!("[pid {}:stderr]: {}", child_pid, line));
Ok(child)
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
tracing_subscriber::fmt::init();
info!("Tracing logging initialised.");
let mut command = Command::new("ls");
command.arg("-l");
let mut child = start_and_log_command(command)?;
// Compose reading waiting concurrently.
let exit_status = child.wait().await.expect("Cannot reap child process");
dbg!(exit_status.success());
Ok(())
}
Removing the call to take() from the indicated lines fails the build, as "child.stdout partially moved due to this method call", which I mostly understand.
I'd like to understand how using take() does not partially move child.stdout.

It's a call to Option::take(), which avoids a "partial move" by leaving None in place of the moved value. As a result, child is left in a valid state and can be returned from the function.
In other words, child.stdout.take() is equivalent to std::mem::replace(&mut child.stdout, None), and means "take the current value out of the option (whatever it is), and leave None in its place."

Related

How to specify an output type for a Future/async block?

use tokio::runtime::Runtime;
// Create the runtime
let rt = Runtime::new().unwrap();
// Execute the future, blocking the current thread until completion
let s = rt.block_on(async {
println!("hello");
});
is it possible to specify an output type for a future block? On the code above, s: (), but I wanted to be Result<(), Error> so I can return some error from inside the block.
I am not quite familiar with async rust, but as far as I know, the return type of an async fn or async block is impl Future<Output=TheRealReturnTypeOfFnBody>. Once blocked on, as you did rt.block_on(async block), the return type will become TheRealReturnTypeOfFnBody, therefore:
To make s have type Result<(), Error>, you have to implement it in the function body(i.e. make TheRealReturnTypeOfFnBody Result<(), Error>)
use tokio::runtime::Runtime;
fn main() {
// Create the runtime
let rt = Runtime::new().unwrap();
// Execute the future, blocking the current thread until completion
let s: () = rt.block_on(async {
println!("hello");
});
let s_with_error_case: Result<(), &str> = rt.block_on(async {
if false {
Err("run into a trouble")
} else {
println!("erverything is fine");
Ok(())
}
});
if let Err(err_info) = s_with_error_case {
eprintln!("{}", err_info);
}
}

Unable to read child stderr twice in Rust

For some reason, I am unable to read from the child process stderr the second time. Here is what I do.
I am spawning a child process for cucumber tests. In the first step I spawn the process, take its stderr, save it, and then read from it. Here is the code:
pub fn wait_process_output(
reader: &mut BufReader<ChildStderr>,
output: Vec<(String, u16)>,
) -> Result<(), String> {
let mut process_output = String::new();
loop {
match reader.read_line(&mut process_output) {
Err(e) => {
return Err(format!("Unable to read output: {}", e));
}
Ok(_) => {
// processing here
}
};
}
}
pub fn step1(world: &mut TestWorld) {
world.app_handler = Some(
Command::new(&world.app_bin)
.stderr(Stdio::piped())
.spawn()
.unwrap(),
);
let app_stderr = world.app_handler.as_mut().unwrap().stderr.take().unwrap();
world.app_reader = Some(BufReader::new(app_stderr));
wait_process_output(world.app_reader.as_mut().unwrap(), /* expected data */).ok();
}
This code works correctly: stderr is being read as expected.
In the third test step, I try to read process output one more time:
pub fn step3(world: &mut TestWorld) {
wait_process_output(world.app_reader.as_mut().unwrap(), /* expected data */).ok();
}
This time reader.read_line hangs infinitely: nothing is being read. I am sure the child process produces some output: I can see it if I run it within the same conditions separately.
Could you please suggest any ideas why the BufReader object becomes corrupted when I try to read from it the second time?
I got to the solution. The issue was that app produced an output earlier than step3 started to read it. I thought that there is some kind of buffer implemented for this kind of situation but it seems I was wrong. So I finally used the following two methods to solve my issue:
pub fn wait_process_output(
receiver: &Receiver<String>,
output: Vec<(String, u16)>,
) -> Result<(), String> {
loop {
match receiver.try_recv() {
// process output
}
}
}
pub fn start_listener(sender: Sender<String>, stream: ChildStderr) {
spawn(move || {
let mut f = BufReader::new(stream);
loop {
let mut buf = String::new();
match f.read_line(&mut buf) {
Ok(_) => {
sender.send(buf).unwrap();
continue;
}
Err(e) => {
println!("Unable to read process stderr: {:?}", e);
break;
}
}
}
});
}

How to add special NotReady logic to tokio-io?

I'm trying to make a Stream that would wait until a specific character is in buffer. I know there's read_until() on BufRead but I actually need a custom solution, as this is a stepping stone to implement waiting until a specific string in in buffer (or, for example, a regexp match happens).
In my project where I first encountered the problem, problem was that future processing just hanged when I get a Ready(_) from inner future and return NotReady from my function. I discovered I shouldn't do that per docs (last paragraph). However, what I didn't get, is what's the actual alternative that is promised in that paragraph. I read all the published documentation on the Tokio site and it doesn't make sense for me at the moment.
So following is my current code. Unfortunately I couldn't make it simpler and smaller as it's already broken. Current result is this:
Err(Custom { kind: Other, error: Error(Shutdown) })
Err(Custom { kind: Other, error: Error(Shutdown) })
Err(Custom { kind: Other, error: Error(Shutdown) })
<ad infinum>
Expected result is getting some Ok(Ready(_)) out of it, while printing W and W', and waiting for specific character in buffer.
extern crate futures;
extern crate tokio_core;
extern crate tokio_io;
extern crate tokio_io_timeout;
extern crate tokio_process;
use futures::stream::poll_fn;
use futures::{Async, Poll, Stream};
use tokio_core::reactor::Core;
use tokio_io::AsyncRead;
use tokio_io_timeout::TimeoutReader;
use tokio_process::CommandExt;
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
struct Process {
child: tokio_process::Child,
stdout: Arc<Mutex<tokio_io_timeout::TimeoutReader<tokio_process::ChildStdout>>>,
}
impl Process {
fn new(
command: &str,
reader_timeout: Option<Duration>,
core: &tokio_core::reactor::Core,
) -> Self {
let mut cmd = Command::new(command);
let cat = cmd.stdout(Stdio::piped());
let mut child = cat.spawn_async(&core.handle()).unwrap();
let stdout = child.stdout().take().unwrap();
let mut timeout_reader = TimeoutReader::new(stdout);
timeout_reader.set_timeout(reader_timeout);
let timeout_reader = Arc::new(Mutex::new(timeout_reader));
Self {
child,
stdout: timeout_reader,
}
}
}
fn work() -> Result<(), ()> {
let window = Arc::new(Mutex::new(Vec::new()));
let mut core = Core::new().unwrap();
let process = Process::new("cat", Some(Duration::from_secs(20)), &core);
let mark = Arc::new(Mutex::new(b'c'));
let read_until_stream = poll_fn({
let window = window.clone();
let timeout_reader = process.stdout.clone();
move || -> Poll<Option<u8>, std::io::Error> {
let mut buf = [0; 8];
let poll;
{
let mut timeout_reader = timeout_reader.lock().unwrap();
poll = timeout_reader.poll_read(&mut buf);
}
match poll {
Ok(Async::Ready(0)) => Ok(Async::Ready(None)),
Ok(Async::Ready(x)) => {
{
let mut window = window.lock().unwrap();
println!("W: {:?}", *window);
println!("buf: {:?}", &buf[0..x]);
window.extend(buf[0..x].into_iter().map(|x| *x));
println!("W': {:?}", *window);
if let Some(_) = window.iter().find(|c| **c == *mark.lock().unwrap()) {
Ok(Async::Ready(Some(1)))
} else {
Ok(Async::NotReady)
}
}
}
Ok(Async::NotReady) => Ok(Async::NotReady),
Err(e) => Err(e),
}
}
});
let _stream_thread = thread::spawn(move || {
for o in read_until_stream.wait() {
println!("{:?}", o);
}
});
match core.run(process.child) {
Ok(_) => {}
Err(e) => {
println!("Child error: {:?}", e);
}
}
Ok(())
}
fn main() {
work().unwrap();
}
This is complete example project.
If you need more data you need to call poll_read again until you either find what you were looking for or poll_read returns NotReady.
You might want to avoid looping in one task for too long, so you can build yourself a yield_task function to call instead if poll_read didn't return NotReady; it makes sure your task gets called again ASAP after other pending tasks were run.
To use it just run return yield_task();.
fn yield_inner() {
use futures::task;
task::current().notify();
}
#[inline(always)]
pub fn yield_task<T, E>() -> Poll<T, E> {
yield_inner();
Ok(Async::NotReady)
}
Also see futures-rs#354: Handle long-running, always-ready futures fairly #354.
With the new async/await API futures::task::current is gone; instead you'll need a std::task::Context reference, which is provided as parameter to the new std::future::Future::poll trait method.
If you're already manually implementing the std::future::Future trait you can simply insert:
context.waker().wake_by_ref();
return std::task::Poll::Pending;
Or build yourself a Future-implementing type that yields exactly once:
pub struct Yield {
ready: bool,
}
impl core::future::Future for Yield {
type Output = ();
fn poll(self: core::pin::Pin<&mut Self>, cx: &mut core::task::Context<'_>) -> core::task::Poll<Self::Output> {
let this = self.get_mut();
if this.ready {
core::task::Poll::Ready(())
} else {
cx.waker().wake_by_ref();
this.ready = true; // ready next round
core::task::Poll::Pending
}
}
}
pub fn yield_task() -> Yield {
Yield { ready: false }
}
And then use it in async code like this:
yield_task().await;

How to cleanly break tokio-core event loop and futures::Stream in Rust

I am dabbling in tokio-core and can figure out how to spawn an event loop. However there are two things i am not sure of - how to gracefully exit the event loop and how to exit a stream running inside an event loop. For e.g consider this simple piece of code which spawns two listeners into the event loop and waits for another thread to indicate an exit condition:
extern crate tokio_core;
extern crate futures;
use tokio_core::reactor::Core;
use futures::sync::mpsc::unbounded;
use tokio_core::net::TcpListener;
use std::net::SocketAddr;
use std::str::FromStr;
use futures::{Stream, Future};
use std::thread;
use std::time::Duration;
use std::sync::mpsc::channel;
fn main() {
let (get_tx, get_rx) = channel();
let j = thread::spawn(move || {
let mut core = Core::new().unwrap();
let (tx, rx) = unbounded();
get_tx.send(tx).unwrap(); // <<<<<<<<<<<<<<< (1)
// Listener-0
{
let l = TcpListener::bind(&SocketAddr::from_str("127.0.0.1:44444").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
// Listener1
{
let l = TcpListener::bind(&SocketAddr::from_str("127.0.0.1:55555").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
let work = rx.for_each(|v| {
if v {
// (3) I want to shut down listener-0 above the release the resources
Ok(())
} else {
Err(()) // <<<<<<<<<<<<<<< (2)
}
});
let _ = core.run(work);
println!("Exiting event loop thread");
});
let tx = get_rx.recv().unwrap();
thread::sleep(Duration::from_secs(2));
println!("Want to terminate listener-0"); // <<<<<< (3)
tx.send(true).unwrap();
thread::sleep(Duration::from_secs(2));
println!("Want to exit event loop");
tx.send(false).unwrap();
j.join().unwrap();
}
So say after the sleep in the main thread i want a clean exit of the event loop thread. Currently I send something to the event loop to make it exit and thus releasing the thread.
However both, (1) and (2) feel hacky - i am forcing an error as an exit condition. My questions are:
1) Am I doing it right ? If not then what is the correct way to gracefully exit the event loop thread.
2) I don't event know how to do (3) - i.e. indicate a condition externally to shutdown listener-0 and free all it's resources. How do i achieve this ?
The event loop (core) is not being turned any more (e.g. by run()) or is forgotten (drop()ed). There is no synchronous exit. core.run() returns and stops turning the loop when the Future passed to it completes.
A Stream completes by yielding None (marked with (3) in the code below).
When e.g. a TCP connection is closed the Stream representing it completes and the other way around.
extern crate tokio_core;
extern crate futures;
use tokio_core::reactor::Core;
use futures::sync::mpsc::unbounded;
use tokio_core::net::TcpListener;
use std::net::SocketAddr;
use std::str::FromStr;
use futures::{Async, Stream, Future, Poll};
use std::thread;
use std::time::Duration;
struct CompletionPact<S, C>
where S: Stream,
C: Stream,
{
stream: S,
completer: C,
}
fn stream_completion_pact<S, C>(s: S, c: C) -> CompletionPact<S, C>
where S: Stream,
C: Stream,
{
CompletionPact {
stream: s,
completer: c,
}
}
impl<S, C> Stream for CompletionPact<S, C>
where S: Stream,
C: Stream,
{
type Item = S::Item;
type Error = S::Error;
fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
match self.completer.poll() {
Ok(Async::Ready(None)) |
Err(_) |
Ok(Async::Ready(Some(_))) => {
// We are done, forget us
Ok(Async::Ready(None)) // <<<<<< (3)
},
Ok(Async::NotReady) => {
self.stream.poll()
},
}
}
}
fn main() {
// unbounded() is the equivalent of a Stream made from a channel()
// directly create it in this thread instead of receiving a Sender
let (tx, rx) = unbounded::<()>();
// A second one to cause forgetting the listener
let (l0tx, l0rx) = unbounded::<()>();
let j = thread::spawn(move || {
let mut core = Core::new().unwrap();
// Listener-0
{
let l = TcpListener::bind(
&SocketAddr::from_str("127.0.0.1:44444").unwrap(),
&core.handle())
.unwrap();
// wrap the Stream of incoming connections (which usually doesn't
// complete) into a Stream that completes when the
// other side is drop()ed or sent on
let fe = stream_completion_pact(l.incoming(), l0rx)
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
// Listener1
{
let l = TcpListener::bind(
&SocketAddr::from_str("127.0.0.1:55555").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
let _ = core.run(rx.into_future());
println!("Exiting event loop thread");
});
thread::sleep(Duration::from_secs(2));
println!("Want to terminate listener-0");
// A drop() will result in the rx side Stream being completed,
// which is indicated by Ok(Async::Ready(None)).
// Our wrapper behaves the same when something is received.
// When the event loop encounters a
// Stream that is complete it forgets about it. Which propagates to a
// drop() that close()es the file descriptor, which closes the port if
// nothing else uses it.
l0tx.send(()).unwrap(); // alternatively: drop(l0tx);
// Note that this is async and is only the signal
// that starts the forgetting.
thread::sleep(Duration::from_secs(2));
println!("Want to exit event loop");
// Same concept. The reception or drop() will cause Stream completion.
// A completed Future will cause run() to return.
tx.send(()).unwrap();
j.join().unwrap();
}
I implemented graceful shutdown via a oneshot channel.
The trick was to use both a oneshot channel to cancel the tcp listener, and use a select! on the two futures. Note I'm using tokio 0.2 and futures 0.3 in the example below.
use futures::channel::oneshot;
use futures::{FutureExt, StreamExt};
use std::thread;
use tokio::net::TcpListener;
pub struct ServerHandle {
// This is the thread in which the server will block
thread: thread::JoinHandle<()>,
// This switch can be used to trigger shutdown of the server.
kill_switch: oneshot::Sender<()>,
}
impl ServerHandle {
pub fn stop(self) {
self.kill_switch.send(()).unwrap();
self.thread.join().unwrap();
}
}
pub fn run_server() -> ServerHandle {
let (kill_switch, kill_switch_receiver) = oneshot::channel::<()>();
let thread = thread::spawn(move || {
info!("Server thread begun!!!");
let mut runtime = tokio::runtime::Builder::new()
.basic_scheduler()
.enable_all()
.thread_name("Tokio-server-thread")
.build()
.unwrap();
runtime.block_on(async {
server_prog(kill_switch_receiver).await.unwrap();
});
info!("Server finished!!!");
});
ServerHandle {
thread,
kill_switch,
}
}
async fn server_prog(kill_switch_receiver: oneshot::Receiver<()>) -> std::io::Result<()> {
let addr = "127.0.0.1:12345";
let addr: std::net::SocketAddr = addr.parse().unwrap();
let mut listener = TcpListener::bind(&addr).await?;
let mut kill_switch_receiver = kill_switch_receiver.fuse();
let mut incoming = listener.incoming().fuse();
loop {
futures::select! {
x = kill_switch_receiver => {
break;
},
optional_new_client = incoming.next() => {
if let Some(new_client) = optional_new_client {
let peer_socket = new_client?;
info!("Client connected!");
let peer = process_client(peer_socket, db.clone());
peers.lock().unwrap().push(peer);
} else {
info!("No more incoming connections.");
break;
}
},
};
}
Ok(())
}
Hopes this helps others (or future me ;)).
My code lives here:
https://github.com/windelbouwman/lognplot/blob/master/lognplot/src/server/server.rs

How can I pass a socket as an argument to a function being called within a thread?

I'm going to have multiple functions that all need access to one main socket.
Would it better to:
Pass this socket to each function that needs access to it
Have a globally accessible socket
Can someone provide an example of the best way to do this?
I come from a Python/Nim background where things like this are easily done.
Edit:
How can I pass a socket as an arg to a function being called within a thread.
Ex.
fn main() {
let mut s = BufferedStream::new((TcpStream::connect(server).unwrap()));
let thread = Thread::spawn(move || {
func1(s, arg1, arg2);
});
while true {
func2(s, arg1);
}
}
Answer for updated question
We can use TcpStream::try_clone:
use std::io::Read;
use std::net::{TcpStream, Shutdown};
use std::thread;
fn main() {
let mut stream = TcpStream::connect("127.0.0.1:34254").unwrap();
let stream2 = stream.try_clone().unwrap();
let _t = thread::spawn(move || {
// close this stream after one second
thread::sleep_ms(1000);
stream2.shutdown(Shutdown::Read).unwrap();
});
// wait for some data, will get canceled after one second
let mut buf = [0];
stream.read(&mut buf).unwrap();
}
Original answer
It's usually (let's say 99.9% of the time) a bad idea to have any global mutable state, if you can help it. Just do as you said: pass the socket to the functions that need it.
use std::io::{self, Write};
use std::net::TcpStream;
fn send_name(stream: &mut TcpStream) -> io::Result<()> {
stream.write(&[42])?;
Ok(())
}
fn send_number(stream: &mut TcpStream) -> io::Result<()> {
stream.write(&[1, 2, 3])?;
Ok(())
}
fn main() {
let mut stream = TcpStream::connect("127.0.0.1:31337").unwrap();
let r = send_name(&mut stream).and_then(|_| send_number(&mut stream));
match r {
Ok(..) => println!("Yay, sent!"),
Err(e) => println!("Boom! {}", e),
}
}
You could also pass the TcpStream to a struct that manages it, and thus gives you a place to put similar methods.
use std::io::{self, Write};
use std::net::TcpStream;
struct GameService {
stream: TcpStream,
}
impl GameService {
fn send_name(&mut self) -> io::Result<()> {
self.stream.write(&[42])?;
Ok(())
}
fn send_number(&mut self) -> io::Result<()> {
self.stream.write(&[1, 2, 3])?;
Ok(())
}
}
fn main() {
let stream = TcpStream::connect("127.0.0.1:31337").unwrap();
let mut service = GameService { stream: stream };
let r = service.send_name().and_then(|_| service.send_number());
match r {
Ok(..) => println!("Yay, sent!"),
Err(e) => println!("Boom! {}", e),
}
}
None of this is really Rust-specific, these are generally-applicable programming practices.

Resources