Stop yielding in nested select!

[Banyc]

I want to encapsulate the inner select! as a single entry point so that the outer select! don't need to explicitly handle every arm in the inner select!. However, if the body of an arm in the inner select! happens to await, the await will yield and let other select! arms interrupt the body. Worse still, the lines after the await won't be called at all.

use std::time::Duration;

use tokio::{select, time};

#[tokio::main]
async fn main() {
    println!("Hello, world!");

    run().await;
}

async fn run() {
    let mut interval_3s = time::interval(Duration::from_secs(3));
    let mut interval_sleep = time::interval(Duration::from_secs(7));
    let mut s = S::new();

    loop {
        select! {
            _ = s.another_select() => {
                println!("run: another_select");  // never called
            }
            _ = interval_3s.tick() => {
                println!("run: interval_3s");
            }
            _ = interval_sleep.tick() => {
                println!("run: interval_sleep");
                // during this time, no other arms interrupt which is ideal
                time::sleep(Duration::from_secs(5)).await;
                println!("run: interval_sleep + 5s");
            }
        }
    }
}

struct S {
    interval_5s: time::Interval,
}

impl S {
    fn new() -> Self {
        S {
            interval_5s: time::interval(Duration::from_secs(5)),
        }
    }

    async fn another_select(&mut self) {
        select! {
            _ = self.interval_5s.tick() => {
                println!("another_select: tick");
                // during this time, it is expected not being interrupted by other arms of in the outer `select!`
                time::sleep(Duration::from_secs(5)).await;
                println!("another_select: tick + 5s");  // never called
            }
        }
    }
}

I don't want any code (interval_3s.tick() in this case) in the same tokio::spawn/task to preempt my encapsulated code (s in this case).

A raw workaround is to break the encapsulation and merge the another_select to the select! under run. But what if I have private members being used in another_select? Should they also be exposed and made public?

[Darksonn]

Perhaps you are looking for something like this?

tokio::join!(
    s.another_select(),
    async {
        loop {
            select! {
                _ = interval_3s.tick() => {
                    println!("run: interval_3s");
                }
                _ = interval_sleep.tick() => {
                    println!("run: interval_sleep");
                    // during this time, no other arms interrupt which is ideal
                    time::sleep(Duration::from_secs(5)).await;
                    println!("run: interval_sleep + 5s");
                }
            }
        }
    }
)

[Banyc]

Thanks for the speedy reply! I really appreciate this!

It doesn't work as expected. I expect no interruptions between another_select: tick and another_select: tick + 5s, but I tested running the solution and found output:

another_select: tick
run: interval_sleep
another_select: tick + 5s
run: interval_sleep + 5s
...

It seems like another_select: tick and run: interval_sleep are interleaved.

After that, another_select never runs again, which should also be in a loop.

Some more output:

Hello, world!
another_select: tick
run: interval_sleep
another_select: tick + 5s
run: interval_sleep + 5s
run: interval_3s
run: interval_3s
run: interval_3s
run: interval_sleep
run: interval_sleep + 5s
run: interval_3s
run: interval_3s
run: interval_sleep
run: interval_sleep + 5s
run: interval_3s
run: interval_3s
run: interval_3s
run: interval_sleep
run: interval_sleep + 5s
run: interval_3s
...

[Darksonn]

Your S struct will need to expose one function for the branch condition, and another function for the body. It is not possible to extend the select! macro itself in a way to support something like this.

[Banyc]

If the inner select! has more than one arm, then the outer select! has to handle those arms, exposing the return items from the branch conditions to the outer select!. If the current tools available do not solve the problem, I want the macro select! adding an special branch type that support what I depicted.

[Darksonn]

You can do this:

struct S {
    interval_5s: time::Interval,
    interval_7s: time::Interval,
}

enum WhichBranch {
    Interval5s,
    Interval7s,
}

impl S {
    async fn condition(&mut self) -> WhichBranch {
        tokio::select! {
            _ = self.interval_5s.tick() => WhichBranch::Interval5s,
            _ = self.interval_7s.tick() => WhichBranch::Interval7s,
        }
    }
    
    async fn body(&mut self, branch: WhichBranch) {
        match branch {
            WhichBranch::Interval5s => println!("Got 5s tick"),
            WhichBranch::Interval7s => println!("Got 7s tick"),
        }
    }
}

loop {
    select! {
        branch = s.condition() => s.body(branch),
        _ = interval_3s.tick() => {
            println!("run: interval_3s");
        }
        _ = interval_sleep.tick() => {
            println!("run: interval_sleep");
            // during this time, no other arms interrupt which is ideal
            time::sleep(Duration::from_secs(5)).await;
            println!("run: interval_sleep + 5s");
        }
    }
}

[Banyc]

This solution is neat!