[runtime/tokio] support for core pinned dedicated tasks

runtime/src/deterministic.rs

@@ -1134,7 +1134,12 @@ impl Context {
 
 impl crate::Spawner for Context {
     fn dedicated(mut self) -> Self {
-        self.execution = Execution::Dedicated;
+        self.execution = Execution::Dedicated(None);
+        self
+    }
+
+    fn pinned(mut self, core: usize) -> Self {
+        self.execution = Execution::Dedicated(Some(core));
         self
     }
 

runtime/src/lib.rs

@@ -179,6 +179,21 @@ stability_scope!(BETA {
         /// This is not the default behavior. See [`Spawner::shared`] for more information.
         fn dedicated(self) -> Self;
 
+        /// Return a [`Spawner`] that runs tasks on a dedicated thread pinned to the given core.
+        ///
+        /// Core pinning is currently Linux only and a no-op on other platforms. Pinning may
+        /// silently fail in restricted environments (e.g. containers with cgroup CPU limits),
+        /// this method will still succeed but the thread will run unpinned.
+        ///
+        /// Use [`available_cores`] to query the number of available CPUs.
+        ///
+        /// Implies [`Spawner::dedicated`].
+        ///
+        /// # Panics
+        ///
+        /// Panics if `core` is greater than or equal to the number of available CPUs.
+        fn pinned(self, core: usize) -> Self;
+
         /// Return a [`Spawner`] that instruments the next spawned task with the label of the spawning context.
         fn instrumented(self) -> Self;
 
@@ -1686,6 +1701,16 @@ mod tests {
         });
     }
 
+    fn test_spawn_pinned<R: Runner>(runner: R)
+    where
+        R::Context: Spawner,
+    {
+        runner.start(|context| async move {
+            let handle = context.pinned(0).spawn(|_| async move { 42 });
+            assert!(matches!(handle.await, Ok(42)));
+        });
+    }
+
     fn test_spawn<R: Runner>(runner: R)
     where
         R::Context: Spawner + Clock,
@@ -3317,6 +3342,12 @@ mod tests {
         test_spawn_dedicated(executor);
     }
 
+    #[test]
+    fn test_deterministic_spawn_pinned() {
+        let executor = deterministic::Runner::default();
+        test_spawn_pinned(executor);
+    }
+
     #[test]
     fn test_deterministic_spawn() {
         let runner = deterministic::Runner::default();
@@ -3666,6 +3697,92 @@ mod tests {
         test_spawn_dedicated(executor);
     }
 
+    #[test]
+    fn test_tokio_spawn_pinned() {
+        let executor = tokio::Runner::default();
+        test_spawn_pinned(executor);
+    }
+
+    #[test]
+    fn test_tokio_spawn_pinned_dedicated_thread() {
+        // Verify that pinned implies dedicated.
+        let executor = tokio::Runner::default();
+        executor.start(|context| async move {
+            let root_thread = std::thread::current().id();
+            let task_thread = context
+                .pinned(0)
+                .spawn(|_| async move { std::thread::current().id() })
+                .await
+                .unwrap();
+            // The task should run on a different thread than the root thread.
+            assert_ne!(root_thread, task_thread);
+        });
+    }
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn test_tokio_spawn_pinned_correct_core() {
+        // Verify that a pinned task is actually running on the expected core,
+        // for every available core.
+        let num_cores = crate::available_cores().unwrap();
+        let executor = tokio::Runner::default();
+        executor.start(|context| async move {
+            for core in 0..num_cores {
+                let actual = context
+                    .clone()
+                    .pinned(core)
+                    .spawn(|_| async move {
+                        // SAFETY: `sched_getcpu` is a read-only query with no
+                        // preconditions.
+                        unsafe { libc::sched_getcpu() as usize }
+                    })
+                    .await
+                    .unwrap();
+                assert_eq!(actual, core);
+            }
+        });
+    }
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn test_tokio_spawn_pinned_same_core() {
+        // Verify that two separate tasks pinned to the same core run on
+        // different threads but report the same CPU.
+        let executor = tokio::Runner::default();
+        executor.start(|context| async move {
+            let core = crate::available_cores().unwrap() - 1;
+            let t1 = context.clone().pinned(core).spawn(|_| async move {
+                // SAFETY: `sched_getcpu` is a read-only query with no
+                // preconditions.
+                (std::thread::current().id(), unsafe { libc::sched_getcpu() })
+            });
+            let t2 = context.clone().pinned(core).spawn(|_| async move {
+                // SAFETY: `sched_getcpu` is a read-only query with no
+                // preconditions.
+                (std::thread::current().id(), unsafe { libc::sched_getcpu() })
+            });
+            let (r1, r2) = futures::future::join(t1, t2).await;
+            let (thread1, cpu1) = r1.unwrap();
+            let (thread2, cpu2) = r2.unwrap();
+            // Different dedicated threads.
+            assert_ne!(thread1, thread2);
+            // Same core.
+            assert_eq!(cpu1, cpu2);
+        });
+    }
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    #[should_panic(expected = "out of range")]
+    fn test_tokio_spawn_pinned_invalid_core() {
+        // Pinning to a core beyond the available count panics eagerly.
+        let num_cores = crate::available_cores().unwrap();
+        let executor = tokio::Runner::default();
+        executor.start(|context| async move {
+            context.pinned(num_cores).spawn(|_| async {});
+        });
+    }
+
     #[test]
     fn test_tokio_spawn() {
         let runner = tokio::Runner::default();

runtime/src/telemetry/metrics/task.rs

@@ -29,7 +29,7 @@ impl Label {
             name,
             kind: Kind::Task,
             execution: match execution {
-                crate::Execution::Dedicated => Execution::Dedicated,
+                crate::Execution::Dedicated(_) => Execution::Dedicated,
                 crate::Execution::Shared(blocking) => {
                     if blocking {
                         Execution::SharedBlocking

runtime/src/tokio/runtime.rs

@@ -548,7 +548,18 @@ impl Context {
 
 impl crate::Spawner for Context {
     fn dedicated(mut self) -> Self {
-        self.execution = Execution::Dedicated;
+        self.execution = Execution::Dedicated(None);
+        self
+    }
+
+    fn pinned(mut self, core: usize) -> Self {
+        if let Some(num_cores) = utils::thread::available_cores() {
+            assert!(
+                core < num_cores,
+                "core {core} out of range ({num_cores} available)"
+            );
+        }
+        self.execution = Execution::Dedicated(Some(core));
         self
     }
 
@@ -601,11 +612,15 @@ impl crate::Spawner for Context {
             Arc::clone(&parent),
         );
 
-        if matches!(past, Execution::Dedicated) {
+        if let Execution::Dedicated(core) = past {
             utils::thread::spawn(executor.thread_stack_size, {
                 // Ensure the task can access the tokio runtime
                 let handle = executor.runtime.handle().clone();
                 move || {
+                    // Pin before running any work on this thread
+                    if let Some(core) = core {
+                        utils::thread::pin_to_core(core);
+                    }
                     handle.block_on(f);
                 }
             });

runtime/src/utils/cell.rs

@@ -110,6 +110,10 @@ where
         Self::Present(self.into_present().dedicated())
     }
 
+    fn pinned(self, core: usize) -> Self {
+        Self::Present(self.into_present().pinned(core))
+    }
+
     fn shared(self, blocking: bool) -> Self {
         Self::Present(self.into_present().shared(blocking))
     }

runtime/src/utils/mod.rs

@@ -16,6 +16,8 @@ commonware_macros::stability_mod!(BETA, pub mod buffer);
 pub mod signal;
 #[cfg(not(target_arch = "wasm32"))]
 pub(crate) mod thread;
+#[cfg(not(target_arch = "wasm32"))]
+pub use thread::available_cores;
 
 mod handle;
 pub use handle::Handle;
@@ -31,8 +33,9 @@ pub(crate) mod supervision;
 /// The execution mode of a task.
 #[derive(Copy, Clone, Debug)]
 pub enum Execution {
-    /// Task runs on a dedicated thread.
-    Dedicated,
+    /// Task runs on a dedicated thread, optionally pinned to a core. Core pinning is
+    /// currently Linux only and a no-op on other platforms.
+    Dedicated(Option<usize>),
     /// Task runs on the shared executor. `true` marks short blocking work that should
     /// use the runtime's blocking-friendly pool.
     Shared(bool),

runtime/src/utils/thread.rs

@@ -1,10 +1,9 @@
 //! Helpers for resolving the configured thread stack size.
 
+#[cfg(target_os = "linux")]
+use commonware_utils::sync::Once;
 use std::{env, sync::OnceLock, thread};
 
-/// Cached configured thread stack size.
-static SYSTEM_THREAD_STACK_SIZE: OnceLock<usize> = OnceLock::new();
-
 /// Rust's default thread stack size.
 ///
 /// See <https://doc.rust-lang.org/std/thread/#stack-size>.
@@ -32,7 +31,10 @@ fn rust_min_stack() -> Option<usize> {
 ///
 /// On other platforms, or if the platform-specific query fails, this falls back
 /// to [RUST_DEFAULT_THREAD_STACK_SIZE].
+///
+/// The result is cached after the first call.
 pub(crate) fn system_thread_stack_size() -> usize {
+    static SYSTEM_THREAD_STACK_SIZE: OnceLock<usize> = OnceLock::new();
     *SYSTEM_THREAD_STACK_SIZE.get_or_init(|| {
         rust_min_stack()
             .or(system_thread_stack_size_impl())
@@ -111,3 +113,94 @@ where
         .spawn(f)
         .expect("failed to spawn thread")
 }
+
+/// Returns the number of available CPUs, or `None` if it cannot be determined.
+///
+/// The result is cached after the first call.
+#[cfg(unix)]
+pub fn available_cores() -> Option<usize> {
+    static CORES: OnceLock<Option<usize>> = OnceLock::new();
+    *CORES.get_or_init(|| {
+        // SAFETY: `sysconf(_SC_NPROCESSORS_ONLN)` is a read-only query with no
+        // preconditions.
+        let n = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
+        if n <= 0 {
+            None
+        } else {
+            Some(n as usize)
+        }
+    })
+}
+
+/// Returns the number of available CPUs, or `None` if it cannot be determined.
+///
+/// Always returns `None` on non-Unix platforms.
+#[cfg(not(unix))]
+pub const fn available_cores() -> Option<usize> {
+    None
+}
+
+/// Pins the current thread to the given core.
+///
+/// If the CPU count cannot be queried or `sched_setaffinity` fails, a warning
+/// is logged once and the thread continues unpinned.
+///
+/// # Panics
+///
+/// Panics if `core` is greater than or equal to the number of available CPUs.
+#[cfg(target_os = "linux")]
+pub(crate) fn pin_to_core(core: usize) {
+    static WARN_CPUS: Once = Once::new();
+    static WARN_AFFINITY: Once = Once::new();
+
+    let Some(num_cores) = available_cores() else {
+        WARN_CPUS.call_once(|| {
+            tracing::warn!("failed to query CPU count, skipping core pinning");
+        });
+        return;
+    };
+    assert!(
+        core < num_cores,
+        "core {core} out of range ({num_cores} available)"
+    );
+
+    // SAFETY: `cpu_set` is zeroed and then a single valid CPU index is set.
+    unsafe {
+        let mut cpu_set: libc::cpu_set_t = std::mem::zeroed();
+        libc::CPU_SET(core, &mut cpu_set);
+        let result = libc::sched_setaffinity(
+            0, // current thread
+            std::mem::size_of::<libc::cpu_set_t>(),
+            &cpu_set,
+        );
+        if result != 0 {
+            WARN_AFFINITY.call_once(|| {
+                tracing::warn!(core, "sched_setaffinity failed, skipping core pinning");
+            });
+        }
+    }
+}
+
+/// Pins the current thread to the given core.
+///
+/// No-op on non-Linux platforms.
+#[cfg(not(target_os = "linux"))]
+pub(crate) const fn pin_to_core(_core: usize) {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[cfg(unix)]
+    #[test]
+    fn test_available_cores() {
+        let n = available_cores().expect("available_cores returned None on Unix");
+        assert!(n >= 1, "expected at least 1 core, got {n}");
+    }
+
+    #[cfg(not(unix))]
+    #[test]
+    fn test_available_cores_non_unix() {
+        assert!(available_cores().is_none());
+    }
+}