使用PyO3时无需全局修改Rust代码处理SIGINT信号的方案

你当前用std::process::exit(2)直接终止进程的方式，会跳过Python的资源清理流程（比如Arrow S3的Finalize调用），同时导致终端冻结——因为进程没有正常释放终端资源。要实现Rust响应SIGINT后让Python抛出常规KeyboardInterrupt的目标，可按以下方案调整：

核心思路 #

1. 用全局原子标志跟踪中断状态，避免直接终止进程；
2. 通过PyO3在信号处理中触发Python的KeyboardInterrupt异常；
3. 在耗时Rust操作中定期检查中断标志，主动返回错误让Python处理异常。

具体实现 #

1. Rust端代码调整 #

首先定义全局原子变量跟踪中断状态，然后修改信号处理逻辑，同时在耗时任务中加入中断检查：

use pyo3::prelude::*;
use pyo3::types::PyKeyboardInterrupt;
use std::sync::atomic::{AtomicBool, Ordering};

// 全局原子标志，标记是否收到中断
static INTERRUPTED: AtomicBool = AtomicBool::new(false);

#[pyfunction]
fn setup_keyboard_interrupt() -> PyResult<()> {
    ctrlc::set_handler(move || {
        // 设置中断标志
        INTERRUPTED.store(true, Ordering::SeqCst);
        
        // 获取Python GIL，抛出KeyboardInterrupt异常
        Python::with_gil(|py| {
            PyErr::new::<PyKeyboardInterrupt, _>(()).restore(py);
        });
    })
    .expect("Failed to set Ctrl-C handler");
    Ok(())
}

// 封装中断检查逻辑，减少重复代码
fn check_interrupted() -> PyResult<()> {
    if INTERRUPTED.load(Ordering::SeqCst) {
        // 返回KeyboardInterrupt对应的PyErr
        Err(PyErr::new::<PyKeyboardInterrupt, _>(()))
    } else {
        Ok(())
    }
}

#[pyfunction]
fn long_running_task() -> PyResult<()> {
    // 模拟耗时操作，这里以循环为例
    for i in 0..10000 {
        // 定期检查中断（频率可根据任务耗时调整）
        check_interrupted()?;
        
        // 你的实际耗时Rust逻辑
        std::thread::sleep(std::time::Duration::from_millis(10));
    }
    Ok(())
}

#[pymodule]
fn your_package(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(setup_keyboard_interrupt, m)?)?;
    m.add_function(wrap_pyfunction!(long_running_task, m)?)?;
    Ok(())
}

2. Python端配合 #

在__init__.py中调用中断设置函数即可，无需额外信号处理：

from .your_package import setup_keyboard_interrupt, long_running_task

# 初始化中断处理
setup_keyboard_interrupt()

关键细节说明 #

• 避免占用GIL过长时间：如果你的Rust任务是计算密集型或阻塞IO，建议用py.allow_threads()释放GIL，让Python的信号处理线程正常运行：

  #[pyfunction]
  fn long_running_task(py: Python) -> PyResult<()> {
      py.allow_threads(|| {
          for i in 0..10000 {
              check_interrupted()?;
              std::thread::sleep(std::time::Duration::from_millis(10));
          }
          Ok(())
      })
  }
  

• 中断检查频率：无需在每一步都检查，可根据任务耗时调整（比如每处理1000条数据检查一次），平衡性能和响应速度。
• 资源清理：通过返回PyErr让Python正常处理异常，会触发所有注册的清理逻辑（包括Arrow S3的Finalize调用），解决你遇到的S3资源未清理警告。

效果验证 #

按Ctrl+C后：

1. Rust的信号处理函数设置中断标志并触发Python异常；
2. 耗时任务检查到标志后立即返回错误；
3. Python捕获到常规KeyboardInterrupt异常，正常退出并清理资源，终端不会冻结。

内容的提问来源于stack exchange，提问作者AlbertGarde