[Rust] [tokio]Futureの実装とasync/await



いわゆる「やってみた系」の記事なので多くは望まないで欲しい。

# 何をするのか

`Future`の自前実装による非同期処理とそれがasync/await導入でどのように変わるかを示す。
なおコードはRust2018で記述している。

# 最初の`Future`

まず簡単なFutureを実装してみる。

```rust
use tokio::prelude::*;

struct MyFuture;

impl Future for MyFuture {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        println!("poll called");
        Ok(Async::Ready(()))
    }
}

fn main() {
    tokio::run(MyFuture{});
    println!("finished");
}
```

`poll()`は正常時には計算終了を表す`Async::Ready(...)`または、まだ計算中であることを示す`Async::NotReady`を返す。**`poll()`は自分の処理がまだ終わっていなくても（例えばスリープ中や電文送信中等）でも~~自分以外の要因で~~(要調査)呼び出される可能性がある**ため、その際には`Async::NotReady`を返してまだ自分の処理が終わっていないことを明示する必要がある。

実行すると`MyFuture`の`poll()`が呼ばれ`Ok(Async::Ready(()))`を返すため即座に処理は終了する。

```bash
$ cargo run
poll called
finished
$
```

`Ok(Async::Ready(()))`を`Ok(Async::NotReady)`に変更するとどうなるか。

```diff
@@ -8,7 +8,7 @@ impl Future for MyFuture {
 
     fn poll(&mut self) -> Poll<(), ()> {
         println!("poll called");
-        Ok(Async::Ready(()))
+        Ok(Async::NotReady)
     }
 }
```

初回に`poll()`が呼ばれて`Ok(Async::NotReady)`を返すため処理は終了しないが、`poll()`を呼び出す要因を誰も発生させないためそのまま`poll()`も呼ばれず、処理も終了しないままとなる。

```bash
$ cargo run
poll called
```

# UDPパケットを待ち受ける

```rust
use tokio::prelude::*;
use tokio::net::UdpSocket;
use std::net::SocketAddr;

struct EchoServer {
    socket: UdpSocket,
    buf: Vec<u8>,
}

impl EchoServer {
    fn new(socket: UdpSocket) -> EchoServer {
        EchoServer{socket: socket, buf: vec![0; 1024]}
    }
}

impl Future for EchoServer {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        match self.socket.poll_recv_from(&mut self.buf) {
            Ok(Async::Ready((_size, _peer))) => {
                println!("recv");
                Ok(Async::Ready(()))
            },
            Ok(Async::NotReady) => {
                println!("not ready");
                Ok(Async::NotReady)
            },
            Err(_) => Err(()),
        }
    }
}

fn main() {
    let addr = "127.0.0.1:9000".to_string().parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&addr).unwrap();
    println!("Listening on: {}", socket.local_addr().unwrap());

    let server = EchoServer::new(socket);

    tokio::run(server);
}
```

`$ cargo run`して`$ echo hoge | socat - UDP:localhost:9000`してやると以下のように出力される。

```bash
$ cargo run
# [...]
Listening on: 127.0.0.1:9000
not ready
recv
$
```

`UdpSocket::poll_recv_from()`は初回の`poll()`では電文が届いていないため`Async::NotReady`を返す。`socat`実行後再度`poll()`が呼ばれその際には`UdpSocket::poll_recv_from()`は`Async::Ready(...)`を返し`EchoServer::poll()`も`Async::Ready(())`を返すため、処理が終了する。

いちいち`poll()`に対して`match`書くのは面倒なので`try_ready!`マクロを使うとちょっとだけ簡単になる。

```rust
use tokio::prelude::*;
use tokio::net::UdpSocket;
use std::net::SocketAddr;

use futures::try_ready;

struct EchoServer {
    socket: UdpSocket,
    buf: Vec<u8>,
}

impl EchoServer {
    fn new(socket: UdpSocket) -> EchoServer {
        EchoServer{socket: socket, buf: vec![0; 1024]}
    }
}

impl Future for EchoServer {
    type Item = ();
    type Error = std::io::Error;

    fn poll(&mut self) -> Poll<(), std::io::Error> {
        let _ = try_ready!(self.socket.poll_recv_from(&mut self.buf));
        println!("connect");
        Ok(Async::Ready(()))
    }
}

fn main() {
    let addr = "127.0.0.1:9000".to_string().parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&addr).unwrap();
    println!("Listening on: {}", socket.local_addr().unwrap());

    let server = EchoServer::new(socket);

    tokio::run(server.map_err(|_| ()));
}
```

`try_ready!()`が`Err`を直接returnするため`EchoServer::Error`が変更されている点と`tokio::run()`の引数型に合わせるため`map_err()`してる点に注意されたし。

# Echoする

もう一歩進めてEchoしてみる。

```rust
struct EchoServer {
    socket: UdpSocket,
    buf: Vec<u8>,
    to_send: Option<(usize, SocketAddr)>,
}

[...]

impl Future for EchoServer {

    [...]

    fn poll(&mut self) -> Poll<(), std::io::Error> {
        loop {
            if self.to_send == None {
                self.to_send = Some(try_ready!(self.socket.poll_recv_from(&mut self.buf)));
            }

            if let Some((size, peer)) = self.to_send {
                let amt = try_ready!(self.socket.poll_send_to(&self.buf[..size], &peer));
                println!("Echoed {}/{} bytes to {}", amt, size, peer);
                self.to_send = None;
            }
        }
    }
}
```

recvしてsendするだけのはずが、loopだのifだのあってややこしい。少し解説したい。

まず`UdpSocket::poll_recv_from()`の部分は前回説明したとおり、受信していれば`self.to_send`に受信バイト数と送信元情報を格納する。

次に受信した電文を`UdpSocket::poll_send_to()`でechoするわけだが、ここで`peer`が送信可能になっていない可能性が有り、その場合は`UdpSocket::poll_send_to()`で`Async::NotReady`が返ってくる。その際`try_ready!()`で`poll()`を抜けるわけだが、次回`poll()`が呼ばれた際に**1つめのif文が無いとechoしてないのに受信待ちを行ってしまう**ことになってしまうのだ。つまり1つめのif文があることで送信が完了するまで`UdpSpcket::poll_recv_from()`をスキップし続けることが可能となる。

またloopの意味だが、例えば以下のようにloop無しになっていると

```rust
    fn poll(&mut self) -> Poll<(), std::io::Error> {
        if self.to_send == None {
            self.to_send = Some(try_ready!(self.socket.poll_recv_from(&mut self.buf)));
        }

        if let Some((size, peer)) = self.to_send {
            let amt = try_ready!(self.socket.poll_send_to(&self.buf[..size], &peer));
            println!("Echoed {}/{} bytes to {}", amt, size, peer);
            self.to_send = None;
        }

        Ok(Async::NotReady)
    }
```

1回はechoされるが、次回の受信待ちができなくなってしまう。イベントループに対して受信時に`poll()`が呼ばれるようにするためには再度`UdpSocket::poll_recv_from()`を呼び出す必要がある。

# async/awaitを使う

以上で見てもらったように`Future`を自前実装する場合は、今は受信中なのか送信中なのかといった状態遷移を考えなければならい。これがasync/awaitでどう簡単になるか見てみよう。

```toml:Cargo.toml
[...]

[dependencies]
tokio = { git = "https://github.com/tokio-rs/tokio", features=["async-await-preview"] }
tokio-async-await = { git = "https://github.com/tokio-rs/tokio" }
tokio-tcp = { git = "https://github.com/tokio-rs/tokio" }
futures-preview = "0.3.0-alpha.9"
```

```rust
#![feature(pin, await_macro, async_await, futures_api)]
use tokio::await;
use tokio::net::UdpSocket;
use std::net::SocketAddr;

fn main() {
    let addr = "127.0.0.1:9000".to_string().parse::<SocketAddr>().unwrap();
    let mut socket = UdpSocket::bind(&addr).unwrap();
    println!("Listening on: {}", socket.local_addr().unwrap());

    tokio::run_async(async move {
        loop {
            let buf = vec![0; 1024];
            let (sock, buf, size, peer) = await!(socket.recv_dgram(buf)).unwrap();
            let (sock, _) = await!(sock.send_dgram(buf, &peer)).unwrap();
            socket = sock;
            println!("Echoed {} to {}", size, peer);
        }
    });
}
```

`UdpSocket`には`Future`(`RecvDgram`, `SendDgram`)を返す`recv_dgram()`,`send_dgram()`があるのでそれらの`Future`に対して`await!()`を実行している。
`RecvDgram`や`SendDgram`は自作`Future`の`poll()`内でも利用可能である。
上記自作`Future`で`RecvDgram`や`SendDgram`を利用するのは読者への宿題としたい。

# `Future`モナドを利用したコード

大変さを強調するため`Future`を自前で実装するということをやってみたが、実はtokioの本来の使い方的にはasync/awaitを使用せずとももう少し簡単に書くことができる。

```rust
use tokio::prelude::*;
use futures::future::{loop_fn, Loop};
use tokio::net::UdpSocket;
use std::net::SocketAddr;

fn main() {
    let addr = "127.0.0.1:9000".to_string().parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&addr).unwrap();
    println!("Listening on: {}", socket.local_addr().unwrap());

    let server = loop_fn(socket, |socket| {
        let buf = vec![0; 1024];
        socket.recv_dgram(buf)
            .and_then(|(socket, buf, _size, peer)| socket.send_dgram(buf, &peer))
            .and_then(|(socket, _)| Ok(Loop::Continue(socket)))
    });

    tokio::run(server.map_err(|_| ()));
}
```

ただこれも複雑な処理になってくるといわゆる「コールバック地獄」になるわけで、やはり早くasync/await来てほしいなという感じである。