Keyboard Interrupt from Python does not abort Rust function (PyO3)

I have a Python library written in Rust with PyO3, and it involves some expensive calculations (up to 10 minutes for a single function call). How can I abort the execution when calling from Python ?

Ctrl+C seems to only be handled after the end of the execution, so is essentially useless.

Minimal reproducible example:

# Cargo.toml

[package]
name = "wait"
version = "0.0.0"
authors = []
edition = "2018"

[lib]
name = "wait"
crate-type = ["cdylib"]

[dependencies.pyo3]
version = "0.10.1"
features = ["extension-module"]

// src/lib.rs

use pyo3::wrap_pyfunction;

#[pyfunction]
pub fn sleep() {
    std::thread::sleep(std::time::Duration::from_millis(10000));
}

#[pymodule]
fn wait(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(sleep))
}

$ rustup override set nightly
$ cargo build --release
$ cp target/release/libwait.so wait.so
$ python3
>>> import wait
>>> wait.sleep()

Immediately after having entered wait.sleep() I type Ctrl + C, and the characters ^C are printed to the screen, but only 10 seconds later do I finally get

>>> wait.sleep()
^CTraceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyboardInterrupt
>>>

The KeyboardInterrupt was detected, but was left unhandled until the end of the call to the Rust function. Is there a way to bypass that ?

The behavior is the same when the Python code is put in a file and executed from outside the REPL.

Solution

One option would be to spawn a separate process to run the Rust function. In the child process, we can set up a signal handler to exit the process on interrupt. Python will then be able to raise a KeyboardInterrupt exception as desired. Here's an example of how to do it:

// src/lib.rs
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use ctrlc;

#[pyfunction]
pub fn sleep() {
    ctrlc::set_handler(|| std::process::exit(2)).unwrap();
    std::thread::sleep(std::time::Duration::from_millis(10000));
}

#[pymodule]
fn wait(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(sleep))
}

# wait.py
import wait
import multiprocessing as mp

def f():
    wait.sleep()

p = mp.Process(target=f)
p.start()
p.join()
print("Done")

Here's the output I get on my machine after pressing CTRL-C:

$ python3 wait.py
^CTraceback (most recent call last):
  File "wait.py", line 9, in <module>
    p.join()
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/process.py", line 140, in join
    res = self._popen.wait(timeout)
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/popen_fork.py", line 48, in wait
    return self.poll(os.WNOHANG if timeout == 0.0 else 0)
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/popen_fork.py", line 28, in poll
    pid, sts = os.waitpid(self.pid, flag)
KeyboardInterrupt