diff --git a/easytier/src/instance/shared_virtual_nic.rs b/easytier/src/instance/shared_virtual_nic.rs index 80bb1fc0..202d0242 100644 --- a/easytier/src/instance/shared_virtual_nic.rs +++ b/easytier/src/instance/shared_virtual_nic.rs @@ -2,6 +2,7 @@ use std::{ collections::{BTreeMap, BTreeSet}, net::{Ipv4Addr, Ipv6Addr}, sync::Arc, + sync::atomic::{AtomicBool, Ordering}, }; use cidr::{Ipv4Inet, Ipv6Inet}; @@ -232,6 +233,7 @@ impl SharedIfConfig { pub struct SharedVirtualNic { nic: Arc>, ifcfg: SharedIfConfig, + valid: Arc, } impl SharedVirtualNic { @@ -239,9 +241,18 @@ impl SharedVirtualNic { Self { nic: Arc::new(Mutex::new(VirtualNic::new(config))), ifcfg: SharedIfConfig::default(), + valid: Arc::new(AtomicBool::new(true)), } } + pub fn mark_invalid(&self) { + self.valid.store(false, Ordering::Release); + } + + pub fn is_valid(&self) -> bool { + self.valid.load(Ordering::Acquire) + } + pub fn ifcfg(&self) -> &SharedIfConfig { &self.ifcfg } @@ -253,6 +264,65 @@ impl SharedVirtualNic { pub fn nic(&self) -> Arc> { self.nic.clone() } + + fn valid_flag(&self) -> Arc { + self.valid.clone() + } +} + +#[derive(Default)] +pub struct SharedVirtualNicRegistry { + nics: BTreeMap, +} + +struct SharedVirtualNicRegistryEntry { + nic: Arc>, + valid: Arc, +} + +impl SharedVirtualNicRegistryEntry { + fn new(nic: SharedVirtualNic) -> Self { + Self { + valid: nic.valid_flag(), + nic: Arc::new(Mutex::new(nic)), + } + } + + fn is_valid(&self) -> bool { + self.valid.load(Ordering::Acquire) + } + + fn nic(&self) -> Arc> { + self.nic.clone() + } +} + +impl SharedVirtualNicRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn get(&self, dev_name: &str) -> Option>> { + self.nics + .get(dev_name) + .filter(|entry| entry.is_valid()) + .map(|entry| entry.nic()) + } + + pub fn get_or_create( + &mut self, + dev_name: String, + config: VirtualNicConfig, + ) -> Arc> { + if let Some(nic) = self.get(&dev_name) { + return nic; + } + + let entry = SharedVirtualNicRegistryEntry::new(SharedVirtualNic::new(config)); + let nic = entry.nic(); + self.nics.insert(dev_name, entry); + nic + } } fn update_owned_items( @@ -485,4 +555,40 @@ mod tests { ); drop(shared_nic.nic()); } + + #[test] + fn registry_reuses_shared_virtual_nic_for_same_dev_name() { + let mut registry = SharedVirtualNicRegistry::new(); + + let first = registry.get_or_create("et0".to_string(), virtual_nic_config()); + let second = registry.get_or_create("et0".to_string(), virtual_nic_config()); + + assert!(Arc::ptr_eq(&first, &second)); + } + + #[test] + fn registry_keeps_different_dev_names_separate() { + let mut registry = SharedVirtualNicRegistry::new(); + + let first = registry.get_or_create("et0".to_string(), virtual_nic_config()); + let second = registry.get_or_create("et1".to_string(), virtual_nic_config()); + + assert!(!Arc::ptr_eq(&first, &second)); + } + + #[test] + fn registry_replaces_invalid_shared_virtual_nic() { + let mut registry = SharedVirtualNicRegistry::new(); + + let first = registry.get_or_create("et0".to_string(), virtual_nic_config()); + first.try_lock().unwrap().mark_invalid(); + let second = registry.get_or_create("et0".to_string(), virtual_nic_config()); + + assert!(!Arc::ptr_eq(&first, &second)); + assert!( + registry + .get("et0") + .is_some_and(|nic| Arc::ptr_eq(&nic, &second)) + ); + } }