Async spawn

crates/polars-io/src/parquet/read_impl.rs

@@ -691,22 +691,16 @@ impl BatchedParquetReader {
                     };
 
                     // Spawn the task and wait on it asynchronously.
-                    let (dfs, rows_read, limit) = if POOL.current_thread_index().is_some() {
+                    if POOL.current_thread_index().is_some() {
                         // We are a rayon thread, so we can't use POOL.spawn as it would mean we spawn a task and block until
                         // another rayon thread executes it - we would deadlock if all rayon threads did this.
-
-                        // Activate another tokio thread to poll futures. There should be at least 1 tokio thread that is
-                        // not a rayon thread.
-                        let handle = tokio::spawn(async { rx.await.unwrap() });
-                        // Now spawn the task onto rayon and participate in executing it. The current thread will no longer
-                        // poll async futures until this rayon task is complete.
-                        POOL.install(f);
-                        handle.await.unwrap()
+                        tokio::task::block_in_place(f);
                     } else {
                         POOL.spawn(f);
-                        rx.await.unwrap()
                     };
 
+                    let (dfs, rows_read, limit) = rx.await.unwrap();
+
                     self.rows_read = rows_read;
                     self.limit = limit;
                     dfs

crates/polars-io/src/pl_async.rs

@@ -2,13 +2,10 @@ use std::error::Error;
 use std::future::Future;
 use std::ops::Deref;
 use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering};
-use std::sync::RwLock;
-use std::thread::ThreadId;
 
 use once_cell::sync::Lazy;
 use polars_core::config::verbose;
 use polars_core::POOL;
-use polars_utils::aliases::PlHashSet;
 use tokio::runtime::{Builder, Runtime};
 use tokio::sync::Semaphore;
 
@@ -220,7 +217,6 @@ where
 
 pub struct RuntimeManager {
     rt: Runtime,
-    blocking_threads: RwLock<PlHashSet<ThreadId>>,
 }
 
 impl RuntimeManager {
@@ -232,10 +228,7 @@ impl RuntimeManager {
             .build()
             .unwrap();
 
-        Self {
-            rt,
-            blocking_threads: Default::default(),
-        }
+        Self { rt }
     }
 
     /// Keep track of rayon threads that drive the runtime. Every thread
@@ -247,15 +240,10 @@ impl RuntimeManager {
         F: Future + Send,
         F::Output: Send,
     {
-        let thread_id = std::thread::current().id();
-
-        if self.blocking_threads.read().unwrap().contains(&thread_id) {
-            std::thread::scope(|s| s.spawn(|| self.rt.block_on(future)).join().unwrap())
+        if tokio::runtime::Handle::try_current().is_err() {
+            self.rt.block_on(future)
         } else {
-            self.blocking_threads.write().unwrap().insert(thread_id);
-            let out = self.rt.block_on(future);
-            self.blocking_threads.write().unwrap().remove(&thread_id);
-            out
+            tokio::task::block_in_place(|| self.rt.block_on(future))
         }
     }