Simplify connection: fetch routing table at the first usage

Author: madchicken

lib/src/graph.rs

@@ -1,9 +1,8 @@
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 use {
-    crate::connection::{Connection, ConnectionInfo, Routing},
+    crate::connection::{ConnectionInfo, Routing},
     crate::graph::ConnectionPoolManager::Routed,
-    crate::routing::{RoundRobinStrategy, RouteBuilder, RoutedConnectionManager},
-    log::info,
+    crate::routing::{RoundRobinStrategy, RoutedConnectionManager},
     std::sync::Arc,
 };
 
@@ -75,21 +74,10 @@ impl Graph {
                 &config.tls_config,
             )?;
             if matches!(info.routing, Routing::Yes(_)) {
-                let mut connection = Connection::new(&info).await?;
-                let mut builder = RouteBuilder::new(info.routing, vec![]);
-                if let Some(db) = config.db.clone() {
-                    builder = builder.with_db(db);
-                }
-                let rt = connection
-                    .route(builder.build(connection.version()))
-                    .await?;
-                connection.reset().await?;
-                info!("Connected to routing server, routing table: {:?}", rt);
                 let pool = Routed(
                     RoutedConnectionManager::new(
                         &config,
-                        Arc::new(rt.clone()),
-                        Arc::new(RoundRobinStrategy::new(&rt.resolve())),
+                        Arc::new(RoundRobinStrategy::default()),
                     )
                     .await?,
                 );

lib/src/routing/connection_registry.rs

@@ -6,6 +6,7 @@ use dashmap::DashMap;
 use futures::lock::Mutex;
 use log::debug;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub(crate) struct BoltServer {
@@ -40,19 +41,15 @@ pub type Registry = DashMap<BoltServer, ConnectionPool>;
 pub(crate) struct ConnectionRegistry {
     config: Config,
     creation_time: Arc<Mutex<u64>>,
-    ttl: u64,
+    ttl: Arc<AtomicU64>,
     pub(crate) connections: Registry,
 }
 
 impl ConnectionRegistry {
     pub(crate) async fn new(
         config: &Config,
-        routing_table: Arc<RoutingTable>,
     ) -> Result<Self, Error> {
-        let ttl = routing_table.ttl;
-        let servers = routing_table.resolve();
-
-        let connections = Self::build_registry(config, &servers).await?;
+        let connections = Self::build_registry(config, &[]).await?;
         Ok(ConnectionRegistry {
             config: config.clone(),
             creation_time: Arc::new(Mutex::new(
@@ -61,7 +58,7 @@ impl ConnectionRegistry {
                     .unwrap()
                     .as_secs(),
             )),
-            ttl,
+            ttl: Arc::new(AtomicU64::new(0)),
             connections,
         })
     }
@@ -89,8 +86,8 @@ impl ConnectionRegistry {
             .as_secs();
         debug!("Checking if routing table is expired...");
         if let Some(mut guard) = self.creation_time.try_lock() {
-            if now - *guard > self.ttl {
-                debug!("Routing table expired, refreshing...");
+            if self.connections.is_empty() || now - *guard > self.ttl.load(Ordering::Relaxed) {
+                debug!("Routing table expired or empty, refreshing...");
                 let routing_table = f().await?;
                 debug!("Routing table refreshed: {:?}", routing_table);
                 let registry = &self.connections;
@@ -102,7 +99,12 @@ impl ConnectionRegistry {
                     registry.insert(server.clone(), create_pool(&self.config).await?);
                 }
                 registry.retain(|k, _| servers.contains(k));
-                debug!("Registry updated. New size is {}", registry.len());
+                let _ = self.ttl.fetch_update(
+                    Ordering::Relaxed,
+                    Ordering::Relaxed,
+                    |_ttl| Some(routing_table.ttl),
+                ).unwrap();
+                debug!("Registry updated. New size is {} with TTL {}s", registry.len(), routing_table.ttl);
                 *guard = now;
             }
         } else {
@@ -177,16 +179,14 @@ mod tests {
             fetch_size: 0,
             tls_config: ConnectionTLSConfig::None,
         };
-        let registry = ConnectionRegistry::new(&config, Arc::new(cluster_routing_table.clone()))
+        let registry = ConnectionRegistry::new(&config)
+            .await
+            .unwrap();
+        registry.update_if_expired(|| async { Ok(cluster_routing_table) })
             .await
             .unwrap();
         assert_eq!(registry.connections.len(), 5);
-        let strategy = RoundRobinStrategy::new(&cluster_routing_table.resolve());
-        let router = strategy.select_router(&registry.servers()).unwrap();
-        assert_eq!(
-            format!("{}:{}", router.address, router.port),
-            routers[0].addresses[0]
-        );
+        let strategy = RoundRobinStrategy::default();
         registry.mark_unavailable(BoltServer::resolve(&writers[0]).first().unwrap());
         assert_eq!(registry.connections.len(), 4);
         let writer = strategy.select_writer(&registry.servers()).unwrap();

lib/src/routing/load_balancing/mod.rs

@@ -5,5 +5,4 @@ use crate::routing::connection_registry::BoltServer;
 pub trait LoadBalancingStrategy: Sync + Send {
     fn select_reader(&self, servers: &[BoltServer]) -> Option<BoltServer>;
     fn select_writer(&self, servers: &[BoltServer]) -> Option<BoltServer>;
-    fn select_router(&self, servers: &[BoltServer]) -> Option<BoltServer>;
 }

lib/src/routing/load_balancing/round_robin_strategy.rs

@@ -2,39 +2,13 @@ use crate::routing::connection_registry::BoltServer;
 use crate::routing::load_balancing::LoadBalancingStrategy;
 use std::sync::atomic::AtomicUsize;
 
+#[derive(Default)]
 pub struct RoundRobinStrategy {
     reader_index: AtomicUsize,
     writer_index: AtomicUsize,
-    router_index: AtomicUsize,
 }
 
 impl RoundRobinStrategy {
-    pub(crate) fn new(servers: &[BoltServer]) -> Self {
-        let readers: Vec<BoltServer> = servers
-            .iter()
-            .filter(|s| s.role == "READ")
-            .cloned()
-            .collect();
-        let writers: Vec<BoltServer> = servers
-            .iter()
-            .filter(|s| s.role == "WRITE")
-            .cloned()
-            .collect();
-        let routers: Vec<BoltServer> = servers
-            .iter()
-            .filter(|s| s.role == "ROUTE")
-            .cloned()
-            .collect();
-        let reader_index = AtomicUsize::new(readers.len());
-        let writer_index = AtomicUsize::new(writers.len());
-        let router_index = AtomicUsize::new(routers.len());
-        RoundRobinStrategy {
-            reader_index,
-            writer_index,
-            router_index,
-        }
-    }
-
     fn select(servers: &[BoltServer], index: &AtomicUsize) -> Option<BoltServer> {
         if servers.is_empty() {
             return None;
@@ -77,15 +51,6 @@ impl LoadBalancingStrategy for RoundRobinStrategy {
             .collect::<Vec<BoltServer>>();
         Self::select(&writers, &self.writer_index)
     }
-
-    fn select_router(&self, servers: &[BoltServer]) -> Option<BoltServer> {
-        let routers = servers
-            .iter()
-            .filter(|s| s.role == "ROUTE")
-            .cloned()
-            .collect::<Vec<BoltServer>>();
-        Self::select(&routers, &self.router_index)
-    }
 }
 
 #[cfg(test)]
@@ -123,7 +88,7 @@ mod tests {
         };
         let all_servers = cluster_routing_table.resolve();
         assert_eq!(all_servers.len(), 4);
-        let strategy = RoundRobinStrategy::new(&cluster_routing_table.resolve());
+        let strategy = RoundRobinStrategy::default();
 
         let reader = strategy.select_reader(&all_servers).unwrap();
         assert_eq!(

lib/src/routing/routed_connection_manager.rs

@@ -9,14 +9,15 @@ use std::sync::Arc;
 use std::time::Duration;
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 use {
-    crate::connection::Routing,
     crate::routing::{RouteBuilder, RoutingTable},
 };
+use crate::connection::{Connection, ConnectionInfo};
 
 #[derive(Clone)]
 pub struct RoutedConnectionManager {
     load_balancing_strategy: Arc<dyn LoadBalancingStrategy>,
     registry: Arc<ConnectionRegistry>,
+    #[allow(dead_code)]
     bookmarks: Arc<Mutex<Vec<String>>>,
     backoff: Arc<ExponentialBackoff>,
     config: Config,
@@ -25,10 +26,9 @@ pub struct RoutedConnectionManager {
 impl RoutedConnectionManager {
     pub async fn new(
         config: &Config,
-        routing_table: Arc<RoutingTable>,
         load_balancing_strategy: Arc<dyn LoadBalancingStrategy>,
     ) -> Result<Self, Error> {
-        let registry = Arc::new(ConnectionRegistry::new(config, routing_table.clone()).await?);
+        let registry = Arc::new(ConnectionRegistry::new(config).await?);
         let backoff = Arc::new(
             ExponentialBackoffBuilder::new()
                 .with_initial_interval(Duration::from_millis(1))
@@ -48,43 +48,20 @@ impl RoutedConnectionManager {
     }
 
     pub async fn refresh_routing_table(&self) -> Result<RoutingTable, Error> {
-        while let Some(router) = self.select_router() {
-            if let Some(pool) = self.registry.get_pool(&router) {
-                if let Ok(mut connection) = pool.get().await {
-                    debug!("Refreshing routing table from router {}", router.address);
-                    let bookmarks = self.bookmarks.lock().await;
-                    let bookmarks = bookmarks.iter().map(|b| b.as_str()).collect();
-                    let route = RouteBuilder::new(Routing::Yes(vec![]), bookmarks)
-                        .with_db(self.config.db.clone().unwrap_or_default())
-                        .build(connection.version());
-                    match connection.route(route).await {
-                        Ok(rt) => {
-                            debug!("Routing table refreshed: {:?}", rt);
-                            return Ok(rt);
-                        }
-                        Err(e) => {
-                            self.registry.mark_unavailable(&router);
-                            error!(
-                                "Failed to refresh routing table from router {}: {}",
-                                router.address, e
-                            );
-                        }
-                    }
-                } else {
-                    self.registry.mark_unavailable(&router);
-                    error!("Failed to create connection to router `{}`", router.address);
-                }
-            } else {
-                error!(
-                    "No connection manager available for router `{}` in the registry. Maybe it was marked as unavailable",
-                    router.address
-                );
-            }
+        let info = ConnectionInfo::new(
+            &self.config.uri,
+            &self.config.user,
+            &self.config.password,
+            &self.config.tls_config,
+        )?;
+        let mut connection = Connection::new(&info).await?;
+        let mut builder = RouteBuilder::new(info.routing, vec![]);
+        if let Some(db) = self.config.db.clone() {
+            builder = builder.with_db(db);
         }
-        // After trying all routers, we still couldn't refresh the routing table: return an error
-        Err(Error::ServerUnavailableError(
-            "No router available".to_string(),
-        ))
+        let rt = connection.route(builder.build(connection.version())).await?;
+        debug!("Fetched a new routing table: {:?}", rt);
+        Ok(rt)
     }
 
     pub(crate) async fn get(
@@ -143,8 +120,4 @@ impl RoutedConnectionManager {
     fn select_writer(&self) -> Option<BoltServer> {
         self.load_balancing_strategy.select_writer(&self.registry.servers())
     }
-
-    fn select_router(&self) -> Option<BoltServer> {
-        self.load_balancing_strategy.select_router(&self.registry.servers())
-    }
 }