Feat: Bolt 5.8 - home db cache & optimistic routing

CHANGELOG.md

@@ -15,6 +15,8 @@
      User-code should not need to create arbitrary `ServerError`s.
      In return, `ServerError` now implements `Clone`.
  - Add support for bolt handshake manifest v1.
+ - Add support for Bolt 5.8 (home database resolution cache)
+   - Includes an optimization where the driver uses a home/default database cache to perform optimistic routing under certain circumstances, saving a full round trip. See the [PR description](https://github.com/robsdedude/neo4j-rust-driver/pull/28) for more details.
 
 **🔧 Fixes**
  - Rework `neo4j::value::graph::Path`

neo4j/src/driver.rs

@@ -14,6 +14,7 @@
 
 pub(crate) mod config;
 pub(crate) mod eager_result;
+mod home_db_cache;
 pub(crate) mod io;
 pub(crate) mod record;
 pub mod record_stream;
@@ -40,10 +41,11 @@ pub use config::{
     InvalidRoutingContextError, KeepAliveConfig, TlsConfigError,
 };
 pub use eager_result::{EagerResult, ScalarError};
+use home_db_cache::HomeDbCache;
 use io::bolt::message_parameters::TelemetryAPI;
 #[cfg(feature = "_internal_testkit_backend")]
 pub use io::ConnectionPoolMetrics;
-use io::{AcquireConfig, Pool, PoolConfig, PooledBolt, SessionAuth, UpdateRtArgs};
+use io::{AcquireConfig, Pool, PoolConfig, PooledBolt, SessionAuth, UpdateRtArgs, UpdateRtDb};
 use notification::NotificationFilter;
 pub use record::Record;
 use record_stream::RecordStream;
@@ -76,8 +78,9 @@ pub mod notification {
 ///  * [`Driver::session()`] for several mechanisms offering more advance patterns.
 #[derive(Debug)]
 pub struct Driver {
-    pub(crate) config: ReducedDriverConfig,
-    pub(crate) pool: Pool,
+    config: ReducedDriverConfig,
+    pool: Pool,
+    home_db_cache: Arc<HomeDbCache>,
     capability_check_config: SessionConfig,
     execute_query_bookmark_manager: Arc<dyn BookmarkManager>,
 }
@@ -120,6 +123,7 @@ impl Driver {
                 idle_time_before_connection_test: config.idle_time_before_connection_test,
             },
             pool: Pool::new(Arc::new(connection_config.address), pool_config),
+            home_db_cache: Default::default(),
             capability_check_config: SessionConfig::default()
                 .with_database(Arc::new(String::from("system"))),
             execute_query_bookmark_manager: Arc::new(bookmark_managers::simple(None)),
@@ -169,7 +173,12 @@ impl Driver {
             idle_time_before_connection_test: self.config.idle_time_before_connection_test,
             eager_begin: true,
         };
-        Session::new(config, &self.pool, &self.config)
+        Session::new(
+            config,
+            &self.pool,
+            Arc::clone(&self.home_db_cache),
+            &self.config,
+        )
     }
 
     fn execute_query_session(
@@ -197,7 +206,12 @@ impl Driver {
             idle_time_before_connection_test: self.config.idle_time_before_connection_test,
             eager_begin: false,
         };
-        Session::new(config, &self.pool, &self.config)
+        Session::new(
+            config,
+            &self.pool,
+            Arc::clone(&self.home_db_cache),
+            &self.config,
+        )
     }
 
     /// Execute a single query inside a transaction.
@@ -329,7 +343,13 @@ impl Driver {
             idle_time_before_connection_test: Some(Duration::ZERO),
             eager_begin: true,
         };
-        Session::new(config, &self.pool, &self.config)
+        let mut session = Session::new(
+            config,
+            &self.pool,
+            Arc::clone(&self.home_db_cache),
+            &self.config,
+        );
+        session
             .acquire_connection(RoutingControl::Read)
             .and_then(|mut con| {
                 con.write_all(None)?;
@@ -380,11 +400,21 @@ impl Driver {
         self.pool.acquire(AcquireConfig {
             mode: RoutingControl::Read,
             update_rt_args: UpdateRtArgs {
-                db: self.capability_check_config.database.as_ref(),
+                db: self
+                    .capability_check_config
+                    .database
+                    .as_ref()
+                    .map(|db| UpdateRtDb {
+                        db: Arc::clone(db),
+                        guess: false,
+                    })
+                    .as_ref(),
                 bookmarks: None,
                 imp_user: None,
+                deadline: self.pool.config.connection_acquisition_deadline(),
                 session_auth: SessionAuth::None,
                 idle_time_before_connection_test: None,
+                db_resolution_cb: None,
             },
         })
     }

neo4j/src/driver/config.rs

@@ -795,7 +795,7 @@ impl ConnectionConfig {
                 }
             }
             Some(query) => {
-                if query == "" {
+                if query.is_empty() {
                     Some(HashMap::new())
                 } else {
                     if !routing {

neo4j/src/driver/home_db_cache.rs

@@ -0,0 +1,249 @@
+// Copyright Rouven Bauer
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use itertools::Itertools;
+use log::debug;
+use parking_lot::Mutex;
+use std::collections::HashMap;
+use std::hash::{Hash, Hasher};
+use std::mem;
+use std::sync::Arc;
+use std::time::Instant;
+
+use super::auth::AuthToken;
+use crate::value::spatial;
+use crate::{value, ValueSend};
+
+#[derive(Debug)]
+pub(super) struct HomeDbCache {
+    cache: Mutex<HashMap<HomeDbCacheKey, HomeDbCacheEntry>>,
+    config: HomeDbCacheConfig,
+}
+
+#[derive(Debug, Copy, Clone)]
+struct HomeDbCacheConfig {
+    max_size: usize,
+    prune_size: usize,
+}
+
+impl Default for HomeDbCache {
+    fn default() -> Self {
+        Self::new(1000)
+    }
+}
+
+impl HomeDbCache {
+    pub(super) fn new(max_size: usize) -> Self {
+        let max_size_f64 = max_size as f64;
+        let prune_size = usize::min(max_size, (max_size_f64 * 0.01).log(max_size_f64) as usize);
+        HomeDbCache {
+            cache: Mutex::new(HashMap::with_capacity(max_size)),
+            config: HomeDbCacheConfig {
+                max_size,
+                prune_size,
+            },
+        }
+    }
+
+    pub(super) fn get(&self, key: &HomeDbCacheKey) -> Option<Arc<String>> {
+        let mut lock = self.cache.lock();
+        let cache: &mut HashMap<HomeDbCacheKey, HomeDbCacheEntry> = &mut lock;
+        let res = cache.get_mut(key).map(|entry| {
+            entry.last_used = Instant::now();
+            Arc::clone(&entry.database)
+        });
+        debug!(
+            "Getting home database cache for key: {} -> {:?}",
+            key.log_str(),
+            res.as_deref(),
+        );
+        res
+    }
+
+    pub(super) fn update(&self, key: HomeDbCacheKey, database: Arc<String>) {
+        let mut lock = self.cache.lock();
+        debug!(
+            "Updating home database cache for key: {} -> {:?}",
+            key.log_str(),
+            database.as_str(),
+        );
+        let cache: &mut HashMap<HomeDbCacheKey, HomeDbCacheEntry> = &mut lock;
+        let previous_val = cache.insert(
+            key,
+            HomeDbCacheEntry {
+                database,
+                last_used: Instant::now(),
+            },
+        );
+        if previous_val.is_none() {
+            // cache grew, prune if necessary
+            Self::prune(cache, self.config);
+        }
+    }
+
+    fn prune(cache: &mut HashMap<HomeDbCacheKey, HomeDbCacheEntry>, config: HomeDbCacheConfig) {
+        if cache.len() <= config.max_size {
+            return;
+        }
+        debug!(
+            "Pruning home database cache to size: {}",
+            config.max_size - config.prune_size
+        );
+        let new_cache = mem::take(cache);
+        *cache = new_cache
+            .into_iter()
+            .sorted_by(|(_, v1), (_, v2)| v2.last_used.cmp(&v1.last_used))
+            .take(config.max_size - config.prune_size)
+            .collect();
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(super) enum HomeDbCacheKey {
+    DriverUser,
+    FixedUser(Arc<String>),
+    SessionAuth(SessionAuthKey),
+}
+
+impl HomeDbCacheKey {
+    fn log_str(&self) -> String {
+        match self {
+            HomeDbCacheKey::DriverUser | HomeDbCacheKey::FixedUser(_) => format!("{:?}", self),
+            HomeDbCacheKey::SessionAuth(SessionAuthKey(auth)) => {
+                let mut auth: AuthToken = (**auth).clone();
+                auth.data
+                    .get_mut("credentials")
+                    .map(|c| *c = value!("**********"));
+                format!("SessionAuth({:?})", auth.data)
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(super) struct SessionAuthKey(Arc<AuthToken>);
+
+impl PartialEq for SessionAuthKey {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+            || self.0.data.len() == other.0.data.len()
+                && self
+                    .0
+                    .data
+                    .iter()
+                    .sorted_by(|(k1, _), (k2, _)| k1.cmp(k2))
+                    .zip(other.0.data.iter().sorted_by(|(k1, _), (k2, _)| k1.cmp(k2)))
+                    .all(|((k1, v1), (k2, v2))| k1 == k2 && v1.eq_data(v2))
+    }
+}
+
+impl Eq for SessionAuthKey {}
+
+impl Hash for SessionAuthKey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0
+            .data
+            .iter()
+            .sorted_by(|(k1, _), (k2, _)| k1.cmp(k2))
+            .for_each(|(k, v)| {
+                k.hash(state);
+                Self::hash(v, state);
+            });
+    }
+}
+
+impl SessionAuthKey {
+    fn hash(v: &ValueSend, state: &mut impl Hasher) {
+        match v {
+            ValueSend::Null => state.write_usize(0),
+            ValueSend::Boolean(v) => v.hash(state),
+            ValueSend::Integer(v) => v.hash(state),
+            ValueSend::Float(v) => v.to_bits().hash(state),
+            ValueSend::Bytes(v) => v.hash(state),
+            ValueSend::String(v) => v.hash(state),
+            ValueSend::List(v) => v.iter().for_each(|v| Self::hash(v, state)),
+            ValueSend::Map(v) => {
+                v.iter()
+                    .sorted_by(|(k1, _), (k2, _)| k1.cmp(k2))
+                    .for_each(|(k, v)| {
+                        k.hash(state);
+                        Self::hash(v, state);
+                    });
+            }
+            ValueSend::Cartesian2D(spatial::Cartesian2D { srid, coordinates }) => {
+                srid.hash(state);
+                coordinates
+                    .iter()
+                    .map(|v| v.to_bits())
+                    .for_each(|v| v.hash(state));
+            }
+            ValueSend::Cartesian3D(spatial::Cartesian3D { srid, coordinates }) => {
+                srid.hash(state);
+                coordinates
+                    .iter()
+                    .map(|v| v.to_bits())
+                    .for_each(|v| v.hash(state));
+            }
+            ValueSend::WGS84_2D(spatial::WGS84_2D { srid, coordinates }) => {
+                srid.hash(state);
+                coordinates
+                    .iter()
+                    .map(|v| v.to_bits())
+                    .for_each(|v| v.hash(state));
+            }
+            ValueSend::WGS84_3D(spatial::WGS84_3D { srid, coordinates }) => {
+                srid.hash(state);
+                coordinates
+                    .iter()
+                    .map(|v| v.to_bits())
+                    .for_each(|v| v.hash(state));
+            }
+            ValueSend::Duration(v) => v.hash(state),
+            ValueSend::LocalTime(v) => v.hash(state),
+            ValueSend::Time(v) => v.hash(state),
+            ValueSend::Date(v) => v.hash(state),
+            ValueSend::LocalDateTime(v) => v.hash(state),
+            ValueSend::DateTime(v) => v.hash(state),
+            ValueSend::DateTimeFixed(v) => v.hash(state),
+        }
+    }
+}
+
+impl HomeDbCacheKey {
+    pub(super) fn new(
+        imp_user: Option<&Arc<String>>,
+        session_auth: Option<&Arc<AuthToken>>,
+    ) -> Self {
+        if let Some(user) = imp_user {
+            HomeDbCacheKey::FixedUser(Arc::clone(user))
+        } else if let Some(auth) = session_auth {
+            if let Some(ValueSend::String(scheme)) = auth.data.get("scheme") {
+                if scheme == "basic" {
+                    if let Some(ValueSend::String(user)) = auth.data.get("principal") {
+                        return HomeDbCacheKey::FixedUser(Arc::new(user.clone()));
+                    }
+                }
+            }
+            HomeDbCacheKey::SessionAuth(SessionAuthKey(Arc::clone(auth)))
+        } else {
+            HomeDbCacheKey::DriverUser
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+struct HomeDbCacheEntry {
+    database: Arc<String>,
+    last_used: Instant,
+}

neo4j/src/driver/io.rs

@@ -19,4 +19,6 @@ mod varint;
 
 #[cfg(feature = "_internal_testkit_backend")]
 pub use pool::ConnectionPoolMetrics;
-pub(crate) use pool::{AcquireConfig, Pool, PoolConfig, PooledBolt, SessionAuth, UpdateRtArgs};
+pub(crate) use pool::{
+    AcquireConfig, Pool, PoolConfig, PooledBolt, SessionAuth, UpdateRtArgs, UpdateRtDb,
+};