From 7aaa26debb3bb875c5e49fbd39dec642cd535f88 Mon Sep 17 00:00:00 2001 From: Carl Dong Date: Fri, 23 Aug 2019 15:07:34 -0400 Subject: [PATCH] We don't need BTreeMaps --- src/bgp_client.rs | 39 +++++++++++++++++++-------------------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/src/bgp_client.rs b/src/bgp_client.rs index 4d393c3..04be136 100644 --- a/src/bgp_client.rs +++ b/src/bgp_client.rs @@ -1,8 +1,7 @@ use std::sync::{Arc, Mutex}; use std::sync::atomic::{AtomicBool, Ordering}; use std::cmp; -use std::ops::Bound::Included; -use std::collections::BTreeMap; +use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::{Duration, Instant}; @@ -30,23 +29,22 @@ struct Route { } struct RoutingTable { - v4_table: BTreeMap<(Ipv4Addr, u8, u32), Arc>, - v6_table: BTreeMap<(Ipv6Addr, u8, u32), Arc>, + v4_table: HashMap<(Ipv4Addr, u8), HashMap>>, + v6_table: HashMap<(Ipv6Addr, u8), HashMap>>, } impl RoutingTable { fn new() -> Self { Self { - v4_table: BTreeMap::new(), - v6_table: BTreeMap::new(), + v4_table: HashMap::new(), + v6_table: HashMap::new(), } } fn get_route_attrs(&self, ip: IpAddr) -> Vec> { macro_rules! lookup_res { ($addrty: ty, $addr: expr, $table: expr, $addr_bits: expr) => { { - let mut res = Vec::new(); - //TODO: Optimize this! + //TODO: Optimize this (probably means making the tables btrees)! for i in (0..$addr_bits).rev() { let mut lookup = $addr.octets(); for b in 0..(i / 8) { @@ -54,12 +52,13 @@ impl RoutingTable { } lookup[lookup.len() - (i/8) - 1] &= !(((1u16 << (i % 8)) - 1) as u8); let lookup_addr = <$addrty>::from(lookup); - for attrs in $table.range((Included((lookup_addr, $addr_bits - i as u8, 0)), Included((lookup_addr, $addr_bits - i as u8, std::u32::MAX)))) { - res.push(Arc::clone(&attrs.1)); + if let Some(routes) = $table.get(&(lookup_addr, $addr_bits - i as u8)).map(|hm| hm.values()) { + if routes.len() > 0 { + return routes.map(|x| Arc::clone(&x)).collect(); + } } - if !res.is_empty() { break; } } - res + vec![] } } } match ip { @@ -73,15 +72,15 @@ impl RoutingTable { NLRIEncoding::IP(p) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, 0)), - IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, 0)), + IpAddr::V4(v4a) => self.v4_table.get_mut(&(v4a, len)).and_then(|hm| hm.remove(&0)), + IpAddr::V6(v6a) => self.v6_table.get_mut(&(v6a, len)).and_then(|hm| hm.remove(&0)), } }, NLRIEncoding::IP_WITH_PATH_ID((p, id)) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, id)), - IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, id)), + IpAddr::V4(v4a) => self.v4_table.get_mut(&(v4a, len)).and_then(|hm| hm.remove(&id)), + IpAddr::V6(v6a) => self.v6_table.get_mut(&(v6a, len)).and_then(|hm| hm.remove(&id)), } }, NLRIEncoding::IP_MPLS(_) => None, @@ -93,15 +92,15 @@ impl RoutingTable { NLRIEncoding::IP(p) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, 0), route), - IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, 0), route), + IpAddr::V4(v4a) => self.v4_table.entry((v4a, len)).or_insert(HashMap::new()).insert(0, route), + IpAddr::V6(v6a) => self.v6_table.entry((v6a, len)).or_insert(HashMap::new()).insert(0, route), } }, NLRIEncoding::IP_WITH_PATH_ID((p, id)) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, id), route), - IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, id), route), + IpAddr::V4(v4a) => self.v4_table.entry((v4a, len)).or_insert(HashMap::new()).insert(id, route), + IpAddr::V6(v6a) => self.v6_table.entry((v6a, len)).or_insert(HashMap::new()).insert(id, route), } }, NLRIEncoding::IP_MPLS(_) => None, -- 2.39.5