diff --git a/veilid-tools/src/bin/virtual_router/main.rs b/veilid-tools/src/bin/virtual_router/main.rs index e9077837..29694b28 100644 --- a/veilid-tools/src/bin/virtual_router/main.rs +++ b/veilid-tools/src/bin/virtual_router/main.rs @@ -64,8 +64,8 @@ fn main() -> Result<(), String> { let args = CmdlineArgs::parse(); - let config = - Config::new(args.config_file).map_err(|e| format!("Error loading config: {}", e))?; + let config = config::Config::new(args.config_file) + .map_err(|e| format!("Error loading config: {}", e))?; if args.dump_config { let cfg_yaml = serde_yaml::to_string(&config) diff --git a/veilid-tools/src/virtual_network/router_server/config.rs b/veilid-tools/src/virtual_network/router_server/config.rs index 7c96eadc..baa8c7c3 100644 --- a/veilid-tools/src/virtual_network/router_server/config.rs +++ b/veilid-tools/src/virtual_network/router_server/config.rs @@ -8,9 +8,11 @@ use validator::{Validate, ValidateArgs, ValidationError, ValidationErrors}; const PREDEFINED_CONFIG: &str = include_str!("predefined_config.yml"); const DEFAULT_CONFIG: &str = include_str!("default_config.yml"); -#[derive(Debug)] +#[derive(Debug, ThisError)] pub enum ConfigError { + #[error("parse error")] ParseError(::config::ConfigError), + #[error("validate error")] ValidateError(validator::ValidationErrors), } @@ -515,7 +517,7 @@ pub struct Allocation { pub subnets: Subnets, } -struct ValidateContext<'a> { +pub struct ValidateContext<'a> { config: &'a Config, } diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs index fb10fea5..b7870385 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs @@ -397,7 +397,7 @@ impl AddressPool { let mut current_scope_end_subnet_index = scope_start_subnet_index + scope_ranges[scope_index].2; - let opt_allocation = loop { + loop { // Get the net at this current subnet index let netbits = u32::from(scope_ranges[current_scope_index].0.network()); let subnetbits = if prefix == 32 { @@ -437,9 +437,7 @@ impl AddressPool { current_scope_end_subnet_index = current_scope_start_subnet_index + scope_ranges[current_scope_index].2; } - }; - - opt_allocation + } } fn find_random_allocation_v6(&self, srng: &mut StableRng, prefix: u8) -> Option { @@ -501,7 +499,7 @@ impl AddressPool { let mut current_scope_end_subnet_index = scope_start_subnet_index + scope_ranges[scope_index].2; - let opt_allocation = loop { + loop { // Get the net at this current subnet index let netbits = u128::from(scope_ranges[current_scope_index].0.network()); let subnetbits = if prefix == 128 { @@ -541,8 +539,6 @@ impl AddressPool { current_scope_end_subnet_index = current_scope_start_subnet_index + scope_ranges[current_scope_index].2; } - }; - - opt_allocation + } } } diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/machine_registry_inner.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/machine_registry_inner.rs index 9a72db54..70d021bb 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/machine_registry_inner.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/machine_registry_inner.rs @@ -115,7 +115,7 @@ impl MachineRegistryInner { .template_states() .get_state(template_state_id) .expect("must exist"); - if !template_state.is_active(self)? { + if !template_state.is_active(self) { Ok(None) } else { Ok(Some(template_state)) diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs index 921b9025..0ba8a68c 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs @@ -1,120 +1,117 @@ use super::*; #[derive(Debug)] -struct BlueprintStateUnlockedInner { - /// The global random number generator - srng: StableRng, +struct BlueprintStateImmutable { /// The unique id of this blueprint id: BlueprintStateId, /// The name of this blueprint state name: String, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct BlueprintStateIpv4Params { pub locations: NetworkLocationsList, pub prefix: WeightedList, pub gateway: Option, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct BlueprintStateIpv6Params { pub locations: NetworkLocationsList, pub prefix: WeightedList, pub gateway: Option, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct BlueprintStateGatewayParams { pub translation: WeightedList, pub upnp: Probability, pub locations: Option, } -#[derive(Debug)] +#[derive(Debug, Clone)] struct BlueprintStateIpv4 { params: BlueprintStateIpv4Params, gateway: Option, } -#[derive(Debug)] +#[derive(Debug, Clone)] struct BlueprintStateIpv6 { params: BlueprintStateIpv6Params, gateway: Option, } -#[derive(Debug)] +#[derive(Debug, Clone)] struct BlueprintStateIpv4Gateway { params: BlueprintStateGatewayParams, } -#[derive(Debug)] +#[derive(Debug, Clone)] struct BlueprintStateIpv6Gateway { params: BlueprintStateGatewayParams, } -#[derive(Debug)] -struct BlueprintStateInner { +#[derive(Debug, Clone)] +struct BlueprintStateFields { limit_network_count: Option, - networks: Vec, - model: Option>, + networks: imbl::Vector, + model: Option>>, ipv4: Option, ipv6: Option, } #[derive(Debug, Clone)] pub struct BlueprintState { - unlocked_inner: Arc, - inner: Arc>, + immutable: Arc, + fields: Arc, } pub type BlueprintStateId = StateId; impl BlueprintState { - pub fn new( - srng: StableRng, - id: BlueprintStateId, - name: String, - ) -> MachineRegistryResult { + pub fn new(id: BlueprintStateId, name: String) -> MachineRegistryResult { Ok(Self { - unlocked_inner: Arc::new(BlueprintStateUnlockedInner { srng, id, name }), - inner: Arc::new(Mutex::new(BlueprintStateInner { + immutable: Arc::new(BlueprintStateImmutable { id, name }), + fields: Arc::new(BlueprintStateFields { limit_network_count: None, - networks: Vec::new(), + networks: imbl::Vector::new(), model: None, ipv4: None, ipv6: None, - })), + }), }) } - pub fn set_limit_network_count(&self, limit_network_count: Option) { - let mut inner = self.inner.lock(); - inner.limit_network_count = limit_network_count; + pub fn set_limit_network_count(&mut self, limit_network_count: Option) { + // Update fields + self.fields = Arc::new(BlueprintStateFields { + limit_network_count, + ..(*self.fields).clone() + }); } - pub fn set_model(&self, model: WeightedList) { - let mut inner = self.inner.lock(); - inner.model = Some(model); + pub fn set_model(&mut self, model: WeightedList) { + let model = Some(model.map(|x| Arc::new(x.clone()))); + // Update fields + self.fields = Arc::new(BlueprintStateFields { + model, + ..(*self.fields).clone() + }); } pub fn set_ipv4( - &self, + &mut self, params: BlueprintStateIpv4Params, gateway_params: Option, ) { - let mut inner = self.inner.lock(); - if inner.ipv4.is_none() { - inner.ipv4 = Some(BlueprintStateIpv4 { + let mut ipv4 = if let Some(ipv4) = self.fields.ipv4.clone() { + BlueprintStateIpv4 { params, ..ipv4 } + } else { + BlueprintStateIpv4 { params, gateway: None, - }); - } else { - inner.ipv4.as_mut().map(|ipv4| { - ipv4.params = params; - }); - } - let ipv4 = inner.ipv4.as_mut().expect("must exist"); + } + }; if ipv4.gateway.is_some() { if let Some(gateway_params) = gateway_params { @@ -127,25 +124,27 @@ impl BlueprintState { params: gateway_params, }) } + + // Update fields + self.fields = Arc::new(BlueprintStateFields { + ipv4: Some(ipv4), + ..(*self.fields).clone() + }); } pub fn set_ipv6( - &self, + &mut self, params: BlueprintStateIpv6Params, gateway_params: Option, ) { - let mut inner = self.inner.lock(); - if inner.ipv6.is_none() { - inner.ipv6 = Some(BlueprintStateIpv6 { + let mut ipv6 = if let Some(ipv6) = self.fields.ipv6.clone() { + BlueprintStateIpv6 { params, ..ipv6 } + } else { + BlueprintStateIpv6 { params, gateway: None, - }); - } else { - inner.ipv6.as_mut().map(|ipv6| { - ipv6.params = params; - }); - } - let ipv6 = inner.ipv6.as_mut().expect("must exist"); + } + }; if ipv6.gateway.is_some() { if let Some(gateway_params) = gateway_params { @@ -158,33 +157,38 @@ impl BlueprintState { params: gateway_params, }) } + + // Update fields + self.fields = Arc::new(BlueprintStateFields { + ipv6: Some(ipv6), + ..(*self.fields).clone() + }); } - pub fn is_active(&self) -> MachineRegistryResult { - let inner = self.inner.lock(); + pub fn is_active(&self, machine_registry_inner: &mut MachineRegistryInner) -> bool { + // Save a backup of the entire state + let backup = machine_registry_inner.clone(); - // 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 { - return Ok(false); - } - } + // Make a copy of this blueprint state + let mut current_state = self.clone(); - todo!("needs better implementation"); + // See what would happen if we try to generate this blueprint + let ok = current_state.generate(machine_registry_inner).is_ok(); - Ok(true) + // Restore the backup + *machine_registry_inner = backup; + + // Return if this worked or not + ok } fn generate_model_inner( - inner: &mut BlueprintStateInner, + &mut self, machine_registry_inner: &mut MachineRegistryInner, - network_state: NetworkState, + network_state: &mut NetworkState, ) -> MachineRegistryResult<()> { - let model_name = match inner.model.clone() { - Some(models) => machine_registry_inner - .srng() - .weighted_choice_ref(&models) - .clone(), + let model_name = match self.fields.model.as_ref() { + Some(models) => (**machine_registry_inner.srng().weighted_choice_ref(models)).clone(), None => machine_registry_inner.config().default_model.clone(), }; let Some(model) = machine_registry_inner.config().models.get(&model_name) else { @@ -196,17 +200,51 @@ impl BlueprintState { distance: model.distance.clone(), loss: model.loss, }; - network_state.with_model(params); + network_state.set_model(params); Ok(()) } + /// Network filter that ensures we can allocate an ipv4 gateway address on a network + fn gateway_network_filter_v4( + &self, + machine_registry_inner: &MachineRegistryInner, + network_state_id: NetworkStateId, + ) -> MachineRegistryResult { + // Get the network state + let network_state = machine_registry_inner + .network_states() + .get_state(network_state_id)?; + + // See if we can allocate on this network + let can_allocate = network_state.can_allocate_address_v4(None); + + Ok(can_allocate) + } + + /// Network filter that ensures we can allocate an ipv4 gateway address on a network + fn gateway_network_filter_v6( + &self, + machine_registry_inner: &MachineRegistryInner, + network_state_id: NetworkStateId, + ) -> MachineRegistryResult { + // Get the network state + let network_state = machine_registry_inner + .network_states() + .get_state(network_state_id)?; + + // See if we can allocate on this network + let can_allocate = network_state.can_allocate_address_v6(None); + + Ok(can_allocate) + } + fn generate_ipv4_inner( - inner: &mut BlueprintStateInner, + &mut self, machine_registry_inner: &mut MachineRegistryInner, - network_state: NetworkState, + network_state: &mut NetworkState, ) -> MachineRegistryResult<()> { - network_state.clear_ipv4(machine_registry_inner); - let Some(ipv4) = inner.ipv4.as_ref() else { + network_state.clear_ipv4(machine_registry_inner)?; + let Some(ipv4) = self.fields.ipv4.as_ref() else { return Ok(()); }; @@ -234,8 +272,7 @@ impl BlueprintState { .subnets .subnet4 .as_ref() - .map(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix)) - .flatten()) + .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) })? else { return Err(MachineRegistryError::NoAllocation); @@ -256,14 +293,16 @@ impl BlueprintState { .filter(|p| *p >= net.prefix_len()) .as_ref() .map(|wl| { - let subnet_prefix = - machine_registry_inner.srng().weighted_choice_ref(wl).clone(); + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); // Use an address pool temporarily to pick a subnet - let mut address_pool = - AddressPool::<()>::new(); + let mut address_pool = AddressPool::<()>::new(); address_pool.add_scope_v4(net); - address_pool.allocate_random_v4(machine_registry_inner.srng(), subnet_prefix, ()) + address_pool.allocate_random_v4( + machine_registry_inner.srng(), + subnet_prefix, + (), + ) }) .transpose()? .flatten(); @@ -276,40 +315,39 @@ impl BlueprintState { // Get networks which have subnets that would fit // our maximum requested prefix let Some(available_networks) = networks.try_filter(|network_id| { - let network_state = machine_registry_inner + let super_network_state = machine_registry_inner .network_states() .get_state(*network_id) .expect("must exist"); - Ok(network_state.can_allocate_subnet_v4(None, max_prefix)) + Ok(super_network_state.can_allocate_subnet_v4(None, max_prefix)) })? else { return Err(MachineRegistryError::NoAllocation); }; // Pick a network - let network_id = *machine_registry_inner + let super_network_id = *machine_registry_inner .srng() .weighted_choice_ref(&available_networks); - let network_state = machine_registry_inner + let mut super_network_state = machine_registry_inner .network_states() - .get_state(network_id) + .get_state(super_network_id) .expect("must exist"); // Pick a prefix that fits in this network and allocate from it let opt_subnet = ipv4 .params .prefix - .filter(|p| network_state.can_allocate_subnet_v4(None, *p)) + .filter(|p| super_network_state.can_allocate_subnet_v4(None, *p)) .as_ref() .map(|wl| { - let subnet_prefix = - machine_registry_inner.srng().weighted_choice_ref(wl).clone(); + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); // Allocate subnet from this network - network_state.allocate_subnet_v4( + super_network_state.allocate_subnet_v4( machine_registry_inner, - OwnerTag::Network(network_state.id()), + OwnerTag::Network(super_network_state.id()), None, subnet_prefix, ) @@ -319,7 +357,12 @@ impl BlueprintState { return Err(MachineRegistryError::NoAllocation); }; - (subnet, Some(network_id)) + // Update network state + machine_registry_inner + .network_states_mut() + .set_state(super_network_state); + + (subnet, Some(super_network_id)) } }; @@ -330,53 +373,266 @@ impl BlueprintState { let gateway_params = match ipv4.gateway.as_ref() { Some(v4gw) => { - let translation = machine_registry_inner + let translation = *machine_registry_inner .srng() - .weighted_choice_ref(&v4gw.params.translation) - .clone(); + .weighted_choice_ref(&v4gw.params.translation); let upnp = machine_registry_inner .srng() .probability_test(v4gw.params.upnp); - let location = match v4gw.params.locations { - Some(locations) => todo!(), - None => todo!(), - }; + let (external_network, external_address) = match v4gw.params.locations.as_ref() { + Some(locations_list) => { + // A external network location was specified, pick one + // Get a network to generate the machine on + let mut gateway_network_state = locations_list.pick( + machine_registry_inner, + |machine_registry_inner, id| { + self.gateway_network_filter_v4(machine_registry_inner, id) + }, + )?; + let gateway_network_state_id = gateway_network_state.id(); -xxx instantiate and clean up on failure + // Allocate an external address on this network + let external_interface_address = gateway_network_state + .allocate_address_v4( + machine_registry_inner, + OwnerTag::Gateway(network_state.id()), + None, + )?; + + // Update the network state + machine_registry_inner + .network_states_mut() + .set_state(gateway_network_state); + + ( + gateway_network_state_id, + Some(external_interface_address.ip), + ) + } + None => { + // No external network specified for gateway machine + // So use the same network as ourselves + (network_state.id(), None) + } + }; Some(NetworkStateIpv4GatewayParams { translation, upnp, - external_network: todo!(), + external_network, internal_address: None, - external_address: None, + external_address, }) } None => None, }; - network_state.set_ipv4(machine_registry_inner, params, gateway_params); + network_state.set_ipv4(machine_registry_inner, params, gateway_params)?; Ok(()) } fn generate_ipv6_inner( - inner: &mut BlueprintStateInner, + &mut self, machine_registry_inner: &mut MachineRegistryInner, - network_state: NetworkState, + network_state: &mut NetworkState, ) -> MachineRegistryResult<()> { - // + network_state.clear_ipv6(machine_registry_inner)?; + let Some(ipv6) = self.fields.ipv6.as_ref() else { + return Ok(()); + }; + + // Get maximum prefix + let max_prefix = ipv6 + .params + .prefix + .iter() + .max() + .copied() + .expect("must have at least one element"); + + // Get addresses for network + let (subnet, super_net) = match &ipv6.params.locations { + NetworkLocationsList::Allocations { allocations } => { + // Get allocations which have subnets that would fit + // our maximum requested prefix + let Some(alloc_subnets) = allocations.try_filter_map(|allocation_name| { + let allocation = machine_registry_inner + .config() + .allocations + .get(allocation_name) + .expect("must exist"); + Ok(allocation + .subnets + .subnet6 + .as_ref() + .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick an allocation + let subnets = machine_registry_inner + .srng() + .weighted_choice_ref(&alloc_subnets); + + // Pick a subnet + let net = *machine_registry_inner.srng().weighted_choice_ref(subnets); + + // Pick a prefix length that would fit in the subnet + let opt_subnet = ipv6 + .params + .prefix + .filter(|p| *p >= net.prefix_len()) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Use an address pool temporarily to pick a subnet + let mut address_pool = AddressPool::<()>::new(); + address_pool.add_scope_v6(net); + address_pool.allocate_random_v6( + machine_registry_inner.srng(), + subnet_prefix, + (), + ) + }) + .transpose()? + .flatten(); + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + (subnet, None) + } + NetworkLocationsList::Networks { networks } => { + // Get networks which have subnets that would fit + // our maximum requested prefix + let Some(available_networks) = networks.try_filter(|network_id| { + let super_network_state = machine_registry_inner + .network_states() + .get_state(*network_id) + .expect("must exist"); + + Ok(super_network_state.can_allocate_subnet_v6(None, max_prefix)) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick a network + let super_network_id = *machine_registry_inner + .srng() + .weighted_choice_ref(&available_networks); + let mut super_network_state = machine_registry_inner + .network_states() + .get_state(super_network_id) + .expect("must exist"); + + // Pick a prefix that fits in this network and allocate from it + let opt_subnet = ipv6 + .params + .prefix + .filter(|p| super_network_state.can_allocate_subnet_v6(None, *p)) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Allocate subnet from this network + super_network_state.allocate_subnet_v6( + machine_registry_inner, + OwnerTag::Network(super_network_state.id()), + None, + subnet_prefix, + ) + }) + .transpose()?; + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Update network state + machine_registry_inner + .network_states_mut() + .set_state(super_network_state); + + (subnet, Some(super_network_id)) + } + }; + + let params = NetworkStateIpv6Params { + allocation: subnet, + super_net, + }; + + let gateway_params = match ipv6.gateway.as_ref() { + Some(v6gw) => { + let translation = *machine_registry_inner + .srng() + .weighted_choice_ref(&v6gw.params.translation); + let upnp = machine_registry_inner + .srng() + .probability_test(v6gw.params.upnp); + + let (external_network, external_address) = match v6gw.params.locations.as_ref() { + Some(locations_list) => { + // A external network location was specified, pick one + // Get a network to generate the machine on + let mut gateway_network_state = locations_list.pick( + machine_registry_inner, + |machine_registry_inner, id| { + self.gateway_network_filter_v6(machine_registry_inner, id) + }, + )?; + let gateway_network_state_id = gateway_network_state.id(); + + // Allocate an external address on this network + let external_interface_address = gateway_network_state + .allocate_address_v6( + machine_registry_inner, + OwnerTag::Gateway(network_state.id()), + None, + )?; + + // Update the network state + machine_registry_inner + .network_states_mut() + .set_state(gateway_network_state); + + ( + gateway_network_state_id, + Some(external_interface_address.ip), + ) + } + None => { + // No external network specified for gateway machine + // So use the same network as ourselves + (network_state.id(), None) + } + }; + + Some(NetworkStateIpv6GatewayParams { + translation, + upnp, + external_network, + internal_address: None, + external_address, + }) + } + None => None, + }; + + network_state.set_ipv6(machine_registry_inner, params, gateway_params)?; + Ok(()) } pub fn generate( - &self, + &mut self, machine_registry_inner: &mut MachineRegistryInner, - ) -> MachineRegistryResult { - let mut inner = self.inner.lock(); - + ) -> MachineRegistryResult { // 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 { + if let Some(limit_network_count) = self.fields.limit_network_count { + if self.fields.networks.len() >= limit_network_count { return Err(MachineRegistryError::BlueprintComplete); } } @@ -385,39 +641,47 @@ xxx instantiate and clean up on failure let network_state_id = machine_registry_inner.network_states_mut().allocate_id(); // Create an anonymous network state - let network_state = NetworkState::new(network_state_id, None); + let mut network_state = + NetworkState::new(network_state_id, None, NetworkOrigin::Blueprint(self.id())); if let Err(e) = (|| { - Self::generate_model_inner(&mut *inner, machine_registry_inner, network_state.clone())?; - Self::generate_ipv4_inner(&mut *inner, machine_registry_inner, network_state.clone())?; - Self::generate_ipv6_inner(&mut *inner, machine_registry_inner, network_state.clone())?; + self.generate_model_inner(machine_registry_inner, &mut network_state)?; + self.generate_ipv4_inner(machine_registry_inner, &mut network_state)?; + self.generate_ipv6_inner(machine_registry_inner, &mut network_state)?; Ok(()) })() { // Release the network state and id if things failed to allocate network_state.release(machine_registry_inner); machine_registry_inner .network_states_mut() - .release_id(network_state_id); + .release_id(network_state_id) + .expect("must succeed"); return Err(e); } // Attach the state to the id machine_registry_inner .network_states_mut() - .attach_state(network_state.clone()); + .attach_state(network_state)?; // Record the newly instantiated network - inner.networks.push(network_state_id); + let mut networks = self.fields.networks.clone(); + networks.push_back(network_state_id); - Ok(network_state) + // Update fields + self.fields = Arc::new(BlueprintStateFields { + networks, + ..(*self.fields).clone() + }); + + Ok(network_state_id) } - pub fn for_each_network_id(&self, callback: F) -> MachineRegistryResult> + pub fn for_each_network_id(&self, mut callback: F) -> MachineRegistryResult> where - F: Fn(NetworkStateId) -> MachineRegistryResult>, + F: FnMut(NetworkStateId) -> MachineRegistryResult>, { - let inner = self.inner.lock(); - for network_id in &inner.networks { + for network_id in &self.fields.networks { if let Some(res) = callback(*network_id)? { return Ok(Some(res)); } @@ -425,23 +689,31 @@ xxx instantiate and clean up on failure Ok(None) } - pub fn on_network_released(&self, network_id: NetworkStateId) { - let mut inner = self.inner.lock(); - let pos = inner + pub fn on_network_released(&mut self, network_id: NetworkStateId) { + // Remove network from list + let pos = self + .fields .networks .iter() .position(|id| *id == network_id) .expect("must exist"); - inner.networks.remove(pos); + let mut networks = self.fields.networks.clone(); + networks.remove(pos); + + // Update fields + self.fields = Arc::new(BlueprintStateFields { + networks, + ..(*self.fields).clone() + }); } } impl State for BlueprintState { fn id(&self) -> StateId { - self.unlocked_inner.id.clone() + self.immutable.id } fn name(&self) -> Option { - Some(self.unlocked_inner.name.clone()) + Some(self.immutable.name.clone()) } } diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_locations_list.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_locations_list.rs index c503b0b9..e7787c05 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_locations_list.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_locations_list.rs @@ -1,5 +1,11 @@ use super::*; +#[derive(Debug, Clone)] +enum BlueprintAvailability { + Existing(NetworkState), + Generate(BlueprintState), +} + /// Locations where a machine can be instantiated #[derive(Debug, Clone)] pub enum MachineLocationsList { @@ -10,3 +16,194 @@ pub enum MachineLocationsList { blueprints: WeightedList, }, } + +impl MachineLocationsList { + pub fn can_pick( + &self, + machine_registry_inner: &mut MachineRegistryInner, + mut network_filter: F, + ) -> MachineRegistryResult + where + F: FnMut(&MachineRegistryInner, NetworkStateId) -> MachineRegistryResult, + { + match self { + MachineLocationsList::Networks { networks } => { + // Filter the weighted list of networks to those that are still active and or not yet started + if networks + .try_filter(|id| { + let network_state = + machine_registry_inner.network_states().get_state(*id)?; + self.is_network_available( + machine_registry_inner, + network_state, + &mut network_filter, + ) + })? + .is_none() + { + return Ok(false); + }; + } + MachineLocationsList::Blueprints { blueprints } => { + // Filter the weighted list of blueprints to those that are still active or not yet started and can allocate + if blueprints + .try_filter(|id| { + let blueprint_state = + machine_registry_inner.blueprint_states().get_state(*id)?; + + self.is_blueprint_available( + machine_registry_inner, + blueprint_state, + &mut network_filter, + ) + .map(|x| x.is_some()) + })? + .is_none() + { + return Ok(false); + }; + } + }; + Ok(true) + } + + pub fn pick( + &self, + machine_registry_inner: &mut MachineRegistryInner, + mut network_filter: F, + ) -> MachineRegistryResult + where + F: FnMut(&MachineRegistryInner, NetworkStateId) -> MachineRegistryResult, + { + // Get a network to generate the machine on + let network_state = match self { + MachineLocationsList::Networks { networks } => { + // Filter the weighted list of networks to those that are still active and or not yet started + let Some(available_networks) = networks.try_filter_map(|id| { + let network_state = machine_registry_inner.network_states().get_state(*id)?; + if self.is_network_available( + machine_registry_inner, + network_state.clone(), + &mut network_filter, + )? { + Ok(Some(network_state)) + } else { + Ok(None) + } + })? + else { + return Err(MachineRegistryError::NetworkComplete); + }; + + // Weighted choice of network now that we have a candidate list + let network_state = machine_registry_inner + .srng() + .weighted_choice(available_networks); + + // Return network state to use + network_state + } + MachineLocationsList::Blueprints { blueprints } => { + // Filter the weighted list of blueprints to those that are still active or not yet started and can allocate + let Some(available_blueprints) = blueprints.try_filter_map(|id| { + let blueprint_state = + machine_registry_inner.blueprint_states().get_state(*id)?; + + self.is_blueprint_available( + machine_registry_inner, + blueprint_state, + &mut network_filter, + ) + })? + else { + return Err(MachineRegistryError::BlueprintComplete); + }; + + // Weighted choice of blueprint now that we have a candidate list + match machine_registry_inner + .srng() + .weighted_choice(available_blueprints) + { + BlueprintAvailability::Existing(network_state) => network_state, + BlueprintAvailability::Generate(mut blueprint_state) => { + // Generate network state from blueprint state + let network_state_id = blueprint_state.generate(machine_registry_inner)?; + + // Update blueprint state + machine_registry_inner + .blueprint_states_mut() + .set_state(blueprint_state); + + // Return network state + machine_registry_inner + .network_states() + .get_state(network_state_id)? + } + } + } + }; + + Ok(network_state) + } + + fn is_network_available( + &self, + machine_registry_inner: &MachineRegistryInner, + network_state: NetworkState, + mut network_filter: F, + ) -> MachineRegistryResult + where + F: FnMut(&MachineRegistryInner, NetworkStateId) -> MachineRegistryResult, + { + // If the network is not active, it is not available + if !network_state.is_active()? { + return Ok(false); + } + + // Check the network filter + if !network_filter(machine_registry_inner, network_state.id())? { + return Ok(false); + } + + Ok(true) + } + + fn is_blueprint_available( + &self, + machine_registry_inner: &mut MachineRegistryInner, + blueprint_state: BlueprintState, + mut network_filter: F, + ) -> MachineRegistryResult> + where + F: FnMut(&MachineRegistryInner, NetworkStateId) -> MachineRegistryResult, + { + // See if the networks generated from this blueprint so far have availability + // in this template + if let Some(available_network_state) = blueprint_state.for_each_network_id(|id| { + // Check the network's availability + let network_state = machine_registry_inner.network_states().get_state(id)?; + if self.is_network_available( + machine_registry_inner, + network_state.clone(), + &mut network_filter, + )? { + // We found one + return Ok(Some(network_state)); + } + // Try next network + Ok(None) + })? { + // We found a usable network + return Ok(Some(BlueprintAvailability::Existing( + available_network_state, + ))); + } + + // If the blueprint is active, it is available because it can make a new network + if blueprint_state.is_active(machine_registry_inner) { + return Ok(Some(BlueprintAvailability::Generate(blueprint_state))); + } + + Ok(None) + } +} diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_state.rs index 5ef25392..213dcc21 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_state.rs @@ -61,14 +61,18 @@ impl MachineState { } pub fn release(mut self, machine_registry_inner: &mut MachineRegistryInner) { - self.release_all_interfaces(machine_registry_inner); + self.release_all_interfaces(machine_registry_inner) + .expect("must succeed"); if let MachineOrigin::Template(generating_template) = self.immutable.origin { - let template_state = machine_registry_inner - .template_states_mut() + let mut template_state = machine_registry_inner + .template_states() .get_state(generating_template) .expect("must exist"); template_state.on_machine_released(self.id()); + machine_registry_inner + .template_states_mut() + .set_state(template_state); } } @@ -108,7 +112,7 @@ impl MachineState { if self.fields.interfaces.contains_key(&interface_key) { return Err(MachineRegistryError::DuplicateName); } - let flags = opt_interface_flags.unwrap_or_else(|| InterfaceFlags { + let flags = opt_interface_flags.unwrap_or(InterfaceFlags { is_loopback: false, is_running: true, is_point_to_point: false, @@ -175,7 +179,7 @@ impl MachineState { .set_state(network_state); // Get address flags - let flags = opt_address_flags.unwrap_or_else(|| AddressFlags { + let flags = opt_address_flags.unwrap_or(AddressFlags { is_dynamic, is_temporary: false, is_preferred: true, @@ -240,7 +244,7 @@ impl MachineState { .set_state(network_state); // Get address flags - let flags = opt_address_flags.unwrap_or_else(|| AddressFlags { + let flags = opt_address_flags.unwrap_or(AddressFlags { is_dynamic, is_temporary: false, is_preferred: true, @@ -536,103 +540,11 @@ impl MachineState { 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::::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 { fn id(&self) -> StateId { - self.immutable.id.clone() + self.immutable.id } fn name(&self) -> Option { diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs index afc5eb4f..8699807c 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs @@ -1,7 +1,7 @@ use super::*; /// Locations where a network can be instantiated -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum NetworkLocationsList { /// Network will be a new allocation Allocations { allocations: WeightedList }, diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs index a32e930a..9a2541db 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs @@ -140,11 +140,14 @@ impl NetworkState { pub fn release(self, machine_registry_inner: &mut MachineRegistryInner) { if let NetworkOrigin::Blueprint(generating_blueprint) = self.immutable.origin { - let blueprint_state = machine_registry_inner - .blueprint_states_mut() + let mut blueprint_state = machine_registry_inner + .blueprint_states() .get_state(generating_blueprint) .expect("must exist"); blueprint_state.on_network_released(self.id()); + machine_registry_inner + .blueprint_states_mut() + .set_state(blueprint_state) } } @@ -673,7 +676,7 @@ impl NetworkState { impl State for NetworkState { fn id(&self) -> StateId { - self.immutable.id.clone() + self.immutable.id } fn name(&self) -> Option { diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/profile_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/profile_state.rs index f4bb3777..1bea01c1 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/profile_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/profile_state.rs @@ -47,7 +47,7 @@ impl ProfileState { impl State for ProfileState { fn id(&self) -> StateId { - self.immutable.id.clone() + self.immutable.id } fn name(&self) -> Option { diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/state_allocator.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/state_allocator.rs index 07794448..e252921f 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/state_allocator.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/state_allocator.rs @@ -181,45 +181,9 @@ impl StateAllocator { self.state_by_id.insert(state.id().0, Some(state)); } - // pub fn update_state< - // R, - // F: FnOnce(&mut StateAllocator, S) -> MachineRegistryResult<(R, Option)>, - // >( - // &mut self, - // id: StateId, - // callback: F, - // ) -> MachineRegistryResult { - // // Get state to update - // let state = { - // // Get the allocator slot - // let Some(opt_state) = self.state_by_id.get(&id.0) else { - // return Err(MachineRegistryError::InvalidId); - // }; - // let Some(state) = opt_state else { - // return Err(MachineRegistryError::NotAttached); - // }; - - // // Make copy of state to update - // state.clone() - // }; - - // // Call the callback - // let id = state.id(); - // let (res, opt_new_state) = callback(self, state)?; - // if let Some(new_state) = opt_new_state { - // assert_eq!(id, new_state.id(), "state id must not change"); - // self.state_by_id.insert(id.0, Some(new_state)); - // } - - // Ok(res) - // } - pub fn get_state_id_by_name(&self, name: &String) -> Option> { // Get the id associated with this name - let Some(id) = self.state_id_by_name.get(name) else { - return None; - }; - + let id = self.state_id_by_name.get(name)?; Some(StateId::new(*id)) } } diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/template_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/template_state.rs index b8e9c90e..b8d3fa3f 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/template_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/template_state.rs @@ -2,7 +2,9 @@ use super::*; #[derive(Debug)] struct TemplateStateImmutable { + /// The unique id of this template id: TemplateStateId, + /// The name of this template state name: String, } @@ -12,12 +14,6 @@ struct PerNetworkInfo { machines: imbl::HashSet, } -#[derive(Debug, Clone)] -enum BlueprintAvailability { - Existing(NetworkState), - Generate(BlueprintState), -} - #[derive(Debug, Clone)] struct TemplateStateFields { limit_machine_count: Option, @@ -110,14 +106,27 @@ impl TemplateState { }); } - fn is_network_available(&self, network_state: NetworkState) -> MachineRegistryResult { - // If the network is not active, it is not available - if !network_state.is_active()? { - return Ok(false); - } + pub fn is_active(&self, machine_registry_inner: &mut MachineRegistryInner) -> bool { + // Save a backup of the entire state + let backup = machine_registry_inner.clone(); + // Make a copy of this template state + let mut current_state = self.clone(); + + // See what would happen if we try to generate this template + let ok = current_state.generate(machine_registry_inner).is_ok(); + + // Restore the backup + *machine_registry_inner = backup; + + // Return if this worked or not + ok + } + + /// Network filter that keeps this template generation within per-network limits + fn network_filter(&self, network_state_id: NetworkStateId) -> MachineRegistryResult { // Get the per network info - let Some(pni) = self.fields.machines_per_network.get(&network_state.id()) else { + let Some(pni) = self.fields.machines_per_network.get(&network_state_id) else { // If we haven't allocated anything in the network yet it is // by definition available return Ok(true); @@ -130,87 +139,6 @@ impl TemplateState { return Ok(false); } } - - Ok(true) - } - - fn is_blueprint_available( - &self, - machine_registry_inner: &MachineRegistryInner, - blueprint_state: BlueprintState, - ) -> MachineRegistryResult> { - // See if the networks generated from this blueprint so far have availability - // in this template - if let Some(available_network_state) = blueprint_state.for_each_network_id(|id| { - // Check the network's availability - let network_state = machine_registry_inner.network_states().get_state(id)?; - if self.is_network_available(network_state.clone())? { - // We found one - return Ok(Some(network_state)); - } - // Try next network - Ok(None) - })? { - // We found a usable network - return Ok(Some(BlueprintAvailability::Existing( - available_network_state, - ))); - } - - // If the blueprint is active, it is available because it can make a new network - if blueprint_state.is_active()? { - return Ok(Some(BlueprintAvailability::Generate(blueprint_state))); - } - - Ok(None) - } - - pub fn is_active( - &self, - machine_registry_inner: &MachineRegistryInner, - ) -> MachineRegistryResult { - // See if we have reached our machine limit - if let Some(limit_machine_count) = self.fields.limit_machine_count { - if self.fields.machines.len() >= limit_machine_count { - return Ok(false); - } - } - - let Some(locations_list) = self.fields.locations_list.as_ref() else { - return Ok(false); - }; - - match locations_list { - MachineLocationsList::Networks { networks } => { - // Filter the weighted list of networks to those that are still active and or not yet started - if networks - .try_filter(|id| { - let network_state = - machine_registry_inner.network_states().get_state(*id)?; - self.is_network_available(network_state) - })? - .is_none() - { - return Ok(false); - }; - } - MachineLocationsList::Blueprints { blueprints } => { - // Filter the weighted list of blueprints to those that are still active or not yet started and can allocate - if blueprints - .try_filter(|id| { - let blueprint_state = - machine_registry_inner.blueprint_states().get_state(*id)?; - - self.is_blueprint_available(machine_registry_inner, blueprint_state) - .map(|x| x.is_some()) - })? - .is_none() - { - return Ok(false); - }; - } - }; - Ok(true) } @@ -220,7 +148,7 @@ impl TemplateState { ) -> MachineRegistryResult { // See if we have reached our machine limit if let Some(limit_machine_count) = self.fields.limit_machine_count { - if self.fields.machines.len() < limit_machine_count.try_into().unwrap_or(usize::MAX) { + if self.fields.machines.len() < limit_machine_count { return Err(MachineRegistryError::TemplateComplete); } } @@ -231,65 +159,26 @@ impl TemplateState { }; // Get a network to generate the machine on - let network_state = match locations_list { - MachineLocationsList::Networks { networks } => { - // Filter the weighted list of networks to those that are still active and or not yet started - let Some(available_networks) = networks.try_filter_map(|id| { - let network_state = machine_registry_inner.network_states().get_state(*id)?; - if self.is_network_available(network_state.clone())? { - Ok(Some(network_state)) - } else { - Ok(None) - } - })? - else { - return Err(MachineRegistryError::NetworkComplete); - }; - - // Weighted choice of network now that we have a candidate list - let network_state = machine_registry_inner - .srng() - .weighted_choice_ref(&available_networks); - - // Return network state to use - network_state.clone() - } - MachineLocationsList::Blueprints { blueprints } => { - // Filter the weighted list of blueprints to those that are still active or not yet started and can allocate - let Some(available_blueprints) = blueprints.try_filter_map(|id| { - let blueprint_state = - machine_registry_inner.blueprint_states().get_state(*id)?; - - self.is_blueprint_available(machine_registry_inner, blueprint_state) - })? - else { - return Err(MachineRegistryError::BlueprintComplete); - }; - - // Weighted choice of blueprint now that we have a candidate list - match machine_registry_inner - .srng() - .weighted_choice_ref(&available_blueprints) - { - BlueprintAvailability::Existing(network_state) => network_state.clone(), - BlueprintAvailability::Generate(blueprint_state) => { - blueprint_state.generate(machine_registry_inner)? - } - } - } - }; + let network_state = + locations_list.pick(machine_registry_inner, |_, x| self.network_filter(x))?; // Allocate a machine id let machine_state_id = machine_registry_inner.machine_states_mut().allocate_id(); -xxx continue here - // Create an anonymous machine state - let machine_state = MachineState::new(machine_state_id, None); + let mut machine_state = + MachineState::new(machine_state_id, None, MachineOrigin::Template(self.id())); + // Scope to release state on error if let Err(e) = (|| { // Build out the machine state from the template - machine_state.set_disable_capabilities(inner.disable_capabilities.clone()); + machine_state.set_disable_capabilities( + self.fields + .disable_capabilities + .iter() + .map(|x| (**x).clone()) + .collect(), + ); machine_state.set_bootstrap(false); // Make the default route interface @@ -307,53 +196,67 @@ xxx continue here machine_state.release(machine_registry_inner); machine_registry_inner .machine_states_mut() - .release_id(machine_state_id); + .release_id(machine_state_id) + .expect("must succeed"); return Err(e); } // Attach the state to the id machine_registry_inner .machine_states_mut() - .attach_state(machine_state.clone()); + .attach_state(machine_state) + .expect("must succeed"); // Record the newly instantiated machine - inner.machines.insert(machine_state_id); - let limit_machines_per_network = inner.limit_machines_per_network.clone(); - let per_network_info = inner - .machines_per_network + let machines = self.fields.machines.update(machine_state_id); + let mut machines_per_network = self.fields.machines_per_network.clone(); + let per_network_info = machines_per_network .entry(network_state.id()) .or_insert_with(|| { - let limit_machine_count = limit_machines_per_network.map(|wl| { - machine_registry_inner - .srng() - .weighted_choice_ref(&wl) - .clone() - }); + let limit_machine_count = self + .fields + .limit_machines_per_network + .as_ref() + .map(|wl| *machine_registry_inner.srng().weighted_choice_ref(wl)); PerNetworkInfo { limit_machine_count, - machines: HashSet::new(), + machines: imbl::HashSet::new(), } }); per_network_info.machines.insert(machine_state_id); - Ok(machine_state) + // Update fields + self.fields = Arc::new(TemplateStateFields { + machines, + machines_per_network, + ..(*self.fields).clone() + }); + + Ok(machine_state_id) } - pub fn on_machine_released(&self, machine_state_id: MachineStateId) { - let mut inner = self.inner.lock(); - inner.machines.remove(&machine_state_id); - for (_network_id, pni) in &mut inner.machines_per_network { + pub fn on_machine_released(&mut self, machine_state_id: MachineStateId) { + let machines = self.fields.machines.without(&machine_state_id); + let mut machines_per_network = self.fields.machines_per_network.clone(); + for (_network_id, pni) in machines_per_network.iter_mut() { pni.machines.remove(&machine_state_id); } + + // Update fields + self.fields = Arc::new(TemplateStateFields { + machines, + machines_per_network, + ..(*self.fields).clone() + }); } } impl State for TemplateState { fn id(&self) -> StateId { - self.unlocked_inner.id.clone() + self.immutable.id } fn name(&self) -> Option { - Some(self.unlocked_inner.name.clone()) + Some(self.immutable.name.clone()) } } diff --git a/veilid-tools/src/virtual_network/router_server/weighted_list.rs b/veilid-tools/src/virtual_network/router_server/weighted_list.rs index 72abb4a5..8e803b72 100644 --- a/veilid-tools/src/virtual_network/router_server/weighted_list.rs +++ b/veilid-tools/src/virtual_network/router_server/weighted_list.rs @@ -56,6 +56,29 @@ impl WeightedList { } } + pub fn map(&self, mut map: F) -> WeightedList + where + F: FnMut(&T) -> S, + S: fmt::Debug + Clone, + { + match self { + WeightedList::Single(v) => WeightedList::Single(map(v)), + WeightedList::List(vec) => { + let mut out = Vec::>::with_capacity(vec.len()); + for v in vec { + out.push(match v { + Weighted::Weighted { item, weight } => Weighted::Weighted { + item: map(item), + weight: *weight, + }, + Weighted::Unweighted(item) => Weighted::Unweighted(map(item)), + }); + } + WeightedList::List(out) + } + } + } + pub fn filter(&self, mut filter: F) -> Option> where F: FnMut(&T) -> bool, @@ -63,9 +86,10 @@ impl WeightedList { match self { WeightedList::Single(v) => { if filter(v) { - return Some(self.clone()); + Some(self.clone()) + } else { + None } - return None; } WeightedList::List(vec) => { let mut out = Vec::>::with_capacity(vec.len()); @@ -90,9 +114,10 @@ impl WeightedList { match self { WeightedList::Single(v) => { if filter(v)? { - return Ok(Some(self.clone())); + Ok(Some(self.clone())) + } else { + Ok(None) } - return Ok(None); } WeightedList::List(vec) => { let mut out = Vec::>::with_capacity(vec.len()); @@ -117,9 +142,10 @@ impl WeightedList { match self { WeightedList::Single(v) => { if let Some(item) = filter(v)? { - return Ok(Some(WeightedList::Single(item))); + Ok(Some(WeightedList::Single(item))) + } else { + Ok(None) } - return Ok(None); } WeightedList::List(vec) => { let mut out = Vec::>::with_capacity(vec.len()); @@ -159,7 +185,7 @@ impl core::ops::Index for WeightedList { fn index(&self, index: usize) -> &Self::Output { match self { - WeightedList::Single(s) => &s, + WeightedList::Single(s) => s, WeightedList::List(vec) => vec[index].item(), } }