use std::sync::OnceLock to store tokio runtime instead of round-tripping to python

Michael-J-Ward · Michael-J-Ward · commit 6fa05bac5714 · 2024-10-03T13:12:55.000-05:00
diff --git a/src/utils.rs b/src/utils.rs
@@ -20,20 +20,16 @@ use crate::TokioRuntime;
 use datafusion::logical_expr::Volatility;
 use pyo3::prelude::*;
 use std::future::Future;
+use std::sync::{Arc, OnceLock};
 use tokio::runtime::Runtime;
 
 /// Utility to get the Tokio Runtime from Python
-pub(crate) fn get_tokio_runtime(py: Python) -> PyRef<TokioRuntime> {
-    let datafusion = py.import_bound("datafusion._internal").unwrap();
-    let tmp = datafusion.getattr("runtime").unwrap();
-    match tmp.extract::<PyRef<TokioRuntime>>() {
-        Ok(runtime) => runtime,
-        Err(_e) => {
-            let rt = TokioRuntime(tokio::runtime::Runtime::new().unwrap());
-            let obj: Bound<'_, TokioRuntime> = Py::new(py, rt).unwrap().into_bound(py);
-            obj.extract().unwrap()
-        }
-    }
+pub(crate) fn get_tokio_runtime(_: Python) -> Arc<TokioRuntime> {
+    static RUNTIME: OnceLock<Arc<TokioRuntime>> = OnceLock::new();
+    RUNTIME.get_or_init(|| {
+        let rt = TokioRuntime(tokio::runtime::Runtime::new().unwrap());
+        Arc::new(rt)
+    }).clone()
 }
 
 /// Utility to collect rust futures with GIL released
