Python 的键盘中断不会中止 Rust 功能（PyO3）

Question

我有一个 Python 库，用 Rust 和 PyO3 编写，它涉及一些昂贵的计算（单个函数调用最多 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

您的问题与this one 非常相似，只是您的代码是用 Rust 而不是 C++ 编写的。

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

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

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

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

这在执行 python 代码时可以正常工作 - 它会在每个代码语句中至少检查一次标志 - 但在执行用 Rust（或任何其他外语）编写的长时间运行的函数时，情况就不同了。在您的 rust 函数返回之前，不会检查该标志。

您可以通过检查 rust 函数中的标志来改善问题。 PyO3 exposes PyErr_CheckSignals 函数正是这样做的。这个函数：

检查信号是否已经发送到进程，如果是， 调用相应的信号处理程序。如果信号模块是 支持，这可以调用用 Python 编写的信号处理程序。在所有 在这种情况下，SIGINT 的默认效果是引发 KeyboardInterrupt 例外。如果引发异常，则设置错误指示器并 函数返回-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