Handle Ctrl+C properly (#2)

Cargo.toml

@@ -12,13 +12,12 @@ chainql-core = { git = "https://github.com/UniqueNetwork/chainql.git", branch =
 either = "1.13.0"
 gag = "1.0.0"
 jrsonnet-evaluator = { version = "0.5.0-pre95" }
+nix = { version = "0.29.0", features = ["signal"] }
 num-bigint = "0.4.6"
+pyo3 = { version = "0.23.3", features = ["abi3-py310", "extension-module", "num-bigint", "either"] }
 ss58-registry = "1.34.0"
 tokio = { version = "1.41.1", default-features = false, features = ["rt", "rt-multi-thread"] }
 
-[dependencies.pyo3]
-version = "0.23.1"
-features = ["abi3-py310", "extension-module", "num-bigint", "either"]
 
 [profile.release]
 opt-level = 3

pyproject.toml

@@ -7,7 +7,7 @@ name = "chainql"
 requires-python = ">=3.10"
 classifiers = []
 dynamic = ["version"]
-dependencies = ["patchelf"]
+dependencies = ["patchelf; platform_system == 'Linux'"]
 
 [tool.maturin]
 python-source = "stubs"

src/chain.rs

@@ -80,15 +80,15 @@ impl Chain {
     #[new]
     #[pyo3(signature = (url, opts=None))]
     pub fn new(url: String, opts: Option<ChainOpts>) -> PyResult<Self> {
-        execute_jsonnet(|| {
-            chainql_core::builtin_chain(url, opts.map(Into::into))
+        execute_jsonnet(|cancel| {
+            chainql_core::chain(url, opts.map(Into::into), cancel)
                 .map(|chain| Self(JsonnetObject(chain)))
                 .map_err(|err| PyBaseException::new_err(err.to_string()))
         })
     }
 
     pub fn latest(&self) -> PyResult<JsonnetObject> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let chain = &self.0 .0;
 
             let latest = chain
@@ -103,7 +103,7 @@ impl Chain {
     }
 
     pub fn block(&self, block: u32) -> PyResult<JsonnetObject> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let chain = &self.0 .0;
 
             let block_func = chain
@@ -131,7 +131,7 @@ pub fn dump(
     data: BTreeMap<Vec<u8>, Vec<u8>>,
     opts: Option<ChainOpts>,
 ) -> PyResult<JsonnetObject> {
-    execute_jsonnet(|| {
+    execute_jsonnet(|_| {
         let meta = match meta {
             Either::Left(l) => jrsonnet_evaluator::typed::Either2::A(l.0),
             Either::Right(r) => jrsonnet_evaluator::typed::Either2::B(Hex(r)),

src/jsonnet.rs

@@ -46,7 +46,7 @@ impl JsonnetObject {
         py: Python<'py>,
         name: &Bound<'_, PyAny>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let key = name
                 .extract::<&str>()
                 .map_err(|err| type_error(py, "key should be a string", err))?;
@@ -94,7 +94,7 @@ impl JsonnetObject {
 
     #[pyo3(signature = (minified=true))]
     fn manifest_json(&self, minified: bool) -> PyResult<String> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let preserve_order = true;
             let fmt = if minified {
                 jrsonnet_evaluator::manifest::JsonFormat::minify(preserve_order)
@@ -144,7 +144,7 @@ impl JsonnetObjectValues {
         slf: PyRefMut<'py, Self>,
         py: Python<'py>,
     ) -> PyResult<Option<Bound<'py, PyAny>>> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let Some(key) = slf.iter.borrow_mut().next() else {
                 return Ok(None);
             };
@@ -176,7 +176,7 @@ impl JsonnetObjectItems {
         slf: PyRefMut<'py, Self>,
         py: Python<'py>,
     ) -> PyResult<Option<(String, Bound<'py, PyAny>)>> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let Some(key) = slf.iter.borrow_mut().next() else {
                 return Ok(None);
             };
@@ -207,7 +207,7 @@ impl JsonnetArray {
     }
 
     fn __contains__(&self, py: Python<'_>, object: Bound<'_, PyAny>) -> PyResult<bool> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let value = py_to_jsonnet(py, object)?;
 
             for c in self.0.iter_lazy() {
@@ -227,7 +227,7 @@ impl JsonnetArray {
         py: Python<'py>,
         key: &Bound<'_, PyAny>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let Ok(idx) = key.extract::<usize>() else {
                 return Err(PyTypeError::new_err("index should be a decimal"));
             };
@@ -249,7 +249,7 @@ impl JsonnetArray {
 
     #[pyo3(signature = (minified=true))]
     fn manifest_json(&self, minified: bool) -> PyResult<String> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let preserve_order = true;
             let fmt = if minified {
                 jrsonnet_evaluator::manifest::JsonFormat::minify(preserve_order)
@@ -282,7 +282,7 @@ impl JsonnetArrayIter {
     ) -> PyResult<Option<Bound<'py, PyAny>>> {
         let slf = &mut slf;
 
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             slf.array
                 .0
                 .get(slf.idx)
@@ -310,7 +310,7 @@ impl JsonnetFunc {
         py: Python<'py>,
         args: &Bound<'_, PyTuple>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        execute_jsonnet(|| {
+        execute_jsonnet(|_| {
             let args = pylist_to_jsonnet(py, args.iter())?;
 
             let out = self

src/jsonnet_tokio.rs

@@ -13,26 +13,45 @@
 // limitations under the License.
 
 use gag::Gag;
+use nix::sys::signal::{signal, SigHandler, Signal};
+use pyo3::PyResult;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::sync::LazyLock;
+use tokio::sync::Notify;
 use tokio::task::block_in_place;
 
-pub fn init() {
-    let runtime = tokio::runtime::Builder::new_multi_thread()
+static RUNTIME: LazyLock<tokio::runtime::Runtime> = LazyLock::new(|| {
+    tokio::runtime::Builder::new_multi_thread()
         .thread_name("chainql-tokio-runtime")
         .enable_all()
         .build()
-        .unwrap();
+        .unwrap()
+});
 
-    crate::RUNTIME.set(runtime).unwrap()
-}
+static CANCELLATION_NOTIFIER: LazyLock<Arc<Notify>> = LazyLock::new(|| Arc::new(Notify::new()));
 
+#[track_caller]
 #[inline(always)]
-pub fn execute_jsonnet<F: FnOnce() -> T, T>(f: F) -> T {
+pub fn execute_jsonnet<F, T>(f: F) -> F::Output
+where
+    F: FnOnce(Arc<Notify>) -> PyResult<T>,
+{
+    let pyhandler = unsafe { signal(Signal::SIGINT, SigHandler::Handler(ctrl_c_handler)).unwrap() };
+
+    // TODO: Remove when chainql_core migrates to tracings
     let use_logger = crate::ENABLE_LOGGER.load(Ordering::SeqCst);
     let _gag = (!use_logger).then(|| (Gag::stdout().unwrap(), Gag::stderr().unwrap()));
 
-    let runtime = crate::RUNTIME.get().unwrap();
-    let _enter_guard = runtime.enter();
+    let _enter_guard = RUNTIME.enter();
+    let result = block_in_place(|| f(Arc::clone(&CANCELLATION_NOTIFIER)));
+
+    // TODO: Defer this
+    unsafe { signal(Signal::SIGINT, pyhandler).unwrap() };
+
+    result
+}
 
-    block_in_place(f)
+extern "C" fn ctrl_c_handler(_signal: core::ffi::c_int) {
+    CANCELLATION_NOTIFIER.notify_last();
 }

src/lib.rs

@@ -23,22 +23,17 @@ mod utils;
 
 use pyo3::prelude::*;
 
-use std::sync::{
-    atomic::{AtomicBool, Ordering},
-    OnceLock,
-};
+use std::sync::{atomic::{AtomicBool, Ordering}, };
 use utils::value_error;
 
 pub(crate) static ENABLE_LOGGER: AtomicBool = AtomicBool::new(false);
-pub(crate) static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
 
 #[pymodule]
 mod chainql {
     use super::*;
 
     #[pymodule_init]
     fn init(_m: &Bound<'_, PyModule>) -> PyResult<()> {
-        jsonnet_tokio::init();
         Ok(())
     }