From 27680cc889a3c407a24b494e71908ac3fb253c66 Mon Sep 17 00:00:00 2001
From: Faizaan Pervaiz <faizaan@byte-trading.com>
Date: Tue, 12 Dec 2023 11:24:29 +0000
Subject: [PATCH] Add ability to pin threads

---
 spdlog/Cargo.toml         |  3 ++-
 spdlog/src/thread_pool.rs | 23 ++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/spdlog/Cargo.toml b/spdlog/Cargo.toml
index 13760a24..bda01cd0 100644
--- a/spdlog/Cargo.toml
+++ b/spdlog/Cargo.toml
@@ -35,13 +35,14 @@ release-level-trace    = []
 source-location = []
 native = []
 libsystemd = ["libsystemd-sys"]
-multi-thread = ["crossbeam"]
+multi-thread = ["crossbeam", "core_affinity"]
 
 [dependencies]
 arc-swap = "1.5.1"
 atomic = "0.5.1"
 cfg-if = "1.0.0"
 chrono = "0.4.22"
+core_affinity = { version = "0.8.1", optional = true }
 crossbeam = { version = "0.8.2", optional = true }
 flexible-string = { version = "0.1.0", optional = true }
 if_chain = "1.0.2"
diff --git a/spdlog/src/thread_pool.rs b/spdlog/src/thread_pool.rs
index 94da96eb..18923bfb 100644
--- a/spdlog/src/thread_pool.rs
+++ b/spdlog/src/thread_pool.rs
@@ -40,6 +40,7 @@ pub struct ThreadPool {
 pub struct ThreadPoolBuilder {
     capacity: usize,
     threads: usize,
+    core_id: Option<usize>,
 }
 
 struct Worker {
@@ -53,6 +54,7 @@ impl ThreadPool {
         ThreadPoolBuilder {
             capacity: 8192,
             threads: 1,
+            core_id: None,
         }
     }
 
@@ -117,7 +119,17 @@ impl ThreadPoolBuilder {
         self
     }
 
+    /// Specify the core ID to which the thread pool should be pinned when built. Only one thread
+    /// will be pinned to the specified core as the thread pool currently only has one thread.
+    pub fn affinity(&mut self, core_id: usize) -> &mut Self {
+        self.core_id = Some(core_id);
+        self
+    }
+
     /// Builds a [`ThreadPool`].
+    ///
+    /// # Panics
+    /// Panics if core affinity has been set using [`ThreadPoolBuilder::affinity`] and pinning the thread to that core failed.
     pub fn build(&self) -> Result<ThreadPool> {
         if self.capacity < 1 {
             return Err(Error::InvalidArgument(
@@ -136,7 +148,16 @@ impl ThreadPoolBuilder {
         let mut threads = Vec::new();
         threads.resize_with(self.threads, || {
             let receiver = receiver.clone();
-            Some(thread::spawn(move || Worker { receiver }.run()))
+            let core_id = self.core_id;
+
+            Some(thread::spawn(move || {
+                if let Some(core_id) = core_id {
+                    if !core_affinity::set_for_current(core_affinity::CoreId { id: core_id }) {
+                        panic!("failed to pin thread to core {}", core_id);
+                    }
+                }
+                Worker { receiver }.run()
+            }))
         });
 
         Ok(ThreadPool {