feat(agent): improve deleting a set of rules into the router

Author: joelwurtz

src/api/rules_message.rs

@@ -1,6 +1,7 @@
 use crate::api::Rule;
 use crate::router::Router;
 use serde::{Deserialize, Serialize};
+use std::collections::HashSet;
 use std::sync::Arc;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -11,9 +12,9 @@ pub struct RulesMessage {
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct RuleChangeSet {
-    pub deleted: Vec<String>,
     pub added: Vec<Rule>,
     pub updated: Vec<Rule>,
+    pub deleted: HashSet<String>,
 }
 
 impl RuleChangeSet {

src/regex_radix_tree/item.rs

@@ -67,6 +67,17 @@ impl<V> Item<V> {
         }
     }
 
+    pub fn retain<F>(self, f: &F) -> Item<V>
+    where
+        F: Fn(&str, &mut V) -> bool,
+    {
+        match self {
+            Item::Empty(_) => self,
+            Item::Node(node) => node.retain(f),
+            Item::Leaf(leaf) => leaf.retain(f),
+        }
+    }
+
     /// Length of node
     pub fn len(&self) -> usize {
         match self {

src/regex_radix_tree/leaf.rs

@@ -99,6 +99,18 @@ impl<V> Leaf<V> {
             }
         }
     }
+    pub fn retain<F>(mut self, f: &F) -> Item<V>
+    where
+        F: Fn(&str, &mut V) -> bool,
+    {
+        self.values.retain(|k, v| f(k, v));
+
+        if self.values.is_empty() {
+            Item::Empty(self.regex.ignore_case)
+        } else {
+            Item::Leaf(self)
+        }
+    }
 
     /// Length of node
     pub fn len(&self) -> usize {

src/regex_radix_tree/node.rs

@@ -140,6 +140,33 @@ impl<V> Node<V> {
         (Item::Node(self), removed)
     }
 
+    pub fn retain<F>(mut self, f: &F) -> Item<V>
+    where
+        F: Fn(&str, &mut V) -> bool,
+    {
+        let mut children = Vec::new();
+
+        for child in self.children {
+            let child = child.retain(f);
+
+            if !child.is_empty() {
+                children.push(child);
+            }
+        }
+
+        if children.is_empty() {
+            return Item::Empty(self.regex.ignore_case);
+        }
+
+        if children.len() == 1 {
+            return children.pop().unwrap();
+        }
+
+        self.children = children;
+
+        Item::Node(self)
+    }
+
     /// Length of node
     pub fn len(&self) -> usize {
         let mut count = 0;

src/regex_radix_tree/tree.rs

@@ -43,6 +43,16 @@ impl<V> RegexTreeMap<V> {
         removed
     }
 
+    pub fn retain<F>(&mut self, f: &F)
+    where
+        F: Fn(&str, &mut V) -> bool,
+    {
+        let mut root = Item::Empty(false);
+        std::mem::swap(&mut self.root, &mut root);
+        let root = root.retain(f);
+        self.root = root;
+    }
+
     pub fn len(&self) -> usize {
         self.root.len()
     }
@@ -114,6 +124,13 @@ impl<V> UniqueRegexTreeMap<V> {
         self.tree.remove(regex)
     }
 
+    pub fn retain<F>(&mut self, f: &F)
+    where
+        F: Fn(&str, &mut V) -> bool,
+    {
+        self.tree.retain(f)
+    }
+
     pub fn len(&self) -> usize {
         self.tree.len()
     }

src/router/mod.rs

@@ -20,7 +20,7 @@ pub use route_weekday::RouteWeekday;
 pub use trace::{RouteTrace, Trace};
 
 use core::cmp::Reverse;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -76,6 +76,11 @@ impl<T> Router<T> {
         }
     }
 
+    pub fn batch_remove(&mut self, ids: &HashSet<String>) {
+        self.routes.retain(|id, _| !ids.contains(id));
+        self.matcher.batch_remove(ids);
+    }
+
     pub fn rebuild_request(&self, request: &Request) -> Request {
         Request::rebuild_with_config(self.config.as_ref(), request)
     }
@@ -154,16 +159,17 @@ where
         self.insert_route(item.into_route(self.config.as_ref()));
     }
 
-    pub fn apply_change_set(&mut self, added: Vec<T>, updated: Vec<T>, removed: Vec<String>) {
-        for id in removed {
-            self.remove(id.as_str());
-        }
+    pub fn apply_change_set(&mut self, added: Vec<T>, updated: Vec<T>, mut removed: HashSet<String>) {
+        let updated_route = updated
+            .into_iter()
+            .map(|item| item.into_route(self.config.as_ref()))
+            .collect::<Vec<Route<T>>>();
 
-        for item in updated {
-            let route = item.into_route(self.config.as_ref());
+        removed.extend(updated_route.iter().map(|item| item.id().to_string()));
+        self.batch_remove(&removed);
 
-            self.remove(route.id());
-            self.insert_route(route);
+        for item in updated_route {
+            self.insert_route(item);
         }
 
         for item in added {

src/router/request_matcher/datetime.rs

@@ -6,8 +6,8 @@ use super::super::trace::{TraceInfo, TraceInfoDateTimeCondition};
 use super::super::{Route, RouterConfig, Trace};
 use crate::http::Request;
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
 use std::collections::BTreeSet;
+use std::collections::{BTreeMap, HashSet};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -106,6 +106,18 @@ impl<T> DateTimeMatcher<T> {
         removed
     }
 
+    pub fn batch_remove(&mut self, ids: &HashSet<String>) -> bool {
+        self.any_datetime.batch_remove(ids);
+
+        self.condition_groups.retain(|_, matcher| {
+            matcher.batch_remove(ids);
+
+            !matcher.is_empty()
+        });
+
+        self.any_datetime.is_empty() && self.condition_groups.is_empty()
+    }
+
     pub fn match_request(&self, request: &Request) -> Vec<Arc<Route<T>>> {
         let mut rules = self.any_datetime.match_request(request);
         let mut execute_conditions = BTreeMap::new();

src/router/request_matcher/header.rs

@@ -5,8 +5,8 @@ use super::super::{Route, RouteHeaderKind, Trace};
 use crate::http::Request;
 use regex::Regex;
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
 use std::collections::BTreeSet;
+use std::collections::{BTreeMap, HashSet};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -117,6 +117,18 @@ impl<T> HeaderMatcher<T> {
         removed
     }
 
+    pub fn batch_remove(&mut self, ids: &HashSet<String>) -> bool {
+        self.any_header.batch_remove(ids);
+
+        self.condition_groups.retain(|_, matcher| {
+            matcher.batch_remove(ids);
+
+            !matcher.is_empty()
+        });
+
+        self.any_header.is_empty() && self.condition_groups.is_empty()
+    }
+
     pub fn match_request(&self, request: &Request) -> Vec<Arc<Route<T>>> {
         let mut rules = self.any_header.match_request(request);
         let mut execute_conditions = BTreeMap::new();

src/router/request_matcher/host.rs

@@ -3,7 +3,7 @@ use super::super::{IpMatcher, Route, RouterConfig, Trace};
 use crate::http::Request;
 use crate::marker::StaticOrDynamic;
 use crate::regex_radix_tree::{Trace as TreeTrace, UniqueRegexTreeMap};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -91,6 +91,24 @@ impl<T> HostMatcher<T> {
         removed
     }
 
+    pub fn batch_remove(&mut self, ids: &HashSet<String>) -> bool {
+        self.any_host.batch_remove(ids);
+
+        self.static_hosts.retain(|_, matcher| {
+            matcher.batch_remove(ids);
+
+            !matcher.is_empty()
+        });
+
+        self.regex_tree_rule.retain(&|_, matcher| {
+            matcher.batch_remove(ids);
+
+            !matcher.is_empty()
+        });
+
+        self.any_host.is_empty() && self.static_hosts.is_empty() && self.regex_tree_rule.is_empty()
+    }
+
     pub fn match_request(&self, request: &Request) -> Vec<Arc<Route<T>>> {
         let mut routes = Vec::new();
 

src/router/request_matcher/ip.rs

@@ -2,7 +2,7 @@ use super::super::route_ip::RouteIp;
 use super::super::trace::TraceInfo;
 use super::super::{MethodMatcher, Route, RouterConfig, Trace};
 use crate::http::Request;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -67,6 +67,18 @@ impl<T> IpMatcher<T> {
         removed
     }
 
+    pub fn batch_remove(&mut self, ids: &HashSet<String>) -> bool {
+        self.no_matcher.batch_remove(ids);
+
+        self.matchers.retain(|_, matcher| {
+            matcher.batch_remove(ids);
+
+            !matcher.is_empty()
+        });
+
+        self.no_matcher.is_empty() && self.matchers.is_empty()
+    }
+
     pub fn match_request(&self, request: &Request) -> Vec<Arc<Route<T>>> {
         let mut routes = self.no_matcher.match_request(request);