rust - 什么是解决 future 需要多长时间的干净方法？

Question

我正在与 Tokio 一起做一些 UDP 的事情。

我想记录我的 UDP 探测 future 花费的时间来解决。我想出了以下函数，time_future()，来包装 future 并给我结果和持续时间。这个函数看起来很幼稚，我认为 Rust 有能力更清晰地表达这个概念。

我的工作代码(Playground):

extern crate futures; // 0.1.25
extern crate tokio; // 0.1.11

use std::time::{Duration, Instant};

use futures::future::{lazy, ok};
use futures::Future;
use tokio::runtime::current_thread::Runtime;
use tokio::timer::Delay;

struct TimedFutureResult<T, E> {
    elapsed: Duration,
    result: Result<T, E>,
}

impl<T, E> TimedFutureResult<T, E> {
    pub fn elapsed_ms(&self) -> i64 {
        return (self.elapsed.as_secs() * 1000 + (self.elapsed.subsec_nanos() / 1000000) as u64)
            as i64;
    }
}

fn time_future<F: Future>(f: F) -> impl Future<Item = TimedFutureResult<F::Item, F::Error>> {
    lazy(|| {
        let start = Instant::now();

        f.then(move |result| {
            ok::<TimedFutureResult<F::Item, F::Error>, ()>(TimedFutureResult {
                elapsed: start.elapsed(),
                result: result,
            })
        })
    })
}

fn main() {
    let when = Instant::now() + Duration::from_millis(100);

    let f = time_future(Delay::new(when)).then(|r| match r {
        Ok(r) => {
            println!("resolved in {}ms", r.elapsed_ms());
            r.result
        }
        _ => unreachable!(),
    });

    let mut runtime = Runtime::new().unwrap();
    runtime.block_on(f).unwrap();
}

我怎样才能改进它并使其更加地道？我能否以某种方式让界面像 inspect() 或 then() 一样工作？

Delay::new(when)
  .timed(|res, elapsed| println!("{}ms!", elapsed))
  .and_then(...);

我尝试创建一个 Timed trait 并为 Future 实现它，但我对自己的做法一点信心都没有。这些类型真的让我大吃一惊。

我至少是在咆哮正确的树吗？

Answer 1

Shepmaster 的回答很棒。但是，他们使用的 futures 版本已经过时并且与 stdlib futures 不兼容。这是我使用 stdlib future 的重写。

use pin_project::pin_project;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use tokio::time; // 0.2

/// A wrapper around a Future which adds timing data.
#[pin_project]
struct Timed<Fut, F>
where
    Fut: Future,
    F: FnMut(&Fut::Output, Duration),
{
    #[pin]
    inner: Fut,
    f: F,
    start: Option<Instant>,
}

impl<Fut, F> Future for Timed<Fut, F>
where
    Fut: Future,
    F: FnMut(&Fut::Output, Duration),
{
    type Output = Fut::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = self.project();
        let start = this.start.get_or_insert_with(Instant::now);

        match this.inner.poll(cx) {
            // If the inner future is still pending, this wrapper is still pending.
            Poll::Pending => Poll::Pending,

            // If the inner future is done, measure the elapsed time and finish this wrapper future.
            Poll::Ready(v) => {
                let elapsed = start.elapsed();
                (this.f)(&v, elapsed);

                Poll::Ready(v)
            }
        }
    }
}

trait TimedExt: Sized + Future {
    fn timed<F>(self, f: F) -> Timed<Self, F>
    where
        F: FnMut(&Self::Output, Duration),
    {
        Timed {
            inner: self,
            f,
            start: None,
        }
    }
} 

// All futures can use the `.timed` method defined above
impl<F: Future> TimedExt for F {}

#[tokio::main]
async fn main() {
    let when = Instant::now() + Duration::from_millis(100);

    let fut = time::delay_until(when.into())
        .timed(|res, elapsed| println!("{:?} elapsed, got {:?}", elapsed, res));
    fut.await;
}