The correct pattern for a trait that starts and stops a future

Hello all

I'd like to define a trait for a process which starts (and continues to run asynchronously), and can later be stopped - an example below:

use std::future::Future;
use std::pin::Pin;
use tokio_util::sync::CancellationToken;
use async_trait::async_trait;

#[tokio::main]
pub async fn main() {

    #[async_trait]
    trait Process {
        async fn start(&mut self);
        fn log(&self, msg: &str);
        async fn stop(&mut self);
    }

    struct Processor {
        running_future: Option<(CancellationToken, Pin<Box<dyn Future<Output = ()> + Send>>)>,
    }

    #[async_trait]
    impl Process for Processor {
        async fn start(&mut self) {
            let cancellation_token = CancellationToken::new();
            let cancellation_token_ = cancellation_token.clone();
            let future = async move {
                for x in 0..20 {
                    if cancellation_token_.is_cancelled() {
                        println!("Process cancelled!");
                        break;
                    };
                    println!("x: {}", x);
                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
                }
            };
            let future = Box::pin(future);
            self.running_future = Some((cancellation_token, future));
        }

        fn log(&self, msg: &str) {
            println!("logging: {}", msg);
        }

        async fn stop(&mut self) {
            if let Some((cancellation_token, future)) = Option::take(&mut self.running_future) {
                cancellation_token.cancel();
                future.await;
                self.running_future = None;
            }
        }
    }

    let mut processor = Processor { running_future: None };
    println!("Starting processor");
    processor.start().await;
    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
    processor.log("logging!");
    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
    processor.stop().await;
}

The issue with this is that the process (within the future variable) never does anything, because it is not .await-ed until stop() is called, at which point the cancellation token has already been cancelled. If I await the future within start(), then obviously it never gets to the stop() call at all.

Is the trait itself a bad design pattern? Or is there some idiomatic way to make the future start during start(), without having to await it?

Thanks for any help!

In order for a future to make progress, something must poll it. await is one way to do that — incorporating that future into a bigger future. But for your job what you want to do is spawn() it — turn it into a task. Your Processor will no longer own the future at all; it might own the JoinHandle for the task, if it needs to, or not if it doesn't. (JoinHandles also allow cancelling a task, though without any cleanup other than drops.)

Is the trait itself a bad design pattern?

The set of operations doesn't seem unreasonable, but:

  • It doesn't obviously to need to be a trait at all. Consider removing the trait Process and change impl Process for Processor into an impl Processor. If you want to have many such stoppable processes, then make Processor take a function for constructing the future, rather than having only one built-in future.

    (But the trait might make sense if you have multiple start-stops with different means of stopping.)

  • As a matter of API design: Right now, if start() is called twice, then the second time it will just drop the old future and cancellation token. You should make it do something sensible, such as one of:

    • cancel the old task before dropping its token,
    • do nothing if the task is already started, or
    • panic!("start() called when task is alredy started"),

    depending on what suits your application's needs, but not simply misbehave when misused. This sort of check helps create applications that are either robust against slight logic errors, or report them promptly so they can be fixed.

2 Likes

Thank you, that's very helpful. Yes, it was spawn that I needed. And I agree with your other comments, I just wanted to keep the minimal example fairly minimal.

1 Like

Which example?

My code example at the top.