wip

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

vortex-io/src/runtime/handle.rs

@@ -21,6 +21,7 @@ use crate::file::IoRequestStream;
 use crate::runtime::AbortHandleRef;
 use crate::runtime::Executor;
 use crate::runtime::IoTask;
+use crate::runtime::LocalExecutor;
 
 /// A handle to an active Vortex runtime.
 ///
@@ -45,6 +46,11 @@ impl Handle {
         })
     }
 
+    /// Returns a [`LocalExecutor`] if the underlying runtime supports spawning `!Send` futures.
+    pub fn as_local_executor(&self) -> Option<Arc<dyn LocalExecutor>> {
+        self.runtime().as_local_executor()
+    }
+
     /// Returns a handle to the current runtime, if such a reasonable choice exists.
     ///
     /// For example, if called from within a Tokio context this will return a

vortex-io/src/runtime/mod.rs

@@ -13,6 +13,10 @@
 //! * Tokio: work is driven on a Tokio runtime provided by the caller.
 //!
 
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 
@@ -34,6 +38,8 @@ pub mod single;
 mod smol;
 #[cfg(feature = "tokio")]
 pub mod tokio;
+#[cfg(not(target_arch = "wasm32"))]
+pub mod uring;
 #[cfg(target_arch = "wasm32")]
 pub mod wasm;
 
@@ -65,8 +71,29 @@ pub(crate) trait Executor: Send + Sync {
     ///
     /// Cancellation is implied by termination of the request stream.
     fn spawn_io(&self, task: IoTask);
+
+    /// Returns a [`LocalExecutor`] view if the runtime supports spawning `!Send` futures.
+    ///
+    /// Default implementation returns `None` for runtimes that only support `Send` futures.
+    fn as_local_executor(&self) -> Option<Arc<dyn LocalExecutor>> {
+        None
+    }
 }
 
+/// Extension trait for runtimes that can build and drive `!Send` futures on a single thread.
+///
+/// The factory is `Send` so it may be sent to the runtime's thread; the produced future can be
+/// `!Send` because it never leaves that thread after creation.
+pub(crate) trait LocalExecutor: Executor {
+    fn spawn_local(&self, f: LocalSpawn) -> AbortHandleRef;
+}
+
+/// A boxed future that may be `!Send`.
+pub(crate) type LocalBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
+
+/// A boxed factory for building a local future on the target runtime thread.
+pub(crate) type LocalSpawn = Box<dyn FnOnce() -> LocalBoxFuture<'static, ()> + Send + 'static>;
+
 /// A handle that may be used to optimistically abort a spawned task.
 ///
 /// If dropped, the task should continue to completion.

vortex-io/src/runtime/uring.rs

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Per-core runtime helpers and dispatching executor.
+
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use futures::future::BoxFuture;
+use vortex_error::vortex_panic;
+
+use crate::runtime::AbortHandleRef;
+use crate::runtime::Executor;
+use crate::runtime::Handle;
+use crate::runtime::IoTask;
+use crate::runtime::LocalExecutor;
+use crate::runtime::LocalSpawn;
+
+/// An executor that dispatches work across a fixed set of underlying executors.
+///
+/// Tasks are assigned round-robin; there is no work stealing. This is intended to pair with
+/// per-core runtimes where each executor owns a single thread/reactor.
+pub(crate) struct HandleSetExecutor {
+    executors: Arc<[Arc<dyn Executor>]>,
+    picker: AtomicUsize,
+}
+
+impl HandleSetExecutor {
+    pub(crate) fn new(executors: Vec<Arc<dyn Executor>>) -> Self {
+        assert!(!executors.is_empty());
+        Self {
+            executors: executors.into(),
+            picker: AtomicUsize::new(0),
+        }
+    }
+
+    fn pick(&self) -> &Arc<dyn Executor> {
+        let idx = self.picker.fetch_add(1, Ordering::Relaxed);
+        // Relaxed is sufficient: we only need uniqueness, not ordering guarantees.
+        &self.executors[idx % self.executors.len()]
+    }
+}
+
+/// A thin wrapper around a set of executors that produces a dispatching [`Handle`].
+///
+/// This is intended to be backed by per-core runtimes (e.g., io_uring reactors), but it can be
+/// constructed from any set of executors for now.
+pub(crate) struct HandleSet {
+    executors: Arc<[Arc<dyn Executor>]>,
+    dispatcher: Arc<HandleSetExecutor>,
+}
+
+impl HandleSet {
+    pub(crate) fn new(executors: Vec<Arc<dyn Executor>>) -> Self {
+        let executors: Arc<[Arc<dyn Executor>]> = executors.into();
+        let dispatcher = Arc::new(HandleSetExecutor::new(
+            executors.iter().cloned().collect(),
+        ));
+        Self {
+            executors,
+            dispatcher,
+        }
+    }
+
+    /// Returns a handle that round-robins spawned work across the underlying executors.
+    pub(crate) fn dispatching_handle(&self) -> Handle {
+        let exec: Arc<dyn Executor> = self.dispatcher.clone();
+        Handle::new(Arc::downgrade(&exec))
+    }
+
+    /// Access to the underlying executors, useful for building per-core pools later.
+    pub(crate) fn executors(&self) -> &[Arc<dyn Executor>] {
+        &self.executors
+    }
+}
+
+impl Executor for HandleSetExecutor {
+    fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef {
+        self.pick().spawn(fut)
+    }
+
+    fn spawn_cpu(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        self.pick().spawn_cpu(task)
+    }
+
+    fn spawn_blocking(&self, task: Box<dyn FnOnce() + Send + 'static>) -> AbortHandleRef {
+        self.pick().spawn_blocking(task)
+    }
+
+    fn spawn_io(&self, task: IoTask) {
+        self.pick().spawn_io(task)
+    }
+
+    fn as_local_executor(&self) -> Option<Arc<dyn LocalExecutor>> {
+        self.pick().as_local_executor()
+    }
+}
+
+impl LocalExecutor for HandleSetExecutor {
+    fn spawn_local(&self, f: LocalSpawn) -> AbortHandleRef {
+        match self.pick().as_local_executor() {
+            Some(exec) => exec.spawn_local(f),
+            None => vortex_panic!("LocalExecutor requested but not supported by any underlying executor"),
+        }
+    }
+}