From ff1e7582a84351e159136eaadaf5af2792c9dd84 Mon Sep 17 00:00:00 2001
From: Michael-J-Ward <ward.michael.j@gmail.com>
Date: Thu, 3 Oct 2024 13:12:55 -0500
Subject: [PATCH] use std::sync::OnceLock to store tokio runtime instead of
 round-tripping to python

---
 src/utils.rs | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/utils.rs b/src/utils.rs
index 0d72eaf75..0500e9a89 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -20,20 +20,18 @@ 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) => {
+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());
-            let obj: Bound<'_, TokioRuntime> = Py::new(py, rt).unwrap().into_bound(py);
-            obj.extract().unwrap()
-        }
-    }
+            Arc::new(rt)
+        })
+        .clone()
 }
 
 /// Utility to collect rust futures with GIL released