Add spawn_blocking to runtime (#4623)

We had hard-coded using the unblock crate, now we can use tokio's
spawn_blocking in that context

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

Cargo.lock

@@ -75,7 +75,6 @@ async-trait = "0.1.88"
 bindgen = "0.72.0"
 bit-vec = "0.8.0"
 bitvec = "1.0.1"
-blocking = "1.6.2"
 bytes = "1.10"
 bzip2 = "0.6.0"
 cbindgen = "0.29.0"

Cargo.toml

@@ -19,7 +19,6 @@ all-features = true
 [dependencies]
 async-compat = { workspace = true }
 async-stream = { workspace = true }
-blocking = { workspace = true }
 bytes = { workspace = true }
 cfg-if = { workspace = true }
 futures = { workspace = true, features = ["std", "executor"] }

vortex-io/Cargo.toml

@@ -6,7 +6,6 @@ use std::os::unix::fs::FileExt;
 use std::sync::Arc;
 
 use async_compat::Compat;
-use blocking::unblock;
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 use futures::{FutureExt, StreamExt};
@@ -101,6 +100,7 @@ impl IoSource for ObjectStoreIoSource {
         let self2 = self.clone();
         requests
             .map(move |req| {
+                let handle = self.handle.clone();
                 let store = self.io.store.clone();
                 let path = self.io.path.clone();
 
@@ -130,12 +130,13 @@ impl IoSource for ObjectStoreIoSource {
                             // The read_exact_at call will either fill the entire buffer or return an error,
                             // ensuring no uninitialized memory is exposed.
                             unsafe { buffer.set_len(len) };
-                            unblock(move || {
-                                file.read_exact_at(&mut buffer, range.start)?;
-                                Ok::<_, io::Error>(buffer)
-                            })
-                            .await
-                            .map_err(io::Error::other)?
+                            handle
+                                .spawn_blocking(move || {
+                                    file.read_exact_at(&mut buffer, range.start)?;
+                                    Ok::<_, io::Error>(buffer)
+                                })
+                                .await
+                                .map_err(io::Error::other)?
                         }
                         object_store::GetResultPayload::Stream(mut byte_stream) => {
                             while let Some(bytes) = byte_stream.next().await {

vortex-io/src/file/object_store.rs

@@ -6,7 +6,6 @@ use std::os::unix::fs::FileExt;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
-use blocking::unblock;
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 use futures::{FutureExt, StreamExt};
@@ -29,16 +28,17 @@ impl IntoIoSource for PathBuf {
 }
 
 impl IntoIoSource for &Path {
-    fn into_io_source(self, _handle: Handle) -> VortexResult<IoSourceRef> {
+    fn into_io_source(self, handle: Handle) -> VortexResult<IoSourceRef> {
         let uri = self.to_string_lossy().to_string().into();
         let file = Arc::new(File::open(self)?);
-        Ok(Arc::new(FileIoSource { uri, file }))
+        Ok(Arc::new(FileIoSource { uri, file, handle }))
     }
 }
 
 pub(crate) struct FileIoSource {
     uri: Arc<str>,
     file: Arc<File>,
+    handle: Handle,
 }
 
 impl IoSource for FileIoSource {
@@ -66,20 +66,22 @@ impl IoSource for FileIoSource {
         requests
             .map(move |req| {
                 let file = self.file.clone();
+                let handle = self.handle.clone();
                 async move {
                     let offset = req.offset();
                     let len = req.len();
                     let alignment = req.alignment();
 
-                    let result = unblock(move || {
-                        let mut buffer = ByteBufferMut::with_capacity_aligned(len, alignment);
-                        unsafe { buffer.set_len(len) };
-                        match file.read_exact_at(&mut buffer, offset) {
-                            Ok(()) => Ok(buffer.freeze()),
-                            Err(e) => Err(VortexError::from(e)),
-                        }
-                    })
-                    .await;
+                    let result = handle
+                        .spawn_blocking(move || {
+                            let mut buffer = ByteBufferMut::with_capacity_aligned(len, alignment);
+                            unsafe { buffer.set_len(len) };
+                            match file.read_exact_at(&mut buffer, offset) {
+                                Ok(()) => Ok(buffer.freeze()),
+                                Err(e) => Err(VortexError::from(e)),
+                            }
+                        })
+                        .await;
                     req.resolve(result);
                 }
             })

vortex-io/src/file/std_file.rs

@@ -91,11 +91,33 @@ impl Handle {
         F: FnOnce() -> R + Send + 'static,
         R: Send + 'static,
     {
-        // TODO(ngates): we want a droppable handle for this.
         let (send, recv) = oneshot::channel();
         let abort_handle = self.runtime.spawn_cpu(Box::new(move || {
-            // Task::detach allows the receiver to be dropped, so we ignore send errors.
-            let _ = send.send(f());
+            // Optimistically avoid the work if the result won't be used.
+            if !send.is_closed() {
+                // Task::detach allows the receiver to be dropped, so we ignore send errors.
+                let _ = send.send(f());
+            }
+        }));
+        Task {
+            recv,
+            abort_handle: Some(abort_handle),
+        }
+    }
+
+    /// Spawn a blocking I/O task for execution on the runtime.
+    pub fn spawn_blocking<F, R>(&self, f: F) -> Task<R>
+    where
+        F: FnOnce() -> R + Send + 'static,
+        R: Send + 'static,
+    {
+        let (send, recv) = oneshot::channel();
+        let abort_handle = self.runtime.spawn_blocking(Box::new(move || {
+            // Optimistically avoid the work if the result won't be used.
+            if !send.is_closed() {
+                // Task::detach allows the receiver to be dropped, so we ignore send errors.
+                let _ = send.send(f());
+            }
         }));
         Task {
             recv,

vortex-io/src/runtime/handle.rs

@@ -50,6 +50,12 @@ pub(crate) trait Runtime: Send + Sync {
     /// yet started executing.
     fn spawn_cpu(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef;
 
+    /// Spawns a blocking I/O task for execution on the runtime.
+    ///
+    /// The returned `AbortHandle` may be used to optimistically cancel the task if it has not
+    /// yet started executing.
+    fn spawn_blocking(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef;
+
     /// Spawns an I/O task for execution on the runtime.
     /// The runtime can choose to invoke the task's `Send` or `!Send` versions.
     ///

vortex-io/src/runtime/mod.rs

@@ -32,15 +32,17 @@ impl Default for SingleThreadRuntime {
 }
 
 struct Sender {
-    scheduling: kanal::Sender<SpawnFuture<'static>>,
-    cpu: kanal::Sender<SpawnCpu<'static>>,
+    scheduling: kanal::Sender<SpawnAsync<'static>>,
+    cpu: kanal::Sender<SpawnSync<'static>>,
+    blocking: kanal::Sender<SpawnSync<'static>>,
     io: kanal::Sender<IoTask>,
 }
 
 impl Sender {
     fn new(local: &Rc<LocalExecutor<'static>>) -> Self {
-        let (scheduling_send, scheduling_recv) = kanal::unbounded::<SpawnFuture>();
-        let (cpu_send, cpu_recv) = kanal::unbounded::<SpawnCpu>();
+        let (scheduling_send, scheduling_recv) = kanal::unbounded::<SpawnAsync>();
+        let (cpu_send, cpu_recv) = kanal::unbounded::<SpawnSync>();
+        let (blocking_send, blocking_recv) = kanal::unbounded::<SpawnSync>();
         let (io_send, io_recv) = kanal::unbounded::<IoTask>();
 
         // We pass weak references to the local executor into the async tasks such that the task's
@@ -68,10 +70,26 @@ impl Sender {
             .spawn(async move {
                 while let Ok(spawn) = cpu_recv.as_async().recv().await {
                     if let Some(local) = weak_local2.upgrade() {
-                        let cpu = spawn.cpu;
+                        let work = spawn.sync;
                         // Ignore send errors since it means the caller immediately detached.
                         let _ = spawn.task_callback.send(SmolAbortHandle::new_handle(
-                            local.spawn(async move { cpu() }),
+                            local.spawn(async move { work() }),
+                        ));
+                    }
+                }
+            })
+            .detach();
+
+        // Drive blocking tasks.
+        let weak_local2 = weak_local.clone();
+        local
+            .spawn(async move {
+                while let Ok(spawn) = blocking_recv.as_async().recv().await {
+                    if let Some(local) = weak_local2.upgrade() {
+                        let work = spawn.sync;
+                        // Ignore send errors since it means the caller immediately detached.
+                        let _ = spawn.task_callback.send(SmolAbortHandle::new_handle(
+                            local.spawn(async move { work() }),
                         ));
                     }
                 }
@@ -93,6 +111,7 @@ impl Sender {
         Self {
             scheduling: scheduling_send,
             cpu: cpu_send,
+            blocking: blocking_send,
             io: io_send,
         }
     }
@@ -105,7 +124,7 @@ impl Sender {
 impl Runtime for Sender {
     fn spawn(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef {
         let (send, recv) = oneshot::channel();
-        if let Err(e) = self.scheduling.send(SpawnFuture {
+        if let Err(e) = self.scheduling.send(SpawnAsync {
             future,
             task_callback: send,
         }) {
@@ -118,8 +137,21 @@ impl Runtime for Sender {
 
     fn spawn_cpu(&self, cpu: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
         let (send, recv) = oneshot::channel();
-        if let Err(e) = self.cpu.send(SpawnCpu {
-            cpu,
+        if let Err(e) = self.cpu.send(SpawnSync {
+            sync: cpu,
+            task_callback: send,
+        }) {
+            vortex_panic!("Executor missing: {}", e);
+        }
+        Box::new(LazyAbortHandle {
+            task: Mutex::new(recv),
+        })
+    }
+
+    fn spawn_blocking(&self, work: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        let (send, recv) = oneshot::channel();
+        if let Err(e) = self.blocking.send(SpawnSync {
+            sync: work,
             task_callback: send,
         }) {
             vortex_panic!("Executor missing: {}", e);
@@ -198,14 +230,14 @@ where
 /// we invert the behaviour of abort and drop. Dropping the abort handle results in the task being
 /// detached, whereas dropping the smol::Task results in the task being canceled. This helps avoid
 /// a race where the caller detaches the LazyAbortHandle before the smol::Task has been launched.
-struct SpawnFuture<'rt> {
+struct SpawnAsync<'rt> {
     future: BoxFuture<'rt, ()>,
     task_callback: oneshot::Sender<AbortHandleRef>,
 }
 
-// A spawn request for a CPU job.
-struct SpawnCpu<'rt> {
-    cpu: Box<dyn FnOnce() + Send + 'rt>,
+// A spawn request for a synchronous job.
+struct SpawnSync<'rt> {
+    sync: Box<dyn FnOnce() + Send + 'rt>,
     task_callback: oneshot::Sender<AbortHandleRef>,
 }
 

vortex-io/src/runtime/single.rs

@@ -16,6 +16,10 @@ impl Runtime for Executor<'static> {
         SmolAbortHandle::new_handle(self.spawn(async move { task() }))
     }
 
+    fn spawn_blocking(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        SmolAbortHandle::new_handle(smol::unblock(task))
+    }
+
     fn spawn_io(&self, task: IoTask) {
         self.spawn(task.source.drive_send(task.stream)).detach()
     }

vortex-io/src/runtime/smol.rs

@@ -41,6 +41,10 @@ impl Runtime for CurrentTokioRuntime {
         )
     }
 
+    fn spawn_blocking(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        Box::new(TokioHandle::current().spawn_blocking(task).abort_handle())
+    }
+
     fn spawn_io(&self, task: IoTask) {
         TokioHandle::current().spawn(task.source.drive_send(task.stream));
     }
@@ -55,6 +59,10 @@ impl Runtime for TokioHandle {
         Box::new(TokioHandle::spawn(self, async move { cpu() }).abort_handle())
     }
 
+    fn spawn_blocking(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        Box::new(TokioHandle::spawn_blocking(self, task).abort_handle())
+    }
+
     fn spawn_io(&self, task: IoTask) {
         TokioHandle::spawn(self, task.source.drive_send(task.stream));
     }