Search code examples
rustconcurrencyfuturerust-tokio

Canceling `read_line` with `^C`


I'm writing a shell interpreter as a hobby in order to learn Rust.

Here are a few facts about this code:

  • In order to catch when users enter ^C and cancel their input, I use the ctrlc crate, as it seems popular.
  • To my knowledge, the loop_exec() need to be a future. This is where my shell awaits for user input with io::stdin().read_line(&mut buf). Still to my knowledge, it can't be canceled.
  • I hence want to run the loop_exec() future, kill it when users hit ^C, and make another with loop {}.
  • In order to do that, I need to register the ^C handler to use the JoinHandle<()> of the future (execution).
  • I obviously can't use it at two places (I have to use it to register the handler AND to await it) as per how Rust works (borrow-checker). However, I can't clone execution: JoinHandle<()> either or use as_bytes(). I'm assuming this is in order to protect devs from shooting themselves in the foot, or do things that are never supposed to happen.

Hence, I'm kind of stuck here. Would you have seen that coming and taken a whole different approach, or do I miss just a bit of information ? Heck I am probably missing a whole lot of knowledge in order to arrive at that point.

How would you do things in this situation or more globally ?

#[tokio::main]
async fn main() {
    let mut kill_execution = false;

    ctrlc::set_handler(move || {
        kill_execution = true;
    })
    .expect("should be able to set Ctrl-C handler");

    loop {
        let execution = tokio::spawn(loop_exec());

        let killer = tokio::spawn(async move {
            while !kill_execution {}
            kill_execution = false;
            execution.abort();
        });

        match execution.await {
            Ok(_) => {}
            Err(error) => {
                println!("Error while executing the main thread: {:?}", error);
            }
        }
    }
}

Edit

With the combined answers below, I simplified the code to the following:

#[tokio::main]
async fn main() {
    loop {
        tokio::select! {
            error = tokio::signal::ctrl_c() => {
                println!("Ctrl-C received, exiting: {:?}", error);
                break;
            },
            res = loop_exec() => match res {
                Ok(_) => {
                    println!("Success")
                },
                Err(e) => {
                    println!("Error in loop_exec: {:?}", e);
                    break;
                }
            }
        }
    }
}

async fn loop_exec() -> Result<(), Error> {
    let stdin = io::stdin();
    let mut input = String::new();

    stdin.read_line(&mut input).expect("Failed to read line");

    // do stuff

    Ok(())
}

The problem is still partially to be identified: ^C won't trigger tokio::signal::ctrl_c() once the listener has been initialized.


Solution

  • There are multiple possible solutions to this problem (for example, wrapping the JoinHandle in Mutex<Option>, or using a select to abort it) but you seem to misunderstand JoinHandle::abort(): it will cancel the task at the next .await point. So if your function does a computationally-intensive task (such as an interpreter), it won't do anything.

    Instead, you need to set an AtomicBool, and check it every few operations (for example, at each bytecode operation). You should also use tokio's own signal handling and not the ctrlc crate, so this will look like this:

    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;
    
    #[tokio::main]
    async fn main() {
        let kill_execution = Arc::new(AtomicBool::new(false));
    
        tokio::spawn({
            let kill_execution = Arc::clone(&kill_execution);
            async move {
                loop {
                    tokio::signal::ctrl_c()
                        .await
                        .expect("failed to listen for Ctrl+C");
                    kill_execution.store(true, Ordering::Relaxed);
                }
            }
        });
    
        loop {
            let execution = tokio::spawn(loop_exec(Arc::clone(&kill_execution)));
    
            match execution.await {
                Ok(_) => {}
                Err(error) => {
                    println!("Error while executing the main thread: {:?}", error);
                }
            }
        }
    }
    
    async fn loop_exec(kill_execution: Arc<AtomicBool>) -> Result<(), String> {
        // Take input from stdin.
    
        kill_execution.store(false, Ordering::Relaxed);
        
        // And every some operations:
        if kill_execution.load(Ordering::Relaxed) {
            return Err("Ctrl+C pressed".to_owned())
        }
    }