awestlake87 / pyo3-asyncio

Other
300 stars 45 forks source link

Calling Python from Rust: no way to catch SIGINT signal in asyncio, causes whole program to crash #107

Open whitevegagabriel opened 11 months ago

whitevegagabriel commented 11 months ago

🌍 Environment

💥 Reproducing

Run this code snippet:

use std::time::Duration;
use pyo3::PyResult;

#[pyo3_asyncio::tokio::main]
async fn main() -> PyResult<()> {
    tokio::signal::ctrl_c().await?;
    println!("ctrl-c received");
    // let the signal propagate in python...
    tokio::time::sleep(Duration::from_millis(10)).await;
    println!("doesn't print because asyncio raises KeyboardInterrupt and the program exits");
    Ok(())
}

No special flags, just cargo run.


When I ctrl-c, I expect to be able to catch the signal and close the program gracefully. However, asyncio seems to be doing its own signal handling.

Ideally I'd like the ability disable signal handling in asyncio entirely and do my own graceful shutdown, but at least I'd like the chance to gracefully shut down my Rust code, which I can't do either.

However, the whole program just crashes. In the logs, I see a KeyboardInterrupt being raised by asyncio.

whitevegagabriel commented 11 months ago
Traceback (most recent call last):
  File ".../asyncio/base_events.py", line 640, in run_until_complete
    self.run_forever()
  File ".../asyncio/base_events.py", line 607, in run_forever
    self._run_once()
  File ".../asyncio/base_events.py", line 1884, in _run_once
    event_list = self._selector.select(timeout)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../selectors.py", line 566, in select
    kev_list = self._selector.control(None, max_ev, timeout)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt
thread 'main' panicked at 'called `Result::unwrap()` on an `Err` value: ()', src/main.rs:5:1
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
awestlake87 commented 11 months ago

Python does do it's own signal handling, but I'm not seeing anything in your example that is catching and handling the CTRL-C. When Python sees a CTRL-C it throws a KeyboardInterrupt exception, which looks like what you're getting there.

AFAICT the panic you're getting originates in the macro expansion when it unwraps the error returned by pyo3_asyncio::tokio::run:

        ...
        #(#attrs)*
        #vis fn main() {
            async fn main() #ret {
                #body
            }

            pyo3::prepare_freethreaded_python();

            let mut builder = #builder;
            #builder_init;

            pyo3_asyncio::tokio::init(builder);

            #rt_init

            pyo3::Python::with_gil(|py| {
                pyo3_asyncio::tokio::run(py, main())
                    .map_err(|e| {
                        e.print_and_set_sys_last_vars(py);
                    })
                    .unwrap();
            });
        }
        ...

If you want to catch errors like this and do some cleanup in tokio, the best way to do that is to include this boilerplate in your code instead of using the #[pyo3_asyncio::tokio::main] attribute. Then you can change the unwrap() to a match and run some cleanup code on Err(e).

whitevegagabriel commented 11 months ago

Thank you for the quick response!

While Python does in general do its own signal handling, it is my understanding that pyo3 specifically initializes Python without a signal handler via ffi::Py_InitializeEx(0).

Python docs on Py_InitializeEx:

... This function works like Py_Initialize() if initsigs is 1. If initsigs is 0, it skips initialization registration of signal handlers, which might be useful when Python is embedded. ...

This can be demonstrated by running this snippet:

#[tokio::main]
async fn main() -> PyResult<()> {
    let handle1 = tokio::spawn(async {
        tokio::signal::ctrl_c().await.unwrap();
        println!("received Ctrl+C!");
    });

    let handle2 = tokio::spawn(async {
        Python::with_gil(|py| {
            let i = 0;
            py_run!(py, i, r#"
    while True:
        i += 1
        if i % 30000000 == 0:
            print(i)
    "#);
        });
    });

    handle1.await.unwrap();
    println!("ctrl+c signal handler returned");
    handle2.await.unwrap();

    println!("never reached, and the python loop continues forever");

    Ok(())
}

Starting the program and triggering ctrl+c gives me these logs:

30000000
60000000
^Creceived Ctrl+C!
ctrl+c signal handler returned
90000000
120000000

I take your point that tokio::signal::ctrl_c() may not actually be setting a signal handler.

And to your point, when I set a signal handler via

ctrlc::set_handler(move || {
        println!("received Ctrl+C!");
    })
    .expect("Error setting Ctrl-C handler");

then it never reaches asyncio, as I would expect.

That being said, since pyo3 initializes Python with signal handling disabled, do you think it would be reasonable to initialize asyncio similarly, with signal handling disabled?

whitevegagabriel commented 10 months ago

After some more investigation, it seems that asyncio sets up its signal handlers during import, regardless of whether it was initialized in a Python env that doesn't register signal handlers (Py_InitializeEx(0)).

Could this be an issue that needs raised directly against asyncio on Github?


The following program terminates with a raised KeyboardInterrupt main.c

#include <Python.h>

int main(int argc, char** argv) {
  Py_Initialize();

  char* exp = "input('press ctrl-c')";
  PyRun_SimpleString(exp);

  Py_Finalize();
}

press ctrl-c^CTraceback (most recent call last): File "", line 1, in KeyboardInterrupt

The following program skips registering Python signal handlers and terminates without a raised KeyboardInterrupt main.c


#include <Python.h>

int main(int argc, char** argv) { Py_InitializeEx(0); // skips registering Python signal handler

char* exp = "input('press ctrl-c')"; PyRun_SimpleString(exp);

Py_FinalizeEx(); }

> press ctrl-c^C
---
The following program skips registering Python signal handlers and terminates _**with**_ a raised KeyboardInterrupt
`main.c`

include

int main(int argc, char** argv) { Py_InitializeEx(0); // skips registering Python signal handler

char* exp = "import asyncio; input('press ctrl-c')"; PyRun_SimpleString(exp);

Py_FinalizeEx(); }


> press ctrl-c^CTraceback (most recent call last):
>   File "<string>", line 1, in <module>
KeyboardInterrupt

---

FYI, I was able to compile the above file with
`gcc main.c $(python3-config --include) $(python3-config --ldflags) -lpython<major.minor> -o main`
whitevegagabriel commented 10 months ago

I have kicked off a conversation in the Python forum as a prelude to officially filing a feature request in cpython: https://discuss.python.org/t/asyncio-skipping-signal-handling-setup-during-import-for-python-embedded-context/37054

whitevegagabriel commented 10 months ago

I have kicked off a Feature Request into CPython: https://github.com/python/cpython/issues/111400

Until then, looks like there are a couple decent options:

  1. (per your reccomendation) Every program using pyo3_asyncio that doesn't want rude KeyboardInterrupts needs to set up its own signal handler
  2. During asyncio_asyncio initialization, reset the signal handler right after first import of asyncio

I think (2) could be tricky, but worth the effort. If I find some time in the next weeks, I'll design and propose a solution.

extraymond commented 8 months ago

After some more investigation, it seems that asyncio sets up its signal handlers during import, regardless of whether it was initialized in a Python env that doesn't register signal handlers (Py_InitializeEx(0)).

Could this be an issue that needs raised directly against asyncio on Github?

The following program terminates with a raised KeyboardInterrupt main.c

#include <Python.h>

int main(int argc, char** argv) {
  Py_Initialize();

  char* exp = "input('press ctrl-c')";
  PyRun_SimpleString(exp);

  Py_Finalize();
}

press ctrl-c^CTraceback (most recent call last): File "", line 1, in KeyboardInterrupt

The following program skips registering Python signal handlers and terminates without a raised KeyboardInterrupt main.c

#include <Python.h>

int main(int argc, char** argv) {
  Py_InitializeEx(0); // skips registering Python signal handler

  char* exp = "input('press ctrl-c')";
  PyRun_SimpleString(exp);

  Py_FinalizeEx();
}

press ctrl-c^C

The following program skips registering Python signal handlers and terminates with a raised KeyboardInterrupt main.c

#include <Python.h>

int main(int argc, char** argv) {
  Py_InitializeEx(0); // skips registering Python signal handler

  char* exp = "import asyncio; input('press ctrl-c')";
  PyRun_SimpleString(exp);

  Py_FinalizeEx();
}

press ctrl-c^CTraceback (most recent call last): File "", line 1, in KeyboardInterrupt

FYI, I was able to compile the above file with gcc main.c $(python3-config --include) $(python3-config --ldflags) -lpython<major.minor> -o main

Shout out to the great investigation. This explains a lot for my rust python interop program which I was never able to figure out who hijacked my signal handling workflow.

BTOdell commented 6 months ago

This is great work!

I assume that if this is fixed in upstream CPython, then pyo3-asyncio wouldn't have to run on the main thread? The docs suggest that the reason it currently has to is due to signal handling. I'm currently working on creating a thin Python plugin SDK for my Rust app and it's quite a headache to have to give up control of the main thread.

awestlake87 commented 6 months ago

@BTOdell to be honest, it's been a long time since I wrote those docs, so I'm not 100% certain that's the only reason Python wants the main thread. If Python doesn't need the main thread, I imagine we can greatly simplify the test runner in pyo3-asyncio as well

whitevegagabriel commented 5 months ago

CPython has unforutnately yet to provide feedback on my proposition. That said, even if a fix is introduced into future versions of python, it would be interesting to think about what we can do for older versions of python.

I know this is a pretty dumb fix, and there's probably a better way, but even something simple like this was able to fix the problem for me.

fn asyncio(py: Python) -> PyResult<&PyAny> {
    ASYNCIO
        .get_or_try_init(|| {
            #[cfg(feature = "reset-signal-handler")]
            {
                let signal = py.import("signal")?;
                let sigint = signal.getattr("SIGINT")?;
                let sig_dfl = signal.getattr("SIG_DFL")?;
                signal.call_method1("signal", (sigint, sig_dfl))?;  
            }
            Ok(py.import("asyncio")?.into())
        })
        .map(|asyncio| asyncio.as_ref(py))
}

I know I can fix this by setting up signal handlers in my rust code, but it can be bothersome and repetitive.

Do you feel like we could do somthing similar to this? The idea is to import signal and undo its signal handler shenanigans before any client code runs. I just happened to do it when importing asyncio. Likely behind a feature flag or similar, to make it opt-in.