Use sched_setaffinity on tokio blocking threads

Author: coolreader18

Cargo.lock

@@ -200,6 +200,7 @@ log = "0.4.17"
 memchr = "2"
 mimalloc = "0.1.39"
 nohash-hasher = "0.2"
+nix = "0.30"
 once_cell = "1.16"
 parking_lot = { version = "0.12.1", features = ["send_guard", "arc_lock"] }
 parse-size = "1.1.0"

Cargo.toml

@@ -120,6 +120,9 @@ core_affinity = "0.8"
 tikv-jemallocator = {workspace = true}
 tikv-jemalloc-ctl = {workspace = true}
 
+[target.'cfg(target_os = "linux")'.dependencies]
+nix = { workspace = true, features = ["sched"] }
+
 [features]
 # Print a warning when doing an unindexed `iter_by_col_range` on a large table.
 unindexed_iter_by_col_range_warn = []

crates/core/Cargo.toml

@@ -155,8 +155,8 @@ fn reload_config<S>(conf_file: &ConfigToml, reload_handle: &reload::Handle<EnvFi
 #[derive(Default)]
 pub struct Cores {
     pub databases: JobCores,
-    pub tokio_workers: ThreadPoolCores,
-    pub rayon: ThreadPoolCores,
+    pub tokio: TokioCores,
+    pub rayon: RayonCores,
 }
 
 impl Cores {
@@ -170,13 +170,25 @@ impl Cores {
 
         let databases = cores.take(frac(1.0 / 8.0)).collect();
 
-        let tokio_workers = ThreadPoolCores(Some(cores.take(frac(4.0 / 8.0)).collect()));
+        let tokio_workers = cores.take(frac(4.0 / 8.0)).collect();
 
-        let rayon = ThreadPoolCores(Some(cores.take(frac(1.0 / 8.0)).collect()));
+        let rayon = RayonCores(Some(cores.take(frac(1.0 / 8.0)).collect()));
+
+        #[cfg(target_os = "linux")]
+        let remaining = cores.try_fold(nix::sched::CpuSet::new(), |mut cpuset, core| {
+            cpuset.set(core.id).ok()?;
+            Some(cpuset)
+        });
+
+        let tokio = TokioCores {
+            workers: Some(tokio_workers),
+            #[cfg(target_os = "linux")]
+            blocking: remaining,
+        };
 
         Some(Self {
             databases,
-            tokio_workers,
+            tokio,
             rayon,
         })
     }
@@ -192,30 +204,40 @@ fn vec_to_queue(cores: Vec<CoreId>) -> CoreQueue {
 }
 
 #[derive(Default)]
-pub struct ThreadPoolCores(Option<Vec<CoreId>>);
+pub struct TokioCores {
+    workers: Option<Vec<CoreId>>,
+    #[cfg(target_os = "linux")]
+    blocking: Option<nix::sched::CpuSet>,
+}
 
-impl ThreadPoolCores {
-    fn into_setup_fn(self) -> Option<(usize, impl Fn())> {
-        self.0.map(|cores| {
+impl TokioCores {
+    pub fn configure(self, builder: &mut tokio::runtime::Builder) {
+        if let Some(cores) = self.workers {
+            let num = cores.len();
             let cores = vec_to_queue(cores);
-            (cores.len(), move || {
-                if let Some(core) = cores.pop() {
-                    core_affinity::set_for_current(core);
-                }
-            })
-        })
-    }
-
-    pub fn configure_tokio(self, builder: &mut tokio::runtime::Builder) {
-        if let Some((num, setup)) = self.into_setup_fn() {
             // `on_thread_start` gets called for both async worker threads and blocking threads,
             // but the first `worker_threads` threads that tokio spawns are worker threads,
             // so this ends up working fine
-            builder.worker_threads(num).on_thread_start(setup);
+            builder.worker_threads(num).on_thread_start(move || {
+                if let Some(core) = cores.pop() {
+                    core_affinity::set_for_current(core);
+                } else {
+                    #[cfg(target_os = "linux")]
+                    if let Some(cpuset) = &self.blocking {
+                        let this = nix::unistd::Pid::from_raw(0);
+                        let _ = nix::sched::sched_setaffinity(this, cpuset);
+                    }
+                }
+            });
         }
     }
+}
+
+#[derive(Default)]
+pub struct RayonCores(Option<Vec<CoreId>>);
 
-    pub fn configure_rayon(self, tokio_handle: &tokio::runtime::Handle) {
+impl RayonCores {
+    pub fn configure(self, tokio_handle: &tokio::runtime::Handle) {
         rayon_core::ThreadPoolBuilder::new()
             .thread_name(|_idx| "rayon-worker".to_string())
             .spawn_handler(thread_spawn_handler(tokio_handle))

crates/core/src/startup.rs

@@ -73,8 +73,8 @@ fn main() -> anyhow::Result<()> {
     // Create a multi-threaded run loop
     let mut builder = Builder::new_multi_thread();
     builder.enable_all();
-    cores.tokio_workers.configure_tokio(&mut builder);
+    cores.tokio.configure(&mut builder);
     let rt = builder.build().unwrap();
-    cores.rayon.configure_rayon(rt.handle());
+    cores.rayon.configure(rt.handle());
     rt.block_on(async_main(cores.databases))
 }