diff --git a/vortex-io/src/runtime/current.rs b/vortex-io/src/runtime/current.rs index 93b698b43a2..6182c7ccb7d 100644 --- a/vortex-io/src/runtime/current.rs +++ b/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 Iterator for ThreadSafeIterator { type Item = T; fn next(&mut self) -> Option { + // 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 ThreadSafeIterator { + // Join current task so panics are re-raised + fn get_task_error(&self) -> Option { + let task = self.driver.lock().take(); + if let Some(task) = task { + block_on(self.executor.run(task)); } + None } }