feat(server): respect CPU affinity and cgroup limits in shards and stats (#3615)

The server assumed it owned the whole host. Shards pinned to
absolute core ids 0..n, which fails with EINVAL (or lands on
forbidden cores) under systemd AllowedCPUs=, container cpusets,
or taskset. Stats reported host-wide CPU and memory, so a
confined instance showed its neighbors' load and a meaningless
host total.

Draw pinned cores from the process's allowed CPU set
(sched_getaffinity) instead, and validate configured ranges
against it at boot. A new pin_cores knob (default true) turns
pinning off for shared-core hosts (cgroup CPU quotas), where
every tenant pinning to the same low-numbered cores would pile
onto one core; unpinned shards let the kernel scheduler place
threads freely.

Stats now scope total_cpu_usage to the allowed cores and report
the effective cgroup memory limit as total_memory. Available
memory adds reclaimable file cache back
(limit - (current - inactive_file - active_file), minimum over
capped ancestors, v1 and v2), since memory.current charges page
cache as used and the naive limit - current trends to zero on a
cache-heavy server long before real OOM pressure.

Author: mmodzelewski

Cargo.lock

@@ -26,13 +26,15 @@ pub struct Stats {
     pub process_id: u32,
     /// The CPU usage of the process.
     pub cpu_usage: f32,
-    /// the total CPU usage of the system.
+    /// The total CPU usage of the system, scoped to the cores this process may run on
+    /// when confined by an affinity/cpuset mask.
     pub total_cpu_usage: f32,
     /// The memory usage of the process.
     pub memory_usage: IggyByteSize,
-    /// The total memory of the system.
+    /// The total memory of the system, or the effective cgroup memory limit when the
+    /// server runs inside a memory-capped cgroup (container, systemd slice).
     pub total_memory: IggyByteSize,
-    /// The available memory of the system.
+    /// The available memory of the system, scoped to the cgroup limit when one applies.
     pub available_memory: IggyByteSize,
     /// The run time of the process.
     pub run_time: IggyDuration,

core/common/src/types/stats/mod.rs

@@ -101,6 +101,10 @@ const fn default_inbox_capacity() -> usize {
     DEFAULT_INBOX_CAPACITY
 }
 
+const fn default_pin_cores() -> bool {
+    true
+}
+
 fn default_shutdown_drain_timeout() -> IggyDuration {
     IggyDuration::new(DEFAULT_SHUTDOWN_DRAIN_TIMEOUT)
 }
@@ -119,6 +123,18 @@ pub struct ShardingConfig {
     #[serde(default)]
     #[config_env(leaf)]
     pub cpu_allocation: CpuAllocation,
+    /// Whether shard threads are pinned to dedicated CPU cores
+    /// (`sched_setaffinity`). Pinning maximizes cache locality when this
+    /// server owns its cores (dedicated host, `numa:` allocations). Set to
+    /// `false` when the server shares cores with other workloads — e.g. a
+    /// multi-tenant host slicing CPU via cgroup quotas — where every process
+    /// pinning to the same low-numbered cores would pile onto one core while
+    /// the rest sit idle; unpinned shards let the kernel scheduler place
+    /// threads freely within the allowed set. With a NUMA-aware allocation,
+    /// `false` drops both the CPU and memory-node bindings (and logs a
+    /// warning, since NUMA placement without pinning is meaningless).
+    #[serde(default = "default_pin_cores")]
+    pub pin_cores: bool,
     /// Per-shard inter-shard inbox channel capacity. Bounded by design.
     /// Drops on full inbox of consensus frames are recovered by VSR
     /// retransmit. Drops of cross-shard client Reply frames are terminal:
@@ -176,6 +192,7 @@ impl Default for ShardingConfig {
     fn default() -> Self {
         Self {
             cpu_allocation: CpuAllocation::default(),
+            pin_cores: default_pin_cores(),
             inbox_capacity: DEFAULT_INBOX_CAPACITY,
             shutdown_drain_timeout: default_shutdown_drain_timeout(),
             shutdown_poll_interval: default_shutdown_poll_interval(),

core/configs/src/server_config/sharding.rs

@@ -28,6 +28,7 @@ use super::sharding::{
 use super::system::SegmentConfig;
 use super::system::{CompressionConfig, LoggingConfig, PartitionConfig};
 use crate::ConfigurationError;
+use cpu_allocation::allowed_cpus;
 use err_trail::ErrContext;
 use iggy_common::CompressionAlgorithm;
 use iggy_common::IggyExpiry;
@@ -488,9 +489,20 @@ impl Validatable<ConfigurationError> for ShardingConfig {
                     );
                     return Err(ConfigurationError::InvalidConfigurationValue);
                 }
-                if *end > available_cpus {
+                if *end - *start > available_cpus {
                     eprintln!(
-                        "Invalid sharding configuration: cpu_allocation range {start}..{end} exceeds available CPU cores (max: {available_cpus})"
+                        "Invalid sharding configuration: cpu_allocation range {start}..{end} yields {} shards, exceeding available CPU cores {available_cpus}",
+                        *end - *start
+                    );
+                    return Err(ConfigurationError::InvalidConfigurationValue);
+                }
+                if !self.pin_cores {
+                    return Ok(());
+                }
+                let allowed = allowed_cpus();
+                if let Some(cpu) = (*start..*end).find(|cpu| !allowed.contains(cpu)) {
+                    eprintln!(
+                        "Invalid sharding configuration: cpu_allocation range {start}..{end} includes CPU {cpu}, which is outside the set of cores allowed for this process (affinity/cpuset mask)"
                     );
                     return Err(ConfigurationError::InvalidConfigurationValue);
                 }
@@ -1008,3 +1020,71 @@ mod sharding_shutdown_knob_tests {
         assert!(cfg.validate().is_err());
     }
 }
+
+#[cfg(test)]
+mod sharding_cpu_range_tests {
+    use super::*;
+
+    #[test]
+    fn inverted_range_is_rejected() {
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(2, 2),
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_err());
+    }
+
+    #[test]
+    fn pinned_range_within_allowed_set_is_accepted() {
+        let first = allowed_cpus()[0];
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(first, first + 1),
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_ok());
+    }
+
+    #[test]
+    fn pinned_range_outside_allowed_set_is_rejected() {
+        let past_last = allowed_cpus().last().copied().unwrap() + 1;
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(past_last, past_last + 1),
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_err());
+    }
+
+    #[test]
+    fn pinned_range_wider_than_parallelism_is_rejected() {
+        // Under a cgroup CPU quota the affinity mask stays full while
+        // `available_parallelism` shrinks, so membership alone would
+        // accept this; the shard-count cap must reject it.
+        let first = allowed_cpus()[0];
+        let available = available_parallelism().unwrap().get();
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(first, first + available + 1),
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_err());
+    }
+
+    #[test]
+    fn unpinned_range_is_capped_by_shard_count_not_core_ids() {
+        // Core ids outside the machine are fine unpinned; only the
+        // resulting shard count matters.
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(1 << 20, (1 << 20) + 1),
+            pin_cores: false,
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_ok());
+
+        let available = available_parallelism().unwrap().get();
+        let cfg = ShardingConfig {
+            cpu_allocation: CpuAllocation::Range(0, available + 1),
+            pin_cores: false,
+            ..ShardingConfig::default()
+        };
+        assert!(cfg.validate().is_err());
+    }
+}

core/configs/src/server_config/validators.rs

@@ -18,13 +18,16 @@
 [package]
 name = "cpu_allocation"
 version = "0.1.0"
-description = "Shard CPU/NUMA allocation config types (CpuAllocation, NumaConfig) parsed from the iggy server config."
+description = "Shard CPU/NUMA allocation config types (CpuAllocation, NumaConfig) parsed from the iggy server config, plus the allowed-CPU-set probe."
 edition = "2024"
 license = "Apache-2.0"
 publish = false
 
 [dependencies]
 serde = { workspace = true }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+nix = { workspace = true }
+
 [dev-dependencies]
 serde_json = { workspace = true }

core/cpu_allocation/Cargo.toml

@@ -0,0 +1,59 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[cfg(target_os = "linux")]
+use nix::{
+    sched::{CpuSet, sched_getaffinity},
+    unistd::Pid,
+};
+use std::thread::available_parallelism;
+
+/// CPUs the calling process is currently allowed to run on, ascending.
+///
+/// Respects restrictions imposed by the parent environment (systemd
+/// `AllowedCPUs=`, container cpusets, `taskset`), which absolute core ids
+/// `0..n` would silently violate: `sched_setaffinity` to a core outside the
+/// allowed set fails with `EINVAL`. Falls back to `0..available_parallelism()`
+/// where the affinity mask is unavailable (non-Linux).
+pub fn allowed_cpus() -> Vec<usize> {
+    #[cfg(target_os = "linux")]
+    {
+        if let Ok(mask) = sched_getaffinity(Pid::from_raw(0)) {
+            let cpus: Vec<usize> = (0..CpuSet::count())
+                .filter(|&cpu| mask.is_set(cpu).unwrap_or(false))
+                .collect();
+            if !cpus.is_empty() {
+                return cpus;
+            }
+        }
+    }
+
+    let fallback = available_parallelism().map(|n| n.get()).unwrap_or(1);
+    (0..fallback).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn allowed_cpus_is_non_empty_and_ascending() {
+        let allowed = allowed_cpus();
+        assert!(!allowed.is_empty());
+        assert!(allowed.windows(2).all(|pair| pair[0] < pair[1]));
+    }
+}

core/cpu_allocation/src/allowed_cpus.rs

@@ -22,11 +22,17 @@
 //! the server config (TOML). The config crate re-exports them, and the
 //! `shard_allocator` crate turns them into a real plan. Kept in their
 //! own little crate so neither side has to pull in the other's heavy
-//! dependencies just to share two small enums.
+//! dependencies just to share two small enums. Also home to
+//! [`allowed_cpus`], the probe for the process's allowed CPU set, which
+//! both sides consult when validating and pinning cores.
 
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::str::FromStr;
 
+mod allowed_cpus;
+
+pub use allowed_cpus::allowed_cpus;
+
 /// Tell server how many CPU cores to grab for shards, and how.
 ///
 /// Server make one shard per core. This say which cores. Pick one:

core/cpu_allocation/src/lib.rs

@@ -630,6 +630,15 @@ ports = { tcp = 8091, quic = 8081, http = 3001, websocket = 8093, tcp_replica =
 # TODO(hubcio): revert to "numa:auto" once multi-shard server-ng is stable.
 cpu_allocation = 1
 
+# Whether shard threads are pinned to dedicated CPU cores (default: true).
+# Pinned cores are drawn from the process's allowed CPU set (affinity/cpuset
+# mask), so the server cooperates with systemd `AllowedCPUs=` and container
+# cpusets. Set to false when the server shares cores with other workloads
+# (e.g. a multi-tenant host slicing CPU via cgroup quotas): unpinned shards
+# let the kernel scheduler place threads freely instead of piling every
+# process onto the same low-numbered cores.
+# pin_cores = true
+
 # Wall-clock budget for a single shard's bus drain on shutdown. Drives
 # the per-shard watchdog and the parallel-join survivor path; sized
 # larger than typical TCP RTT times in-flight write-batch so writers

core/server-ng/config.toml

@@ -451,6 +451,35 @@ pub async fn load_config(logging: &mut Logging) -> Result<ServerNgConfig, Server
     Ok(config)
 }
 
+/// Resolve the operator's `cpu_allocation` into concrete shard
+/// assignments plus the checked `u16` shard count.
+///
+/// Shard ids index `ReplicaOwnerTable` slots as `u16`. `OWNER_NONE`
+/// (`u16::MAX`) is reserved as the empty-slot sentinel, so a server
+/// configured with `u16::MAX` shards would mint a shard id that
+/// collides with the sentinel and an owner-table lookup could never
+/// tell that shard apart from an unowned slot. Reject at boot so the
+/// invariant is held by the type system, not by hoping the operator
+/// never configures 65535 cores worth of shards.
+fn resolve_shard_assignments(
+    sharding: &configs::sharding::ShardingConfig,
+) -> Result<(Vec<ShardInfo>, u16), ServerNgError> {
+    let allocator = ShardAllocator::new(&sharding.cpu_allocation, sharding.pin_cores)
+        .map_err(ServerNgError::ShardAllocator)?;
+    let assignments = allocator
+        .to_shard_assignments()
+        .map_err(ServerNgError::ShardAllocator)?;
+    if assignments.is_empty() {
+        return Err(ServerNgError::ShardsCountZero);
+    }
+    match u16::try_from(assignments.len()) {
+        Ok(count) if count < message_bus::OWNER_NONE => Ok((assignments, count)),
+        _ => Err(ServerNgError::ShardsCountOverflow {
+            count: assignments.len(),
+        }),
+    }
+}
+
 /// Re-validate the runtime sharding knobs that the per-shard runtime
 /// consumes directly. Mirrors `ShardingConfig::validate` so a caller
 /// that built the config without running it (e.g. tests, embedded
@@ -523,30 +552,8 @@ pub fn bootstrap(
     current_replica_id: Option<u8>,
 ) -> Result<ShardHandles, ServerNgError> {
     warm_dummy_password_hash();
-    let allocator = ShardAllocator::new(&config.system.sharding.cpu_allocation)
-        .map_err(ServerNgError::ShardAllocator)?;
-    let assignments = allocator
-        .to_shard_assignments()
-        .map_err(ServerNgError::ShardAllocator)?;
+    let (assignments, total_shards) = resolve_shard_assignments(&config.system.sharding)?;
     let shards_count = assignments.len();
-    if shards_count == 0 {
-        return Err(ServerNgError::ShardsCountZero);
-    }
-    // Shard ids index `ReplicaOwnerTable` slots as `u16`. `OWNER_NONE`
-    // (`u16::MAX`) is reserved as the empty-slot sentinel, so a server
-    // configured with `u16::MAX` shards would mint a shard id that
-    // collides with the sentinel and an owner-table lookup could never
-    // tell that shard apart from an unowned slot. Reject at boot so the
-    // invariant is held by the type system above this line, not by hoping
-    // the operator never configures 65535 cores worth of shards.
-    let total_shards = match u16::try_from(shards_count) {
-        Ok(count) if count < message_bus::OWNER_NONE => count,
-        _ => {
-            return Err(ServerNgError::ShardsCountOverflow {
-                count: shards_count,
-            });
-        }
-    };
 
     // Re-check the full valid range, not just the zero floor: a caller
     // that built the config without running `ShardingConfig::validate`

core/server-ng/src/bootstrap.rs

@@ -48,6 +48,7 @@ chrono = { workspace = true }
 clap = { workspace = true }
 compio = { workspace = true }
 configs = { workspace = true }
+cpu_allocation = { workspace = true }
 ctrlc = { workspace = true }
 cyper = { workspace = true }
 cyper-axum = { workspace = true }