You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iggy.apache.org by ma...@apache.org on 2026/07/06 09:05:06 UTC
(iggy) 01/01: feat(server): respect allowed CPU set and add pin_cores toggle
This is an automated email from the ASF dual-hosted git repository.
mmodzelewski pushed a commit to branch cpu_pinning_config
in repository https://gitbox.apache.org/repos/asf/iggy.git
commit cb9859fadf18fbeef365acb8c125ba33118f52ae
Author: Maciej Modzelewski <ma...@modzelewski.dev>
AuthorDate: Mon Jul 6 09:59:21 2026 +0200
feat(server): respect allowed CPU set and add pin_cores toggle
Shard threads pinned to absolute cores 0..n, silently violating any
affinity/cpuset restriction from the parent environment (systemd
AllowedCPUs=, container cpusets, taskset): sched_setaffinity to a
core outside the allowed set fails with EINVAL. On shared hosts it
also piled every process onto the same low-numbered cores while the
rest sat idle.
Draw pinned cores from the process's allowed CPU set instead, and
fail fast when an explicit range names a core outside it. Add
pin_cores (default true); set false on multi-tenant hosts to drop
pinning entirely so the kernel scheduler places shard threads
freely. With a NUMA-aware allocation, false also drops the memory-
node bindings and warns, since NUMA placement without pinning is
meaningless.
Also report cgroup memory limits in shard stats when tighter than
the host, so a quota-restricted server reports its real ceiling.
---
core/common/src/types/stats/mod.rs | 5 +-
core/configs/src/server_config/sharding.rs | 17 +++
core/configs/src/server_config/validators.rs | 105 ++++++++++++-
core/server-ng/config.toml | 9 ++
core/server-ng/src/bootstrap.rs | 3 +-
core/server/config.toml | 8 +
core/server/src/main.rs | 5 +-
core/server/src/shard/system/stats.rs | 8 +
core/server/src/shard_allocator.rs | 191 ++++++++++++++++++++++--
core/server_common/src/sharding/allowed_cpus.rs | 59 ++++++++
core/server_common/src/sharding/mod.rs | 2 +
11 files changed, 390 insertions(+), 22 deletions(-)
diff --git a/core/common/src/types/stats/mod.rs b/core/common/src/types/stats/mod.rs
index fa848858d..8aee40a73 100644
--- a/core/common/src/types/stats/mod.rs
+++ b/core/common/src/types/stats/mod.rs
@@ -30,9 +30,10 @@ pub struct Stats {
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,
diff --git a/core/configs/src/server_config/sharding.rs b/core/configs/src/server_config/sharding.rs
index ec332462b..9602ee208 100644
--- a/core/configs/src/server_config/sharding.rs
+++ b/core/configs/src/server_config/sharding.rs
@@ -96,6 +96,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)
}
@@ -114,6 +118,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:
@@ -171,6 +187,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(),
diff --git a/core/configs/src/server_config/validators.rs b/core/configs/src/server_config/validators.rs
index 93c61fb9b..151c7516f 100644
--- a/core/configs/src/server_config/validators.rs
+++ b/core/configs/src/server_config/validators.rs
@@ -33,6 +33,7 @@ use iggy_common::CompressionAlgorithm;
use iggy_common::IggyExpiry;
use iggy_common::MaxTopicSize;
use iggy_common::Validatable;
+use server_common::sharding::allowed_cpus;
use std::thread::available_parallelism;
use tracing::warn;
@@ -488,9 +489,41 @@ impl Validatable<ConfigurationError> for ShardingConfig {
);
return Err(ConfigurationError::InvalidConfigurationValue);
}
- if *end > available_cpus {
+ if !self.pin_cores {
+ // Unpinned shards never bind to the named cores; the
+ // range only sets the shard count, so cap it like `Count`.
+ if *end - *start > available_cpus {
+ eprintln!(
+ "Invalid sharding configuration: cpu_allocation range {start}..{end} yields {} shards, exceeding available CPU cores {available_cpus}",
+ *end - *start
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ }
+ return Ok(());
+ }
+ // A pinned range still must not spawn more shards than the
+ // effective parallelism: a cgroup CPU quota shrinks
+ // `available_parallelism` without shrinking the affinity
+ // mask, so membership alone would let e.g. 0..64 pass on a
+ // 2-CPU-quota host. Compare counts, not core ids, so a
+ // high-numbered mask (`taskset -c 8-15` with range 8..16)
+ // still validates.
+ 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);
+ }
+ // Pinned cores are named exactly: check membership in the
+ // allowed CPU set (affinity/cpuset mask) rather than against
+ // `available_parallelism`, which under-counts when the mask
+ // names high-numbered cores (e.g. `taskset -c 8-15` yields a
+ // parallelism of 8 while core ids up to 15 are valid).
+ 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 +1041,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());
+ }
+}
diff --git a/core/server-ng/config.toml b/core/server-ng/config.toml
index 11de70d34..70af981f5 100644
--- a/core/server-ng/config.toml
+++ b/core/server-ng/config.toml
@@ -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
diff --git a/core/server-ng/src/bootstrap.rs b/core/server-ng/src/bootstrap.rs
index 46bd45c29..164f81d99 100644
--- a/core/server-ng/src/bootstrap.rs
+++ b/core/server-ng/src/bootstrap.rs
@@ -514,7 +514,8 @@ pub fn bootstrap(
config: ServerNgConfig,
current_replica_id: Option<u8>,
) -> Result<ShardHandles, ServerNgError> {
- let allocator = ShardAllocator::new(&config.system.sharding.cpu_allocation)
+ let sharding = &config.system.sharding;
+ let allocator = ShardAllocator::new(&sharding.cpu_allocation, sharding.pin_cores)
.map_err(ServerNgError::ShardAllocator)?;
let assignments = allocator
.to_shard_assignments()
diff --git a/core/server/config.toml b/core/server/config.toml
index d6d0bdf24..b31bf8a1c 100644
--- a/core/server/config.toml
+++ b/core/server/config.toml
@@ -592,6 +592,14 @@ ports = { tcp = 8091, quic = 8081, http = 3001, websocket = 8093, tcp_replica =
# + "numa:auto": Use all available numa node, cores
# + "numa:nodes=0,1;cores=4;no_ht=true": Use NUMA node 0 and 1, each nodes use 4 cores, and no hyperthreads
cpu_allocation = "numa:auto"
+# 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
[websocket]
enabled = true
diff --git a/core/server/src/main.rs b/core/server/src/main.rs
index d9baaac85..9bd4758ec 100644
--- a/core/server/src/main.rs
+++ b/core/server/src/main.rs
@@ -301,7 +301,10 @@ fn main() -> Result<(), ServerError> {
);
// ELEVENTH DISCRETE LOADING STEP.
- let shard_allocator = ShardAllocator::new(&config.system.sharding.cpu_allocation)?;
+ let shard_allocator = ShardAllocator::new(
+ &config.system.sharding.cpu_allocation,
+ config.system.sharding.pin_cores,
+ )?;
let shard_assignment = shard_allocator.to_shard_assignments()?;
#[cfg(feature = "disable-mimalloc")]
diff --git a/core/server/src/shard/system/stats.rs b/core/server/src/shard/system/stats.rs
index dc1e92e07..cb542256b 100644
--- a/core/server/src/shard/system/stats.rs
+++ b/core/server/src/shard/system/stats.rs
@@ -88,6 +88,14 @@ impl IggyShard {
stats.written_bytes = disk_usage.total_written_bytes.into();
stats.threads_count = process.tasks().map(|t| t.len() as u32).unwrap_or(0);
+
+ if let Some(limits) = process
+ .cgroup_limits()
+ .filter(|limits| limits.total_memory < sys.total_memory())
+ {
+ stats.total_memory = limits.total_memory.into();
+ stats.available_memory = limits.free_memory.into();
+ }
}
let (streams_count, topics_count, partitions_count, consumer_groups_count, stream_ids) =
diff --git a/core/server/src/shard_allocator.rs b/core/server/src/shard_allocator.rs
index 26a1f0e38..94fe6221e 100644
--- a/core/server/src/shard_allocator.rs
+++ b/core/server/src/shard_allocator.rs
@@ -23,6 +23,7 @@ use hwlocality::memory::binding::{MemoryBindingFlags, MemoryBindingPolicy};
use hwlocality::object::types::ObjectType::{self, NUMANode};
#[cfg(target_os = "linux")]
use nix::{sched::sched_setaffinity, unistd::Pid};
+use server_common::sharding::allowed_cpus;
use std::collections::HashSet;
use std::sync::Arc;
use std::thread::available_parallelism;
@@ -49,6 +50,16 @@ pub enum ShardingError {
node: usize,
},
+ #[error(
+ "Requested {requested} shard(s) but only {available} CPU core(s) are allowed for this process (affinity/cpuset mask)"
+ )]
+ AllowedCpusExceeded { requested: usize, available: usize },
+
+ #[error(
+ "Configured CPU {cpu} is outside the set of cores allowed for this process (affinity/cpuset mask)"
+ )]
+ CpuNotAllowed { cpu: usize },
+
#[error("Invalid NUMA node: requested {requested}, only available {available} node")]
InvalidNode { requested: usize, available: usize },
@@ -230,13 +241,50 @@ impl ShardInfo {
}
}
+/// One shard per core drawn from `allowed`, each pinned to its own core.
+fn pinned_from_allowed(allowed: &[usize], count: usize) -> Result<Vec<ShardInfo>, ShardingError> {
+ if count > allowed.len() {
+ return Err(ShardingError::AllowedCpusExceeded {
+ requested: count,
+ available: allowed.len(),
+ });
+ }
+
+ Ok(allowed
+ .iter()
+ .take(count)
+ .map(|&cpu_id| ShardInfo {
+ cpu_set: HashSet::from([cpu_id]),
+ numa_node: None,
+ })
+ .collect())
+}
+
+/// `count` shards with no CPU affinity: `bind_cpu` becomes a no-op and the
+/// kernel scheduler places shard threads freely. This is the right mode when
+/// the process shares its cores with other workloads (e.g. a multi-tenant
+/// host slicing CPU via cgroup quotas), where per-process pinning to the same
+/// low-numbered cores would pile every tenant onto one core.
+fn unpinned(count: usize) -> Vec<ShardInfo> {
+ (0..count)
+ .map(|_| ShardInfo {
+ cpu_set: HashSet::new(),
+ numa_node: None,
+ })
+ .collect()
+}
+
pub struct ShardAllocator {
allocation: CpuAllocation,
+ pin_cores: bool,
topology: Option<Arc<NumaTopology>>,
}
impl ShardAllocator {
- pub fn new(allocation: &CpuAllocation) -> Result<ShardAllocator, ShardingError> {
+ pub fn new(
+ allocation: &CpuAllocation,
+ pin_cores: bool,
+ ) -> Result<ShardAllocator, ShardingError> {
let topology = if matches!(allocation, CpuAllocation::NumaAware(_)) {
let numa_topology = NumaTopology::detect()?;
@@ -247,6 +295,7 @@ impl ShardAllocator {
Ok(Self {
allocation: allocation.clone(),
+ pin_cores,
topology,
})
}
@@ -254,39 +303,66 @@ impl ShardAllocator {
pub fn to_shard_assignments(&self) -> Result<Vec<ShardInfo>, ShardingError> {
match &self.allocation {
CpuAllocation::All => {
+ // `available_parallelism` already accounts for both the
+ // affinity mask and any cgroup CPU quota, so a
+ // quota-restricted process gets proportionally fewer shards.
let available_cpus = available_parallelism()
.map_err(|err| ShardingError::Other {
msg: format!("Failed to get available_parallelism: {:?}", err),
})?
.get();
- let shard_assignments: Vec<_> = (0..available_cpus)
- .map(|cpu_id| ShardInfo {
- cpu_set: HashSet::from([cpu_id]),
- numa_node: None,
- })
- .collect();
+ if !self.pin_cores {
+ info!("Using all available CPU cores ({available_cpus} shards, unpinned)");
+ return Ok(unpinned(available_cpus));
+ }
+
+ let allowed = allowed_cpus();
+ let shard_assignments =
+ pinned_from_allowed(&allowed, available_cpus.min(allowed.len()))?;
info!(
- "Using all available CPU cores ({} shards with affinity)",
- shard_assignments.len()
+ "Using all available CPU cores ({} shards pinned within allowed set {:?})",
+ shard_assignments.len(),
+ allowed
);
Ok(shard_assignments)
}
CpuAllocation::Count(count) => {
- let shard_assignments = (0..*count)
- .map(|cpu_id| ShardInfo {
- cpu_set: HashSet::from([cpu_id]),
- numa_node: None,
- })
+ if !self.pin_cores {
+ info!("Using {count} shard(s), unpinned");
+ return Ok(unpinned(*count));
+ }
+
+ let allowed = allowed_cpus();
+ let shard_assignments = pinned_from_allowed(&allowed, *count)?;
+ let cores: Vec<usize> = shard_assignments
+ .iter()
+ .flat_map(|shard| shard.cpu_set.iter().copied())
.collect();
- info!("Using {count} shards with affinity to cores 0..{count}");
+ info!("Using {count} shard(s) with affinity to cores {cores:?}");
Ok(shard_assignments)
}
CpuAllocation::Range(start, end) => {
+ if !self.pin_cores {
+ info!(
+ "Using {} shard(s) for range {start}..{end}, unpinned",
+ end - start
+ );
+ return Ok(unpinned(end - start));
+ }
+
+ // An explicit range is operator intent for exact cores; honor
+ // it as-is but fail fast when a core is outside the allowed
+ // set instead of letting `sched_setaffinity` EINVAL later.
+ let allowed = allowed_cpus();
+ if let Some(cpu) = (*start..*end).find(|cpu| !allowed.contains(cpu)) {
+ return Err(ShardingError::CpuNotAllowed { cpu });
+ }
+
let shard_assignments = (*start..*end)
.map(|cpu_id| ShardInfo {
cpu_set: HashSet::from([cpu_id]),
@@ -303,7 +379,20 @@ impl ShardAllocator {
}
CpuAllocation::NumaAware(numa_config) => {
let topology = self.topology.as_ref().ok_or(ShardingError::NoTopology)?;
- self.compute_numa_assignments(topology, numa_config)
+ let assignments = self.compute_numa_assignments(topology, numa_config)?;
+
+ if !self.pin_cores {
+ // NUMA-aware placement is meaningless without pinning: an
+ // unpinned thread migrates across nodes and a bound memory
+ // policy would then hurt. Keep the shard count, drop the
+ // bindings.
+ tracing::warn!(
+ "pin_cores = false with a NUMA-aware cpu_allocation: NUMA bindings are ignored"
+ );
+ return Ok(unpinned(assignments.len()));
+ }
+
+ Ok(assignments)
}
}
}
@@ -375,3 +464,73 @@ impl ShardAllocator {
Ok(shard_infos)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn pinned_from_allowed_draws_cores_from_the_allowed_set() {
+ let allowed = vec![2, 3, 6, 7];
+ let shards = pinned_from_allowed(&allowed, 2).unwrap();
+ assert_eq!(shards.len(), 2);
+ assert_eq!(shards[0].cpu_set, HashSet::from([2]));
+ assert_eq!(shards[1].cpu_set, HashSet::from([3]));
+ assert!(shards.iter().all(|shard| shard.numa_node.is_none()));
+ }
+
+ #[test]
+ fn pinned_from_allowed_rejects_more_shards_than_allowed_cores() {
+ let allowed = vec![0, 1];
+ let err = pinned_from_allowed(&allowed, 3).unwrap_err();
+ assert!(matches!(
+ err,
+ ShardingError::AllowedCpusExceeded {
+ requested: 3,
+ available: 2,
+ }
+ ));
+ }
+
+ #[test]
+ fn unpinned_shards_have_empty_cpu_sets() {
+ let shards = unpinned(4);
+ assert_eq!(shards.len(), 4);
+ assert!(shards.iter().all(|shard| shard.cpu_set.is_empty()));
+ assert!(shards.iter().all(|shard| shard.numa_node.is_none()));
+ }
+
+ #[test]
+ fn count_without_pinning_yields_unpinned_shards() {
+ let allocator = ShardAllocator::new(&CpuAllocation::Count(1), false).unwrap();
+ let shards = allocator.to_shard_assignments().unwrap();
+ assert_eq!(shards.len(), 1);
+ assert!(shards[0].cpu_set.is_empty());
+ }
+
+ #[test]
+ fn count_with_pinning_stays_within_the_allowed_set() {
+ let allocator = ShardAllocator::new(&CpuAllocation::Count(1), true).unwrap();
+ let shards = allocator.to_shard_assignments().unwrap();
+ assert_eq!(shards.len(), 1);
+ let allowed = allowed_cpus();
+ assert!(shards[0].cpu_set.iter().all(|cpu| allowed.contains(cpu)));
+ }
+
+ #[test]
+ fn range_without_pinning_keeps_shard_count() {
+ let allocator = ShardAllocator::new(&CpuAllocation::Range(0, 1), false).unwrap();
+ let shards = allocator.to_shard_assignments().unwrap();
+ assert_eq!(shards.len(), 1);
+ assert!(shards[0].cpu_set.is_empty());
+ }
+
+ #[test]
+ fn bind_cpu_with_empty_set_is_a_no_op() {
+ let shard = ShardInfo {
+ cpu_set: HashSet::new(),
+ numa_node: None,
+ };
+ assert!(shard.bind_cpu().is_ok());
+ }
+}
diff --git a/core/server_common/src/sharding/allowed_cpus.rs b/core/server_common/src/sharding/allowed_cpus.rs
new file mode 100644
index 000000000..df7fe5490
--- /dev/null
+++ b/core/server_common/src/sharding/allowed_cpus.rs
@@ -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]));
+ }
+}
diff --git a/core/server_common/src/sharding/mod.rs b/core/server_common/src/sharding/mod.rs
index 761e41d04..76285d981 100644
--- a/core/server_common/src/sharding/mod.rs
+++ b/core/server_common/src/sharding/mod.rs
@@ -15,11 +15,13 @@
// specific language governing permissions and limitations
// under the License.
+mod allowed_cpus;
mod local_idx;
mod namespace;
mod partition_location;
mod shard_id;
+pub use allowed_cpus::allowed_cpus;
pub use local_idx::LocalIdx;
pub use namespace::{
IggyNamespace, MAX_PARTITIONS, MAX_STREAMS, MAX_TOPICS, METADATA_CONSENSUS_NAMESPACE,