Tokio select macro conditions and arm evaluation

Question

I have a code that needs to be read from several TcpStream. One of these TcpStream will always be there, but another may or may not.

I wrote the code before and (naively) used conditions on tokio select macro. Unfortunately, I quickly learned that the macro will first evaluate the arm (to get the future) and only after that will check a condition and based on this condition will either start the future or skip it.

As a result, as I run this code I will get: 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
        }
    }
    
}

The question which I have is how to handle this?

I stumbled on following ideas and I am not happy about any of these

Add if statement. If stream2.is_some() then have select with two arms, otherwise select with one mark. This works but results in quite ugly code duplication
I looked into FutureUnordered direction. However, taking into account that the above code will be executed in the loop, I bailed out from this approach, because it didn't like that I try to borrow things as mutable in the loop.
I was thinking to have a "fake" TcpStream, just to remove Option<> part. However, it also looks ugly.

I am looking for some way that will allow me to keep this select! over Option

Looks like this behavior changed. An old version of the docs: _"Additionally, each branch may include an optional `if` precondition. This precondition is evaluated __before__ the ``. If the precondition returns `false`, the branch is entirely disabled."_ But the current version: _"Additionally, each branch may include an optional `if` precondition. If the precondition returns `false`, then the branch is disabled. The provided `` is still evaluated but the resulting future is never polled."_ — cdhowie, May 16 '22 at 22:52
This is a little weird to me, and I'm curious why they made this change. The old version's behavior makes a lot more sense to me. — cdhowie, May 16 '22 at 22:53
@cdhowie I have never used the old version. However, I agree, that old behavior made way more sense. — Victor Ronin, May 17 '22 at 01:37

Netwave · Accepted Answer · 2022-05-16T19:30:49.043

You can wrap the handling of the optiona stream into another future, maybe returning a custom error. Then pattern match on successful reads:

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

If you do not like the custom error, you can rewrap everything in an Option<Result> and pattern match everything:

...
    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

Thank you @Netwave. This is exactly what I was looking for. I went for a walk and had a similar idea and it was amazing to come back and see you have a complete solution. — Victor Ronin, May 16 '22 at 21:39

Tokio select macro conditions and arm evaluation

1 Answers1