Python的键盘中断不会中止Rust函数（PyO3）

Question

我有一个用PyO3编写的Python库，它涉及一些昂贵的计算（单个函数调用最多10分钟）。从Python调用时如何中止执行

Ctrl+C似乎只在执行结束后处理，因此基本上是无用的

最小可复制示例：

# 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()

输入wait.sleep()后，我立即键入Ctrl + C，字符^C被打印到屏幕上，但仅10秒后，我最终获得

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

已检测到KeyboardInterrupt，但在调用Rust函数结束之前未进行处理。有没有办法绕过这个问题

当Python代码放在文件中并从REPL外部执行时，行为是相同的

Answer 1

一种选择是生成一个单独的进程来运行Rust函数。在子进程中，我们可以设置一个信号处理程序，在中断时退出进程。Python将能够根据需要引发键盘中断异常。下面是一个如何执行此操作的示例：

// 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")

以下是我在按下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

Answer 2

<>你的问题与this one非常相似，除了你的代码是用锈写的，而不是用C++写的。p>

您没有说您使用的是哪一个平台，我将假定它与unix类似。对于Windows，此答案的某些方面可能不正确

在类unix系统中，Ctrl+C将导致向进程发送SIGINT信号。在C库的极低级别上，应用程序可以注册函数，当接收到这些信号时将调用这些函数。有关信号的更详细说明，请参见man signal(7)

因为在任何时候都可以调用信号处理程序（甚至通过一些通常认为是原子的操作），所以信号处理程序实际上可以做很大的限制。这与编程语言或环境无关。大多数程序只是在接收到信号并返回时设置一个标志，然后检查该标志并对其采取行动

Python也不例外——它为SIGINT信号设置了一个信号处理程序，该信号处理程序设置了一些标志，用于检查（在安全的情况下）并对其执行操作

这在执行python代码时可以正常工作——每个代码语句至少检查一次标志——但在执行用Rust（或任何其他外语）编写的长时间运行的函数时则是另一回事。直到rust函数返回，才会检查该标志

您可以通过检查rust函数中的标志来改善问题。PyO3exposes是PyErr_CheckSignals函数，它正是这样做的。此功能：

checks whether a signal has been sent to the processes and if so, invokes the corresponding signal handler. If the signal module is supported, this can invoke a signal handler written in Python. In all cases, the default effect for SIGINT is to raise the KeyboardInterrupt exception. If an exception is raised the error indicator is set and the function returns -1; otherwise the function returns 0

因此，您可以在Rust函数内的适当时间间隔调用此函数，并检查返回值。如果是-1，您应该立即从Rust函数返回；否则继续

如果您的锈迹代码是多线程的，则情况会更复杂。您只能从python解释器调用您的同一线程调用PyErr_CheckSignals；如果返回-1，则必须清除返回之前启动的所有其他线程。具体如何做到这一点超出了本答案的范围