use std::borrow::Cow;
use super::super::{AffinityIntent, WorkType};
use super::{MemPolicy, MpolFlags, SchedPolicy};
pub(crate) fn validate_task_comm_string(field: &str, name: &str) {
assert!(
!name.is_empty(),
"{field}: empty string rejected — use `None` (default) for no override, not an empty value",
);
assert!(
!name.contains('\0'),
"{field}: string {name:?} contains an interior NUL byte; \
prctl(PR_SET_NAME) treats it as a C string and would \
truncate at the NUL — strip it before calling .{field}()",
);
assert!(
name.len() <= 15,
"{field}: name {name:?} is {} bytes; kernel TASK_COMM_LEN \
limit is 15 bytes (TASK_COMM_LEN-1=15 in include/linux/sched.h; \
`__set_task_comm` truncates at that cap) — shorten before \
calling .{field}()",
name.len(),
);
}
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(bound(deserialize = ""))]
pub struct WorkSpec {
pub work_type: WorkType,
pub sched_policy: SchedPolicy,
pub num_workers: Option<usize>,
pub affinity: AffinityIntent,
pub mem_policy: MemPolicy,
pub mpol_flags: MpolFlags,
pub nice: Option<i32>,
pub comm: Option<Cow<'static, str>>,
pub pcomm: Option<Cow<'static, str>>,
pub uid: Option<u32>,
pub gid: Option<u32>,
pub numa_node: Option<u32>,
pub workers_pct: Option<f64>,
}
impl Default for WorkSpec {
fn default() -> Self {
Self {
work_type: WorkType::SpinWait,
sched_policy: SchedPolicy::Normal,
num_workers: None,
affinity: AffinityIntent::Inherit,
mem_policy: MemPolicy::Default,
mpol_flags: MpolFlags::NONE,
nice: None,
comm: None,
pcomm: None,
uid: None,
gid: None,
numa_node: None,
workers_pct: None,
}
}
}
impl WorkSpec {
#[must_use = "builder methods consume self; bind the result"]
pub fn workers(mut self, n: usize) -> Self {
self.num_workers = Some(n);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn workers_pct(mut self, pct: f64) -> Self {
assert!(
pct.is_finite() && pct > 0.0,
"WorkSpec::workers_pct({pct}): pct must be finite and > 0.0",
);
self.workers_pct = Some(pct);
self
}
pub(crate) fn resolve_workers_pct(
mut self,
cpuset_cpus: usize,
cgroup_name: &str,
) -> anyhow::Result<Self> {
let Some(pct) = self.workers_pct else {
return Ok(self);
};
if let Some(n) = self.num_workers {
anyhow::bail!(
"cgroup '{}': WorkSpec sets BOTH workers({n}) and \
workers_pct({pct}); pick one — workers_pct resolves the \
cpuset fraction at apply-setup time and is incompatible \
with an explicit count",
cgroup_name,
);
}
let scaled = (cpuset_cpus as f64 * pct).ceil() as usize;
if scaled == 0 {
anyhow::bail!(
"cgroup '{cgroup_name}': workers_pct({pct}) on a cpuset of \
{cpuset_cpus} CPU(s) resolved to 0 workers \
(ceil({cpuset_cpus} * {pct}) = 0); the cgroup would \
have no workers and downstream assertions would \
vacuously pass — narrow the cpuset, raise the fraction, \
or use `workers(N)` instead",
);
}
self.num_workers = Some(scaled);
self.workers_pct = None;
Ok(self)
}
#[must_use = "builder methods consume self; bind the result"]
pub fn work_type(mut self, wt: WorkType) -> Self {
self.work_type = wt;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn sched_policy(mut self, p: SchedPolicy) -> Self {
self.sched_policy = p;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn affinity(mut self, a: AffinityIntent) -> Self {
self.affinity = a;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn mem_policy(mut self, p: MemPolicy) -> Self {
self.mem_policy = p;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn mpol_flags(mut self, f: MpolFlags) -> Self {
self.mpol_flags = f;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn nice(mut self, n: i32) -> Self {
self.nice = Some(n);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn comm(mut self, name: impl Into<Cow<'static, str>>) -> Self {
let name: Cow<'static, str> = name.into();
validate_task_comm_string("WorkSpec::comm", &name);
self.comm = Some(name);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn uid(mut self, uid: u32) -> Self {
self.uid = Some(uid);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn gid(mut self, gid: u32) -> Self {
self.gid = Some(gid);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn numa_node(mut self, node: u32) -> Self {
self.numa_node = Some(node);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn pcomm(mut self, name: impl Into<Cow<'static, str>>) -> Self {
let name: Cow<'static, str> = name.into();
validate_task_comm_string("WorkSpec::pcomm", &name);
self.pcomm = Some(name);
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[should_panic(expected = "WorkSpec::comm: empty string rejected")]
fn work_spec_comm_rejects_empty() {
let _ = WorkSpec::default().comm("");
}
#[test]
#[should_panic(expected = "interior NUL byte")]
fn work_spec_comm_rejects_interior_nul() {
let _ = WorkSpec::default().comm("foo\0bar");
}
#[test]
#[should_panic(expected = "WorkSpec::pcomm: empty string rejected")]
fn work_spec_pcomm_rejects_empty() {
let _ = WorkSpec::default().pcomm("");
}
#[test]
#[should_panic(expected = "interior NUL byte")]
fn work_spec_pcomm_rejects_interior_nul() {
let _ = WorkSpec::default().pcomm("foo\0bar");
}
#[test]
fn work_spec_comm_accepts_15_byte_boundary() {
let fifteen = "a".repeat(15);
let spec = WorkSpec::default().comm(fifteen.clone());
assert_eq!(spec.comm.as_deref(), Some(fifteen.as_str()));
}
#[test]
#[should_panic(expected = "WorkSpec::comm: name")]
fn work_spec_comm_rejects_16_byte_overflow() {
let _ = WorkSpec::default().comm("a".repeat(16));
}
#[test]
fn work_spec_pcomm_accepts_15_byte_boundary() {
let fifteen = "a".repeat(15);
let spec = WorkSpec::default().pcomm(fifteen.clone());
assert_eq!(spec.pcomm.as_deref(), Some(fifteen.as_str()));
}
#[test]
#[should_panic(expected = "WorkSpec::pcomm: name")]
fn work_spec_pcomm_rejects_16_byte_overflow() {
let _ = WorkSpec::default().pcomm("a".repeat(16));
}
#[test]
fn work_spec_comm_accepts_10_byte_utf8_within_cap() {
let cyr10 = "приве";
assert_eq!(
cyr10.len(),
10,
"test fixture: cyrillic 5-char must be 10 bytes"
);
let spec = WorkSpec::default().comm(cyr10);
assert_eq!(spec.comm.as_deref(), Some(cyr10));
}
#[test]
#[should_panic(expected = "16 bytes")]
fn work_spec_comm_rejects_16_byte_utf8_overflow() {
let cyr16 = "приветик";
assert_eq!(
cyr16.len(),
16,
"test fixture: cyrillic 8-char must be 16 bytes"
);
let _ = WorkSpec::default().comm(cyr16);
}
}