cpu pinning: set affinity to multiple hyperthreads, fix issues

aquatic_common/Cargo.toml

@@ -14,7 +14,7 @@ name = "aquatic_common"
 ahash = "0.7"
 anyhow = "1"
 arc-swap = "1"
-core_affinity = "0.5"
+affinity = "0.1"
 hashbrown = "0.11.2"
 hex = "0.4"
 indexmap-amortized = "1"

aquatic_common/src/cpu_pinning.rs

@@ -17,8 +17,8 @@ impl Default for CpuPinningMode {
 pub struct CpuPinningConfig {
     pub active: bool,
     pub mode: CpuPinningMode,
-    pub offset: usize,
-    pub multiple: usize,
+    pub virtual_per_physical_cpu: usize,
+    pub offset_cpus: usize,
 }
 
 impl Default for CpuPinningConfig {
@@ -26,8 +26,8 @@ impl Default for CpuPinningConfig {
         Self {
             active: false,
             mode: Default::default(),
-            offset: 0,
-            multiple: 1,
+            virtual_per_physical_cpu: 2,
+            offset_cpus: 0,
         }
     }
 }
@@ -49,41 +49,58 @@ pub enum WorkerIndex {
 }
 
 impl WorkerIndex {
-    pub fn get_cpu_index(self, config: &CpuPinningConfig, socket_workers: usize) -> usize {
-        let index = match self {
-            Self::Other => config.offset,
-            Self::SocketWorker(index) => config.multiple * (config.offset + 1 + index),
+    fn get_cpu_indices(self, config: &CpuPinningConfig, socket_workers: usize) -> Vec<usize> {
+        let offset = match self {
+            Self::Other => config.virtual_per_physical_cpu * config.offset_cpus,
+            Self::SocketWorker(index) => {
+                config.virtual_per_physical_cpu * (config.offset_cpus + 1 + index)
+            }
             Self::RequestWorker(index) => {
-                config.multiple * (config.offset + 1 + socket_workers + index)
+                config.virtual_per_physical_cpu * (config.offset_cpus + 1 + socket_workers + index)
             }
         };
 
-        let index = match config.mode {
-            CpuPinningMode::Ascending => index,
+        let virtual_cpus = (0..config.virtual_per_physical_cpu).map(|i| offset + i);
+
+        let virtual_cpus: Vec<usize> = match config.mode {
+            CpuPinningMode::Ascending => virtual_cpus.collect(),
             CpuPinningMode::Descending => {
-                let max = core_affinity::get_core_ids()
-                    .map(|ids| ids.iter().map(|id| id.id).max())
-                    .flatten()
-                    .unwrap_or(0);
+                let max_index = affinity::get_core_num() - 1;
 
-                max - index
+                virtual_cpus
+                    .map(|i| max_index.checked_sub(i).unwrap_or(0))
+                    .collect()
             }
         };
 
-        ::log::info!("Calculated CPU pin index {} for {:?}", index, self);
+        ::log::info!(
+            "Calculated virtual CPU pin indices {:?} for {:?}",
+            virtual_cpus,
+            self
+        );
 
-        index
+        virtual_cpus
     }
 }
 
+/// Note: don't call this when affinities were already set in the current or in
+/// a parent thread. Doing so limits the number of cores that are seen and
+/// messes up setting affinities.
 pub fn pin_current_if_configured_to(
     config: &CpuPinningConfig,
     socket_workers: usize,
     worker_index: WorkerIndex,
 ) {
     if config.active {
-        core_affinity::set_for_current(core_affinity::CoreId {
-            id: worker_index.get_cpu_index(config, socket_workers),
-        });
+        let indices = worker_index.get_cpu_indices(config, socket_workers);
+
+        if let Err(err) = affinity::set_thread_affinity(indices.clone()) {
+            ::log::error!(
+                "Failed setting thread affinities {:?} for {:?}: {:#?}",
+                indices,
+                worker_index,
+                err
+            );
+        }
     }
 }