block_on_stream_thread_safe: increase channel size to number of cores

vortex-io/src/runtime/current.rs

@@ -9,6 +9,7 @@ use futures::StreamExt;
 use futures::stream::BoxStream;
 use parking_lot::Mutex;
 use smol::block_on;
+use vortex_utils::parallelism::get_available_parallelism;
 
 use crate::runtime::BlockingRuntime;
 use crate::runtime::Executor;
@@ -67,8 +68,10 @@ impl CurrentThreadRuntime {
 
         // We create an MPMC result channel and spawn a task to drive the stream and send results.
         // This allows multiple worker threads to drive the execution while all waiting for results
-        // on the channel.
-        let (result_tx, result_rx) = kanal::bounded_async(1);
+        // on the channel. Channel buffers up to one item per core so calling threads can get a
+        // ready item without every one of them having to become an executor.
+        let capacity = get_available_parallelism().unwrap_or(1).max(1);
+        let (result_tx, result_rx) = kanal::bounded_async(capacity);
         let driver = self.executor.spawn(async move {
             futures::pin_mut!(stream);
             while let Some(item) = stream.next().await {
@@ -154,19 +157,28 @@ impl<T> Iterator for ThreadSafeIterator<T> {
     type Item = T;
 
     fn next(&mut self) -> Option<Self::Item> {
+        // If driver already has an item, take it without driving the executor
+        match self.results.try_recv() {
+            Ok(Some(item)) => return Some(item),
+            Ok(None) => {}
+            Err(_) => return self.get_task_error(),
+        }
+
         match block_on(self.executor.run(self.results.recv())) {
             Ok(item) => Some(item),
-            // The result channel closes when the driver task finishes. Join the task so a panic
-            // raised while driving the stream is re-raised here instead of being lost. The first
-            // consumer to observe closure joins it; any later consumer just sees the stream end.
-            Err(_) => {
-                let task = self.driver.lock().take();
-                if let Some(task) = task {
-                    block_on(self.executor.run(task));
-                }
-                None
-            }
+            Err(_) => self.get_task_error(),
+        }
+    }
+}
+
+impl<T> ThreadSafeIterator<T> {
+    // Join current task so panics are re-raised
+    fn get_task_error(&self) -> Option<T> {
+        let task = self.driver.lock().take();
+        if let Some(task) = task {
+            block_on(self.executor.run(task));
         }
+        None
     }
 }