review feedback

oxidecomputer · Jan 13, 2024 · 73d4669 · 73d4669
1 parent e3657fe
commit 73d4669
Showing 1 changed file with 25 additions and 26 deletions.
diff --git a/lib/oxide-vpc/src/engine/overlay.rs b/lib/oxide-vpc/src/engine/overlay.rs
@@ -26,6 +26,7 @@ use opte::api::Ipv4Cidr;
 use opte::api::MacAddr;
 use opte::api::OpteError;
 use opte::ddi::sync::KMutex;
+use opte::ddi::sync::KMutexGuard;
 use opte::ddi::sync::KMutexType;
 use opte::ddi::sync::KRwLock;
 use opte::ddi::sync::KRwLockType;
@@ -562,19 +563,19 @@ pub const INTERNET_GATEWAY_MAC: [u8; 6] = [0xA8, 0x40, 0x25, 0x77, 0x77, 0x77];
 
 impl Virt2Boundary {
     pub fn dump_ip4(&self) -> Vec<(Ipv4Cidr, BTreeSet<TunnelEndpoint>)> {
-        let mut ip4 = Vec::new();
-        for (vip, baddrs) in self.ip4.lock().iter() {
-            ip4.push((*vip, baddrs.clone()));
-        }
-        ip4
+        self.ip4
+            .lock()
+            .iter()
+            .map(|(vip, baddr)| (*vip, baddr.clone()))
+            .collect()
     }
 
     pub fn dump_ip6(&self) -> Vec<(Ipv6Cidr, BTreeSet<TunnelEndpoint>)> {
-        let mut ip6 = Vec::new();
-        for (vip, baddrs) in self.ip6.lock().iter() {
-            ip6.push((*vip, baddrs.clone()));
-        }
-        ip6
+        self.ip6
+            .lock()
+            .iter()
+            .map(|(vip, baddr)| (*vip, baddr.clone()))
+            .collect()
     }
 
     pub fn dump(&self) -> DumpVirt2BoundaryResp {
@@ -640,8 +641,7 @@ impl Virt2Boundary {
                 if clear {
                     tbl.remove(&ip4);
                 }
-                drop(tbl);
-                self.update_poptrie_v4();
+                self.update_poptrie_v4(&tbl);
                 orig
             }
             IpCidr::Ip6(ip6) => {
@@ -659,8 +659,7 @@ impl Virt2Boundary {
                 if clear {
                     tbl.remove(&ip6);
                 }
-                drop(tbl);
-                self.update_poptrie_v6();
+                self.update_poptrie_v6(&tbl);
                 orig
             }
         }
@@ -682,8 +681,7 @@ impl Virt2Boundary {
                     }
                     None => tbl.insert(ip4, tep.into_iter().collect()),
                 };
-                drop(tbl);
-                self.update_poptrie_v4();
+                self.update_poptrie_v4(&tbl);
                 result
             }
             IpCidr::Ip6(ip6) => {
@@ -696,29 +694,30 @@ impl Virt2Boundary {
                     }
                     None => tbl.insert(ip6, tep.into_iter().collect()),
                 };
-                drop(tbl);
-                self.update_poptrie_v6();
+                self.update_poptrie_v6(&tbl);
                 result
             }
         }
     }
 
-    fn update_poptrie_v4(&self) {
+    fn update_poptrie_v4(
+        &self,
+        tree: &KMutexGuard<BTreeMap<Ipv4Cidr, BTreeSet<TunnelEndpoint>>>,
+    ) {
         let table = poptrie::Ipv4RoutingTable(
-            self.ip4
-                .lock()
-                .iter()
+            tree.iter()
                 .map(|(k, v)| ((u32::from(k.ip()), k.prefix_len()), v.clone()))
                 .collect(),
         );
         *self.pt4.write() = poptrie::Poptrie::from(table);
     }
 
-    fn update_poptrie_v6(&self) {
+    fn update_poptrie_v6(
+        &self,
+        tree: &KMutexGuard<BTreeMap<Ipv6Cidr, BTreeSet<TunnelEndpoint>>>,
+    ) {
         let table = poptrie::Ipv6RoutingTable(
-            self.ip6
-                .lock()
-                .iter()
+            tree.iter()
                 .map(|(k, v)| ((u128::from(k.ip()), k.prefix_len()), v.clone()))
                 .collect(),
         );