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

Question

我有一个用 PyO3 用 Rust 编写的 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"]

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

已进入后立即wait.sleep()I型Ctrl + C，字符^C打印到屏幕上，但只有10秒钟后我终于得到

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

该KeyboardInterrupt被检测到，但留下未处理，直到调用防锈功能的结束。有没有办法绕过它？

当 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 的默认效果是引发 KeyboardInterrupt 异常。如果引发异常，则设置错误指示符并且函数返回 -1；否则函数返回 0

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

如果您的 Rust 代码是多线程的，情况会更加复杂。您只能PyErr_CheckSignals从与 python 解释器调用您相同的线程进行调用；如果它返回 -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