Rename HashedTokenRouter -> TokenRouter (googleforgames#980)

goalsgame · Jun 12, 2024 · 50d91e4 · 50d91e4
1 parent 64d4289
commit 50d91e4
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 219 deletions.
diff --git a/benches/token_router.rs b/benches/token_router.rs
@@ -1,21 +1,17 @@
 use divan::Bencher;
-use quilkin::filters::token_router::{HashedTokenRouter, Router, TokenRouter};
+use quilkin::filters::token_router::TokenRouter;
 use rand::SeedableRng;
 
 mod shared;
 
-#[divan::bench(types = [TokenRouter, HashedTokenRouter], args = ["single:duplicates", "single:unique", "multi:2..128:duplicates", "multi:2..128:unique"])]
-fn token_router<T>(b: Bencher, token_kind: &str)
-where
-    T: Router + Sync,
-{
-    let filter = <T as Router>::new();
+#[divan::bench(args = ["single:duplicates", "single:unique", "multi:2..128:duplicates", "multi:2..128:unique"])]
+fn token_router(b: Bencher, token_kind: &str) {
+    let filter = TokenRouter::default();
     let gc = shared::gen_cluster_map::<42>(token_kind.parse().unwrap());
 
     let mut tokens = Vec::new();
 
     let cm = std::sync::Arc::new(gc.cm);
-    cm.build_token_maps();
 
     // Calculate the amount of bytes for all the tokens
     for eps in cm.iter() {

diff --git a/crates/test/tests/mesh.rs b/crates/test/tests/mesh.rs
@@ -52,7 +52,7 @@ trace_test!(relay_routing, {
                         }),
                     })
                     .unwrap(),
-                    HashedTokenRouter::as_filter_config(None).unwrap(),
+                    TokenRouter::as_filter_config(None).unwrap(),
                 ])
                 .unwrap(),
                 ..Default::default()
@@ -226,7 +226,7 @@ trace_test!(filter_update, {
                         }),
                     })
                     .unwrap(),
-                    HashedTokenRouter::as_filter_config(None).unwrap(),
+                    TokenRouter::as_filter_config(None).unwrap(),
                 ])
                 .unwrap(),
                 ..Default::default()
@@ -279,7 +279,7 @@ trace_test!(filter_update, {
                     token.len().to_string(),
                 )
                 .unwrap(),
-                HashedTokenRouter::as_filter_config(None).unwrap(),
+                TokenRouter::as_filter_config(None).unwrap(),
             ])
             .unwrap();
         });

diff --git a/src/filters.rs b/src/filters.rs
@@ -66,7 +66,7 @@ pub use self::{
     registry::FilterRegistry,
     set::{FilterMap, FilterSet},
     timestamp::Timestamp,
-    token_router::{HashedTokenRouter, TokenRouter},
+    token_router::TokenRouter,
     write::WriteContext,
 };
 

diff --git a/src/filters/set.rs b/src/filters/set.rs
@@ -51,7 +51,6 @@ impl FilterSet {
                 filters::Debug::factory(),
                 filters::Drop::factory(),
                 filters::Firewall::factory(),
-                filters::HashedTokenRouter::factory(),
                 filters::LoadBalancer::factory(),
                 filters::LocalRateLimit::factory(),
                 filters::Match::factory(),

diff --git a/src/filters/token_router.rs b/src/filters/token_router.rs
@@ -21,51 +21,24 @@ use crate::{
     net::endpoint::metadata,
 };
 
-use crate::generated::quilkin::filters::token_router::v1alpha1 as proto;
+use xds::generated::quilkin::filters::token_router::v1alpha1 as proto;
 
 /// Filter that only allows packets to be passed to Endpoints that have a matching
 /// connection_id to the token stored in the Filter's dynamic metadata.
+#[derive(Default)]
 pub struct TokenRouter {
     config: Config,
 }
 
 impl TokenRouter {
-    fn new(config: Config) -> Self {
-        Self { config }
-    }
-}
-
-impl StaticFilter for TokenRouter {
-    const NAME: &'static str = "quilkin.filters.token_router.v1alpha1.TokenRouter";
-    type Configuration = Config;
-    type BinaryConfiguration = proto::TokenRouter;
-
-    fn try_from_config(config: Option<Self::Configuration>) -> Result<Self, CreationError> {
-        Ok(TokenRouter::new(config.unwrap_or_default()))
-    }
-}
-
-#[async_trait::async_trait]
-impl Filter for TokenRouter {
-    async fn read(&self, ctx: &mut ReadContext) -> Result<(), FilterError> {
-        self.sync_read(ctx)
-    }
-}
-
-impl Router for TokenRouter {
-    fn sync_read(&self, ctx: &mut ReadContext) -> Result<(), FilterError> {
+    /// Non-async version of [`Filter::read`], as this filter does no actual async
+    /// operations. Used in benchmarking.
+    pub fn sync_read(&self, ctx: &mut ReadContext) -> Result<(), FilterError> {
         match ctx.metadata.get(&self.config.metadata_key) {
             Some(metadata::Value::Bytes(token)) => {
-                let destinations = ctx.endpoints.filter_endpoints(|endpoint| {
-                    if endpoint.metadata.known.tokens.contains(&**token) {
-                        tracing::trace!(%endpoint.address, token = &*crate::codec::base64::encode(token), "Endpoint matched");
-                        true
-                    } else {
-                        false
-                    }
-                });
+                let tok = crate::net::cluster::Token::new(token);
 
-                ctx.destinations = destinations.into_iter().map(|ep| ep.address).collect();
+                ctx.destinations = ctx.endpoints.addresses_for_token(tok);
 
                 if ctx.destinations.is_empty() {
                     Err(FilterError::new(Error::NoEndpointMatch(
@@ -85,76 +58,27 @@ impl Router for TokenRouter {
             ))),
         }
     }
-
-    fn new() -> Self {
-        Self::from_config(None)
-    }
-}
-
-pub struct HashedTokenRouter {
-    config: Config,
-}
-
-impl HashedTokenRouter {
-    fn new(config: Config) -> Self {
-        Self { config }
-    }
 }
 
-impl StaticFilter for HashedTokenRouter {
-    const NAME: &'static str = "quilkin.filters.token_router.v1alpha1.HashedTokenRouter";
+impl StaticFilter for TokenRouter {
+    const NAME: &'static str = "quilkin.filters.token_router.v1alpha1.TokenRouter";
     type Configuration = Config;
     type BinaryConfiguration = proto::TokenRouter;
 
     fn try_from_config(config: Option<Self::Configuration>) -> Result<Self, CreationError> {
-        Ok(Self::new(config.unwrap_or_default()))
+        Ok(Self {
+            config: config.unwrap_or_default(),
+        })
     }
 }
 
 #[async_trait::async_trait]
-impl Filter for HashedTokenRouter {
+impl Filter for TokenRouter {
     async fn read(&self, ctx: &mut ReadContext) -> Result<(), FilterError> {
         self.sync_read(ctx)
     }
 }
 
-impl Router for HashedTokenRouter {
-    fn sync_read(&self, ctx: &mut ReadContext) -> Result<(), FilterError> {
-        match ctx.metadata.get(&self.config.metadata_key) {
-            Some(metadata::Value::Bytes(token)) => {
-                let tok = crate::net::cluster::Token::new(token);
-
-                ctx.destinations = ctx.endpoints.addresses_for_token(tok);
-
-                if ctx.destinations.is_empty() {
-                    Err(FilterError::new(Error::NoEndpointMatch(
-                        self.config.metadata_key,
-                        crate::codec::base64::encode(token),
-                    )))
-                } else {
-                    Ok(())
-                }
-            }
-            Some(value) => Err(FilterError::new(Error::InvalidType(
-                self.config.metadata_key,
-                value.clone(),
-            ))),
-            None => Err(FilterError::new(Error::NoTokenFound(
-                self.config.metadata_key,
-            ))),
-        }
-    }
-
-    fn new() -> Self {
-        Self::from_config(None)
-    }
-}
-
-pub trait Router {
-    fn sync_read(&self, ctx: &mut ReadContext) -> Result<(), FilterError>;
-    fn new() -> Self;
-}
-
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error("no routing token found for `{0}`")]

diff --git a/src/net/cluster.rs b/src/net/cluster.rs
@@ -505,21 +505,6 @@ where
         ret
     }
 
-    /// Builds token maps for every locality. Only used by testing/benching
-    #[doc(hidden)]
-    pub fn build_token_maps(&self) {
-        self.token_map.clear();
-
-        for mut eps in self.map.iter_mut() {
-            eps.build_token_map();
-
-            for (token_hash, addrs) in &eps.token_map {
-                self.token_map
-                    .insert(*token_hash, addrs.iter().cloned().collect());
-            }
-        }
-    }
-
     pub fn addresses_for_token(&self, token: Token) -> Vec<EndpointAddress> {
         self.token_map
             .get(&token.0)

diff --git a/tests/hashed_token_router.rs b/tests/hashed_token_router.rs