use std::{borrow::Cow, fmt};
pub type StrCow = Cow<'static, str>;
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[repr(C)]
pub enum Resource {
LocalMachine,
CpuPackage { id: u32 },
CpuCore { id: u32 },
Dram { pkg_id: u32 },
Gpu { bus_id: StrCow },
Custom { kind: StrCow, id: StrCow },
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[repr(C)]
pub enum ResourceConsumer {
LocalMachine,
Process { pid: u32 },
ControlGroup { path: StrCow },
Custom { kind: StrCow, id: StrCow },
}
impl Resource {
pub fn custom(kind: impl Into<StrCow>, id: impl Into<StrCow>) -> Resource {
Resource::Custom {
kind: kind.into(),
id: id.into(),
}
}
pub fn kind(&self) -> &str {
match self {
Resource::LocalMachine => "local_machine",
Resource::CpuPackage { .. } => "cpu_package",
Resource::CpuCore { .. } => "cpu_core",
Resource::Dram { .. } => "dram",
Resource::Gpu { .. } => "gpu",
Resource::Custom { kind, id: _ } => kind,
}
}
pub fn id_string(&self) -> Option<String> {
match self {
Resource::LocalMachine => None,
r => Some(r.id_display().to_string()),
}
}
pub fn id_display(&self) -> impl fmt::Display + '_ {
match self {
Resource::LocalMachine => LazyDisplayable::Str(""),
Resource::CpuPackage { id } => LazyDisplayable::U32(*id),
Resource::CpuCore { id } => LazyDisplayable::U32(*id),
Resource::Dram { pkg_id } => LazyDisplayable::U32(*pkg_id),
Resource::Gpu { bus_id } => LazyDisplayable::Str(bus_id),
Resource::Custom { kind: _, id } => LazyDisplayable::Str(id),
}
}
pub fn parse(kind: impl Into<StrCow>, id: impl Into<StrCow>) -> Result<Self, InvalidResourceError> {
Resource::custom(kind, id).normalize()
}
pub fn normalize(self) -> Result<Self, InvalidResourceError> {
match self {
Resource::Custom { kind, id } => match kind.as_ref() {
"local_machine" => {
if id.is_empty() {
Ok(Resource::LocalMachine)
} else {
Err(InvalidResourceError::InvalidId(kind))
}
}
"cpu_package" => {
let id = id.parse().map_err(|_| InvalidResourceError::InvalidId(kind))?;
Ok(Resource::CpuPackage { id })
}
"cpu_core" => {
let id = id.parse().map_err(|_| InvalidResourceError::InvalidId(kind))?;
Ok(Resource::CpuCore { id })
}
"dram" => {
let pkg_id = id.parse().map_err(|_| InvalidResourceError::InvalidId(kind))?;
Ok(Resource::Dram { pkg_id })
}
"gpu" => Ok(Resource::Gpu { bus_id: id }),
_ => Ok(Resource::Custom { kind, id }),
},
r => Ok(r),
}
}
}
impl ResourceConsumer {
pub fn custom(kind: impl Into<StrCow>, id: impl Into<StrCow>) -> ResourceConsumer {
ResourceConsumer::Custom {
kind: kind.into(),
id: id.into(),
}
}
pub fn kind(&self) -> &str {
match self {
ResourceConsumer::LocalMachine => "local_machine",
ResourceConsumer::Process { .. } => "process",
ResourceConsumer::ControlGroup { .. } => "cgroup",
ResourceConsumer::Custom { kind, id: _ } => kind,
}
}
pub fn id_string(&self) -> Option<String> {
match self {
ResourceConsumer::LocalMachine => None,
c => Some(c.id_display().to_string()),
}
}
pub fn id_display(&self) -> impl fmt::Display + '_ {
match self {
ResourceConsumer::LocalMachine => LazyDisplayable::Str(""),
ResourceConsumer::Process { pid } => LazyDisplayable::U32(*pid),
ResourceConsumer::ControlGroup { path } => LazyDisplayable::Str(path),
ResourceConsumer::Custom { kind: _, id } => LazyDisplayable::Str(id),
}
}
pub fn parse(kind: impl Into<StrCow>, id: impl Into<StrCow>) -> Result<Self, InvalidConsumerError> {
ResourceConsumer::custom(kind, id).normalize()
}
pub fn normalize(self) -> Result<Self, InvalidConsumerError> {
match self {
ResourceConsumer::Custom { kind, id } => match kind.as_ref() {
"process" => {
let pid = id.parse().map_err(|_| InvalidConsumerError::InvalidId(kind))?;
Ok(ResourceConsumer::Process { pid })
}
"cgroup" => Ok(ResourceConsumer::ControlGroup { path: id }),
_ => Ok(ResourceConsumer::Custom { kind, id }),
},
r => Ok(r),
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum InvalidResourceError {
#[error("invalid resource identifier for kind {0}")]
InvalidId(StrCow),
}
#[derive(Debug, thiserror::Error)]
pub enum InvalidConsumerError {
#[error("invalid consumer identifier for kind {0}")]
InvalidId(StrCow),
}
enum LazyDisplayable<'a> {
U32(u32),
Str(&'a str),
}
impl<'a> fmt::Display for LazyDisplayable<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LazyDisplayable::U32(id) => write!(f, "{id}"),
LazyDisplayable::Str(id) => write!(f, "{id}"),
}
}
}