fix bug

Author: jayshrivastava

src/planning.rs

@@ -3,6 +3,7 @@ use std::{
     env,
     sync::{Arc, LazyLock},
 };
+use tokio::sync::OnceCell;
 
 use anyhow::{anyhow, Context};
 use arrow_flight::Action;
@@ -89,17 +90,20 @@ impl DDStage {
     }
 }
 
-static STATE: LazyLock<Result<SessionState>> = LazyLock::new(|| {
-    let wait_result = wait_for(make_state(), "make_state");
-    match wait_result {
-        Ok(Ok(state)) => Ok(state),
-        Ok(Err(e)) => Err(anyhow!("Failed to initialize state: {}", e).into()),
-        Err(e) => Err(anyhow!("Failed to initialize state: {}", e).into()),
-    }
-});
+// This is thead safe: https://docs.rs/tokio/latest/tokio/sync/struct.OnceCell.html
+static STATE: OnceCell<Result<SessionState>> = OnceCell::const_new();
+
+pub async fn get_ctx() -> Result<SessionContext> {
+    let result = STATE
+        .get_or_init(|| async {
+            match make_state().await {
+                Ok(state) => Ok(state),
+                Err(e) => Err(anyhow!("Failed to initialize state: {}", e).into()),
+            }
+        })
+        .await;
 
-pub fn get_ctx() -> Result<SessionContext> {
-    match &*STATE {
+    match result {
         Ok(state) => Ok(SessionContext::new_with_state(state.clone())),
         Err(e) => Err(anyhow!("Context initialization failed: {}", e).into()),
     }

src/proxy_service.rs

@@ -121,8 +121,9 @@ impl DDProxyHandler {
         stage_id: u64,
         addrs: Addrs,
     ) -> Result<Response<crate::flight::DoGetStream>, Status> {
-        let mut ctx =
-            get_ctx().map_err(|e| Status::internal(format!("Could not create context {e:?}")))?;
+        let mut ctx = get_ctx()
+            .await
+            .map_err(|e| Status::internal(format!("Could not create context {e:?}")))?;
 
         add_ctx_extentions(&mut ctx, &self.host, &query_id, stage_id, addrs, vec![])
             .map_err(|e| Status::internal(format!("Could not add context extensions {e:?}")))?;

src/query_planner.rs

@@ -91,7 +91,9 @@ impl QueryPlanner {
     /// distributed plan, and distributed stages, but it does not yet contain
     /// worker addresses or tasks, as they are filled in later by `distribute_plan()`.
     pub async fn prepare(&self, sql: &str) -> Result<QueryPlan> {
-        let mut ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?;
+        let mut ctx = get_ctx()
+            .await
+            .map_err(|e| anyhow!("Could not create context: {e}"))?;
         if let Some(customizer) = &self.customizer {
             customizer
                 .customize(&mut ctx)
@@ -119,7 +121,9 @@ impl QueryPlanner {
     /// distributed plan, and distributed stages, but it does not yet contain
     /// worker addresses or tasks, as they are filled in later by `distribute_plan()`.
     pub async fn prepare_substrait(&self, substrait_plan: Plan) -> Result<QueryPlan> {
-        let mut ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?;
+        let mut ctx = get_ctx()
+            .await
+            .map_err(|e| anyhow!("Could not create context: {e}"))?;
         if let Some(customizer) = &self.customizer {
             customizer
                 .customize(&mut ctx)

src/util.rs

@@ -9,7 +9,9 @@ use std::{
     time::Duration,
 };
 
-use anyhow::{anyhow, Context as anyhowctx};
+use tokio::time::timeout;
+
+use anyhow::{anyhow, Context as anyhowctx, Error};
 use arrow::{
     array::RecordBatch,
     datatypes::SchemaRef,
@@ -53,6 +55,7 @@ use tokio::{
 use tonic::transport::Channel;
 use url::Url;
 
+use crate::result::DDError;
 use crate::{
     logging::{debug, error, trace},
     protobuf::StageAddrs,
@@ -76,43 +79,47 @@ impl Spawner {
         Self { runtime }
     }
 
+    // wait_for_future can be dc
     fn wait_for_future<F>(&self, f: F, name: &str) -> Result<F::Output>
     where
         F: Future + Send + 'static,
         F::Output: Send,
     {
+        if Handle::try_current().is_ok() {
+            panic!("cannot call wait_for_future within an async runtime")
+        }
+
         let name_c = name.to_owned();
         trace!("Spawner::wait_for {name_c}");
-        let (tx, rx) = std::sync::mpsc::channel::<F::Output>();
+        let (tx, mut rx) = tokio::sync::mpsc::channel::<F::Output>(1);
 
-        let func = move || {
+        // Let's not drop the tx.
+        let func = async move || {
             trace!("spawned fut start {name_c}");
 
-            let out = Handle::current().block_on(f);
+            let out = f.await;
             trace!("spawned fut stop {name_c}");
-            tx.send(out).inspect_err(|e| {
-                error!("ERROR sending future reesult over channel!!!! {e:?}");
-            })
+            let result = tx.send(out).await;
+            if let Err(e) = result {
+                error!("ERROR sending future result over channel!!!! {e:?}");
+            }
+            // tx is dropped
         };
 
         {
             let _guard = self.runtime.enter();
-            let handle = Handle::current();
-
-            trace!("Spawner spawning {name}");
-            handle.spawn_blocking(func);
-            trace!("Spawner spawned {name}");
+            trace!("Spawner spawning {name} (sync)");
+            tokio::spawn(func());
+            trace!("Spawner spawned {name} (sync)");
         }
 
-        let out = rx
-            .recv_timeout(Duration::from_secs(5))
-            .inspect_err(|e| {
-                error!("Spawner::wait_for {name} timed out waiting for future result: {e:?}");
-            })
-            .context("Spawner::wait_for failed to receive future result")?;
+        let result = rx.blocking_recv();
+        if result.is_none() {
+            error!("Spawner::wait_for {name} timed out waiting for future result");
+            return Err(DDError::Other(anyhow!("blah")));
+        }
 
-        debug!("Spawner::wait_for {name} returning");
-        Ok(out)
+        Ok(result.unwrap())
     }
 }
 
@@ -652,21 +659,21 @@ mod test {
         assert_eq!(out, 5);
     }
 
-    #[test]
-    fn test_wait_for_nested() {
-        println!("test_wait_for_nested");
-        let fut = async || {
-            println!("in outter fut");
-            let fut5 = async || {
-                println!("in inner fut");
-                5
-            };
-            wait_for(fut5(), "inner").unwrap()
-        };
-
-        let out = wait_for(fut(), "outer").unwrap();
-        assert_eq!(out, 5);
-    }
+    // #[test]
+    // fn test_wait_for_nested() {
+    //     println!("test_wait_for_nested");
+    //     let fut = async || {
+    //         println!("in outter fut");
+    //         let fut5 = async || {
+    //             println!("in inner fut");
+    //             5
+    //         };
+    //         wait_for(fut5(), "inner").unwrap()
+    //     };
+    //
+    //     let out = wait_for(fut(), "outer").unwrap();
+    //     assert_eq!(out, 5);
+    // }
 
     #[test(tokio::test)]
     async fn test_max_rows_stream() {

src/worker_service.rs

@@ -261,7 +261,7 @@ impl DDWorkerHandler {
         stage_addrs: Addrs,
         partition_group: Vec<u64>,
     ) -> Result<SessionContext> {
-        let mut ctx = get_ctx()?;
+        let mut ctx = get_ctx().await?;
         let host = Host {
             addr: self.addr.clone(),
             name: self.name.clone(),