wip

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

Author: onursatici

vortex-io/src/runtime/handle.rs

@@ -40,7 +40,7 @@ impl Handle {
         Self { runtime }
     }
 
-    fn runtime(&self) -> Arc<dyn Executor> {
+    pub(crate) fn runtime(&self) -> Arc<dyn Executor> {
         self.runtime.upgrade().unwrap_or_else(|| {
             vortex_panic!("Attempted to use a Handle after its runtime was dropped")
         })

vortex-io/src/runtime/mod.rs

@@ -47,7 +47,7 @@ pub mod wasm;
 mod tests;
 
 /// Trait used to abstract over different async runtimes.
-pub(crate) trait Executor: Send + Sync {
+pub trait Executor: Send + Sync {
     /// Spawns a future to be executed on the runtime.
     ///
     /// The future should continue to be polled in the background by the runtime.
@@ -84,7 +84,7 @@ pub(crate) trait Executor: Send + Sync {
 ///
 /// 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 {
+pub trait LocalExecutor: Executor {
     fn spawn_local(&self, f: LocalSpawn) -> AbortHandleRef;
 }
 
@@ -98,11 +98,11 @@ pub(crate) type LocalSpawn = Box<dyn FnOnce() -> LocalBoxFuture<'static, ()> + S
 ///
 /// If dropped, the task should continue to completion.
 /// If explicitly aborted, the task should be cancelled if it has not yet started executing.
-pub(crate) trait AbortHandle: Send + Sync {
+pub trait AbortHandle: Send + Sync {
     fn abort(self: Box<Self>);
 }
 
-pub(crate) type AbortHandleRef = Box<dyn AbortHandle>;
+pub type AbortHandleRef = Box<dyn AbortHandle>;
 
 /// A task for driving I/O requests against a source.
 ///
@@ -112,7 +112,7 @@ pub(crate) type AbortHandleRef = Box<dyn AbortHandle>;
 ///
 // NOTE(ngates): We could in theory make IoSource support as_any if we wanted each runtime to implement the
 // actual read logic themselves? Not sure yet...
-pub(crate) struct IoTask {
+pub struct IoTask {
     pub(crate) source: ReadSourceRef,
     pub(crate) stream: BoxStream<'static, IoRequest>,
 }

vortex-io/src/runtime/uring.rs

@@ -10,13 +10,17 @@ use std::sync::atomic::Ordering;
 use futures::future::BoxFuture;
 use vortex_error::vortex_panic;
 
+use crate::runtime::blocking::BlockingRuntime;
 use crate::runtime::AbortHandleRef;
+use crate::runtime::current::CurrentThreadRuntime;
+use crate::runtime::current::CurrentThreadWorkerPool;
 use crate::runtime::Executor;
 use crate::runtime::Handle;
 use crate::runtime::IoTask;
 use crate::runtime::LocalExecutor;
 use crate::runtime::LocalSpawn;
 
+#[allow(dead_code)]
 /// 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
@@ -26,6 +30,7 @@ pub(crate) struct HandleSetExecutor {
     picker: AtomicUsize,
 }
 
+#[allow(dead_code)]
 impl HandleSetExecutor {
     pub(crate) fn new(executors: Vec<Arc<dyn Executor>>) -> Self {
         assert!(!executors.is_empty());
@@ -42,6 +47,7 @@ impl HandleSetExecutor {
     }
 }
 
+#[allow(dead_code)]
 /// 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
@@ -51,6 +57,7 @@ pub(crate) struct HandleSet {
     dispatcher: Arc<HandleSetExecutor>,
 }
 
+#[allow(dead_code)]
 impl HandleSet {
     pub(crate) fn new(executors: Vec<Arc<dyn Executor>>) -> Self {
         let executors: Arc<[Arc<dyn Executor>]> = executors.into();
@@ -75,6 +82,71 @@ impl HandleSet {
     }
 }
 
+/// Create a [`Handle`] that dispatches work round-robin across the provided handles.
+///
+/// This is useful for thread-per-core runtimes where each handle is tied to a single reactor.
+pub fn dispatching_handle(handles: &[Handle]) -> Handle {
+    let executors = handles
+        .iter()
+        .map(|h| h.runtime())
+        .collect::<Vec<_>>();
+    let set = HandleSet::new(executors);
+    set.dispatching_handle()
+}
+
+/// A lightweight per-core pool using current-thread runtimes and background workers.
+///
+/// This is a stopgap until a true io_uring-backed runtime is wired in. Each core owns its own
+/// executor driven by a single worker thread, and the exposed handle dispatches round-robin
+/// across them.
+#[allow(dead_code)]
+pub struct PerCoreRuntimePool {
+    cores: Vec<CurrentThreadCore>,
+    handle: Handle,
+}
+
+#[allow(dead_code)]
+impl PerCoreRuntimePool {
+    /// Build a pool with `cores` runtimes (defaults to available_parallelism if None).
+    pub fn new(cores: Option<usize>) -> Self {
+        let core_count = cores
+            .or_else(|| std::thread::available_parallelism().ok().map(|n| n.get()))
+            .unwrap_or(1);
+
+        let cores: Vec<_> = (0..core_count).map(|_| CurrentThreadCore::new()).collect();
+        let handles: Vec<_> = cores.iter().map(|c| c.handle()).collect();
+        let handle = dispatching_handle(&handles);
+
+        Self { cores, handle }
+    }
+
+    /// A handle that spreads work across the per-core runtimes.
+    pub fn handle(&self) -> Handle {
+        self.handle.clone()
+    }
+}
+
+struct CurrentThreadCore {
+    runtime: CurrentThreadRuntime,
+    _pool: CurrentThreadWorkerPool,
+}
+
+impl CurrentThreadCore {
+    fn new() -> Self {
+        let runtime = CurrentThreadRuntime::new();
+        let pool = runtime.new_pool();
+        pool.set_workers(1);
+        Self {
+            runtime,
+            _pool: pool,
+        }
+    }
+
+    fn handle(&self) -> Handle {
+        self.runtime.handle()
+    }
+}
+
 impl Executor for HandleSetExecutor {
     fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef {
         self.pick().spawn(fut)

vortex-io/src/session.rs

@@ -52,5 +52,13 @@ pub trait RuntimeSessionExt: SessionExt {
         self.get_mut::<RuntimeSession>().handle = Some(handle);
         self
     }
+
+    /// Configure the runtime session to use a dispatching handle that round-robins across the
+    /// provided handles. Useful for thread-per-core runtimes.
+    fn with_dispatching_handles(self, handles: Vec<Handle>) -> Self {
+        self.get_mut::<RuntimeSession>().handle =
+            Some(crate::runtime::uring::dispatching_handle(&handles));
+        self
+    }
 }
 impl<S: SessionExt> RuntimeSessionExt for S {}