Git-Mirrors/veilid
1
0
Fork 0
mirror of https://gitlab.com/veilid/veilid.git synced 2025-01-27 06:47:16 -05:00
Code Issues Packages Projects Releases Wiki Activity

[skip ci] start conversion to using persistent/immutable states

Browse Source
This commit is contained in:
Christien Rioux 2024-12-28 21:13:40 -05:00
parent 9c68df4274
commit 751d27c30d
19 changed files with 1363 additions and 588 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
Cargo.lock generated
View File

@ -728,6 +728,12 @@ version = "2.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
[[package]]
name = "bitmaps"
version = "3.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d084b0137aaa901caf9f1e8b21daa6aa24d41cd806e111335541eff9683bd6"
[[package]]
name = "blake2"
version = "0.10.6"
@ -2791,6 +2797,29 @@ dependencies = [
"libc",
]
[[package]]
name = "imbl"
version = "3.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bc3be8d8cd36f33a46b1849f31f837c44d9fa87223baee3b4bd96b8f11df81eb"
dependencies = [
"bitmaps",
"imbl-sized-chunks",
"rand_core",
"rand_xoshiro",
"serde",
"version_check",
]
[[package]]
name = "imbl-sized-chunks"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "144006fb58ed787dcae3f54575ff4349755b00ccc99f4b4873860b654be1ed63"
dependencies = [
"bitmaps",
]
[[package]]
name = "indent"
version = "0.1.1"
@ -4467,6 +4496,15 @@ dependencies = [
"getrandom",
]
[[package]]
name = "rand_xoshiro"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
dependencies = [
"rand_core",
]
[[package]]
name = "range-set-blaze"
version = "0.1.16"
@ -6590,6 +6628,7 @@ dependencies = [
"futures_codec",
"getrandom",
"ifstructs",
"imbl",
"ipnet",
"jni",
"jni-sys",

1
veilid-tools/Cargo.toml
View File

@ -84,6 +84,7 @@ backtrace = "0.3.71"
fn_name = "0.1.0"
range-set-blaze = "0.1.16"
flume = { version = "0.11.0", features = ["async"] }
imbl = { version = "3.0.0", features = ["serde"] }
# Dependencies for native builds only
# Linux, Windows, Mac, iOS, Android

0
veilid-tools/src/virtual_network/router_server/machine_registry/weighted_list.rs → veilid-tools/src/rollback_guard.rs
View File

75
veilid-tools/src/virtual_network/router_server/config.rs
View File

@ -312,10 +312,40 @@ fn validate_blueprint_limits(limits: &BlueprintLimits) -> Result<(), ValidationE
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlueprintIpv4 {
#[serde(untagged)]
pub enum BlueprintLocation {
Allocation {
allocation: WeightedList<String>,
},
Network {
#[serde(default)]
pub allocation: Option<String>,
pub additional_prefix: u8,
network: Option<WeightedList<String>>,
},
}
fn validate_blueprint_location(
value: &BlueprintLocation,
context: &ValidateContext,
) -> Result<(), ValidationError> {
match value {
BlueprintLocation::Allocation { allocation } => {
allocation.try_for_each(|a| validate_allocation_exists(a, context))?;
}
BlueprintLocation::Network { network } => {
if let Some(network) = network {
network.try_for_each(|n| validate_network_exists(n, context))?;
}
}
}
Ok(())
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlueprintIpv4 {
#[serde(flatten)]
pub location: BlueprintLocation,
pub prefix: WeightedList<u8>,
#[serde(default)]
pub gateway: Option<BlueprintGateway>,
}
@ -324,14 +354,15 @@ fn validate_blueprint_ipv4(
blueprint_ipv4: &BlueprintIpv4,
context: &ValidateContext,
) -> Result<(), ValidationError> {
if let Some(allocation) = &blueprint_ipv4.allocation {
validate_allocation_exists(allocation, context)?;
}
if blueprint_ipv4.additional_prefix > 32 {
validate_blueprint_location(&blueprint_ipv4.location, context)?;
blueprint_ipv4.prefix.validate_once()?;
blueprint_ipv4.prefix.try_for_each(|x| {
if *x > 32 {
return Err(ValidationError::new("badprefix")
.with_message("ipv4 blueprint additional prefix too long".into()));
.with_message("ipv4 blueprint prefix too long".into()));
}
Ok(())
})?;
if let Some(gateway) = &blueprint_ipv4.gateway {
validate_blueprint_gateway(gateway, context)?;
@ -341,9 +372,9 @@ fn validate_blueprint_ipv4(
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlueprintIpv6 {
#[serde(default)]
pub allocation: Option<String>,
pub additional_prefix: u8,
#[serde(flatten)]
pub allocation: BlueprintLocation,
pub prefix: WeightedList<u8>,
#[serde(default)]
pub gateway: Option<BlueprintGateway>,
}
@ -352,14 +383,15 @@ fn validate_blueprint_ipv6(
blueprint_ipv6: &BlueprintIpv6,
context: &ValidateContext,
) -> Result<(), ValidationError> {
if let Some(allocation) = &blueprint_ipv6.allocation {
validate_allocation_exists(allocation, context)?;
}
if blueprint_ipv6.additional_prefix > 128 {
validate_blueprint_location(&blueprint_ipv6.allocation, context)?;
blueprint_ipv6.prefix.validate_once()?;
blueprint_ipv6.prefix.try_for_each(|x| {
if *x > 128 {
return Err(ValidationError::new("badprefix")
.with_message("ipv6 blueprint additional prefix too long".into()));
.with_message("ipv6 blueprint prefix too long".into()));
}
Ok(())
})?;
if let Some(gateway) = &blueprint_ipv6.gateway {
validate_blueprint_gateway(gateway, context)?;
@ -371,7 +403,8 @@ fn validate_blueprint_ipv6(
pub struct BlueprintGateway {
pub translation: WeightedList<Translation>,
pub upnp: Probability,
pub network: Option<String>,
#[serde(flatten)]
pub location: TemplateLocation,
}
fn validate_blueprint_gateway(
@ -379,9 +412,7 @@ fn validate_blueprint_gateway(
context: &ValidateContext,
) -> Result<(), ValidationError> {
gateway.translation.validate_once()?;
if let Some(network) = &gateway.network {
validate_network_exists(network, context)?;
}
validate_template_location(&gateway.location, context)?;
Ok(())
}
////////////////////////////////////////////////////////////////

16
veilid-tools/src/virtual_network/router_server/default_config.yml
View File

@ -124,22 +124,23 @@ networks:
# Blueprints are used to generate Networks for use with Machines
blueprints:
# A subnet of the internet directly attached with no translation
# * A subnet of the internet directly attached with no translation
# with both ipv4 and ipv6 networking
direct:
ipv4:
prefix: 24
ipv6:
prefix: 64
# An ipv4-only subnet of the internet directly attached with no translation
# * An ipv4-only subnet of the internet directly attached with no translation
direct_ipv4_no_ipv6:
ipv4:
prefix: 24
# An ipv6-only subnet of the internet directly attached with no translation
# * An ipv6-only subnet of the internet directly attached with no translation
direct_ipv6_no_ipv4:
ipv6:
prefix: 64
# An ipv4-only subnet of the internet attached via NAT
# * An ipv4-only subnet of the internet attached via NAT to an
# an ipv4-only subnet of the internet directly attached with no translation
nat_ipv4_no_ipv6:
ipv4:
allocation: "$private"
@ -147,8 +148,10 @@ blueprints:
gateway:
translation: "port_restricted"
upnp: 0.25
# An ipv4 subnet of the internet attached via NAT and
# an ipv6 subnet of the internet directly attached with no translation
blueprint: "direct_ipv4_no_ipv6"
# * An ipv4 subnet of the internet attached via NAT to an
# an ipv4-only subnet of the internet directly attached with no translation
# * An ipv6 subnet of the internet directly attached with no translation
nat_ipv4_direct_ipv6:
ipv4:
allocation: "$private"
@ -156,6 +159,7 @@ blueprints:
gateway:
translation: "port_restricted"
upnp: 0.25
blueprint: "direct_ipv4_no_ipv6"
ipv6:
prefix: 56

282
veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs
View File

@ -1,38 +1,44 @@
use super::*;
use ipnet::*;
#[derive(Debug)]
pub struct AddressPool {
srng: StableRng,
#[derive(Debug, Clone)]
pub struct AddressPool<T: fmt::Debug + Clone> {
scope_v4: imbl::Vector<Ipv4Net>,
scope_v6: imbl::Vector<Ipv6Net>,
scope_v4: Vec<Ipv4Net>,
scope_v6: Vec<Ipv6Net>,
allocated_v4: imbl::Vector<Ipv4Net>,
allocated_v6: imbl::Vector<Ipv6Net>,
allocated_v4: Vec<Ipv4Net>,
allocated_v6: Vec<Ipv6Net>,
owner_tags_v4: imbl::HashMap<Ipv4Net, Option<T>>,
owner_tags_v6: imbl::HashMap<Ipv6Net, Option<T>>,
}
impl AddressPool {
pub fn new(srng: StableRng) -> Self {
impl<T: fmt::Debug + Clone> AddressPool<T> {
pub fn new() -> Self {
Self {
srng,
scope_v4: Vec::new(),
scope_v6: Vec::new(),
allocated_v4: Vec::new(),
allocated_v6: Vec::new(),
scope_v4: imbl::Vector::new(),
scope_v6: imbl::Vector::new(),
allocated_v4: imbl::Vector::new(),
allocated_v6: imbl::Vector::new(),
owner_tags_v4: imbl::HashMap::new(),
owner_tags_v6: imbl::HashMap::new(),
}
}
/////////////////////////////////////////////////////////////////////
pub fn add_scope_v4(&mut self, allocation: Ipv4Net) {
self.scope_v4.push(allocation);
self.scope_v4 = Ipv4Net::aggregate(&self.scope_v4);
let mut scopes = self.scope_v4.iter().copied().collect::<Vec<_>>();
scopes.push(allocation);
scopes = Ipv4Net::aggregate(&scopes);
self.scope_v4 = scopes.into();
}
pub fn add_scope_v6(&mut self, allocation: Ipv6Net) {
self.scope_v6.push(allocation);
self.scope_v6 = Ipv6Net::aggregate(&self.scope_v6);
let mut scopes = self.scope_v6.iter().copied().collect::<Vec<_>>();
scopes.push(allocation);
scopes = Ipv6Net::aggregate(&scopes);
self.scope_v6 = scopes.into();
}
pub fn is_in_scope_v4(&self, allocation: Ipv4Net) -> bool {
@ -53,30 +59,66 @@ impl AddressPool {
false
}
/////////////////////////////////////////////////////////////////////
pub fn can_allocate_v6(&self, prefix: u8) -> MachineRegistryResult<bool> {
if prefix > 128 {
return Err(MachineRegistryError::InvalidPrefix);
}
pub fn allocate_v4(&mut self, allocation: Ipv4Net) -> MachineRegistryResult<()> {
let mut srng = StableRng::new(0);
let opt_allocation = self.find_random_allocation_v6(&mut srng, prefix);
Ok(opt_allocation.is_some())
}
pub fn can_allocate_v4(&self, prefix: u8) -> MachineRegistryResult<bool> {
if prefix > 32 {
return Err(MachineRegistryError::InvalidPrefix);
}
let mut srng = StableRng::new(0);
let opt_allocation = self.find_random_allocation_v4(&mut srng, prefix);
Ok(opt_allocation.is_some())
}
pub fn reserve_allocation_v4(
&mut self,
allocation: Ipv4Net,
opt_tag: Option<T>,
) -> MachineRegistryResult<()> {
// Ensure the allocation is in our scope
if !self.is_in_scope_v4(allocation) {
return Err(MachineRegistryError::NoAllocation);
}
// Only reserve if it's not overlapping an allocation
if !self.get_overlaps_v4(allocation).is_empty() {
return Err(MachineRegistryError::NoAllocation);
}
// Add to our allocated pool
self.allocated_v4.push(allocation);
self.allocated_v4 = Ipv4Net::aggregate(&self.allocated_v4);
self.allocated_v4.insert_ord(allocation);
self.owner_tags_v4.insert(allocation, opt_tag);
Ok(())
}
pub fn allocate_v6(&mut self, allocation: Ipv6Net) -> MachineRegistryResult<()> {
pub fn reserve_allocation_v6(
&mut self,
allocation: Ipv6Net,
opt_tag: Option<T>,
) -> MachineRegistryResult<()> {
// Ensure the allocation is in our scope
if !self.is_in_scope_v6(allocation) {
return Err(MachineRegistryError::NoAllocation);
}
// Only reserve if it's not overlapping an allocation
if !self.get_overlaps_v6(allocation).is_empty() {
return Err(MachineRegistryError::NoAllocation);
}
// Add to our allocated pool
self.allocated_v6.push(allocation);
self.allocated_v6 = Ipv6Net::aggregate(&self.allocated_v6);
self.allocated_v6.insert_ord(allocation);
self.owner_tags_v6.insert(allocation, opt_tag);
Ok(())
}
@ -103,6 +145,166 @@ impl AddressPool {
overlaps
}
pub fn allocate_random_v4(
&mut self,
srng: &mut StableRng,
prefix: u8,
tag: T,
) -> MachineRegistryResult<Option<Ipv4Net>> {
if prefix > 32 {
return Err(MachineRegistryError::InvalidPrefix);
}
let opt_allocation = self.find_random_allocation_v4(srng, prefix);
// If we found a free subnet, add it to our allocations
if let Some(allocation) = opt_allocation {
// Add to our allocated pool
self.allocated_v4.insert_ord(allocation);
self.owner_tags_v4.insert(allocation, Some(tag));
return Ok(Some(allocation));
}
// No allocation
Ok(None)
}
pub fn allocate_random_v6(
&mut self,
srng: &mut StableRng,
prefix: u8,
tag: T,
) -> MachineRegistryResult<Option<Ipv6Net>> {
if prefix > 128 {
return Err(MachineRegistryError::InvalidPrefix);
}
let opt_allocation = self.find_random_allocation_v6(srng, prefix);
// If we found a free subnet, add it to our allocations
if let Some(allocation) = opt_allocation {
// Add to our allocated pool
self.allocated_v6.insert_ord(allocation);
self.owner_tags_v6.insert(allocation, Some(tag));
return Ok(Some(allocation));
}
// No allocation
Ok(None)
}
pub fn release_allocation_v4(
&mut self,
allocation: Ipv4Net,
) -> MachineRegistryResult<Option<T>> {
let Some(pos) = self.allocated_v4.iter().position(|x| *x == allocation) else {
return Err(MachineRegistryError::NoAllocation);
};
let Some(opt_tag) = self.owner_tags_v4.remove(&allocation) else {
return Err(MachineRegistryError::NoAllocation);
};
self.allocated_v4.remove(pos);
Ok(opt_tag)
}
pub fn release_allocation_v6(
&mut self,
allocation: Ipv6Net,
) -> MachineRegistryResult<Option<T>> {
let Some(pos) = self.allocated_v6.iter().position(|x| *x == allocation) else {
return Err(MachineRegistryError::NoAllocation);
};
let Some(opt_tag) = self.owner_tags_v6.remove(&allocation) else {
return Err(MachineRegistryError::NoAllocation);
};
self.allocated_v4.remove(pos);
Ok(opt_tag)
}
pub fn is_ipv4(&self) -> bool {
!self.scope_v4.is_empty()
}
pub fn is_ipv4_allocated(&self) -> bool {
self.is_ipv4() && !self.allocated_v4.is_empty()
}
pub fn is_ipv6(&self) -> bool {
!self.scope_v6.is_empty()
}
pub fn is_ipv6_allocated(&self) -> bool {
self.is_ipv6() && !self.allocated_v6.is_empty()
}
pub fn is_in_use<F: FnMut(IpNet, &T) -> bool>(&self, mut check: F) -> bool {
for (netv4, opt_tag) in self.owner_tags_v4.iter() {
if let Some(tag) = opt_tag.as_ref() {
if check(IpNet::V4(*netv4), tag) {
return true;
}
}
}
for (netv6, opt_tag) in self.owner_tags_v6.iter() {
if let Some(tag) = opt_tag.as_ref() {
if check(IpNet::V6(*netv6), tag) {
return true;
}
}
}
false
}
pub fn clear_ipv4<F: FnMut(Ipv4Net, &T) -> bool>(
&mut self,
mut check: F,
) -> MachineRegistryResult<()> {
if !self.is_ipv4() {
return Ok(());
}
if self.is_in_use(|n, t| match n {
IpNet::V4(ipv4_net) => check(ipv4_net, t),
IpNet::V6(_ipv6_net) => false,
}) {
return Err(MachineRegistryError::ResourceInUse);
}
assert!(self.owner_tags_v4.is_empty(), "tags should be empty");
self.scope_v4.clear();
self.allocated_v4.clear();
self.owner_tags_v4.clear();
Ok(())
}
pub fn clear_ipv6<F: FnMut(Ipv6Net, &T) -> bool>(
&mut self,
mut check: F,
) -> MachineRegistryResult<()> {
if !self.is_ipv6() {
return Ok(());
}
if self.is_in_use(|n, t| match n {
IpNet::V4(_ipv4_net) => false,
IpNet::V6(ipv6_net) => check(ipv6_net, t),
}) {
return Err(MachineRegistryError::ResourceInUse);
}
assert!(self.owner_tags_v6.is_empty(), "tags should be empty");
self.scope_v6.clear();
self.allocated_v6.clear();
self.owner_tags_v6.clear();
Ok(())
}
/////////////////////////////////////////////////////////////////////
fn range_in_prefix_32(scope_prefix: u8, iterable_prefix_bits: u8) -> u32 {
// If we're allocating addresses, exclude scope's network and broadcast address
if scope_prefix + iterable_prefix_bits == 32 {
@ -137,7 +339,7 @@ impl AddressPool {
}
}
pub fn allocate_random_v4(&mut self, prefix: u8) -> Option<Ipv4Net> {
fn find_random_allocation_v4(&self, srng: &mut StableRng, prefix: u8) -> Option<Ipv4Net> {
// Scope ranges to iterate
let mut scope_ranges = Vec::<(Ipv4Net, u8, u32)>::new();
let mut total_subnets = 0u32;
@ -164,7 +366,7 @@ impl AddressPool {
}
// Choose a random subnet to start with
let chosen_subnet_index = self.srng.next_u32(0, total_subnets - 1);
let chosen_subnet_index = srng.next_u32(0, total_subnets - 1);
// Find the starting scope and starting subnet index within
// the scope of the chosen subnet index
@ -237,19 +439,10 @@ impl AddressPool {
}
};
// If we found a free subnet, add it to our allocations
if let Some(allocation) = opt_allocation {
// Add to our allocated pool
self.allocated_v4.push(allocation);
self.allocated_v4 = Ipv4Net::aggregate(&self.allocated_v4);
return Some(allocation);
opt_allocation
}
// No allocation
None
}
pub fn allocate_random_v6(&mut self, prefix: u8) -> Option<Ipv6Net> {
fn find_random_allocation_v6(&self, srng: &mut StableRng, prefix: u8) -> Option<Ipv6Net> {
// Scope ranges to iterate
let mut scope_ranges = Vec::<(Ipv6Net, u8, u128)>::new();
let mut total_subnets = 0u128;
@ -277,7 +470,7 @@ impl AddressPool {
}
// Choose a random subnet to start with
let chosen_subnet_index = self.srng.next_u128(0, total_subnets - 1);
let chosen_subnet_index = srng.next_u128(0, total_subnets - 1);
// Find the starting scope and starting subnet index within
// the scope of the chosen subnet index
@ -350,15 +543,6 @@ impl AddressPool {
}
};
// If we found a free subnet, add it to our allocations
if let Some(allocation) = opt_allocation {
// Add to our allocated pool
self.allocated_v6.push(allocation);
self.allocated_v6 = Ipv6Net::aggregate(&self.allocated_v6);
return Some(allocation);
}
// No allocation
None
opt_allocation
}
}

65
veilid-tools/src/virtual_network/router_server/machine_registry/machine_registry_inner.rs
View File

@ -1,15 +1,15 @@
use super::*;
#[derive(Debug)]
#[derive(Debug, Clone)]
pub(super) struct MachineRegistryInner {
unlocked_inner: Arc<MachineRegistryUnlockedInner>,
allocated_machines: HashSet<MachineStateId>,
srng: StableRng,
allocated_machines: imbl::HashSet<MachineStateId>,
profile_state_allocator: StateAllocator<ProfileState>,
machine_state_allocator: StateAllocator<MachineState>,
template_state_allocator: StateAllocator<TemplateState>,
network_state_allocator: StateAllocator<NetworkState>,
blueprint_state_allocator: StateAllocator<BlueprintState>,
address_pool: AddressPool,
}
impl MachineRegistryInner {
@ -17,20 +17,21 @@ impl MachineRegistryInner {
/// Public Interface
pub fn new(unlocked_inner: Arc<MachineRegistryUnlockedInner>) -> Self {
let srng = unlocked_inner.srng.clone();
let srng = StableRng::new(unlocked_inner.config.seed.unwrap_or_default());
MachineRegistryInner {
unlocked_inner,
allocated_machines: HashSet::new(),
srng,
allocated_machines: imbl::HashSet::new(),
profile_state_allocator: StateAllocator::new(),
machine_state_allocator: StateAllocator::new(),
template_state_allocator: StateAllocator::new(),
network_state_allocator: StateAllocator::new(),
blueprint_state_allocator: StateAllocator::new(),
address_pool: AddressPool::new(srng),
}
}
pub fn srng(&self) -> StableRng {
self.unlocked_inner.srng.clone()
pub fn srng(&mut self) -> &mut StableRng {
&mut self.srng
}
pub fn config(&self) -> &config::Config {
&self.unlocked_inner.config
@ -55,7 +56,7 @@ impl MachineRegistryInner {
};
// Get current profile state, creating one if we have not yet started executing the profile
let profile_state = self
let profile_state_id = self
.profile_state_allocator
.get_or_create_by_name(profile, |id, name| {
ProfileState::new(id, name, profile_def.clone())
@ -63,52 +64,57 @@ impl MachineRegistryInner {
// Get the next instance from the definition
loop {
// Move to the next profile instance
let mut profile_state = self.profile_states().get_state(profile_state_id)?;
let Some(instance_def) = profile_state.next_instance() else {
return Err(MachineRegistryError::ProfileComplete);
};
self.profile_states_mut().set_state(profile_state);
let machine_state = match instance_def {
let machine_state_id = 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)
let opt_machine_states_ids = machine_names.try_filter_map(|name| {
let Some(machine_state_id) =
self.machine_states().get_state_id_by_name(name)
else {
return Err(MachineRegistryError::MachineNotFound);
};
if self.allocated_machines.contains(&machine_state.id()) {
if self.allocated_machines.contains(&machine_state_id) {
Ok(None)
} else {
Ok(Some(machine_state))
Ok(Some(machine_state_id))
}
})?;
let Some(machine_states) = opt_machine_states else {
let Some(machine_state_ids) = opt_machine_states_ids 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();
let machine_state_id = self.srng.weighted_choice(machine_state_ids);
// Activate it
self.allocated_machines.insert(machine_state.id());
self.allocated_machines.insert(machine_state_id);
machine_state
machine_state_id
}
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)
let Some(template_state_id) =
self.template_states().get_state_id_by_name(name)
else {
return Err(MachineRegistryError::TemplateNotFound);
};
let template_state = self
.template_states()
.get_state(template_state_id)
.expect("must exist");
if !template_state.is_active(self)? {
Ok(None)
} else {
@ -120,13 +126,20 @@ impl MachineRegistryInner {
continue;
};
let template_state = self.unlocked_inner.srng.weighted_choice(&template_states);
// Chose a template
let mut template_state = self.srng.weighted_choice(template_states);
template_state.generate(self)?
// Generate a machine from the template
let machine_state_id = template_state.generate(self)?;
// Save the updated template
self.template_states_mut().set_state(template_state);
machine_state_id
}
};
break Ok(machine_state.external_id());
break Ok(machine_state_id.0);
}
}

24
veilid-tools/src/virtual_network/router_server/machine_registry/mod.rs
View File

@ -14,13 +14,13 @@ struct Machine {}
#[derive(Debug)]
struct MachineRegistryUnlockedInner {
config: config::Config,
srng: StableRng,
}
#[derive(Debug, Clone)]
pub enum MachineRegistryError {
InvalidId,
InvalidName,
InvalidPrefix,
AlreadyAttached,
NotAttached,
DuplicateName,
@ -35,6 +35,7 @@ pub enum MachineRegistryError {
BlueprintNotFound,
ModelNotFound,
NoAllocation,
ResourceInUse,
}
pub type MachineRegistryResult<T> = Result<T, MachineRegistryError>;
@ -49,8 +50,7 @@ impl MachineRegistry {
///////////////////////////////////////////////////////////
/// Public Interface
pub fn new(config: config::Config) -> Self {
let srng = StableRng::new(config.seed.unwrap_or_default());
let unlocked_inner = Arc::new(MachineRegistryUnlockedInner { srng, config });
let unlocked_inner = Arc::new(MachineRegistryUnlockedInner { config });
Self {
inner: Arc::new(Mutex::new(MachineRegistryInner::new(
unlocked_inner.clone(),
@ -61,11 +61,25 @@ impl MachineRegistry {
pub fn allocate(&self, profile: String) -> MachineRegistryResult<MachineId> {
let mut inner = self.inner.lock();
inner.allocate(profile)
let saved_state = (*inner).clone();
match inner.allocate(profile) {
Ok(v) => Ok(v),
Err(e) => {
*inner = saved_state;
Err(e)
}
}
}
pub fn release(&self, machine_id: MachineId) -> MachineRegistryResult<()> {
let mut inner = self.inner.lock();
inner.release(machine_id)
let saved_state = (*inner).clone();
match inner.release(machine_id) {
Ok(v) => Ok(v),
Err(e) => {
*inner = saved_state;
Err(e)
}
}
}
}

195
veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs
View File

@ -1,23 +1,26 @@
use super::*;
use ipnet::*;
#[derive(Debug)]
struct BlueprintStateUnlockedInner {
/// The global random number generator
srng: StableRng,
/// The unique id of this blueprint
id: BlueprintStateId,
/// The name of this blueprint state
name: String,
}
#[derive(Debug)]
pub struct BlueprintStateIpv4Params {
pub allocation: Option<String>,
pub prefix: u8,
pub locations: NetworkLocationsList,
pub prefix: WeightedList<u8>,
pub gateway: Option<BlueprintStateGatewayParams>,
}
#[derive(Debug)]
pub struct BlueprintStateIpv6Params {
pub allocation: Option<String>,
pub prefix: u8,
pub locations: NetworkLocationsList,
pub prefix: WeightedList<u8>,
pub gateway: Option<BlueprintStateGatewayParams>,
}
@ -25,7 +28,7 @@ pub struct BlueprintStateIpv6Params {
pub struct BlueprintStateGatewayParams {
pub translation: WeightedList<config::Translation>,
pub upnp: Probability,
pub network: Option<NetworkStateId>,
pub locations: Option<MachineLocationsList>,
}
#[derive(Debug)]
@ -68,9 +71,13 @@ pub struct BlueprintState {
pub type BlueprintStateId = StateId<BlueprintState>;
impl BlueprintState {
pub fn new(id: BlueprintStateId, name: String) -> MachineRegistryResult<BlueprintState> {
pub fn new(
srng: StableRng,
id: BlueprintStateId,
name: String,
) -> MachineRegistryResult<BlueprintState> {
Ok(Self {
unlocked_inner: Arc::new(BlueprintStateUnlockedInner { id, name }),
unlocked_inner: Arc::new(BlueprintStateUnlockedInner { srng, id, name }),
inner: Arc::new(Mutex::new(BlueprintStateInner {
limit_network_count: None,
networks: Vec::new(),
@ -163,6 +170,8 @@ impl BlueprintState {
}
}
todo!("needs better implementation");
Ok(true)
}
@ -174,7 +183,7 @@ impl BlueprintState {
let model_name = match inner.model.clone() {
Some(models) => machine_registry_inner
.srng()
.weighted_choice(&models)
.weighted_choice_ref(&models)
.clone(),
None => machine_registry_inner.config().default_model.clone(),
};
@ -187,7 +196,7 @@ impl BlueprintState {
distance: model.distance.clone(),
loss: model.loss,
};
network_state.set_model(params);
network_state.with_model(params);
Ok(())
}
@ -201,43 +210,153 @@ impl BlueprintState {
return Ok(());
};
let allocation = ipv4
// Get maximum prefix
let max_prefix = ipv4
.params
.allocation
.clone()
.map(|x| {
.prefix
.iter()
.max()
.copied()
.expect("must have at least one element");
// Get addresses for network
let (subnet, super_net) = match &ipv4.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(&x)
.ok_or(MachineRegistryError::InvalidName)?;
if let Some(subnet4) = &allocation.subnets.subnet4 {
let ipv4net = machine_registry_inner
.srng()
.weighted_choice(subnet4)
.clone();
Ok(Some(ipv4net))
} else {
Ok(None)
}
})
.unwrap_or_else(|| {
Ok(Some(
Ipv4Net::new(Ipv4Addr::UNSPECIFIED, 0).expect("must be valid"),
))
})?;
let Some(allocation) = allocation else {
return Ok(());
.get(allocation_name)
.expect("must exist");
Ok(allocation
.subnets
.subnet4
.as_ref()
.map(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))
.flatten())
})?
else {
return Err(MachineRegistryError::NoAllocation);
};
// Do suballocation
xxx figure out how to do allocation inside internet vs private network
// Pick an allocation
let subnets = machine_registry_inner
.srng()
.weighted_choice_ref(&alloc_subnets);
let params = NetworkStateIpv4Params { allocation };
// Pick a subnet
let net = *machine_registry_inner.srng().weighted_choice_ref(subnets);
let gateway_params = inner.
// Pick a prefix length that would fit in the subnet
let opt_subnet = ipv4
.params
.prefix
.filter(|p| *p >= net.prefix_len())
.as_ref()
.map(|wl| {
let subnet_prefix =
machine_registry_inner.srng().weighted_choice_ref(wl).clone();
network_state.set_ipv4(, gateway_params);
// Use an address pool temporarily to pick a subnet
let mut address_pool =
AddressPool::<()>::new();
address_pool.add_scope_v4(net);
address_pool.allocate_random_v4(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 network_state = machine_registry_inner
.network_states()
.get_state(*network_id)
.expect("must exist");
Ok(network_state.can_allocate_subnet_v4(None, max_prefix))
})?
else {
return Err(MachineRegistryError::NoAllocation);
};
// Pick a network
let network_id = *machine_registry_inner
.srng()
.weighted_choice_ref(&available_networks);
let network_state = machine_registry_inner
.network_states()
.get_state(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))
.as_ref()
.map(|wl| {
let subnet_prefix =
machine_registry_inner.srng().weighted_choice_ref(wl).clone();
// Allocate subnet from this network
network_state.allocate_subnet_v4(
machine_registry_inner,
OwnerTag::Network(network_state.id()),
None,
subnet_prefix,
)
})
.transpose()?;
let Some(subnet) = opt_subnet else {
return Err(MachineRegistryError::NoAllocation);
};
(subnet, Some(network_id))
}
};
let params = NetworkStateIpv4Params {
allocation: subnet,
super_net,
};
let gateway_params = match ipv4.gateway.as_ref() {
Some(v4gw) => {
let translation = machine_registry_inner
.srng()
.weighted_choice_ref(&v4gw.params.translation)
.clone();
let upnp = machine_registry_inner
.srng()
.probability_test(v4gw.params.upnp);
let location = match v4gw.params.locations {
Some(locations) => todo!(),
None => todo!(),
};
xxx instantiate and clean up on failure
Some(NetworkStateIpv4GatewayParams {
translation,
upnp,
external_network: todo!(),
internal_address: None,
external_address: None,
})
}
None => None,
};
network_state.set_ipv4(machine_registry_inner, params, gateway_params);
Ok(())
}

12
veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_locations_list.rs Normal file
View File

@ -0,0 +1,12 @@
use super::*;
/// Locations where a machine can be instantiated
#[derive(Debug)]
pub enum MachineLocationsList {
Networks {
networks: WeightedList<NetworkStateId>,
},
Blueprints {
blueprints: WeightedList<BlueprintStateId>,
},
}

19
veilid-tools/src/virtual_network/router_server/machine_registry/state/machine_state.rs
View File

@ -30,6 +30,7 @@ struct MachineStateUnlockedInner {
#[derive(Debug, Clone)]
pub struct MachineState {
xxx convert to immutable state
unlocked_inner: Arc<MachineStateUnlockedInner>,
inner: Arc<Mutex<MachineStateInner>>,
}
@ -151,8 +152,11 @@ impl MachineState {
// 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)?;
let ifv4_addr = network_state.allocate_address_v4(
machine_registry_inner,
OwnerTag::Machine(self.id()),
opt_address,
)?;
// Get address flags
let flags = opt_address_flags.unwrap_or_else(|| AddressFlags {
@ -191,8 +195,11 @@ impl MachineState {
// 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)?;
let ifv6_addr = network_state.allocate_address_v6(
machine_registry_inner,
OwnerTag::Machine(self.id()),
opt_address,
)?;
// Get address flags
let flags = opt_address_flags.unwrap_or_else(|| AddressFlags {
@ -274,7 +281,7 @@ impl MachineState {
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
@ -329,7 +336,7 @@ impl MachineState {
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

6
veilid-tools/src/virtual_network/router_server/machine_registry/state/mod.rs
View File

@ -1,16 +1,18 @@
mod blueprint_state;
mod machine_locations_list;
mod machine_state;
mod network_locations_list;
mod network_state;
mod profile_state;
mod state_allocator;
mod template_state;
use std::marker::PhantomData;
use super::*;
pub use blueprint_state::*;
pub use machine_locations_list::*;
pub use machine_state::*;
pub use network_locations_list::*;
pub use network_state::*;
pub use profile_state::*;
pub use state_allocator::*;

12
veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs Normal file
View File

@ -0,0 +1,12 @@
use super::*;
/// Locations where a network can be instantiated
#[derive(Debug)]
pub enum NetworkLocationsList {
/// Network will be a new allocation
Allocations { allocations: WeightedList<String> },
/// Network will be allocated as a subnet of an existing network
Networks {
networks: WeightedList<NetworkStateId>,
},
}

815
veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs
View File

File diff suppressed because it is too large Load Diff

31
veilid-tools/src/virtual_network/router_server/machine_registry/state/profile_state.rs
View File

@ -1,12 +1,12 @@
use super::*;
#[derive(Debug)]
struct ProfileStateInner {
struct ProfileStateFields {
next_instance_index: usize,
}
#[derive(Debug)]
struct ProfileStateUnlockedInner {
struct ProfileStateImmutable {
id: ProfileStateId,
name: String,
def: config::Profile,
@ -14,8 +14,8 @@ struct ProfileStateUnlockedInner {
#[derive(Debug, Clone)]
pub struct ProfileState {
unlocked_inner: Arc<ProfileStateUnlockedInner>,
inner: Arc<Mutex<ProfileStateInner>>,
immutable: Arc<ProfileStateImmutable>,
fields: Arc<ProfileStateFields>,
}
pub type ProfileStateId = StateId<ProfileState>;
@ -23,33 +23,34 @@ pub type ProfileStateId = StateId<ProfileState>;
impl ProfileState {
pub fn new(id: ProfileStateId, name: String, def: config::Profile) -> Self {
Self {
unlocked_inner: Arc::new(ProfileStateUnlockedInner { id, name, def }),
inner: Arc::new(Mutex::new(ProfileStateInner {
immutable: Arc::new(ProfileStateImmutable { id, name, def }),
fields: Arc::new(ProfileStateFields {
next_instance_index: 0,
})),
}),
}
}
pub fn next_instance(&self) -> Option<config::Instance> {
pub fn next_instance(&mut self) -> Option<config::Instance> {
let instance_index = {
let mut inner = self.inner.lock();
let instance_index = inner.next_instance_index;
if instance_index >= self.unlocked_inner.def.instances.len() {
let instance_index = self.fields.next_instance_index;
if instance_index >= self.immutable.def.instances.len() {
return None;
}
inner.next_instance_index += 1;
self.fields = Arc::new(ProfileStateFields {
next_instance_index: instance_index + 1,
});
instance_index
};
Some(self.unlocked_inner.def.instances[instance_index].clone())
Some(self.immutable.def.instances[instance_index].clone())
}
}
impl State for ProfileState {
fn id(&self) -> StateId<Self> {
self.unlocked_inner.id.clone()
self.immutable.id.clone()
}
fn name(&self) -> Option<String> {
Some(self.unlocked_inner.name.clone())
Some(self.immutable.name.clone())
}
}

87
veilid-tools/src/virtual_network/router_server/machine_registry/state/state_allocator.rs
View File

@ -1,4 +1,5 @@
use super::*;
use std::marker::PhantomData;
pub trait State: fmt::Debug + Clone {
fn id(&self) -> StateId<Self>;
@ -38,21 +39,21 @@ impl<S: State> core::hash::Hash for StateId<S> {
}
}
#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct StateAllocator<S: State> {
state_id_by_name: HashMap<String, StateIdInternal>,
state_by_id: HashMap<StateIdInternal, Option<S>>,
state_id_by_name: imbl::HashMap<Arc<String>, StateIdInternal>,
state_by_id: imbl::HashMap<StateIdInternal, Option<S>>,
next_state_id: StateIdInternal,
free_state_ids: Vec<StateIdInternal>,
free_state_ids: imbl::Vector<StateIdInternal>,
}
impl<S: State> StateAllocator<S> {
pub fn new() -> Self {
Self {
state_id_by_name: HashMap::new(),
state_by_id: HashMap::new(),
state_id_by_name: imbl::HashMap::new(),
state_by_id: imbl::HashMap::new(),
next_state_id: 0,
free_state_ids: Vec::new(),
free_state_ids: imbl::Vector::new(),
}
}
@ -60,20 +61,19 @@ impl<S: State> StateAllocator<S> {
&mut self,
name: String,
create: F,
) -> S {
if let Some(v) = self.get_state_by_name(&name) {
return v.clone();
) -> StateId<S> {
if let Some(id) = self.get_state_id_by_name(&name) {
return id;
}
let id = self.allocate_id();
let state = create(id, name);
self.attach_state(state.clone())
.expect("should always attach");
state
self.attach_state(state).expect("should always attach");
id
}
pub fn allocate_id(&mut self) -> StateId<S> {
// Allocate new internal id
let state_id = self.free_state_ids.pop().unwrap_or_else(|| {
let state_id = self.free_state_ids.pop_back().unwrap_or_else(|| {
let x = self.next_state_id;
self.next_state_id += 1;
x
@ -103,7 +103,7 @@ impl<S: State> StateAllocator<S> {
}
// Keep old id in the free list
self.free_state_ids.push(id.0);
self.free_state_ids.push_back(id.0);
Ok(())
}
@ -129,7 +129,7 @@ impl<S: State> StateAllocator<S> {
}
// Register the named state
self.state_id_by_name
.insert(name, id.0)
.insert(Arc::new(name), id.0)
.expect("should not have a duplicated name here");
}
@ -164,32 +164,59 @@ impl<S: State> StateAllocator<S> {
pub fn get_state(&self, id: StateId<S>) -> MachineRegistryResult<S> {
// Get the allocator slot
let Some(opt_state) = self.state_by_id.get(&id.0).cloned() else {
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);
};
Ok(state)
Ok(state.clone())
}
pub fn get_state_by_name(&self, name: &String) -> Option<S> {
pub fn set_state(&mut self, state: S) {
self.state_by_id.insert(state.id().0, Some(state));
}
// pub fn update_state<
// R,
// F: FnOnce(&mut StateAllocator<S>, S) -> MachineRegistryResult<(R, Option<S>)>,
// >(
// &mut self,
// id: StateId<S>,
// callback: F,
// ) -> MachineRegistryResult<R> {
// // 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<StateId<S>> {
// Get the id associated with this name
let Some(id) = self.state_id_by_name.get(name) else {
return None;
};
// Get the allocator slot
let opt_state = self
.state_by_id
.get(&id)
.cloned()
.expect("id should always be valid");
// The state should be attached otherwise we screwed up
let state = opt_state.expect("named states should always be attached");
Some(state)
Some(StateId::new(*id))
}
}

36
veilid-tools/src/virtual_network/router_server/machine_registry/state/template_state.rs
View File

@ -12,16 +12,6 @@ struct PerNetworkInfo {
machines: HashSet<MachineStateId>,
}
#[derive(Debug)]
enum LocationsList {
Networks {
networks: WeightedList<NetworkStateId>,
},
Blueprints {
blueprints: WeightedList<BlueprintStateId>,
},
}
#[derive(Debug, Clone)]
enum BlueprintAvailability {
Existing(NetworkState),
@ -32,7 +22,7 @@ enum BlueprintAvailability {
struct TemplateStateInner {
limit_machine_count: Option<usize>,
limit_machines_per_network: Option<WeightedList<usize>>,
locations_list: Option<LocationsList>,
locations_list: Option<MachineLocationsList>,
machines: HashSet<MachineStateId>,
machines_per_network: HashMap<NetworkStateId, PerNetworkInfo>,
disable_capabilities: Vec<String>,
@ -68,12 +58,12 @@ impl TemplateState {
pub fn set_networks_list(&self, networks: WeightedList<NetworkStateId>) {
let mut inner = self.inner.lock();
inner.locations_list = Some(LocationsList::Networks { networks })
inner.locations_list = Some(MachineLocationsList::Networks { networks })
}
pub fn set_blueprints_list(&self, blueprints: WeightedList<BlueprintStateId>) {
let mut inner = self.inner.lock();
inner.locations_list = Some(LocationsList::Blueprints { blueprints })
inner.locations_list = Some(MachineLocationsList::Blueprints { blueprints })
}
pub fn clear_locations_list(&self) {
@ -153,7 +143,7 @@ impl TemplateState {
pub fn is_active(
&self,
machine_registry_inner: &mut MachineRegistryInner,
machine_registry_inner: &MachineRegistryInner,
) -> MachineRegistryResult<bool> {
let inner = self.inner.lock();
@ -169,7 +159,7 @@ impl TemplateState {
};
match locations_list {
LocationsList::Networks { networks } => {
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| {
@ -182,7 +172,7 @@ impl TemplateState {
return Ok(false);
};
}
LocationsList::Blueprints { blueprints } => {
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| {
@ -207,9 +197,9 @@ impl TemplateState {
}
pub fn generate(
&self,
&mut self,
machine_registry_inner: &mut MachineRegistryInner,
) -> MachineRegistryResult<MachineState> {
) -> MachineRegistryResult<MachineStateId> {
let mut inner = self.inner.lock();
// See if we have reached our machine limit
@ -226,7 +216,7 @@ impl TemplateState {
// Get a network to generate the machine on
let network_state = match locations_list {
LocationsList::Networks { networks } => {
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)?;
@ -243,12 +233,12 @@ impl TemplateState {
// Weighted choice of network now that we have a candidate list
let network_state = machine_registry_inner
.srng()
.weighted_choice(&available_networks);
.weighted_choice_ref(&available_networks);
// Return network state to use
network_state.clone()
}
LocationsList::Blueprints { blueprints } => {
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 =
@ -267,7 +257,7 @@ impl TemplateState {
// Weighted choice of blueprint now that we have a candidate list
match machine_registry_inner
.srng()
.weighted_choice(&available_blueprints)
.weighted_choice_ref(&available_blueprints)
{
BlueprintAvailability::Existing(network_state) => network_state.clone(),
BlueprintAvailability::Generate(blueprint_state) => {
@ -320,7 +310,7 @@ impl TemplateState {
.entry(network_state.id())
.or_insert_with(|| {
let limit_machine_count = limit_machines_per_network
.map(|wl| machine_registry_inner.srng().weighted_choice(&wl).clone());
.map(|wl| machine_registry_inner.srng().weighted_choice_ref(&wl).clone());
PerNetworkInfo {
limit_machine_count,
machines: HashSet::new(),

88
veilid-tools/src/virtual_network/router_server/stable_rng.rs
View File

@ -3,12 +3,13 @@ use super::*;
use rand::{seq::SliceRandom, Rng, SeedableRng};
use rand_chacha::ChaCha20Rng;
struct StableRngInner {
#[derive(Clone)]
pub struct StableRngState {
srng: ChaCha20Rng,
count: usize,
}
impl fmt::Debug for StableRngInner {
impl fmt::Debug for StableRngState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StableRngInner")
.field("count", &self.count)
@ -18,7 +19,7 @@ impl fmt::Debug for StableRngInner {
#[derive(Clone, Debug)]
pub struct StableRng {
inner: Arc<Mutex<StableRngInner>>,
state: StableRngState,
}
impl StableRng {
@ -27,14 +28,33 @@ impl StableRng {
pub fn new(seed: u64) -> Self {
Self {
inner: Arc::new(Mutex::new(StableRngInner {
state: StableRngState {
srng: ChaCha20Rng::seed_from_u64(seed),
count: 0,
})),
},
}
}
pub fn weighted_choice<'a, T: fmt::Debug + Clone>(
&self,
pub fn save_state(&self) -> StableRngState {
self.state.clone()
}
pub fn restore_state(&mut self, state: StableRngState) {
self.state = state;
}
pub fn probability_test(&mut self, probability: Probability) -> bool {
if probability == 1.0 {
return true;
} else if probability == 0.0 {
return false;
}
let num = self.next_f32(0.0, 1.0);
num < probability
}
pub fn weighted_choice_ref<'a, T: fmt::Debug + Clone>(
&mut self,
weighted_list: &'a WeightedList<T>,
) -> &'a T {
match weighted_list {
@ -59,27 +79,49 @@ impl StableRng {
}
}
}
pub fn shuffle_vec<T>(&self, v: &mut Vec<T>) {
let mut inner = self.inner.lock();
inner.count += 1;
v.shuffle(&mut inner.srng);
pub fn weighted_choice<T: fmt::Debug + Clone>(&mut self, weighted_list: WeightedList<T>) -> T {
match weighted_list {
WeightedList::Single(x) => x,
WeightedList::List(mut vec) => {
let total_weight = vec
.iter()
.map(|x| x.weight())
.reduce(|acc, x| acc + x)
.expect("config validation broken");
let r = self.next_f32(0.0, total_weight);
let mut current_weight = 0.0f32;
let last = vec.pop().expect("config validation broken").into_item();
for x in vec {
current_weight += x.weight();
if r < current_weight {
return x.into_item();
}
}
// Catch f32 imprecision
last
}
}
}
pub fn next_u32(&self, min: u32, max: u32) -> u32 {
let mut inner = self.inner.lock();
inner.count += 1;
inner.srng.gen_range(min..=max)
pub fn shuffle_vec<T>(&mut self, v: &mut Vec<T>) {
self.state.count += 1;
v.shuffle(&mut self.state.srng);
}
pub fn next_u128(&self, min: u128, max: u128) -> u128 {
let mut inner = self.inner.lock();
inner.count += 1;
inner.srng.gen_range(min..=max)
pub fn next_u32(&mut self, min: u32, max: u32) -> u32 {
self.state.count += 1;
self.state.srng.gen_range(min..=max)
}
pub fn next_f32(&self, min: f32, max: f32) -> f32 {
let mut inner = self.inner.lock();
inner.count += 1;
inner.srng.gen_range(min..=max)
pub fn next_u128(&mut self, min: u128, max: u128) -> u128 {
self.state.count += 1;
self.state.srng.gen_range(min..=max)
}
pub fn next_f32(&mut self, min: f32, max: f32) -> f32 {
self.state.count += 1;
self.state.srng.gen_range(min..=max)
}
}

118
veilid-tools/src/virtual_network/router_server/weighted_list.rs
View File

@ -4,6 +4,9 @@ use validator::{Validate, ValidationError, ValidationErrors};
pub type Probability = f32;
//////////////////////////////////////////////////////////////////////////
/// WeightedList
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WeightedList<T: fmt::Debug + Clone> {
@ -15,32 +18,18 @@ impl<T: fmt::Debug + Clone> Default for WeightedList<T> {
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 len(&self) -> usize {
match self {
WeightedList::Single(_) => 1,
WeightedList::List(vec) => vec.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn validate_once(&self) -> Result<(), ValidationError> {
match self {
Self::List(v) => {
@ -153,8 +142,81 @@ impl<T: fmt::Debug + Clone> WeightedList<T> {
}
}
}
pub fn iter(&self) -> WeightedListIter<'_, T> {
WeightedListIter {
values: self,
index: 0,
}
}
}
//////////////////////////////////////////////////////////////////////////
/// Index
impl<T: fmt::Debug + Clone> core::ops::Index<usize> for WeightedList<T> {
type Output = T;
fn index(&self, index: usize) -> &Self::Output {
match self {
WeightedList::Single(s) => &s,
WeightedList::List(vec) => vec[index].item(),
}
}
}
//////////////////////////////////////////////////////////////////////////
/// Iterator
pub struct WeightedListIter<'a, T: fmt::Debug + Clone> {
values: &'a WeightedList<T>,
index: usize,
}
impl<'a, T: fmt::Debug + Clone> Iterator for WeightedListIter<'a, T> {
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
if self.index >= self.values.len() {
return None;
}
self.index += 1;
Some(&self.values[self.index - 1])
}
}
//////////////////////////////////////////////////////////////////////////
/// Validate
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)
}
}
}
//////////////////////////////////////////////////////////////////////////
/// Weighted
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Weighted<T: fmt::Debug + Clone> {
@ -190,6 +252,12 @@ impl<T: fmt::Debug + Clone> Weighted<T> {
Weighted::Unweighted(item) => item,
}
}
pub fn into_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,