ADR 031: Homedb Cache (Bolt 5.8) (#1115)

Home database cache/optimistic routing (Bolt 5.8)

Introducing support for Bolt 5.8, which (among other things) makes the server
echo back the resolved home/default db on starting a transaction as well as a
connection hint whether server-side routing (SSR) is enabled. If so, the
echoed db can be used together with an optimistic home db cache to guess a
client-side route (optimistic routing - falling back to SSR on wrong guesses).
This feature saves a round-trip under the following conditions
(*all* must be fulfilled):
 - Server has SSR enabled
 - Sever supports Bolt 5.8
 - A fresh (yet unused) session without an explicitly configured target database starts its first transaction (auto-commit or explicit)

Co-authored-by: MaxAake <[email protected]>

Author: robsdedude

.gitattributes

@@ -1,5 +1,5 @@
 # configure github not to display generated files
 /src/neo4j/_sync/** linguist-generated=true
-/tests/unit/sync_/** linguist-generated=true
-/tests/integration/sync_/** linguist-generated=true
+/tests/unit/sync/** linguist-generated=true
+/tests/integration/sync/** linguist-generated=true
 /testkitbackend/_sync/** linguist-generated=true

docs/source/api.rst

@@ -260,7 +260,8 @@ Closing a driver will immediately shut down all connections in the pool.
         :param database\_:
             Database to execute the query against.
 
-            None (default) uses the database configured on the server side.
+            :data:`None` (default) uses the database configured on the server
+            side.
 
             .. Note::
                 It is recommended to always specify the database explicitly
@@ -1034,7 +1035,7 @@ Specifically, the following applies:
   all queries within that session are executed with the explicit database
   name 'movies' supplied. Any change to the user’s home database is
   reflected only in sessions created after such change takes effect. This
-  behavior requires additional network communication. In clustered
+  behavior may require additional network communication. In clustered
   environments, it is strongly recommended to avoid a single point of
   failure. For instance, by ensuring that the connection URI resolves to
   multiple endpoints. For older Bolt protocol versions the behavior is the

docs/source/async_api.rst

@@ -247,7 +247,8 @@ Closing a driver will immediately shut down all connections in the pool.
         :param database\_:
             Database to execute the query against.
 
-            None (default) uses the database configured on the server side.
+            :data:`None` (default) uses the database configured on the server
+            side.
 
             .. Note::
                 It is recommended to always specify the database explicitly

src/neo4j/_async/driver.py

@@ -799,7 +799,8 @@ async def example(driver: neo4j.AsyncDriver) -> int:
         :param database_:
             Database to execute the query against.
 
-            None (default) uses the database configured on the server side.
+            :data:`None` (default) uses the database configured on the server
+            side.
 
             .. Note::
                 It is recommended to always specify the database explicitly

src/neo4j/_async/home_db_cache.py

@@ -0,0 +1,150 @@
+# Copyright (c) "Neo4j"
+# Neo4j Sweden AB [https://neo4j.com]
+#
+# This file is part of Neo4j.
+#
+# 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.
+
+
+from __future__ import annotations
+
+import math
+import typing as t
+from time import monotonic
+
+from .._async_compat.concurrency import AsyncCooperativeLock
+
+
+if t.TYPE_CHECKING:
+    import typing_extensions as te
+
+    TKey: te.TypeAlias = t.Union[
+        str,
+        t.Tuple[t.Tuple[str, t.Hashable], ...],
+        t.Tuple[None],
+    ]
+    TVal: te.TypeAlias = t.Tuple[float, str]
+
+
+class AsyncHomeDbCache:
+    _ttl: float
+    _enabled: bool
+    _max_size: int | None
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        ttl: float = float("inf"),
+        max_size: int | None = None,
+    ) -> None:
+        if math.isnan(ttl) or ttl <= 0:
+            raise ValueError(f"home db cache ttl must be greater 0, got {ttl}")
+        self._enabled = enabled
+        self._ttl = ttl
+        self._cache: dict[TKey, TVal] = {}
+        self._lock = AsyncCooperativeLock()
+        self._oldest_entry = monotonic()
+        if max_size is not None and max_size <= 0:
+            raise ValueError(
+                f"home db cache max_size must be greater 0 or None, "
+                f"got {max_size}"
+            )
+        self._max_size = max_size
+        self._truncate_size = (
+            min(max_size, int(0.01 * max_size * math.log(max_size)))
+            if max_size is not None
+            else None
+        )
+
+    def compute_key(
+        self,
+        imp_user: str | None,
+        auth: dict | None,
+    ) -> TKey:
+        if not self._enabled:
+            return (None,)
+        if imp_user is not None:
+            return imp_user
+        if auth is not None:
+            return _consolidate_auth_token(auth)
+        return (None,)
+
+    def get(self, key: TKey) -> str | None:
+        if not self._enabled:
+            return None
+        with self._lock:
+            self._clean(monotonic())
+            val = self._cache.get(key)
+            if val is None:
+                return None
+            return val[1]
+
+    def set(self, key: TKey, value: str | None) -> None:
+        if not self._enabled:
+            return
+        with self._lock:
+            now = monotonic()
+            self._clean(now)
+            if value is None:
+                self._cache.pop(key, None)
+            else:
+                self._cache[key] = (now, value)
+
+    def clear(self) -> None:
+        if not self._enabled:
+            return
+        with self._lock:
+            self._cache = {}
+            self._oldest_entry = monotonic()
+
+    def _clean(self, now: float | None = None) -> None:
+        now = monotonic() if now is None else now
+        if now - self._oldest_entry > self._ttl:
+            self._cache = {
+                k: v
+                for k, v in self._cache.items()
+                if now - v[0] < self._ttl * 0.9
+            }
+            self._oldest_entry = min(
+                (v[0] for v in self._cache.values()), default=now
+            )
+        if self._max_size and len(self._cache) > self._max_size:
+            self._cache = dict(
+                sorted(
+                    self._cache.items(),
+                    key=lambda item: item[1][0],
+                    reverse=True,
+                )[: self._truncate_size]
+            )
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+    @property
+    def enabled(self) -> bool:
+        return self._enabled
+
+
+def _consolidate_auth_token(auth: dict) -> tuple | str:
+    if auth.get("scheme") == "basic" and isinstance(
+        auth.get("principal"), str
+    ):
+        return auth["principal"]
+    return _hashable_dict(auth)
+
+
+def _hashable_dict(d: dict) -> tuple:
+    return tuple(
+        (k, _hashable_dict(v) if isinstance(v, dict) else v)
+        for k, v in sorted(d.items())
+    )

src/neo4j/_async/io/__init__.py

@@ -22,7 +22,8 @@
 """
 
 __all__ = [
-    "AcquireAuth",
+    "AcquisitionAuth",
+    "AcquisitionDatabase",
     "AsyncBolt",
     "AsyncBoltPool",
     "AsyncNeo4jPool",
@@ -37,7 +38,8 @@
     ConnectionErrorHandler,
 )
 from ._pool import (
-    AcquireAuth,
+    AcquisitionAuth,
+    AcquisitionDatabase,
     AsyncBoltPool,
     AsyncNeo4jPool,
 )

src/neo4j/_async/io/_bolt.py

@@ -25,6 +25,7 @@
 
 from ..._async_compat.network import AsyncBoltSocket
 from ..._async_compat.util import AsyncUtil
+from ..._auth_management import to_auth_dict
 from ..._codec.hydration import (
     HydrationHandlerABC,
     v1 as hydration_v1,
@@ -39,12 +40,10 @@
 from ..._sync.config import PoolConfig
 from ...addressing import ResolvedAddress
 from ...api import (
-    Auth,
     ServerInfo,
     Version,
 )
 from ...exceptions import (
-    AuthError,
     ConfigurationError,
     DriverError,
     IncompleteCommit,
@@ -158,10 +157,7 @@ def __init__(
             ),
             self.PROTOCOL_VERSION,
         )
-        # so far `connection.recv_timeout_seconds` is the only available
-        # configuration hint that exists. Therefore, all hints can be stored at
-        # connection level. This might change in the future.
-        self.configuration_hints = {}
+        self.connection_hints = {}
         self.patch = {}
         self.outbox = AsyncOutbox(
             self.socket,
@@ -187,7 +183,7 @@ def __init__(
             self.user_agent = USER_AGENT
 
         self.auth = auth
-        self.auth_dict = self._to_auth_dict(auth)
+        self.auth_dict = to_auth_dict(auth)
         self.auth_manager = auth_manager
         self.telemetry_disabled = telemetry_disabled
 
@@ -206,26 +202,14 @@ def _get_server_state_manager(self) -> ServerStateManagerBase: ...
     @abc.abstractmethod
     def _get_client_state_manager(self) -> ClientStateManagerBase: ...
 
-    @classmethod
-    def _to_auth_dict(cls, auth):
-        # Determine auth details
-        if not auth:
-            return {}
-        elif isinstance(auth, tuple) and 2 <= len(auth) <= 3:
-            return vars(Auth("basic", *auth))
-        else:
-            try:
-                return vars(auth)
-            except (KeyError, TypeError) as e:
-                # TODO: 6.0 - change this to be a DriverError (or subclass)
-                raise AuthError(
-                    f"Cannot determine auth details from {auth!r}"
-                ) from e
-
     @property
     def connection_id(self):
         return self.server_info._metadata.get("connection_id", "<unknown id>")
 
+    @property
+    @abc.abstractmethod
+    def ssr_enabled(self) -> bool: ...
+
     @property
     @abc.abstractmethod
     def supports_multiple_results(self):
@@ -308,6 +292,7 @@ def protocol_handlers(cls, protocol_version=None):
             AsyncBolt5x5,
             AsyncBolt5x6,
             AsyncBolt5x7,
+            AsyncBolt5x8,
         )
 
         handlers = {
@@ -325,6 +310,7 @@ def protocol_handlers(cls, protocol_version=None):
             AsyncBolt5x5.PROTOCOL_VERSION: AsyncBolt5x5,
             AsyncBolt5x6.PROTOCOL_VERSION: AsyncBolt5x6,
             AsyncBolt5x7.PROTOCOL_VERSION: AsyncBolt5x7,
+            AsyncBolt5x8.PROTOCOL_VERSION: AsyncBolt5x8,
         }
 
         if protocol_version is None:
@@ -461,7 +447,10 @@ async def open(
 
         # avoid new lines after imports for better readability and conciseness
         # fmt: off
-        if protocol_version == (5, 7):
+        if protocol_version == (5, 8):
+            from ._bolt5 import AsyncBolt5x8
+            bolt_cls = AsyncBolt5x8
+        elif protocol_version == (5, 7):
             from ._bolt5 import AsyncBolt5x7
             bolt_cls = AsyncBolt5x7
         elif protocol_version == (5, 6):
@@ -626,7 +615,7 @@ def re_auth(
 
         :returns: whether the auth was changed
         """
-        new_auth_dict = self._to_auth_dict(auth)
+        new_auth_dict = to_auth_dict(auth)
         if not force and new_auth_dict == self.auth_dict:
             self.auth_manager = auth_manager
             self.auth = auth

src/neo4j/_async/io/_bolt3.py

@@ -148,6 +148,8 @@ class AsyncBolt3(AsyncBolt):
 
     PROTOCOL_VERSION = Version(3, 0)
 
+    ssr_enabled = False
+
     supports_multiple_results = False
 
     supports_multiple_databases = False

src/neo4j/_async/io/_bolt4.py

@@ -64,6 +64,8 @@ class AsyncBolt4x0(AsyncBolt):
 
     PROTOCOL_VERSION = Version(4, 0)
 
+    ssr_enabled = False
+
     supports_multiple_results = True
 
     supports_multiple_databases = True
@@ -614,10 +616,10 @@ async def hello(self, dehydration_hooks=None, hydration_hooks=None):
         )
 
         def on_success(metadata):
-            self.configuration_hints.update(metadata.pop("hints", {}))
+            self.connection_hints.update(metadata.pop("hints", {}))
             self.server_info.update(metadata)
-            if "connection.recv_timeout_seconds" in self.configuration_hints:
-                recv_timeout = self.configuration_hints[
+            if "connection.recv_timeout_seconds" in self.connection_hints:
+                recv_timeout = self.connection_hints[
                     "connection.recv_timeout_seconds"
                 ]
                 if isinstance(recv_timeout, int) and recv_timeout > 0: