Tokio select 宏观条件和手臂评估

Question

我有一段代码需要从多个TcpStream中读取。其中一个 TcpStream 将永远存在，但另一个可能存在也可能不存在。

我之前写过代码并（天真地）在 tokio select 宏上使用了条件。不幸的是，我很快了解到宏将首先评估手臂（以获得未来），然后才会检查条件，并根据该条件启动未来或跳过它。

因此，当我运行这段代码时，我将得到： thread 'main' panicked at 'called `Option::unwrap()` on a `None` value', src/main.rs:16:35

use tokio::io::AsyncReadExt;
use tokio::net::TcpStream;

#[tokio::main]
async fn main() {
    let mut stream1 : TcpStream = TcpStream::connect("info.cern.ch:80").await.unwrap();
    let mut stream2 : Option<TcpStream> = None;
    let mut buf = vec![0_u8; 1024];
    let mut buf2 = vec![0_u8; 1024];
    
    tokio::select! {
        result = stream1.read(&mut buf) => {
            // do something here
        }
        
        result = stream2.as_mut().unwrap().read(&mut buf2), if stream2.is_some() => {
            // do somethihng here
        }
    }
    
}

我的问题是如何处理？

我偶然发现了以下想法，但我对其中任何一个都不满意

添加 if 语句。如果 stream2.is_some() 则 select 有两个手臂，否则 select 有一个标记。这可行，但会导致非常难看的代码重复
我查看了 FutureUnordered 的方向。但是，考虑到上面的代码将在循环中执行，我放弃了这种方法，因为它不喜欢我尝试在循环中借用可变的东西。
我想有一个“假的”TcpStream，只是为了删除 Option<> 部分。然而，它也很难看。

我正在寻找可以让我保留这个 select 的方法！在选项

Answer 1

您可以将 optiona 流的处理包装到另一个未来，也许会返回一个自定义错误。然后在成功读取时进行模式匹配：

use tokio::io::AsyncReadExt;
use tokio::net::TcpStream;
use std::io::{Error, ErrorKind};

#[tokio::main]
async fn main() {
    let mut stream1 : TcpStream = TcpStream::connect("info.cern.ch:80").await.unwrap();
    let mut stream2 : Option<TcpStream> = None;
    let mut buf = vec![0_u8; 1024];
    let mut buf2 = vec![0_u8; 1024];
    
    let read2 = async move {
        match stream2 {
            Some(mut stream) => stream.read(&mut buf2).await,
            _ => Err(Error::new(ErrorKind::Other, "No available stream")),
        }
    };
    
    tokio::select! {
        result = stream1.read(&mut buf) => {
            // do something here
        }
        
        Ok(result) = read2 => {
            // do somethihng here
        }
    }
    
}

Playground

如果您不喜欢自定义错误，您可以在 Option<Result> 中重新包装所有内容并匹配所有内容：

...
    let read2 = async move {
        match stream2 {
            Some(mut stream) => Some(stream.read(&mut buf2).await),
            None => None,
        }
    };
    
    tokio::select! {
        result = stream1.read(&mut buf) => {
            // do something here
        }
        
        Some(result) = read2 => {
            // do somethihng here
        }
    }
...

Playground

Tokio select 宏观条件和手臂评估

Tokio select macro conditions and arm evaluation

rust

rust-tokio