Move bookmarks handling into session

Author: madchicken

ci/Cargo.lock.min

@@ -42,11 +42,10 @@ pastey = "0.1.0"
 pin-project-lite = "0.2.9"
 rustls-native-certs = "0.7.3"
 rustls-pemfile = "2.1.2"
-scc = "3.0.2"
 serde = { version = "1.0.185", features = ["derive"] }    # TODO: eliminate derive
 serde_json = { version = "1.0.0", optional = true }
 thiserror = "1.0.7"
-time = { version = "0.3.22", optional = true }
+time = { version = "=0.3.22", optional = true }
 tokio = { version = "1.5.0", features = ["full"] }
 url = "2.0.0"
 

ci/Cargo.lock.msrv

@@ -1,20 +1,26 @@
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 use {
     crate::connection::{ConnectionInfo, Routing},
-    crate::graph::ConnectionPoolManager::Routed,
     crate::routing::{ClusterRoutingTableProvider, RoutedConnectionManager},
     crate::summary::ResultSummary,
     log::debug,
     std::sync::Arc,
 };
 
 use crate::config::ImpersonateUser;
-use crate::graph::ConnectionPoolManager::Direct;
 use crate::pool::ManagedConnection;
 use crate::query::RetryableQuery;
 use crate::retry::Retry;
 use crate::RunResult;
-use crate::{config::{Config, ConfigBuilder, Database, LiveConfig}, errors::Result, pool::{create_pool, ConnectionPool}, query::Query, stream::DetachedRowStream, txn::Txn, Operation};
+use crate::{
+    config::{Config, ConfigBuilder, Database, LiveConfig},
+    errors::Result,
+    pool::{create_pool, ConnectionPool},
+    query::Query,
+    stream::DetachedRowStream,
+    txn::Txn,
+    Operation,
+};
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 use crate::{
     session::{Session, SessionConfig},
@@ -41,16 +47,18 @@ impl ConnectionPoolManager {
     ) -> Result<ManagedConnection> {
         match self {
             #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
-            Routed(manager) => manager.get(operation, db, imp_user, bookmarks).await,
-            Direct(pool) => pool.get().await.map_err(crate::Error::from),
+            ConnectionPoolManager::Routed(manager) => {
+                manager.get(operation, db, imp_user, bookmarks).await
+            }
+            ConnectionPoolManager::Direct(pool) => pool.get().await.map_err(crate::Error::from),
         }
     }
 
     fn backoff(&self) -> ExponentialBuilder {
         match self {
             #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
-            Routed(manager) => manager.backoff(),
-            Direct(pool) => pool.manager().backoff(),
+            ConnectionPoolManager::Routed(manager) => manager.backoff(),
+            ConnectionPoolManager::Direct(pool) => pool.manager().backoff(),
         }
     }
 }
@@ -84,7 +92,7 @@ impl Graph {
             )?;
             if matches!(info.init.routing, Routing::Yes(_)) {
                 debug!("Routing enabled, creating a routed connection manager");
-                let pool = Routed(RoutedConnectionManager::new(
+                let pool = ConnectionPoolManager::Routed(RoutedConnectionManager::new(
                     &config,
                     Arc::new(ClusterRoutingTableProvider::new(config.clone())),
                 )?);
@@ -93,7 +101,7 @@ impl Graph {
                     pool,
                 })
             } else {
-                let pool = Direct(create_pool(&config)?);
+                let pool = ConnectionPoolManager::Direct(create_pool(&config)?);
                 Ok(Graph {
                     config: config.into_live_config(),
                     pool,
@@ -102,7 +110,7 @@ impl Graph {
         }
         #[cfg(not(feature = "unstable-bolt-protocol-impl-v2"))]
         {
-            let pool = Direct(create_pool(&config)?);
+            let pool = ConnectionPoolManager::Direct(create_pool(&config)?);
             Ok(Graph {
                 config: config.into_live_config(),
                 pool,
@@ -194,7 +202,7 @@ impl Graph {
         {
             Txn::new(
                 db,
-                fetch_size.or(Some(self.config.fetch_size)).unwrap(),
+                fetch_size.unwrap_or(self.config.fetch_size),
                 connection,
                 operation,
                 imp_user,
@@ -247,14 +255,26 @@ impl Graph {
         db: impl Into<Database>,
         q: impl Into<Query>,
     ) -> Result<ResultSummary> {
-        self.impl_run_on(Some(db.into()), self.config.imp_user.clone(), &[], Some(self.config.fetch_size), q.into())
-            .await
+        self.impl_run_on(
+            Some(db.into()),
+            self.config.imp_user.clone(),
+            &[],
+            Some(self.config.fetch_size),
+            q.into(),
+        )
+        .await
     }
 
     #[cfg(not(feature = "unstable-bolt-protocol-impl-v2"))]
     pub async fn run_on(&self, db: impl Into<Database>, q: impl Into<Query>) -> Result<()> {
-        self.impl_run_on(Some(db.into()), self.config.imp_user.clone(), &[], Some(self.config.fetch_size), q.into())
-            .await
+        self.impl_run_on(
+            Some(db.into()),
+            self.config.imp_user.clone(),
+            &[],
+            Some(self.config.fetch_size),
+            q.into(),
+        )
+        .await
     }
 
     #[allow(unused_variables)]
@@ -284,40 +304,20 @@ impl Graph {
             .await;
 
         match result {
-            Ok(result) => {
-                #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
-                {
-                    if let Some(bookmark) = result.bookmark.as_deref() {
-                        match &self.pool {
-                            Routed(routed) => {
-                                routed.add_bookmark(bookmark).await;
-                            }
-                            Direct(_) => {}
-                        }
-                    } else if query.is_read() {
-                        match &self.pool {
-                            Routed(routed) => {
-                                debug!("No bookmark received after a read operation, discarding all bookmarks");
-                                routed.clear_bookmarks().await;
-                            }
-                            Direct(_) => {}
-                        }
-                    }
-                }
-                Ok(result)
-            }
+            Ok(result) => Ok(result),
             Err(e) => Err(e.into_inner()),
         }
     }
 
-    /// Executes a READ query on the configured database and returns a [`DetachedRowStream`]
+    /// Executes a READ|WRITE query on the configured database and returns a [`DetachedRowStream`]
     ///
     /// This operation retires the query on certain failures.
     /// All errors with the `Transient` error class as well as a few other error classes are considered retryable.
     /// This includes errors during a leader election or when the transaction resources on the server (memory, handles, ...) are exhausted.
     /// Retries happen with an exponential backoff until a retry delay exceeds 60s, at which point the query fails with the last error as it would without any retry.
     pub async fn execute(&self, q: impl Into<Query>) -> Result<DetachedRowStream> {
         self.impl_execute_on(
+            Operation::Write,
             self.config.db.clone(),
             self.config.imp_user.clone(),
             &[],
@@ -335,6 +335,7 @@ impl Graph {
     /// Retries happen with an exponential backoff until a retry delay exceeds 60s, at which point the query fails with the last error as it would without any retry.
     pub async fn execute_read(&self, q: impl Into<Query>) -> Result<DetachedRowStream> {
         self.impl_execute_on(
+            Operation::Read,
             self.config.db.clone(),
             self.config.imp_user.clone(),
             &[],
@@ -353,11 +354,19 @@ impl Graph {
     #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
     pub async fn execute_on(
         &self,
+        operation: Operation,
         db: impl Into<Database>,
         q: impl Into<Query>,
     ) -> Result<DetachedRowStream> {
-        self.impl_execute_on(Some(db.into()), self.config.imp_user.clone(), &[], Some(self.config.fetch_size), q.into())
-            .await
+        self.impl_execute_on(
+            operation,
+            Some(db.into()),
+            self.config.imp_user.clone(),
+            &[],
+            Some(self.config.fetch_size),
+            q.into(),
+        )
+        .await
     }
 
     /// Executes a query on the provided database and returns a [`DetachedRowStream`]
@@ -372,6 +381,7 @@ impl Graph {
         q: impl Into<Query>,
     ) -> Result<DetachedRowStream> {
         self.impl_execute_on(
+            Operation::Write,
             Some(db.into()),
             self.config.imp_user.clone(),
             &[],
@@ -384,6 +394,7 @@ impl Graph {
     #[allow(unused_variables)]
     pub(crate) async fn impl_execute_on(
         &self,
+        operation: Operation,
         db: Option<Database>,
         imp_user: Option<ImpersonateUser>,
         bookmarks: &[String],
@@ -393,7 +404,7 @@ impl Graph {
         let query = query.into_retryable(
             db,
             imp_user,
-            Operation::Read,
+            operation,
             &self.pool,
             fetch_size.or(Some(self.config.fetch_size)),
             bookmarks,
@@ -439,7 +450,7 @@ impl Graph {
                 .config
                 .db
                 .clone()
-                .map_or(Ok(Some("".into())), |db| Ok(Some(db.into()))),
+                .map_or(Ok(Some("".into())), |db| Ok(Some(db))),
         }
     }
 }

lib/Cargo.toml

@@ -515,7 +515,7 @@ mod version;
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 pub use {
     session::{Session, SessionConfig, SessionConfigBuilder},
-    utils::ConcurrentHashMap
+    utils::ConcurrentHashMap,
 };
 
 pub use crate::auth::ClientCertificate;

lib/src/graph.rs

@@ -259,11 +259,6 @@ pub(crate) struct RetryableQuery<'a> {
 }
 
 impl<'a> RetryableQuery<'a> {
-    #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
-    pub(crate) fn is_read(&self) -> bool {
-        self.operation.is_read()
-    }
-
     pub(crate) async fn retry_run(self) -> (Self, QueryResult<RunResult>) {
         let result = self.run().await;
         (self, result)

lib/src/lib.rs

@@ -42,18 +42,19 @@ impl ConnectionRegistry {
 
     /// Retrieve the pool for a specific server and database.
     pub(crate) fn get_server_pool(&self, server: &BoltServer) -> Option<ConnectionPool> {
-        self.pool_registry.get(&server.to_neo_url(self.scheme.as_str()))
+        self.pool_registry
+            .get(&server.to_neo_url(self.scheme.as_str()))
     }
 
     /// Mark a server as available for a specific database.
-    pub(crate) fn mark_unavailable(&self, server: &BoltServer, db: Option<Database>) {
-        let db_name = db.map_or(String::new(), |d| d.to_string());
-        if self.databases.contains_key(&db_name) {
+    pub(crate) fn mark_unavailable(&self, server: &BoltServer) {
+        for db_name in self.databases.keys() {
             debug!("Marking server as available: {server:?}");
             let mut table = self.databases.get(&db_name).unwrap();
             if table.mark_server_unavailable(server) {
                 self.databases.insert(db_name, table);
-                self.pool_registry.remove(&server.to_neo_url(self.scheme.as_str()));
+                self.pool_registry
+                    .remove(&server.to_neo_url(self.scheme.as_str()));
             } else {
                 debug!("Server not found in the registry: {server:?}");
             }
@@ -128,7 +129,11 @@ impl ConnectionRegistry {
     }
 
     pub fn all_servers(&self) -> Vec<BoltServer> {
-        self.databases.values().iter().flat_map(|dt| dt.resolve()).collect()
+        self.databases
+            .values()
+            .iter()
+            .flat_map(|dt| dt.resolve())
+            .collect()
     }
 
     pub fn update(&self, config: &Config, routing_table: &RoutingTable) -> Result<(), Error> {
@@ -298,10 +303,10 @@ mod tests {
         assert_eq!(servers.len(), 5);
         assert_eq!(registry.pool_registry.keys().len(), 3);
         registry.mark_unavailable(&BoltServer {
-            address:"host1".to_string(),
+            address: "host1".to_string(),
             port: 7687,
             role: "WRITE".to_string(),
-        }, db.clone());
+        });
         let servers = registry.servers(db, None, &[]).await;
         assert_eq!(servers.len(), 3);
         assert_eq!(registry.pool_registry.keys().len(), 2);
@@ -393,7 +398,7 @@ mod tests {
             ])),
         ));
 
-        let db = None;
+        let db = Some(Database::from(""));
         let db1 = Some(Database::from("db1"));
         let db2 = Some(Database::from("db2"));
 
@@ -406,20 +411,20 @@ mod tests {
         let servers = registry.servers(db2.clone(), None, &[]).await;
         assert_eq!(servers.len(), 5);
 
-        assert_eq!(registry.pool_registry.keys().len(), 7);
+        assert_eq!(registry.pool_registry.keys().len(), 9);
 
         registry.mark_unavailable(&BoltServer {
-            address:"host1".to_string(),
+            address: "host1".to_string(),
             port: 7687,
-            role: "WRITE".to_string(),
-        }, db.clone());
+            role: "READ".to_string(),
+        });
 
         let servers = registry.servers(db, None, &[]).await;
         assert_eq!(servers.len(), 4);
         let servers = registry.servers(db1, None, &[]).await;
         assert_eq!(servers.len(), 4);
         let servers = registry.servers(db2, None, &[]).await;
         assert_eq!(servers.len(), 5);
-        assert_eq!(registry.pool_registry.keys().len(), 6);
+        assert_eq!(registry.pool_registry.keys().len(), 8);
     }
 }

lib/src/query.rs

@@ -9,15 +9,13 @@ use crate::Database;
 #[cfg(feature = "unstable-bolt-protocol-impl-v2")]
 use crate::{Config, Error, Operation};
 use backon::ExponentialBuilder;
-use futures::lock::Mutex;
 use log::{debug, error};
 use std::sync::Arc;
 
 #[derive(Clone)]
 pub struct RoutedConnectionManager {
     load_balancing_strategy: Arc<dyn LoadBalancingStrategy>,
     connection_registry: Arc<ConnectionRegistry>,
-    bookmarks: Arc<Mutex<Vec<String>>>,
     backoff: ExponentialBuilder,
 }
 
@@ -27,7 +25,6 @@ impl RoutedConnectionManager {
         let connection_registry = Arc::new(ConnectionRegistry::new(config, provider));
         Ok(RoutedConnectionManager {
             load_balancing_strategy: Arc::new(RoundRobinStrategy::new(connection_registry.clone())),
-            bookmarks: Arc::new(Mutex::new(vec![])),
             connection_registry,
             backoff,
         })
@@ -78,8 +75,7 @@ impl RoutedConnectionManager {
                     Ok(conn) => return Ok(conn),
                     Err(e) => {
                         error!("Failed to get connection from pool for server {selected_server:?}: {e}");
-                        self.connection_registry
-                            .mark_unavailable(&selected_server, db.clone());
+                        self.connection_registry.mark_unavailable(&selected_server);
                         continue; // Try selecting another server
                     }
                 }
@@ -113,17 +109,4 @@ impl RoutedConnectionManager {
     fn select_writer(&self, servers: &[BoltServer]) -> Option<BoltServer> {
         self.load_balancing_strategy.select_writer(servers)
     }
-
-    #[allow(dead_code)]
-    pub(crate) async fn add_bookmark(&self, bookmark: &str) {
-        let mut guard = self.bookmarks.lock().await;
-        if !guard.contains(&bookmark.to_string()) {
-            guard.push(bookmark.to_string());
-        }
-    }
-
-    #[allow(dead_code)]
-    pub(crate) async fn clear_bookmarks(&self) {
-        self.bookmarks.lock().await.clear();
-    }
 }