From 8e11c7d55c3eafde373a8cdba7911289f28eace6 Mon Sep 17 00:00:00 2001
From: SilasMarvin <19626586+SilasMarvin@users.noreply.github.com>
Date: Fri, 22 Mar 2024 11:50:40 -0700
Subject: [PATCH] Added once_cell crate to make RUNTIME lazy

---
 pgml-sdks/pgml/Cargo.lock |  1 +
 pgml-sdks/pgml/Cargo.toml |  1 +
 pgml-sdks/pgml/src/lib.rs | 24 ++++++++----------------
 3 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/pgml-sdks/pgml/Cargo.lock b/pgml-sdks/pgml/Cargo.lock
index e651e5969..11128b907 100644
--- a/pgml-sdks/pgml/Cargo.lock
+++ b/pgml-sdks/pgml/Cargo.lock
@@ -1547,6 +1547,7 @@ dependencies = [
  "lopdf",
  "md5",
  "neon",
+ "once_cell",
  "parking_lot",
  "pyo3",
  "pyo3-asyncio",
diff --git a/pgml-sdks/pgml/Cargo.toml b/pgml-sdks/pgml/Cargo.toml
index 633c9d30d..6141eb600 100644
--- a/pgml-sdks/pgml/Cargo.toml
+++ b/pgml-sdks/pgml/Cargo.toml
@@ -44,6 +44,7 @@ colored = "2"
 ctrlc = "3"
 inquire = "0.6"
 parking_lot = "0.12.1"
+once_cell = "1.19.0"
 
 [features]
 default = []
diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs
index 8f6039b52..ee5469f0f 100644
--- a/pgml-sdks/pgml/src/lib.rs
+++ b/pgml-sdks/pgml/src/lib.rs
@@ -4,6 +4,7 @@
 //!
 //! With this SDK, you can seamlessly manage various database tables related to documents, text chunks, text splitters, LLM (Language Model) models, and embeddings. By leveraging the SDK's capabilities, you can efficiently index LLM embeddings using PgVector for fast and accurate queries.
 
+use once_cell::sync::Lazy;
 use parking_lot::RwLock;
 use sqlx::{postgres::PgPoolOptions, PgPool};
 use std::collections::HashMap;
@@ -123,24 +124,15 @@ fn internal_init_logger(level: Option<String>, format: Option<String>) -> anyhow
 
 // Normally the global async runtime is handled by tokio but because we are a library being called
 // by javascript and other langauges, we occasionally need to handle it ourselves
-#[allow(dead_code)]
-static mut RUNTIME: Option<Runtime> = None;
+static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
+    Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("Error creating tokio runtime")
+});
 
-#[allow(dead_code)]
 fn get_or_set_runtime<'a>() -> &'a Runtime {
-    unsafe {
-        if let Some(r) = &RUNTIME {
-            r
-        } else {
-            // Need to use multi thread for JavaScript
-            let runtime = Builder::new_multi_thread()
-                .enable_all()
-                .build()
-                .expect("Error creating tokio runtime");
-            RUNTIME = Some(runtime);
-            get_or_set_runtime()
-        }
-    }
+    &RUNTIME
 }
 
 #[cfg(feature = "python")]