来自 Python 的键盘中断不会中止 Rust 函数 (PyO3)

Question

我有一个用 Rust 和 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 而不是C++编写的。

你没有说你正在使用哪个平台 - 我将假设它是类似 unix 的。此答案的某些方面对于Windows可能不正确。

在类似 unix 的系统中，Ctrl+C 会导致将SIGINT信号发送到您的进程。在 C 库的非常低级别，应用程序可以注册在接收到这些信号时将调用的函数。有关信号的更详细描述，请参见 man signal(7(。

由于信号处理程序可以随时调用(即使是您通常认为是原子操作的中途(，因此信号处理程序的实际功能有很大的限制。这与编程语言或环境无关。大多数程序只是在收到信号时设置一个标志，然后返回，然后检查该标志并对其采取行动。

Python 也不例外 - 它为SIGINT信号设置一个信号处理程序，该信号设置一些标志，它会检查(在安全的情况下(并对其采取行动。

这在执行 python 代码时工作正常 - 它将在每个代码语句中至少检查一次标志 - 但是在执行用 Rust(或任何其他外语(编写的长时间运行的函数时，这是另一回事。在 rust 函数返回之前，不会检查该标志。

您可以通过检查 rust 函数中的标志来改进问题。PyO3 公开了正是这样做的 PyErr_CheckSignals 函数。此函数：

检查信号是否已发送到进程，如果是，则 调用相应的信号处理程序。如果信号模块是 支持，这可以调用用 Python 编写的信号处理程序。一起 在这种情况下，SIGINT的默认效果是引发键盘中断 例外。如果引发异常，则设置错误指示器并 该函数返回 -1;否则函数返回 0

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

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

Answer 2

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

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