Fix load balancing strategy for multi database support

madchicken · madchicken · commit 342c638eacc0 · 2025-05-30T15:28:30.000+02:00
diff --git a/lib/src/pool.rs b/lib/src/pool.rs
@@ -63,7 +63,8 @@ pub fn create_pool(config: &Config) -> Result<ConnectionPool> {
         &config.tls_config,
     )?;
     info!(
-        "creating connection pool with max size {}",
+        "creating connection pool for node {} with max size {}",
+        config.uri,
         config.max_connections
     );
     Ok(ConnectionPool::builder(mgr)
diff --git a/lib/src/routing/connection_registry.rs b/lib/src/routing/connection_registry.rs
@@ -5,13 +5,14 @@ use crate::routing::{RoutingTable, Server};
 use crate::{Config, Database, Error};
 use dashmap::DashMap;
 use log::debug;
+use std::fmt::Debug;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::Sender;
 
 /// Represents a Bolt server, with its address, port and role.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub(crate) struct BoltServer {
     pub(crate) address: String,
     pub(crate) port: u16,
@@ -35,23 +36,12 @@ impl BoltServer {
             })
             .collect()
     }
-}
-
-impl Eq for BoltServer {}
-
-impl PartialEq for BoltServer {
-    fn eq(&self, other: &Self) -> bool {
+    
+    pub fn has_same_address(&self, other: &Self) -> bool {
         self.address == other.address && self.port == other.port
     }
 }
 
-impl std::hash::Hash for BoltServer {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.address.hash(state);
-        self.port.hash(state);
-    }
-}
-
 /// A registry of connection pools, indexed by the Bolt server they connect to.
 pub type PoolRegistry = DashMap<BoltServer, ConnectionPool>;
 /// A map of registries, indexed by the database name.
@@ -138,6 +128,17 @@ async fn refresh_routing_tables(
             "No servers available in the routing table".to_string(),
         ));
     }
+
+    // purge the pool registry of servers that are no longer in the routing tables
+    let all_servers: Vec<BoltServer> = connection_registry
+        .databases
+        .iter()
+        .flat_map(|kv| kv.value().clone())
+        .collect();
+    connection_registry
+        .pool_registry
+        .retain(|server, _| all_servers.contains(server));
+
     Ok(*ttls.iter().min().unwrap())
 }
 
@@ -179,9 +180,9 @@ async fn refresh_routing_table(
             })?,
         );
     }
-    registry.retain(|k, _| servers.contains(k));
     debug!(
-        "Registry updated. New size is {} with TTL {}s",
+        "Registry updated for database {}. New size is {} with TTL {}s",
+        db.as_ref().map_or("default".to_string(), |d| d.to_string()),
         registry.len(),
         routing_table.ttl
     );
@@ -271,7 +272,7 @@ impl ConnectionRegistry {
             if let Some(index) = self
                 .databases
                 .get(db_name.as_str())
-                .and_then(|vec| vec.iter().position(|s| s == server))
+                .and_then(|vec| vec.iter().position(|s| server.has_same_address(s)))
             {
                 debug!("Marking server as available: {:?}", server);
                 self.databases
@@ -300,6 +301,13 @@ impl ConnectionRegistry {
         }
     }
 
+    pub fn all_servers(&self) -> Vec<BoltServer> {
+        self.pool_registry
+            .iter()
+            .map(|kv| kv.key().clone())
+            .collect::<Vec<BoltServer>>()
+    }
+
     fn get_db_name(&self, db: Option<Database>) -> String {
         db.as_ref()
             .map(|d| d.to_string())
@@ -419,7 +427,7 @@ mod tests {
         let servers = registry.servers(None);
         assert_eq!(servers.len(), 5);
 
-        let strategy = RoundRobinStrategy::default();
+        let strategy = RoundRobinStrategy::new(registry.clone());
         registry.mark_unavailable(BoltServer::resolve(&writers[0]).first().unwrap(), None);
         let servers = registry.servers(None);
         assert_eq!(servers.len(), 4);
@@ -540,7 +548,7 @@ mod tests {
         assert_eq!(servers2.len(), 5); // 2 readers, 2 writers, 1 router
         assert_eq!(servers2.first().unwrap().address, "host5");
 
-        let strategy = RoundRobinStrategy::default();
+        let strategy = RoundRobinStrategy::new(registry.clone());
         let writer = strategy
             .select_writer(&registry.servers(Some("db1".into())))
             .unwrap();
diff --git a/lib/src/routing/load_balancing/round_robin_strategy.rs b/lib/src/routing/load_balancing/round_robin_strategy.rs
@@ -1,32 +1,46 @@
-use crate::routing::connection_registry::BoltServer;
+use crate::routing::connection_registry::{BoltServer, ConnectionRegistry};
 use crate::routing::load_balancing::LoadBalancingStrategy;
-use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
-#[derive(Default)]
 pub struct RoundRobinStrategy {
+    connection_registry: Arc<ConnectionRegistry>,
     reader_index: AtomicUsize,
     writer_index: AtomicUsize,
 }
 
 impl RoundRobinStrategy {
-    fn select(servers: &[BoltServer], index: &AtomicUsize) -> Option<BoltServer> {
+    pub fn new(connection_registry: Arc<ConnectionRegistry>) -> Self {
+        RoundRobinStrategy {
+            connection_registry,
+            reader_index: AtomicUsize::new(0),
+            writer_index: AtomicUsize::new(0),
+        }
+    }
+
+    fn select(all_servers: &[BoltServer], servers: &[BoltServer], index: &AtomicUsize) -> Option<BoltServer> {
         if servers.is_empty() {
             return None;
         }
 
-        let _ = index.compare_exchange(
-            0,
-            servers.len(),
-            std::sync::atomic::Ordering::Relaxed,
-            std::sync::atomic::Ordering::Relaxed,
-        );
-        let i = index.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
-        if let Some(server) = servers.get(i - 1) {
-            Some(server.clone())
-        } else {
-            //reset index
-            index.store(servers.len(), std::sync::atomic::Ordering::Relaxed);
-            servers.last().cloned()
+        let mut used = vec![];
+        loop {
+            if used.len() >= all_servers.len() {
+                return None; // All servers have been used
+            }
+            let _ = index.compare_exchange(
+                0,
+                all_servers.len(),
+                Ordering::Relaxed,
+                Ordering::Relaxed,
+            );
+            let i = index.fetch_sub(1, Ordering::Relaxed);
+            if let Some(server) = all_servers.get(i - 1) {
+                if servers.contains(server) {
+                    return Some(server.clone());
+                }
+                used.push(server.clone());
+            }
         }
     }
 }
@@ -38,7 +52,13 @@ impl LoadBalancingStrategy for RoundRobinStrategy {
             .filter(|s| s.role == "READ")
             .cloned()
             .collect::<Vec<BoltServer>>();
-        Self::select(&readers, &self.reader_index)
+        let all_readers = self.connection_registry.all_servers()
+            .iter()
+            .filter(|s| s.role == "READ")
+            .cloned()
+            .collect::<Vec<BoltServer>>();
+
+        Self::select(&all_readers, &readers, &self.reader_index)
     }
 
     fn select_writer(&self, servers: &[BoltServer]) -> Option<BoltServer> {
@@ -47,7 +67,13 @@ impl LoadBalancingStrategy for RoundRobinStrategy {
             .filter(|s| s.role == "WRITE")
             .cloned()
             .collect::<Vec<BoltServer>>();
-        Self::select(&writers, &self.writer_index)
+        let all_writers = self.connection_registry.all_servers()
+            .iter()
+            .filter(|s| s.role == "WRITE")
+            .cloned()
+            .collect::<Vec<BoltServer>>();
+
+        Self::select(&all_writers, &writers, &self.writer_index)
     }
 }
 
@@ -59,54 +85,94 @@ mod tests {
     #[test]
     fn should_get_next_server() {
         let routers = vec![Server {
-            addresses: vec!["192.168.0.1:7688".to_string()],
+            addresses: vec!["server1:7687".to_string()],
+            role: "ROUTE".to_string(),
+        }];
+        let readers1 = vec![Server {
+            addresses: vec!["server1:7687".to_string()],
+            role: "READ".to_string(),
+        }, Server {
+            addresses: vec!["server2:7687".to_string()],
+            role: "READ".to_string(),
+        }];
+        let writers1 = vec![Server {
+            addresses: vec!["server4:7687".to_string()],
             role: "WRITE".to_string(),
         }];
-        let readers = vec![Server {
-            addresses: vec![
-                "192.168.0.2:7687".to_string(),
-                "192.168.0.3:7687".to_string(),
-            ],
+        let readers2 = vec![Server {
+            addresses: vec!["server1:7687".to_string()],
+            role: "READ".to_string(),
+        }, Server {
+            addresses: vec!["server3:7687".to_string()],
             role: "READ".to_string(),
         }];
-        let writers = vec![Server {
-            addresses: vec!["192.168.0.4:7688".to_string()],
+
+        let writers2 = vec![Server {
+            addresses: vec!["server4:7687".to_string()],
             role: "WRITE".to_string(),
         }];
 
-        let cluster_routing_table = RoutingTable {
+        let routing_table_1 = RoutingTable {
+            ttl: 300,
+            db: Some("db-1".into()),
+            servers: routers
+                .clone()
+                .into_iter()
+                .chain(readers1.clone())
+                .chain(writers1.clone())
+                .collect(),
+        };
+        let routing_table_2 = RoutingTable {
             ttl: 300,
-            db: Some("neo4j".into()),
+            db: Some("db-2".into()),
             servers: routers
                 .clone()
                 .into_iter()
-                .chain(readers.clone())
-                .chain(writers.clone())
+                .chain(readers2.clone())
+                .chain(writers2.clone())
                 .collect(),
         };
-        let all_servers = cluster_routing_table.resolve();
-        assert_eq!(all_servers.len(), 4);
-        let strategy = RoundRobinStrategy::default();
 
-        let reader = strategy.select_reader(&all_servers).unwrap();
-        assert_eq!(
-            format!("{}:{}", reader.address, reader.port),
-            readers[0].addresses[1]
-        );
-        let reader = strategy.select_reader(&all_servers).unwrap();
+        let registry = Arc::new(ConnectionRegistry::default());
+
+        let mut servers1 = routing_table_1.resolve();
+        servers1.retain(|s| s.role == "READ");
+        let mut servers2 = routing_table_2.resolve();
+        servers2.retain(|s| s.role == "READ");
+
+        let mut all_readers: Vec<BoltServer> = Vec::new();
+        for s in servers1.iter() {
+            if !all_readers.iter().any(|x| x == s) {
+                all_readers.push(s.clone());
+            }
+        }
+        for s in servers2.iter() {
+            if !all_readers.iter().any(|x| x == s) {
+                all_readers.push(s.clone());
+            }
+        }
+        all_readers.retain(|s| s.role == "READ");
+
+        assert_eq!(all_readers.len(), 3);
+        let strategy = RoundRobinStrategy::new(registry.clone());
+
+        // select a reader for db-1
+        let reader = RoundRobinStrategy::select(&all_readers, &servers1, &strategy.reader_index).unwrap();
         assert_eq!(
-            format!("{}:{}", reader.address, reader.port),
-            readers[0].addresses[0]
+            reader.address,
+            "server2"
         );
-        let reader = strategy.select_reader(&all_servers).unwrap();
+        // select a reader for db-2
+        let reader = RoundRobinStrategy::select(&all_readers, &servers2, &strategy.reader_index).unwrap();
         assert_eq!(
-            format!("{}:{}", reader.address, reader.port),
-            readers[0].addresses[1]
+            reader.address,
+            "server1"
         );
-        let writer = strategy.select_writer(&all_servers).unwrap();
+        // select another reader for db-1
+        let reader = RoundRobinStrategy::select(&all_readers, &servers1, &strategy.reader_index).unwrap();
         assert_eq!(
-            format!("{}:{}", writer.address, writer.port),
-            writers[0].addresses[0]
+            reader.address,
+            "server2"
         );
     }
 }
diff --git a/lib/src/routing/routed_connection_manager.rs b/lib/src/routing/routed_connection_manager.rs
@@ -23,6 +23,8 @@ pub struct RoutedConnectionManager {
     channel: Sender<RegistryCommand>,
 }
 
+const ROUTING_TABLE_MAX_WAIT_TIME_MS: i32 = 5000;
+
 impl RoutedConnectionManager {
     pub fn new(config: &Config, provider: Arc<dyn RoutingTableProvider>) -> Result<Self, Error> {
         let backoff = Arc::new(
@@ -37,7 +39,7 @@ impl RoutedConnectionManager {
         let connection_registry = Arc::new(ConnectionRegistry::default());
         let channel = start_background_updater(config, connection_registry.clone(), provider);
         Ok(RoutedConnectionManager {
-            load_balancing_strategy: Arc::new(RoundRobinStrategy::default()),
+            load_balancing_strategy: Arc::new(RoundRobinStrategy::new(connection_registry.clone())),
             bookmarks: Arc::new(Mutex::new(vec![])),
             connection_registry,
             backoff,
@@ -76,15 +78,17 @@ impl RoutedConnectionManager {
             if servers.is_empty() {
                 // the first time we need to wait until we get the routing table
                 tokio::time::sleep(Duration::from_millis(10)).await;
-                attempts += 1;
-                if attempts > 500 {
-                    // 5 seconds max wait time
+                attempts += 10;
+                if attempts > ROUTING_TABLE_MAX_WAIT_TIME_MS {
+                    // 5 seconds max wait time by default
                     error!(
-                        "Failed to get a connection after 5 seconds, routing table is still empty"
+                        "Failed to get a connection after {} seconds, routing table is still empty",
+                        ROUTING_TABLE_MAX_WAIT_TIME_MS / 1000
                     );
-                    return Err(Error::ServerUnavailableError(
-                        "Routing table is still empty after 5 seconds".to_string(),
-                    ));
+                    return Err(Error::ServerUnavailableError(format!(
+                        "Routing table is still empty after {} seconds",
+                        ROUTING_TABLE_MAX_WAIT_TIME_MS / 1000
+                    )));
                 }
                 continue;
             }
