rust-future-wrap

一个最小的crate，允许您包装一个future来跟踪每次poll和修改结果。

它只包含一个名为 WrapFuture 的单个trait，它为所有实现 std::future::Future 的类型添加了一个 wrap 函数。该 wrap 函数接受一个闭包作为参数，该闭包将在future被poll时被调用。闭包接收包装后的 Pin<Future> 并负责使用接收到的waker上下文调用其上的 poll。闭包可以返回任何类型 T 的 std::task::Poll<T>，因此可以修改包装future的返回类型（类似于 Future.map）。

示例

此示例展示了如何跟踪和限制在poll future上花费的时间。

use std::future::Future;
use std::task::Poll;
use std::time::{Duration, Instant};

use tokio::runtime::Builder;
use tokio::time::sleep;

use future_wrap::WrapFuture;

async fn some_async_fn() {
    sleep(Duration::from_secs(1)).await;
    std::thread::sleep(Duration::from_millis(3));
    sleep(Duration::from_secs(1)).await;
    std::thread::sleep(Duration::from_millis(3));
    sleep(Duration::from_secs(1)).await;
    std::thread::sleep(Duration::from_millis(3));
    sleep(Duration::from_secs(1)).await;
    std::thread::sleep(Duration::from_millis(3));
    sleep(Duration::from_secs(1)).await;
}

fn main() {
    let runtime = Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();

    runtime.block_on(async move {
        let fut = some_async_fn();

        let mut remaining_time = Duration::from_millis(10);
        fut.wrap(|fut, cx| {
            let poll_start = Instant::now();

            println!("Poll start");
            let res = fut.poll(cx);
            println!("Poll end");

            remaining_time = remaining_time.saturating_sub(poll_start.elapsed());
            if remaining_time.is_zero() {
                println!("Too much time spent on polls :(");
                Poll::Ready(None)
            } else {
                res.map(|v| Some(v))
            }
        }).await;
    });
}