use std::collections::BTreeMap;
use std::fmt;
use std::str::FromStr;
use mongreldb_types::ids::{NodeId, QueryId};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub enum WorkloadClass {
Control,
Replication,
Oltp,
InteractiveSql,
AiRetrieval,
Analytics,
Maintenance,
Backup,
}
impl WorkloadClass {
pub const ALL: [WorkloadClass; 8] = [
WorkloadClass::Control,
WorkloadClass::Replication,
WorkloadClass::Oltp,
WorkloadClass::InteractiveSql,
WorkloadClass::AiRetrieval,
WorkloadClass::Analytics,
WorkloadClass::Maintenance,
WorkloadClass::Backup,
];
pub fn name(self) -> &'static str {
match self {
WorkloadClass::Control => "control",
WorkloadClass::Replication => "replication",
WorkloadClass::Oltp => "oltp",
WorkloadClass::InteractiveSql => "interactive_sql",
WorkloadClass::AiRetrieval => "ai_retrieval",
WorkloadClass::Analytics => "analytics",
WorkloadClass::Maintenance => "maintenance",
WorkloadClass::Backup => "backup",
}
}
pub fn has_reserved_capacity(self) -> bool {
matches!(self, WorkloadClass::Control | WorkloadClass::Replication)
}
pub fn is_low_priority(self) -> bool {
matches!(
self,
WorkloadClass::Analytics | WorkloadClass::Maintenance | WorkloadClass::Backup
)
}
}
impl fmt::Display for WorkloadClass {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ResourceError {
#[error("resource group name must not be empty")]
EmptyName,
#[error("resource group {name:?}: cpu_weight must be nonzero (§4.9 bounded CPU/work)")]
ZeroCpuWeight {
name: String,
},
#[error("resource group {name:?}: work_units must be nonzero (§4.9 bounded CPU/work)")]
ZeroWorkUnits {
name: String,
},
#[error(
"resource group {name:?} is pinned: control and replication capacity is reserved (§13.1)"
)]
Pinned {
name: String,
},
#[error("resource group {name:?} not found")]
NotFound {
name: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceGroup {
pub name: String,
pub max_concurrency: usize,
pub max_queue: usize,
pub memory_bytes: u64,
pub temporary_disk_bytes: u64,
pub work_units: u64,
pub cpu_weight: u32,
pub priority: u8,
pub max_result_bytes: u64,
}
impl ResourceGroup {
pub fn validate(&self) -> Result<(), ResourceError> {
if self.name.is_empty() {
return Err(ResourceError::EmptyName);
}
if self.cpu_weight == 0 {
return Err(ResourceError::ZeroCpuWeight {
name: self.name.clone(),
});
}
if self.work_units == 0 {
return Err(ResourceError::ZeroWorkUnits {
name: self.name.clone(),
});
}
Ok(())
}
pub fn for_class(class: WorkloadClass) -> Self {
let (name, max_concurrency, max_queue, memory_mib, work_units, cpu_weight, priority) =
match class {
WorkloadClass::Control => {
("control", 8usize, 64usize, 256u64, 1 << 20, 256u32, 255u8)
}
WorkloadClass::Replication => ("replication", 8, 64, 512, 1 << 20, 256, 254),
WorkloadClass::Oltp => ("oltp", 64, 256, 1024, 1 << 20, 128, 200),
WorkloadClass::InteractiveSql => {
("interactive_sql", 16, 64, 1024, 1 << 19, 64, 180)
}
WorkloadClass::AiRetrieval => ("ai_retrieval", 8, 32, 1024, 1 << 18, 32, 150),
WorkloadClass::Analytics => ("analytics", 4, 16, 2048, 1 << 17, 16, 100),
WorkloadClass::Maintenance => ("maintenance", 2, 8, 512, 1 << 16, 8, 50),
WorkloadClass::Backup => ("backup", 1, 4, 256, 1 << 16, 4, 40),
};
Self {
name: name.to_string(),
max_concurrency,
max_queue,
memory_bytes: memory_mib * 1024 * 1024,
temporary_disk_bytes: 1024 * 1024 * 1024,
work_units,
cpu_weight,
priority,
max_result_bytes: 1024 * 1024 * 1024,
}
}
}
#[repr(transparent)]
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct TenantId(pub u64);
impl TenantId {
pub const ZERO: Self = Self(0);
pub const fn new(value: u64) -> Self {
Self(value)
}
pub const fn get(self) -> u64 {
self.0
}
}
impl fmt::Display for TenantId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl fmt::Debug for TenantId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "TenantId({})", self.0)
}
}
impl FromStr for TenantId {
type Err = std::num::ParseIntError;
fn from_str(text: &str) -> Result<Self, Self::Err> {
text.parse::<u64>().map(Self)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeHandle {
id: NodeId,
}
impl NodeHandle {
pub fn new(id: NodeId) -> Self {
Self { id }
}
pub fn id(self) -> NodeId {
self.id
}
pub fn tenant(self, tenant: TenantId) -> TenantHandle {
TenantHandle {
node: self,
id: tenant,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TenantHandle {
node: NodeHandle,
id: TenantId,
}
impl TenantHandle {
pub fn node(self) -> NodeHandle {
self.node
}
pub fn id(self) -> TenantId {
self.id
}
pub fn group(self, name: impl Into<String>) -> ResourceGroupHandle {
ResourceGroupHandle {
tenant: self,
name: name.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ResourceGroupHandle {
tenant: TenantHandle,
name: String,
}
impl ResourceGroupHandle {
pub fn tenant(&self) -> TenantHandle {
self.tenant
}
pub fn node(&self) -> NodeHandle {
self.tenant.node()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn query(&self, id: QueryId, class: WorkloadClass) -> QueryHandle {
QueryHandle {
group: self.clone(),
id,
class,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct QueryHandle {
group: ResourceGroupHandle,
id: QueryId,
class: WorkloadClass,
}
impl QueryHandle {
pub fn group(&self) -> &ResourceGroupHandle {
&self.group
}
pub fn id(&self) -> QueryId {
self.id
}
pub fn class(&self) -> WorkloadClass {
self.class
}
pub fn tenant(&self) -> TenantHandle {
self.group.tenant()
}
pub fn node(&self) -> NodeHandle {
self.group.node()
}
}
pub struct ResourceGroupRegistry {
groups: parking_lot::RwLock<BTreeMap<String, ResourceGroup>>,
}
const PINNED_GROUP_NAMES: [&str; 2] = ["control", "replication"];
impl ResourceGroupRegistry {
pub fn new() -> Self {
Self {
groups: parking_lot::RwLock::new(BTreeMap::new()),
}
}
pub fn with_defaults() -> Self {
let registry = Self::new();
for class in WorkloadClass::ALL {
let group = ResourceGroup::for_class(class);
registry.groups.write().insert(group.name.clone(), group);
}
registry
}
pub fn register(&self, group: ResourceGroup) -> Result<(), ResourceError> {
group.validate()?;
if PINNED_GROUP_NAMES.contains(&group.name.as_str())
&& (group.max_concurrency == 0 || group.memory_bytes == 0)
{
return Err(ResourceError::Pinned {
name: group.name.clone(),
});
}
self.groups.write().insert(group.name.clone(), group);
Ok(())
}
pub fn get(&self, name: &str) -> Option<ResourceGroup> {
self.groups.read().get(name).cloned()
}
pub fn lookup(&self, name: &str) -> Result<ResourceGroup, ResourceError> {
self.get(name).ok_or_else(|| ResourceError::NotFound {
name: name.to_string(),
})
}
pub fn remove(&self, name: &str) -> Result<bool, ResourceError> {
if PINNED_GROUP_NAMES.contains(&name) {
return Err(ResourceError::Pinned {
name: name.to_string(),
});
}
Ok(self.groups.write().remove(name).is_some())
}
pub fn names(&self) -> Vec<String> {
self.groups.read().keys().cloned().collect()
}
pub fn len(&self) -> usize {
self.groups.read().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl Default for ResourceGroupRegistry {
fn default() -> Self {
Self::with_defaults()
}
}
impl fmt::Debug for ResourceGroupRegistry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ResourceGroupRegistry")
.field("groups", &*self.groups.read())
.finish()
}
}
impl Serialize for ResourceGroupRegistry {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.groups.read().serialize(serializer)
}
}
impl<'de> Deserialize<'de> for ResourceGroupRegistry {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let groups = BTreeMap::<String, ResourceGroup>::deserialize(deserializer)?;
let registry = Self::new();
for (_, group) in groups {
registry.register(group).map_err(serde::de::Error::custom)?;
}
for pinned in PINNED_GROUP_NAMES {
if registry.get(pinned).is_none() {
return Err(serde::de::Error::custom(format!(
"pinned resource group {pinned:?} missing (§13.1 reserved capacity)"
)));
}
}
Ok(registry)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn valid_group(name: &str) -> ResourceGroup {
ResourceGroup {
name: name.to_string(),
max_concurrency: 4,
max_queue: 16,
memory_bytes: 1 << 20,
temporary_disk_bytes: 1 << 20,
work_units: 100,
cpu_weight: 1,
priority: 128,
max_result_bytes: 1 << 20,
}
}
#[test]
fn workload_class_set_matches_spec() {
assert_eq!(WorkloadClass::ALL.len(), 8);
let names: Vec<_> = WorkloadClass::ALL.iter().map(|c| c.name()).collect();
assert_eq!(
names,
vec![
"control",
"replication",
"oltp",
"interactive_sql",
"ai_retrieval",
"analytics",
"maintenance",
"backup"
]
);
let mut deduped = names.clone();
deduped.sort_unstable();
deduped.dedup();
assert_eq!(deduped.len(), 8);
}
#[test]
fn control_and_replication_are_the_reserved_classes() {
for class in WorkloadClass::ALL {
assert_eq!(
class.has_reserved_capacity(),
matches!(class, WorkloadClass::Control | WorkloadClass::Replication),
"reserved capacity: {class}"
);
}
}
#[test]
fn low_priority_classes_are_deferrable() {
for class in WorkloadClass::ALL {
assert_eq!(
class.is_low_priority(),
matches!(
class,
WorkloadClass::Analytics | WorkloadClass::Maintenance | WorkloadClass::Backup
),
"low priority: {class}"
);
}
}
#[test]
fn workload_class_serde_round_trip() {
for class in WorkloadClass::ALL {
let json = serde_json::to_string(&class).unwrap();
let back: WorkloadClass = serde_json::from_str(&json).unwrap();
assert_eq!(back, class);
}
assert_eq!(
serde_json::to_string(&WorkloadClass::AiRetrieval).unwrap(),
"\"AiRetrieval\""
);
}
#[test]
fn validation_accepts_a_valid_group() {
assert!(valid_group("g").validate().is_ok());
for priority in [0u8, 255] {
let mut g = valid_group("g");
g.priority = priority;
assert!(g.validate().is_ok());
}
}
#[test]
fn validation_rejects_empty_name() {
let mut g = valid_group("");
g.name.clear();
assert_eq!(g.validate(), Err(ResourceError::EmptyName));
}
#[test]
fn validation_rejects_zero_weights() {
let mut g = valid_group("g");
g.cpu_weight = 0;
assert_eq!(
g.validate(),
Err(ResourceError::ZeroCpuWeight { name: "g".into() })
);
let mut g = valid_group("g");
g.work_units = 0;
assert_eq!(
g.validate(),
Err(ResourceError::ZeroWorkUnits { name: "g".into() })
);
}
#[test]
fn resource_group_serde_round_trip_and_priority_range() {
let g = valid_group("oltp");
let json = serde_json::to_string(&g).unwrap();
let back: ResourceGroup = serde_json::from_str(&json).unwrap();
assert_eq!(back, g);
let too_big = json.replace("\"priority\":128", "\"priority\":256");
assert!(serde_json::from_str::<ResourceGroup>(&too_big).is_err());
}
#[test]
fn registry_defaults_cover_every_class() {
let registry = ResourceGroupRegistry::with_defaults();
assert_eq!(registry.len(), WorkloadClass::ALL.len());
for class in WorkloadClass::ALL {
let group = registry
.get(ResourceGroup::for_class(class).name.as_str())
.unwrap_or_else(|| panic!("default group for {class}"));
assert!(group.validate().is_ok());
assert!(group.max_concurrency > 0);
assert!(group.memory_bytes > 0);
}
assert!(registry.lookup("no-such-group").is_err());
}
#[test]
fn control_and_replication_defaults_have_top_priority() {
let registry = ResourceGroupRegistry::with_defaults();
let control = registry.lookup("control").unwrap();
let replication = registry.lookup("replication").unwrap();
assert_eq!(control.priority, 255);
assert_eq!(replication.priority, 254);
for name in registry.names() {
if name != "control" && name != "replication" {
let g = registry.lookup(&name).unwrap();
assert!(
g.priority < replication.priority,
"{name} outranks replication"
);
}
}
}
#[test]
fn pinned_groups_cannot_be_removed_or_starved() {
let registry = ResourceGroupRegistry::with_defaults();
for pinned in ["control", "replication"] {
assert_eq!(
registry.remove(pinned),
Err(ResourceError::Pinned {
name: pinned.into()
})
);
let mut starved = registry.lookup(pinned).unwrap();
starved.max_concurrency = 0;
assert_eq!(
registry.register(starved),
Err(ResourceError::Pinned {
name: pinned.into()
})
);
let mut starved = registry.lookup(pinned).unwrap();
starved.memory_bytes = 0;
assert_eq!(
registry.register(starved),
Err(ResourceError::Pinned {
name: pinned.into()
})
);
let mut kept = registry.lookup(pinned).unwrap();
kept.max_concurrency += 1;
assert!(registry.register(kept.clone()).is_ok());
assert_eq!(registry.lookup(pinned).unwrap(), kept);
}
registry.register(valid_group("scratch")).unwrap();
assert_eq!(registry.remove("scratch"), Ok(true));
assert_eq!(registry.remove("scratch"), Ok(false));
}
#[test]
fn register_validates_groups() {
let registry = ResourceGroupRegistry::new();
let mut bad = valid_group("bad");
bad.cpu_weight = 0;
assert!(matches!(
registry.register(bad),
Err(ResourceError::ZeroCpuWeight { .. })
));
assert!(registry.get("bad").is_none());
}
#[test]
fn registry_serde_round_trip_preserves_groups() {
let registry = ResourceGroupRegistry::with_defaults();
registry.register(valid_group("tenant_a")).unwrap();
let json = serde_json::to_string(®istry).unwrap();
let back: ResourceGroupRegistry = serde_json::from_str(&json).unwrap();
assert_eq!(back.names(), registry.names());
assert_eq!(back.lookup("tenant_a").unwrap(), valid_group("tenant_a"));
assert_eq!(serde_json::to_string(&back).unwrap(), json);
}
#[test]
fn registry_deserialization_re_enforces_pinned_defaults() {
let bare = serde_json::json!({ "oltp": valid_group("oltp") });
assert!(serde_json::from_value::<ResourceGroupRegistry>(bare).is_err());
let mut starved = valid_group("control");
starved.memory_bytes = 0;
let doc =
serde_json::json!({ "control": starved, "replication": valid_group("replication") });
assert!(serde_json::from_value::<ResourceGroupRegistry>(doc).is_err());
}
#[test]
fn hierarchy_handles_walk_up_the_chain() {
let node = NodeHandle::new(NodeId::new_random());
let tenant = node.tenant(TenantId::new(7));
let group = tenant.group("oltp");
let query_id = QueryId::new_random();
let query = group.query(query_id, WorkloadClass::Oltp);
assert_eq!(tenant.node(), node);
assert_eq!(tenant.id(), TenantId::new(7));
assert_eq!(group.tenant(), tenant);
assert_eq!(group.node(), node);
assert_eq!(group.name(), "oltp");
assert_eq!(query.group(), &group);
assert_eq!(query.id(), query_id);
assert_eq!(query.class(), WorkloadClass::Oltp);
assert_eq!(query.tenant(), tenant);
assert_eq!(query.node(), node);
assert_eq!(node.id(), node.id());
}
#[test]
fn hierarchy_handles_serde_round_trip() {
let node = NodeHandle::new(NodeId::from_bytes([0xAB; 16]));
let query = node
.tenant(TenantId::new(42))
.group("analytics")
.query(QueryId::from_bytes([0xCD; 16]), WorkloadClass::Analytics);
let json = serde_json::to_string(&query).unwrap();
let back: QueryHandle = serde_json::from_str(&json).unwrap();
assert_eq!(back, query);
assert_eq!(back.node().id(), NodeId::from_bytes([0xAB; 16]));
assert_eq!(back.tenant().id(), TenantId::new(42));
}
#[test]
fn tenant_id_forms() {
let id = TenantId::new(123);
assert_eq!(id.get(), 123);
assert_eq!(id.to_string(), "123");
assert_eq!(format!("{id:?}"), "TenantId(123)");
assert_eq!("123".parse::<TenantId>().unwrap(), id);
assert_eq!(TenantId::ZERO.get(), 0);
}
}