How to add special NotReady logic to tokio-io?

Question

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.

Stefan Stefan · Accepted Answer · 2018-05-17T20:47:39

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 add special NotReady logic to tokio-io?

1 Answers