[skip ci] refactor and cleanup

This commit is contained in:
Christien Rioux 2024-12-15 20:09:30 -05:00
parent 68fc6f97eb
commit 826f1cc782
12 changed files with 909 additions and 577 deletions

View File

@ -14,202 +14,6 @@ pub enum ConfigError {
ValidateError(validator::ValidationErrors),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WeightedList<T: fmt::Debug + Clone> {
Single(T),
List(Vec<Weighted<T>>),
}
impl<T: fmt::Debug + Clone> Default for WeightedList<T> {
fn default() -> Self {
Self::List(Vec::new())
}
}
impl<T: fmt::Debug + Clone> Validate for WeightedList<T> {
fn validate(&self) -> Result<(), ValidationErrors> {
let mut errors = ValidationErrors::new();
match self {
Self::List(v) => {
if v.is_empty() {
errors.add(
"List",
ValidationError::new("len")
.with_message("weighted list must not be empty".into()),
)
}
errors.merge_self("List", v.validate());
}
Self::Single(_addr) => {}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
}
impl<T: fmt::Debug + Clone> WeightedList<T> {
fn validate_once(&self) -> Result<(), ValidationError> {
match self {
Self::List(v) => {
if v.is_empty() {
return Err(ValidationError::new("len")
.with_message("weighted list must not be empty".into()));
}
}
Self::Single(_addr) => {}
}
Ok(())
}
pub fn try_for_each<E, F: FnMut(&T) -> Result<(), E>>(&self, mut f: F) -> Result<(), E> {
match self {
WeightedList::Single(v) => f(v),
WeightedList::List(vec) => vec
.iter()
.map(|v| match v {
Weighted::Weighted { item, weight: _ } => item,
Weighted::Unweighted(item) => item,
})
.try_for_each(f),
}
}
pub fn filter<F>(&self, mut filter: F) -> Option<WeightedList<T>>
where
F: FnMut(&T) -> bool,
{
match self {
WeightedList::Single(v) => {
if filter(v) {
return Some(self.clone());
}
return None;
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<T>>::with_capacity(vec.len());
for v in vec {
if filter(v.item()) {
out.push(v.clone());
}
}
if out.is_empty() {
None
} else {
Some(WeightedList::List(out))
}
}
}
}
pub fn try_filter<F, E>(&self, mut filter: F) -> Result<Option<WeightedList<T>>, E>
where
F: FnMut(&T) -> Result<bool, E>,
{
match self {
WeightedList::Single(v) => {
if filter(v)? {
return Ok(Some(self.clone()));
}
return Ok(None);
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<T>>::with_capacity(vec.len());
for v in vec {
if filter(v.item())? {
out.push(v.clone());
}
}
if out.is_empty() {
Ok(None)
} else {
Ok(Some(WeightedList::List(out)))
}
}
}
}
pub fn try_filter_map<F, S, E>(&self, mut filter: F) -> Result<Option<WeightedList<S>>, E>
where
F: FnMut(&T) -> Result<Option<S>, E>,
S: fmt::Debug + Clone,
{
match self {
WeightedList::Single(v) => {
if let Some(item) = filter(v)? {
return Ok(Some(WeightedList::Single(item)));
}
return Ok(None);
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<S>>::with_capacity(vec.len());
for v in vec {
if let Some(item) = filter(v.item())? {
out.push(match v {
Weighted::Weighted { item: _, weight } => Weighted::Weighted {
item,
weight: *weight,
},
Weighted::Unweighted(_) => Weighted::Unweighted(item),
});
}
}
if out.is_empty() {
Ok(None)
} else {
Ok(Some(WeightedList::List(out)))
}
}
}
}
}
pub type Probability = f32;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Weighted<T: fmt::Debug + Clone> {
Weighted { item: T, weight: f32 },
Unweighted(T),
}
impl<T: fmt::Debug + Clone> Validate for Weighted<T> {
fn validate(&self) -> Result<(), ValidationErrors> {
let mut errors = ValidationErrors::new();
if let Self::Weighted { item: _, weight } = self {
if *weight <= 0.0 {
errors.add(
"Weighted",
ValidationError::new("len")
.with_message("weight must be a positive value".into()),
)
}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
}
impl<T: fmt::Debug + Clone> Weighted<T> {
pub fn item(&self) -> &T {
match self {
Weighted::Weighted { item, weight: _ } => item,
Weighted::Unweighted(item) => item,
}
}
pub fn weight(&self) -> f32 {
match self {
Weighted::Weighted { item: _, weight } => *weight,
Weighted::Unweighted(_) => 1.0f32,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
#[validate(context = "ValidateContext<'v_a>")]
pub struct Profile {
@ -324,7 +128,7 @@ pub struct TemplateLimits {
#[serde(default)]
pub machine_count: Option<WeightedList<u32>>,
#[validate(nested)]
pub machines_per_network: WeightedList<u32>,
pub machines_per_network: Option<WeightedList<u32>>,
}
fn validate_template_limits(limits: &TemplateLimits) -> Result<(), ValidationError> {
@ -337,13 +141,15 @@ fn validate_template_limits(limits: &TemplateLimits) -> Result<(), ValidationErr
Ok(())
})?;
}
limits.machines_per_network.try_for_each(|x| {
if *x == 0 {
return Err(ValidationError::new("badcount")
.with_message("template limits has zero machines per network count".into()));
}
Ok(())
})?;
if let Some(machines_per_network) = &limits.machines_per_network {
machines_per_network.try_for_each(|x| {
if *x == 0 {
return Err(ValidationError::new("badcount")
.with_message("template limits has zero machines per network count".into()));
}
Ok(())
})?;
}
Ok(())
}

View File

@ -1,8 +1,11 @@
use std::marker::PhantomData;
use super::*;
#[derive(Debug)]
pub(super) struct MachineRegistryInner {
unlocked_inner: Arc<MachineRegistryUnlockedInner>,
allocated_machines: HashSet<MachineStateId>,
profile_state_allocator: StateAllocator<ProfileState>,
machine_state_allocator: StateAllocator<MachineState>,
template_state_allocator: StateAllocator<TemplateState>,
@ -19,6 +22,7 @@ impl MachineRegistryInner {
let srng = unlocked_inner.srng.clone();
MachineRegistryInner {
unlocked_inner,
allocated_machines: HashSet::new(),
profile_state_allocator: StateAllocator::new(),
machine_state_allocator: StateAllocator::new(),
template_state_allocator: StateAllocator::new(),
@ -27,6 +31,21 @@ impl MachineRegistryInner {
address_pool: AddressPool::new(srng),
}
}
pub fn srng(&self) -> StableRng {
self.unlocked_inner.srng.clone()
}
pub fn execute_config(&self, cfg: config::Config) -> MachineRegistryResult<()> {
// Create all networks
// Create all blueprints
// Create all templates
// Create all machines
Ok(())
}
pub fn allocate(&mut self, profile: String) -> MachineRegistryResult<MachineId> {
// Get profile definition
@ -42,40 +61,89 @@ impl MachineRegistryInner {
});
// Get the next instance from the definition
let Some(instance_def) = profile_state.next_instance() else {
return Err(MachineRegistryError::ProfileComplete);
};
loop {
let Some(instance_def) = profile_state.next_instance() else {
return Err(MachineRegistryError::ProfileComplete);
};
let machine_state = match instance_def {
config::Instance::Machine { machine } => {
let machine = self.unlocked_inner.srng.weighted_choice(&machine);
let unlocked_inner = self.unlocked_inner.clone();
let machine_def = unlocked_inner
.config
.machines
.get(machine)
.cloned()
.expect("config validation is broken");
self.get_or_create_machine_state(machine.clone(), machine_def)?
}
config::Instance::Template { template } => {
let template = self.unlocked_inner.srng.weighted_choice(&template);
let unlocked_inner = self.unlocked_inner.clone();
let template_def = unlocked_inner
.config
.templates
.get(template)
.cloned()
.expect("config validation is broken");
self.get_or_create_machine_state_from_template(template.clone(), template_def)?
}
};
Ok(machine_state.id())
let machine_state = match instance_def {
config::Instance::Machine {
machine: machine_names,
} => {
// Filter out machines that are already allocated
let opt_machine_states = machine_names.try_filter_map(|name| {
let Some(machine_state) = self.machine_states().get_state_by_name(name)
else {
return Err(MachineRegistryError::MachineNotFound);
};
if self.allocated_machines.contains(&machine_state.id()) {
Ok(None)
} else {
Ok(Some(machine_state))
}
})?;
let Some(machine_states) = opt_machine_states else {
// All machines in this instance are allocated
continue;
};
// Choose a machine state to activate
let machine_state = self
.unlocked_inner
.srng
.weighted_choice(&machine_states)
.clone();
// Activate it
self.allocated_machines.insert(machine_state.id());
machine_state
}
config::Instance::Template {
template: template_names,
} => {
// Filter out templates that are no longer active
let opt_template_states = template_names.try_filter_map(|name| {
let Some(template_state) = self.template_states().get_state_by_name(name)
else {
return Err(MachineRegistryError::TemplateNotFound);
};
if !template_state.is_active(self)? {
Ok(None)
} else {
Ok(Some(template_state))
}
})?;
let Some(template_states) = opt_template_states else {
// No templates in this instance are still active
continue;
};
let template_state = self.unlocked_inner.srng.weighted_choice(&template_states);
template_state.generate(self)?
}
};
break Ok(machine_state.external_id());
}
}
pub fn release(&mut self, machine_id: MachineId) -> MachineRegistryResult<()> {
// xxx
// xxx remember machines and networks may not be 'named' if they are generated by templates and blueprints
let id = StateId::<MachineState>::new(machine_id);
if self.allocated_machines.contains(&id) {
// Was a fixed machine, so we leave the machine state so it can
// be reallocated later
self.allocated_machines.remove(&id);
} else {
// Was a templated machine, so remove the machine state
let Some(machine_state) = self.machine_states().get_state(id)? else {
return Err(MachineRegistryError::InvalidId);
};
machine_state.release(self);
self.machine_states().release_id(id)?;
}
Ok(())
}
@ -98,59 +166,59 @@ impl MachineRegistryInner {
&mut self.blueprint_state_allocator
}
pub(super) fn get_or_create_machine_state(
&mut self,
opt_name: Option<String>,
params: config::Machine,
) -> MachineRegistryResult<MachineState> {
// Ensure we don't already have this machine created (name must be unique)
if let Some(name) = &opt_name {
if let Some(machine_id) = self.resolve_to_manager_machine.add(name.clone()).get() {
return Ok(self
.machine_state_by_id
.get(&machine_id)
.cloned()
.expect("must exist"));
}
}
// pub(super) fn get_or_create_machine_state(
// &mut self,
// opt_name: Option<String>,
// params: config::Machine,
// ) -> MachineRegistryResult<MachineState> {
// // Ensure we don't already have this machine created (name must be unique)
// if let Some(name) = &opt_name {
// if let Some(machine_id) = self.resolve_to_manager_machine.add(name.clone()).get() {
// return Ok(self
// .machine_state_by_id
// .get(&machine_id)
// .cloned()
// .expect("must exist"));
// }
// }
// Allocate a machine id
let machine_id = self.free_machine_ids.pop().unwrap_or_else(|| {
let x = self.next_machine_id;
self.next_machine_id += 1;
x
});
// // Allocate a machine id
// let machine_id = self.free_machine_ids.pop().unwrap_or_else(|| {
// let x = self.next_machine_id;
// self.next_machine_id += 1;
// x
// });
// Create a new machine state
let machine_state = match MachineState::try_new(
self,
machine_id,
MachineStateName::Machine(name.clone()),
machine_def.clone(),
) {
Ok(v) => v,
Err(e) => {
// Release the machine id
self.free_machine_ids.push(machine_id);
return Err(e);
}
};
// // Create a new machine state
// let machine_state = match MachineState::try_new(
// self,
// machine_id,
// MachineStateName::Machine(name.clone()),
// machine_def.clone(),
// ) {
// Ok(v) => v,
// Err(e) => {
// // Release the machine id
// self.free_machine_ids.push(machine_id);
// return Err(e);
// }
// };
// Store the machine state with its unique id
self.machine_state_by_id.insert(machine_id, machine_state);
// // Store the machine state with its unique id
// self.machine_state_by_id.insert(machine_id, machine_state);
// Bind the name to the id
self.resolve_to_manager_machine
.resolve(&name, machine_id)
.expect("must resolve");
// // Bind the name to the id
// self.resolve_to_manager_machine
// .resolve(&name, machine_id)
// .expect("must resolve");
// Return the state
Ok(self
.machine_state_by_id
.get(&machine_id)
.cloned()
.expect("must exist"))
}
// // Return the state
// Ok(self
// .machine_state_by_id
// .get(&machine_id)
// .cloned()
// .expect("must exist"))
// }
// pub(super) fn get_machine_state_by_id(
// &mut self,

View File

@ -20,13 +20,16 @@ struct MachineRegistryUnlockedInner {
#[derive(Debug, Clone)]
pub enum MachineRegistryError {
InvalidMachineId,
InvalidAllocationName,
ProfileNotFound,
InvalidId,
InvalidName,
AlreadyAttached,
AlreadyDetached,
DuplicateName,
ProfileComplete,
TemplateComplete,
NetworkComplete,
BlueprintComplete,
ProfileNotFound,
MachineNotFound,
NetworkNotFound,
TemplateNotFound,

View File

@ -2,14 +2,14 @@ use super::*;
#[derive(Debug)]
struct BlueprintStateUnlockedInner {
id: BlueprintStateId,
name: String,
blueprint_def: config::Blueprint,
}
#[derive(Debug)]
struct BlueprintStateInner {
limit_network_count: Option<u32>,
networks: Vec<NetworkId>,
limit_network_count: Option<usize>,
networks: Vec<NetworkStateId>,
}
#[derive(Debug, Clone)]
@ -18,33 +18,22 @@ pub struct BlueprintState {
inner: Arc<Mutex<BlueprintStateInner>>,
}
impl BlueprintState {
pub fn try_new(
machine_registry_inner: &mut MachineRegistryInner,
name: String,
blueprint_def: config::Blueprint,
) -> MachineRegistryResult<BlueprintState> {
let limit_network_count = blueprint_def.limits.network_count.as_ref().map(|nc| {
*machine_registry_inner
.unlocked_inner
.srng
.weighted_choice(nc)
});
pub type BlueprintStateId = StateId<BlueprintState>;
impl BlueprintState {
pub fn new(id: BlueprintStateId, name: String) -> MachineRegistryResult<BlueprintState> {
Ok(Self {
unlocked_inner: Arc::new(BlueprintStateUnlockedInner {
name,
blueprint_def,
}),
unlocked_inner: Arc::new(BlueprintStateUnlockedInner { id, name }),
inner: Arc::new(Mutex::new(BlueprintStateInner {
limit_network_count,
limit_network_count: None,
networks: Vec::new(),
})),
})
}
pub fn def(&self) -> &config::Blueprint {
&self.unlocked_inner.blueprint_def
pub fn set_limit_network_count(&self, limit_network_count: Option<usize>) {
let mut inner = self.inner.lock();
inner.limit_network_count = limit_network_count;
}
pub fn is_active(&self) -> MachineRegistryResult<bool> {
@ -52,7 +41,7 @@ impl BlueprintState {
// See if there's room for another network
if let Some(limit_network_count) = inner.limit_network_count {
if inner.networks.len() >= limit_network_count.try_into().unwrap_or(usize::MAX) {
if inner.networks.len() >= limit_network_count {
return Ok(false);
}
}
@ -63,13 +52,13 @@ impl BlueprintState {
pub fn generate(
&self,
machine_registry_inner: &MachineRegistryInner,
) -> MachineRegistryResult<config::Network> {
) -> MachineRegistryResult<NetworkState> {
//
}
pub fn for_each_network_id<F, R>(&self, mut callback: F) -> MachineRegistryResult<Option<R>>
where
F: FnMut(NetworkId) -> MachineRegistryResult<Option<R>>,
F: FnMut(NetworkStateId) -> MachineRegistryResult<Option<R>>,
{
let inner = self.inner.lock();
for network_id in &inner.networks {
@ -80,3 +69,13 @@ impl BlueprintState {
Ok(None)
}
}
impl State for BlueprintState {
fn id(&self) -> StateId<Self> {
self.unlocked_inner.id.clone()
}
fn name(&self) -> Option<String> {
Some(self.unlocked_inner.name.clone())
}
}

View File

@ -3,7 +3,7 @@ use super::*;
#[derive(Debug)]
struct MachineStateInner {
/// The current network interfaces definition
interfaces: Vec<MachineStateInterface>,
interfaces: HashMap<String, MachineStateInterface>,
}
#[derive(Debug)]
@ -36,140 +36,339 @@ impl MachineState {
Self {
unlocked_inner: Arc::new(MachineStateUnlockedInner { id, opt_name }),
inner: Arc::new(Mutex::new(MachineStateInner {
interfaces: Vec::new(),
interfaces: HashMap::new(),
})),
}
}
pub fn release(&self, machine_registry_inner: &mut MachineRegistryInner) {
let interfaces = {
let mut inner = self.inner.lock();
core::mem::take(&mut inner.interfaces)
};
for intf in interfaces {
let network_state = machine_registry_inner
.network_states()
.get_state(intf.network_id)
.expect("must exist")
.expect("must be bound");
let addrs = &intf.network_interface.addrs;
network_state
.release_all_addresses(addrs.iter().map(|x| x.if_addr.ip()))
.expect("must succeed");
}
self.release_all_interfaces(machine_registry_inner);
}
pub fn external_id(&self) -> MachineId {
self.unlocked_inner.id.0
}
fn next_free_interface_name(&self) -> String {}
fn next_free_interface_name_inner(inner: &MachineStateInner) -> String {
let mut inum = 0usize;
loop {
let name = format!("vin{}", inum);
if !inner.interfaces.contains_key(&name) {
return name;
}
inum += 1;
}
}
xxx implement this
pub fn allocate_interface(
&self,
machine_registry_inner: &mut MachineRegistryInner,
network_id: NetworkStateId,
address4: Option<Ipv4Addr>,
address6: Option<Ipv6Addr>,
opt_name: Option<String>,
opt_interface_flags: Option<InterfaceFlags>,
) -> MachineRegistryResult<String> {
// Find existing network or create a new one from network or blueprint definition
let network_state = match params {
MachineParameters::Direct {
network_id,
disable_capabilities: _,
bootstrap: _,
} => todo!(),
MachineParameters::Config { name, def } => {
machine_registry_inner
.get_or_create_network_state_from_machine_location(&def.location)?;
}
};
let srng = machine_registry_inner.unlocked_inner.srng.clone();
// Build list of default route interface addresses
let mut addrs = Vec::<InterfaceAddress>::new();
// Make the default route interface
let machine_location = machine_def.location.clone();
let (allocate_v4, opt_address4, allocate_v6, opt_address6) = match machine_location {
config::MachineLocation::Network {
network: _,
address4,
address6,
} => (
network_state.is_ipv4() && address4.is_some(),
address4,
network_state.is_ipv6() && address6.is_some(),
address6,
),
config::MachineLocation::Blueprint { blueprint: _ } => {
(network_state.is_ipv4(), None, network_state.is_ipv6(), None)
}
};
if allocate_v4 {
let if_addr4 = match network_state.allocate_address_v4(srng.clone(), id, opt_address4) {
Ok(v) => v,
Err(e) => {
network_state
.release_all_addresses(addrs.iter().map(|x| x.if_addr.ip()))
.expect("must succeed");
return Err(e);
}
};
addrs.push(InterfaceAddress {
if_addr: IfAddr::V4(if_addr4),
flags: AddressFlags {
is_dynamic: false,
is_temporary: false,
is_preferred: true,
},
});
let mut inner = self.inner.lock();
let name = opt_name.unwrap_or_else(|| Self::next_free_interface_name_inner(&*inner));
if inner.interfaces.contains_key(&name) {
return Err(MachineRegistryError::DuplicateName);
}
if allocate_v6 {
let if_addr6 = match network_state.allocate_address_v6(srng.clone(), id, opt_address6) {
Ok(v) => v,
Err(e) => {
network_state
.release_all_addresses(addrs.iter().map(|x| x.if_addr.ip()))
.expect("must succeed");
return Err(e);
}
};
addrs.push(InterfaceAddress {
if_addr: IfAddr::V6(if_addr6),
flags: AddressFlags {
is_dynamic: false,
is_temporary: false,
is_preferred: true,
},
});
}
// Allocate an address on the network and make an veilid-style interface record for it
let network_interface = NetworkInterface {
name: "vin0".to_owned(),
flags: InterfaceFlags {
is_loopback: false,
is_running: true,
is_point_to_point: false,
has_default_route: true,
},
addrs,
};
interfaces.push(MachineStateInterface {
network_id: network_state.id(),
network_interface,
let flags = opt_interface_flags.unwrap_or_else(|| InterfaceFlags {
is_loopback: false,
is_running: true,
is_point_to_point: false,
has_default_route: true,
});
inner.interfaces.insert(
name.clone(),
MachineStateInterface {
network_id,
network_interface: NetworkInterface {
name: name.clone(),
flags,
addrs: Vec::new(),
},
},
);
Ok(name)
}
pub fn release_interface(&self) -> () {
//
pub fn interfaces(&self) -> Vec<String> {
let mut intfs: Vec<String> = self.inner.lock().interfaces.keys().cloned().collect();
intfs.sort();
intfs
}
pub fn allocate_address_ipv4(
&self,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
opt_address: Option<Ipv4Addr>,
opt_address_flags: Option<AddressFlags>,
) -> MachineRegistryResult<Ifv4Addr> {
let mut inner = self.inner.lock();
let Some(intf) = inner.interfaces.get_mut(interface) else {
return Err(MachineRegistryError::InvalidName);
};
// Get the network state
let Some(network_state) = machine_registry_inner
.network_states()
.get_state(intf.network_id)?
else {
return Err(MachineRegistryError::NetworkNotFound);
};
// Allocate interface address
let is_dynamic = opt_address.is_none();
let ifv4_addr =
network_state.allocate_address_v4(machine_registry_inner, self.id(), opt_address)?;
// Get address flags
let flags = opt_address_flags.unwrap_or_else(|| AddressFlags {
is_dynamic,
is_temporary: false,
is_preferred: true,
});
intf.network_interface.addrs.push(InterfaceAddress {
if_addr: IfAddr::V4(ifv4_addr.clone()),
flags,
});
Ok(ifv4_addr)
}
pub fn allocate_address_ipv6(
&self,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
opt_address: Option<Ipv6Addr>,
opt_address_flags: Option<AddressFlags>,
) -> MachineRegistryResult<Ifv6Addr> {
let mut inner = self.inner.lock();
let Some(intf) = inner.interfaces.get_mut(interface) else {
return Err(MachineRegistryError::InvalidName);
};
// Get the network state
let Some(network_state) = machine_registry_inner
.network_states()
.get_state(intf.network_id)?
else {
return Err(MachineRegistryError::NetworkNotFound);
};
// Allocate interface address
let is_dynamic = opt_address.is_none();
let ifv6_addr =
network_state.allocate_address_v6(machine_registry_inner, self.id(), opt_address)?;
// Get address flags
let flags = opt_address_flags.unwrap_or_else(|| AddressFlags {
is_dynamic,
is_temporary: false,
is_preferred: true,
});
intf.network_interface.addrs.push(InterfaceAddress {
if_addr: IfAddr::V6(ifv6_addr.clone()),
flags,
});
Ok(ifv6_addr)
}
pub fn release_address(
&self,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
address: IpAddr,
) -> MachineRegistryResult<()> {
let mut inner = self.inner.lock();
let Some(intf) = inner.interfaces.get_mut(interface) else {
return Err(MachineRegistryError::InvalidName);
};
// Get the network state
let Some(network_state) = machine_registry_inner
.network_states()
.get_state(intf.network_id)?
else {
return Err(MachineRegistryError::NetworkNotFound);
};
// Release the address from the network
match address {
IpAddr::V4(ipv4_addr) => network_state.release_address_v4(ipv4_addr)?,
IpAddr::V6(ipv6_addr) => network_state.release_address_v6(ipv6_addr)?,
}
// Remove the address from the interface
intf.network_interface
.addrs
.retain(|x| x.if_addr().ip() != address);
Ok(())
}
pub fn release_all_addresses(
&self,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
) -> MachineRegistryResult<()> {
let mut inner = self.inner.lock();
Self::release_all_addresses_inner(&mut *inner, machine_registry_inner, interface)
}
fn release_all_addresses_inner(
inner: &mut MachineStateInner,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
) -> MachineRegistryResult<()> {
let Some(intf) = inner.interfaces.get_mut(interface) else {
return Err(MachineRegistryError::InvalidName);
};
// Get the network state
let Some(network_state) = machine_registry_inner
.network_states()
.get_state(intf.network_id)?
else {
return Err(MachineRegistryError::NetworkNotFound);
};
// Release the addresses from the network
for addr in &intf.network_interface.addrs {
match addr.if_addr.ip() {
IpAddr::V4(ipv4_addr) => network_state.release_address_v4(ipv4_addr)?,
IpAddr::V6(ipv6_addr) => network_state.release_address_v6(ipv6_addr)?,
}
}
// Remove the addresses from the interface
intf.network_interface.addrs.clear();
Ok(())
}
pub fn release_interface(
&self,
machine_registry_inner: &mut MachineRegistryInner,
interface: &str,
) -> MachineRegistryResult<()> {
let mut inner = self.inner.lock();
Self::release_all_addresses_inner(&mut *inner, machine_registry_inner, interface)?;
inner
.interfaces
.remove(interface)
.expect("interface must exist");
Ok(())
}
pub fn release_all_interfaces(
&self,
machine_registry_inner: &mut MachineRegistryInner,
) -> MachineRegistryResult<()> {
let mut inner = self.inner.lock();
let interfaces: Vec<String> = inner.interfaces.keys().cloned().collect();
for interface in interfaces {
Self::release_all_addresses_inner(&mut *inner, machine_registry_inner, &interface)?;
}
inner.interfaces.clear();
Ok(())
}
// let network_state = match params {
// MachineParameters::Direct {
// network_id,
// disable_capabilities: _,
// bootstrap: _,
// } => todo!(),
// MachineParameters::Config { name, def } => {
// machine_registry_inner
// .get_or_create_network_state_from_machine_location(&def.location)?;
// }
// };
// let srng = machine_registry_inner.unlocked_inner.srng.clone();
// // Build list of default route interface addresses
// let mut addrs = Vec::<InterfaceAddress>::new();
// // Make the default route interface
// let machine_location = machine_def.location.clone();
// let (allocate_v4, opt_address4, allocate_v6, opt_address6) = match machine_location {
// config::MachineLocation::Network {
// network: _,
// address4,
// address6,
// } => (
// network_state.is_ipv4() && address4.is_some(),
// address4,
// network_state.is_ipv6() && address6.is_some(),
// address6,
// ),
// config::MachineLocation::Blueprint { blueprint: _ } => {
// (network_state.is_ipv4(), None, network_state.is_ipv6(), None)
// }
// };
// if allocate_v4 {
// let if_addr4 = match network_state.allocate_address_v4(srng.clone(), id, opt_address4) {
// Ok(v) => v,
// Err(e) => {
// network_state
// .release_all_addresses(addrs.iter().map(|x| x.if_addr.ip()))
// .expect("must succeed");
// return Err(e);
// }
// };
// addrs.push(InterfaceAddress {
// if_addr: IfAddr::V4(if_addr4),
// flags: AddressFlags {
// is_dynamic: false,
// is_temporary: false,
// is_preferred: true,
// },
// });
// }
// if allocate_v6 {
// let if_addr6 = match network_state.allocate_address_v6(srng.clone(), id, opt_address6) {
// Ok(v) => v,
// Err(e) => {
// network_state
// .release_all_addresses(addrs.iter().map(|x| x.if_addr.ip()))
// .expect("must succeed");
// return Err(e);
// }
// };
// addrs.push(InterfaceAddress {
// if_addr: IfAddr::V6(if_addr6),
// flags: AddressFlags {
// is_dynamic: false,
// is_temporary: false,
// is_preferred: true,
// },
// });
// }
// // Allocate an address on the network and make an veilid-style interface record for it
// let network_interface = NetworkInterface {
// name: "vin0".to_owned(),
// flags: InterfaceFlags {
// is_loopback: false,
// is_running: true,
// is_point_to_point: false,
// has_default_route: true,
// },
// addrs,
// };
// interfaces.push(MachineStateInterface {
// network_id: network_state.id(),
// network_interface,
// });
// }
}
impl State for MachineState {

View File

@ -16,7 +16,7 @@ struct NetworkStateInner {
/// Distance simulation metric
distance: Option<config::Distance>,
/// Packet loss probability
loss: config::Probability,
loss: Probability,
/// IPv4 state if it is enabled
ipv4: Option<NetworkStateIpv4>,
/// IPv6 state if it is enabled
@ -27,14 +27,14 @@ struct NetworkStateInner {
struct NetworkStateIpv4 {
allocation: Ipv4Net,
gateway: Option<NetworkGatewayState>,
machine_addresses: HashMap<Ipv4Addr, MachineId>,
machine_addresses: HashMap<Ipv4Addr, MachineStateId>,
}
#[derive(Debug)]
struct NetworkStateIpv6 {
allocation: Ipv6Net,
gateway: Option<NetworkGatewayState>,
machine_addresses: HashMap<Ipv6Addr, MachineId>,
machine_addresses: HashMap<Ipv6Addr, MachineStateId>,
}
#[derive(Debug)]
@ -181,8 +181,8 @@ impl NetworkState {
pub fn allocate_address_v4(
&self,
srng: StableRng,
machine_id: MachineId,
machine_registry_inner: &mut MachineRegistryInner,
machine_id: MachineStateId,
opt_address: Option<Ipv4Addr>,
) -> MachineRegistryResult<Ifv4Addr> {
let mut inner = self.inner.lock();
@ -217,7 +217,9 @@ impl NetworkState {
ip_addr
} else {
// Any address will do
let addr_end = srng.next_u32(first_host_bits, last_host_bits);
let addr_end = machine_registry_inner
.srng()
.next_u32(first_host_bits, last_host_bits);
// Find a free address starting from here
let mut addr = addr_end;
@ -269,8 +271,8 @@ impl NetworkState {
pub fn allocate_address_v6(
&self,
srng: StableRng,
machine_id: MachineId,
machine_registry_inner: &mut MachineRegistryInner,
machine_id: MachineStateId,
opt_address: Option<Ipv6Addr>,
) -> MachineRegistryResult<Ifv6Addr> {
let mut inner = self.inner.lock();
@ -305,7 +307,9 @@ impl NetworkState {
ip_addr
} else {
// Any address will do
let addr_end = srng.next_u128(first_host_bits, last_host_bits);
let addr_end = machine_registry_inner
.srng()
.next_u128(first_host_bits, last_host_bits);
// Find a free address starting from here
let mut addr = addr_end;

View File

@ -1,46 +1,5 @@
use super::*;
#[derive(ThisError, Debug)]
pub enum StateAllocatorReleaseError {
#[error("invalid state id")]
InvalidId,
}
pub type StateAllocatorReleaseResult<T> = Result<T, StateAllocatorReleaseError>;
#[derive(ThisError, Debug)]
pub enum StateAllocatorAttachError {
#[error("invalid state id")]
InvalidId,
#[error("state already attached")]
AlreadyAttached,
#[error("duplicate name")]
DuplicateName,
}
pub type StateAllocatorAttachResult<T> = Result<T, StateAllocatorAttachError>;
#[derive(ThisError, Debug)]
pub enum StateAllocatorDetachError {
#[error("invalid state id")]
InvalidId,
#[error("state already detached")]
AlreadyDetached,
}
pub type StateAllocatorDetachResult<T> = Result<T, StateAllocatorDetachError>;
#[derive(ThisError, Debug)]
pub enum StateAllocatorGetStateError {
#[error("invalid state id")]
InvalidId,
}
pub type StateAllocatorGetStateResult<T> = Result<T, StateAllocatorGetStateError>;
#[derive(ThisError, Debug)]
pub enum StateAllocatorGetOrCreateByNameError {
#[error("duplicate name")]
DuplicateName,
}
pub type StateAllocatorGetOrCreateByNameResult<T> = Result<T, StateAllocatorGetOrCreateByNameError>;
pub trait State: fmt::Debug + Clone {
fn id(&self) -> StateId<Self>;
fn name(&self) -> Option<String>;
@ -48,8 +7,36 @@ pub trait State: fmt::Debug + Clone {
type StateIdInternal = u64;
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[derive(Debug, Clone)]
pub struct StateId<S: State>(pub StateIdInternal, core::marker::PhantomData<S>);
impl<S: State> StateId<S> {
pub fn new(external_id: u64) -> Self {
Self(external_id, PhantomData {})
}
}
impl<S: State> Copy for StateId<S> {}
impl<S: State> PartialEq for StateId<S> {
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
}
}
impl<S: State> Eq for StateId<S> {}
impl<S: State> PartialOrd for StateId<S> {
fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
Some(self.0.cmp(&other.0))
}
}
impl<S: State> Ord for StateId<S> {
fn cmp(&self, other: &Self) -> core::cmp::Ordering {
self.0.cmp(&other.0)
}
}
impl<S: State> core::hash::Hash for StateId<S> {
fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
self.0.hash(state)
}
}
#[derive(Debug)]
pub struct StateAllocator<S: State> {
@ -99,10 +86,10 @@ impl<S: State> StateAllocator<S> {
StateId(state_id, PhantomData {})
}
pub fn release_id(&mut self, id: StateId<S>) -> StateAllocatorReleaseResult<()> {
pub fn release_id(&mut self, id: StateId<S>) -> MachineRegistryResult<()> {
// Remove id to state mapping
let Some(old_opt_state) = self.state_by_id.remove(&id.0) else {
return Err(StateAllocatorReleaseError::InvalidId);
return Err(MachineRegistryError::InvalidId);
};
// Release state if it is attached
@ -121,24 +108,24 @@ impl<S: State> StateAllocator<S> {
Ok(())
}
pub fn attach_state(&mut self, state: S) -> StateAllocatorAttachResult<()> {
pub fn attach_state(&mut self, state: S) -> MachineRegistryResult<()> {
// Get the id from the state
let id = state.id();
// Get the allocator slot
let Some(opt_state) = self.state_by_id.get_mut(&id.0) else {
return Err(StateAllocatorAttachError::InvalidId);
return Err(MachineRegistryError::InvalidId);
};
// Ensure the state slot isn't attached already
if opt_state.is_some() {
return Err(StateAllocatorAttachError::AlreadyAttached);
return Err(MachineRegistryError::AlreadyAttached);
}
// Ensure the name isn't duplicated
if let Some(name) = state.name() {
if self.state_id_by_name.contains_key(&name) {
return Err(StateAllocatorAttachError::DuplicateName);
return Err(MachineRegistryError::DuplicateName);
}
// Register the named state
self.state_id_by_name
@ -152,15 +139,15 @@ impl<S: State> StateAllocator<S> {
Ok(())
}
pub fn detach_state(&mut self, id: StateId<S>) -> StateAllocatorDetachResult<S> {
pub fn detach_state(&mut self, id: StateId<S>) -> MachineRegistryResult<S> {
// Get the allocator slot
let Some(opt_state) = self.state_by_id.get_mut(&id.0) else {
return Err(StateAllocatorDetachError::InvalidId);
return Err(MachineRegistryError::InvalidId);
};
// Take the state out of the slot and ensure the state slot isn't detached already
let Some(state) = opt_state.take() else {
return Err(StateAllocatorDetachError::AlreadyDetached);
return Err(MachineRegistryError::AlreadyDetached);
};
// Release the name if it exists
@ -175,10 +162,10 @@ impl<S: State> StateAllocator<S> {
Ok(state)
}
pub fn get_state(&self, id: StateId<S>) -> StateAllocatorGetStateResult<Option<S>> {
pub fn get_state(&self, id: StateId<S>) -> MachineRegistryResult<Option<S>> {
// Get the allocator slot
let Some(opt_state) = self.state_by_id.get(&id.0).cloned() else {
return Err(StateAllocatorGetStateError::InvalidId);
return Err(MachineRegistryError::InvalidId);
};
Ok(opt_state)

View File

@ -2,21 +2,33 @@ use super::*;
#[derive(Debug)]
struct TemplateStateUnlockedInner {
id: TemplateStateId,
name: String,
template_def: config::Template,
}
#[derive(Debug)]
struct PerNetworkInfo {
limit_machine_count: u32,
machines: HashSet<MachineId>,
limit_machine_count: Option<usize>,
machines: HashSet<MachineStateId>,
}
#[derive(Debug)]
enum LocationsList {
Networks {
networks: WeightedList<NetworkStateId>,
},
Blueprints {
blueprints: WeightedList<BlueprintStateId>,
},
}
#[derive(Debug)]
struct TemplateStateInner {
limit_machine_count: Option<u32>,
machines: HashSet<MachineId>,
machines_per_network: HashMap<NetworkId, PerNetworkInfo>,
limit_machine_count: Option<usize>,
limit_machines_per_network: Option<WeightedList<usize>>,
locations_list: Option<LocationsList>,
machines: HashSet<MachineStateId>,
machines_per_network: HashMap<NetworkStateId, PerNetworkInfo>,
}
#[derive(Debug, Clone)]
@ -25,35 +37,48 @@ pub struct TemplateState {
inner: Arc<Mutex<TemplateStateInner>>,
}
impl TemplateState {
pub fn try_new(
machine_registry_inner: &mut MachineRegistryInner,
name: String,
template_def: config::Template,
) -> MachineRegistryResult<TemplateState> {
let limit_machine_count = template_def.limits.machine_count.as_ref().map(|mc| {
*machine_registry_inner
.unlocked_inner
.srng
.weighted_choice(mc)
});
pub type TemplateStateId = StateId<TemplateState>;
Ok(Self {
unlocked_inner: Arc::new(TemplateStateUnlockedInner { name, template_def }),
impl TemplateState {
pub fn new(id: TemplateStateId, name: String) -> Self {
Self {
unlocked_inner: Arc::new(TemplateStateUnlockedInner { id, name }),
inner: Arc::new(Mutex::new(TemplateStateInner {
limit_machine_count,
limit_machine_count: None,
limit_machines_per_network: None,
locations_list: None,
machines: HashSet::new(),
machines_per_network: HashMap::new(),
})),
})
}
}
pub fn name(&self) -> String {
self.unlocked_inner.name.clone()
pub fn set_networks_list(&self, networks: WeightedList<NetworkStateId>) {
let mut inner = self.inner.lock();
inner.locations_list = Some(LocationsList::Networks { networks })
}
pub fn def(&self) -> &config::Template {
&self.unlocked_inner.template_def
pub fn set_blueprints_list(&self, blueprints: WeightedList<BlueprintStateId>) {
let mut inner = self.inner.lock();
inner.locations_list = Some(LocationsList::Blueprints { blueprints })
}
pub fn clear_locations_list(&self) {
let mut inner = self.inner.lock();
inner.locations_list = None;
}
pub fn set_limit_machine_count(&self, limit_machine_count: Option<usize>) {
let mut inner = self.inner.lock();
inner.limit_machine_count = limit_machine_count;
}
pub fn set_limit_machines_per_network(
&self,
limit_machines_per_network: Option<WeightedList<usize>>,
) {
let mut inner = self.inner.lock();
inner.limit_machines_per_network = limit_machines_per_network;
}
fn is_network_available_inner(
@ -71,28 +96,35 @@ impl TemplateState {
return Ok(true);
};
// If this template has not yet allocated the maximum number of machines per-network
// for this network, then it is available
if pni.machines.len() < pni.limit_machine_count.try_into().unwrap_or(usize::MAX) {
return Ok(true);
// If this template has allocated the maximum number of machines per-network
// for this network, then it is not available
if let Some(limit_machine_count) = pni.limit_machine_count {
if pni.machines.len() >= limit_machine_count {
return Ok(false);
}
}
Ok(false)
Ok(true)
}
xxx should this be sensitive to already generated blueprint networks?
fn is_blueprint_available_inner(
inner: &TemplateStateInner,
machine_registry_inner: &MachineRegistryInner,
blueprint_state: BlueprintState,
) -> MachineRegistryResult<Availability<NetworkId>> {
) -> MachineRegistryResult<Availability<NetworkStateId>> {
// See if the networks generated from this blueprint so far have availability
// in this template
if let Some(available_network_id) = blueprint_state.for_each_network_id(|network_id| {
if let Some(available_network_id) = blueprint_state.for_each_network_id(|id| {
// Check the network's availability
let network_state = machine_registry_inner.get_network_state_by_id(network_id)?;
let network_state = machine_registry_inner
.network_states()
.get_state(id)?
.expect("must exist");
if Self::is_network_available_inner(inner, network_state)? {
// We found one
return Ok(Some(network_id));
return Ok(Some(id));
}
// Try next network
Ok(None)
@ -117,28 +149,23 @@ impl TemplateState {
// See if we have reached our machine limit
if let Some(limit_machine_count) = inner.limit_machine_count {
if inner.machines.len() < limit_machine_count.try_into().unwrap_or(usize::MAX) {
if inner.machines.len() >= limit_machine_count {
return Ok(false);
}
}
// See if any of our existing networks have room to allocate (machines could have been removed)
for (_network_id, pni) in &inner.machines_per_network {
// If this template has not yet allocated the maximum number of machines per-network
// for this network, then it is available
if pni.machines.len() < pni.limit_machine_count.try_into().unwrap_or(usize::MAX) {
return Ok(true);
}
}
let Some(locations_list) = inner.locations_list.as_ref() else {
return Ok(false);
};
// If existing networks are all full, we'd have to allocate one, see if we'd be able to do that
match self.def().location.clone() {
config::TemplateLocation::Network { network } => {
match locations_list {
LocationsList::Networks { networks } => {
// Filter the weighted list of networks to those that are still active and or not yet started
if network
.try_filter(|n| {
if networks
.try_filter(|id| {
machine_registry_inner
.get_network_state_by_name(&n)
.network_states()
.get_state(*id)?
.clone()
.map(|ns| Self::is_network_available_inner(&*inner, ns))
.unwrap_or(Ok(true))
@ -148,12 +175,13 @@ impl TemplateState {
return Ok(false);
};
}
config::TemplateLocation::Blueprint { blueprint } => {
LocationsList::Blueprints { blueprints } => {
// Filter the weighted list of blueprints to those that are still active or not yet started and can allocate
if blueprint
.try_filter(|b| {
if blueprints
.try_filter(|id| {
machine_registry_inner
.get_blueprint_state(&b)
.blueprint_states()
.get_state(*id)?
.clone()
.map(|bs| {
Self::is_blueprint_available_inner(
@ -180,19 +208,27 @@ impl TemplateState {
machine_registry_inner: &mut MachineRegistryInner,
) -> MachineRegistryResult<MachineState> {
let mut inner = self.inner.lock();
// See if we have reached our machine limit
if let Some(limit_machine_count) = inner.limit_machine_count {
if inner.machines.len() < limit_machine_count.try_into().unwrap_or(usize::MAX) {
return Err(MachineRegistryError::TemplateComplete);
}
}
// Pick or instantiate an available network
let network_state = match self.def().location.clone() {
config::TemplateLocation::Network { network } => {
// Filter the weighted list of networks to those that are still active or not yet started and can allocate
let Some(active_networks) = network.try_filter(|n| {
// If existing networks are all full, we'd have to allocate one, see if we'd be able to do that
let Some(locations_list) = inner.locations_list.as_ref() else {
return Err(MachineRegistryError::TemplateComplete);
};
// Get a network to generate the machine on
let network_state = match locations_list {
LocationsList::Networks { networks } => {
// Filter the weighted list of networks to those that are still active and or not yet started
let Some(active_networks) = networks.try_filter(|id| {
machine_registry_inner
.get_network_state_by_name(&n)
.network_states()
.get_state(*id)?
.clone()
.map(|ns| Self::is_network_available_inner(&*inner, ns))
.unwrap_or(Ok(true))
@ -202,56 +238,85 @@ impl TemplateState {
};
// Weighted choice of network now that we have a candidate list
let network = machine_registry_inner
.unlocked_inner
.srng
let network_id = machine_registry_inner
.srng()
.weighted_choice(&active_networks);
// Instantiate the network if it doesn't yet exist
let network_state = machine_registry_inner.get_or_create_network_state(network.clone())?;
// Get the fixed network
let network_state = machine_registry_inner
.network_states()
.get_state(*network_id)?
.expect("must exist");
// Return network state to use
network_state
}
config::TemplateLocation::Blueprint { blueprint } => {
LocationsList::Blueprints { blueprints } => {
// Filter the weighted list of blueprints to those that are still active or not yet started and can allocate
let Some(active_blueprints) = blueprint.try_filter_map(|b| {
let Some(active_blueprints) = blueprints.try_filter(|id| {
machine_registry_inner
.get_blueprint_state(&b)
.blueprint_states()
.get_state(*id)?
.clone()
.map(|bs| {
Self::is_blueprint_available_inner(inner, machine_registry_inner, bs)
Self::is_blueprint_available_inner(&*inner, machine_registry_inner, bs)
.map(|x| !matches!(x, Availability::None))
})
.unwrap_or(Ok(Some()))
.unwrap_or(Ok(true))
})?
else {
return Err(MachineRegistryError::BlueprintComplete);
};
// Weighted choice of blueprint now that we have a candidate list
let blueprint_name = machine_registry_inner
.unlocked_inner
.srng
let blueprint_id = machine_registry_inner
.srng()
.weighted_choice(&active_blueprints);
config::MachineLocation::Blueprint {
blueprint: blueprint_name.clone(),
}
xxx do not always generate... use most recent network for this blueprint in this template.
// Instantiate a blueprint network
let blueprint_state = machine_registry_inner
.blueprint_states()
.get_state(*blueprint_id)?
.expect("must exist");
blueprint_state.generate(machine_registry_inner)?
}
};
xxx
// Add to machines for this template
{
let template_state = self.get_template_state(&name).expect("must exist");
template_state.machines.insert(machine_id);
}
// Return the unique id
Ok(machine_id)
// Allocate a machine id
let machine_id = machine_registry_inner.machine_states().allocate_id();
Ok(MachineParameters::Direct {
disable_capabilities: self.def().disable_capabilities.clone(),
bootstrap: false,
})
// Create an anonymous machine state
let mut machine_state = MachineState::new(machine_id, None);
// Build out the machine state from the template
//inner.
// Attach the state to the id
machine_registry_inner
.machine_states()
.attach_state(machine_state.clone());
// Record the newly instantiated machine
inner.machines.insert(machine_id);
let per_network_info = inner.machines_per_network.entry(network_state).or_insert_with(|| {
let limit_machine_count = inner.limit_machines_per_network.map(|wl| machine_registry_inner.srng().weighted_choice(&wl)).copied();
PerNetworkInfo{ limit_machine_count, machines: HashSet::new() }
});
per_network_info.machines.insert(machine_id);
Ok(machine_state)
}
}
impl State for TemplateState {
fn id(&self) -> StateId<Self> {
self.unlocked_inner.id.clone()
}
fn name(&self) -> Option<String> {
Some(self.unlocked_inner.name.clone())
}
}

View File

@ -2,12 +2,14 @@ pub mod config;
mod machine_registry;
mod server_processor;
mod stable_rng;
mod weighted_list;
use super::*;
use machine_registry::*;
use server_processor::*;
use stable_rng::*;
use weighted_list::*;
use async_tungstenite::accept_async;
use futures_codec::{Bytes, BytesCodec, FramedRead, FramedWrite};

View File

@ -35,11 +35,11 @@ impl StableRng {
}
pub fn weighted_choice<'a, T: fmt::Debug + Clone>(
&self,
weighted_list: &'a config::WeightedList<T>,
weighted_list: &'a WeightedList<T>,
) -> &'a T {
match weighted_list {
config::WeightedList::Single(x) => x,
config::WeightedList::List(vec) => {
WeightedList::Single(x) => x,
WeightedList::List(vec) => {
let total_weight = vec
.iter()
.map(|x| x.weight())

View File

@ -0,0 +1,199 @@
use super::*;
use serde::*;
use validator::{Validate, ValidationError, ValidationErrors};
pub type Probability = f32;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WeightedList<T: fmt::Debug + Clone> {
Single(T),
List(Vec<Weighted<T>>),
}
impl<T: fmt::Debug + Clone> Default for WeightedList<T> {
fn default() -> Self {
Self::List(Vec::new())
}
}
impl<T: fmt::Debug + Clone> Validate for WeightedList<T> {
fn validate(&self) -> Result<(), ValidationErrors> {
let mut errors = ValidationErrors::new();
match self {
Self::List(v) => {
if v.is_empty() {
errors.add(
"List",
ValidationError::new("len")
.with_message("weighted list must not be empty".into()),
)
}
errors.merge_self("List", v.validate());
}
Self::Single(_addr) => {}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
}
impl<T: fmt::Debug + Clone> WeightedList<T> {
pub fn validate_once(&self) -> Result<(), ValidationError> {
match self {
Self::List(v) => {
if v.is_empty() {
return Err(ValidationError::new("len")
.with_message("weighted list must not be empty".into()));
}
}
Self::Single(_addr) => {}
}
Ok(())
}
pub fn try_for_each<E, F: FnMut(&T) -> Result<(), E>>(&self, mut f: F) -> Result<(), E> {
match self {
WeightedList::Single(v) => f(v),
WeightedList::List(vec) => vec
.iter()
.map(|v| match v {
Weighted::Weighted { item, weight: _ } => item,
Weighted::Unweighted(item) => item,
})
.try_for_each(f),
}
}
pub fn filter<F>(&self, mut filter: F) -> Option<WeightedList<T>>
where
F: FnMut(&T) -> bool,
{
match self {
WeightedList::Single(v) => {
if filter(v) {
return Some(self.clone());
}
return None;
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<T>>::with_capacity(vec.len());
for v in vec {
if filter(v.item()) {
out.push(v.clone());
}
}
if out.is_empty() {
None
} else {
Some(WeightedList::List(out))
}
}
}
}
pub fn try_filter<F, E>(&self, mut filter: F) -> Result<Option<WeightedList<T>>, E>
where
F: FnMut(&T) -> Result<bool, E>,
{
match self {
WeightedList::Single(v) => {
if filter(v)? {
return Ok(Some(self.clone()));
}
return Ok(None);
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<T>>::with_capacity(vec.len());
for v in vec {
if filter(v.item())? {
out.push(v.clone());
}
}
if out.is_empty() {
Ok(None)
} else {
Ok(Some(WeightedList::List(out)))
}
}
}
}
pub fn try_filter_map<F, S, E>(&self, mut filter: F) -> Result<Option<WeightedList<S>>, E>
where
F: FnMut(&T) -> Result<Option<S>, E>,
S: fmt::Debug + Clone,
{
match self {
WeightedList::Single(v) => {
if let Some(item) = filter(v)? {
return Ok(Some(WeightedList::Single(item)));
}
return Ok(None);
}
WeightedList::List(vec) => {
let mut out = Vec::<Weighted<S>>::with_capacity(vec.len());
for v in vec {
if let Some(item) = filter(v.item())? {
out.push(match v {
Weighted::Weighted { item: _, weight } => Weighted::Weighted {
item,
weight: *weight,
},
Weighted::Unweighted(_) => Weighted::Unweighted(item),
});
}
}
if out.is_empty() {
Ok(None)
} else {
Ok(Some(WeightedList::List(out)))
}
}
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Weighted<T: fmt::Debug + Clone> {
Weighted { item: T, weight: f32 },
Unweighted(T),
}
impl<T: fmt::Debug + Clone> Validate for Weighted<T> {
fn validate(&self) -> Result<(), ValidationErrors> {
let mut errors = ValidationErrors::new();
if let Self::Weighted { item: _, weight } = self {
if *weight <= 0.0 {
errors.add(
"Weighted",
ValidationError::new("len")
.with_message("weight must be a positive value".into()),
)
}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
}
impl<T: fmt::Debug + Clone> Weighted<T> {
pub fn item(&self) -> &T {
match self {
Weighted::Weighted { item, weight: _ } => item,
Weighted::Unweighted(item) => item,
}
}
pub fn weight(&self) -> f32 {
match self {
Weighted::Weighted { item: _, weight } => *weight,
Weighted::Unweighted(_) => 1.0f32,
}
}
}