RDBC-833 Implement session.advanced.get_current_session_node

Author: poissoncorp

ravendb/documents/session/document_session.py

@@ -102,6 +102,7 @@
     from ravendb.documents.operations.lazy.definition import LazyOperation
     from ravendb.http.request_executor import RequestExecutor
     from ravendb.documents.store.definition import DocumentStore
+    from ravendb.http.server_node import ServerNode
 
 _T = TypeVar("_T")
 _TIndex = TypeVar("_TIndex", bound=AbstractCommonApiForIndexes)
@@ -753,6 +754,9 @@ def refresh(self, entity: object) -> object:
         def raw_query(self, query: str, object_type: Optional[Type[_T]] = None) -> RawDocumentQuery[_T]:
             return RawDocumentQuery(object_type, self._session, query)
 
+        def get_current_session_node(self) -> ServerNode:
+            return self.session_info.get_current_session_node(self.request_executor)
+
         @property
         def lazily(self) -> LazySessionOperations:
             return self._session._lazily

ravendb/documents/session/misc.py

@@ -6,16 +6,17 @@
 from enum import Enum
 from typing import Union, Optional, TYPE_CHECKING, List, Dict, Generic, TypeVar
 
-from ravendb.http.misc import LoadBalanceBehavior
+from ravendb.http.misc import LoadBalanceBehavior, ReadBalanceBehavior
 
 if TYPE_CHECKING:
     from ravendb.http.request_executor import RequestExecutor
-    from ravendb.documents.session.query import Query
+    from ravendb.documents.queries.misc import Query
     from ravendb.documents.session.operations.query import QueryOperation
     from ravendb.documents.session.document_session_operations.in_memory_document_session_operations import (
         InMemoryDocumentSessionOperations,
     )
-    from ravendb.documents import DocumentStore
+    from ravendb.documents.store.definition import DocumentStore
+    from ravendb.http.server_node import ServerNode
 
 
 _T_Key = TypeVar("_T_Key")
@@ -53,7 +54,7 @@ def __init__(
         self.__session_id: Union[None, int] = None
         self.__session_id_used: Union[None, bool] = None
         self.__load_balancer_context_seed = session._request_executor.conventions.load_balancer_context_seed
-        self.__can_use_load_balance_behavior = (
+        self._can_use_load_balance_behavior = (
             session.conventions.load_balance_behavior == LoadBalanceBehavior.USE_SESSION_CONTEXT
             and session.conventions.load_balancer_per_session_context_selector is not None
         )
@@ -64,7 +65,7 @@ def __init__(
 
     @property
     def can_use_load_balance_behavior(self) -> bool:
-        return self.__can_use_load_balance_behavior
+        return self._can_use_load_balance_behavior
 
     @property
     def session_id(self) -> int:
@@ -95,6 +96,25 @@ def __set_context_internal(self, session_key: str) -> None:
     def increment_request_count(self) -> None:
         self.__session.increment_requests_count()
 
+    def get_current_session_node(self, request_executor: RequestExecutor) -> ServerNode:
+        if request_executor.conventions.load_balance_behavior == LoadBalanceBehavior.USE_SESSION_CONTEXT:
+            if self._can_use_load_balance_behavior:
+                result = request_executor.get_node_by_session_id(self.session_id)
+                return result.current_node
+
+        read_balance_behavior = request_executor.conventions.read_balance_behavior
+
+        if read_balance_behavior == ReadBalanceBehavior.NONE:
+            result = request_executor.preferred_node
+        elif read_balance_behavior == ReadBalanceBehavior.ROUND_ROBIN:
+            result = request_executor.get_node_by_session_id(self.session_id)
+        elif read_balance_behavior == ReadBalanceBehavior.FASTEST_NODE:
+            result = request_executor.get_fastest_node()
+        else:
+            raise ValueError(f"Unsupported read balance behavior '{str(read_balance_behavior)}'")
+
+        return result.current_node
+
 
 class SessionOptions:
     def __init__(

ravendb/http/request_executor.py

@@ -197,11 +197,11 @@ def default_timeout(self, value: datetime.timedelta) -> None:
 
     @property
     def preferred_node(self) -> CurrentIndexAndNode:
-        self.__ensure_node_selector()
+        self._ensure_node_selector()
         return self._node_selector.get_preferred_node()
 
     def get_requested_node(self, node_tag: str, throw_if_contains_failures: bool = False) -> CurrentIndexAndNode:
-        self.__ensure_node_selector()
+        self._ensure_node_selector()
         current_index_and_node = self._node_selector.get_requested_node(node_tag)
 
         if throw_if_contains_failures and not self._node_selector.node_is_available(
@@ -213,6 +213,16 @@ def get_requested_node(self, node_tag: str, throw_if_contains_failures: bool = F
 
         return current_index_and_node
 
+    def get_fastest_node(self) -> CurrentIndexAndNode:
+        self._ensure_node_selector()
+
+        return self._node_selector.get_fastest_node()
+
+    def get_node_by_session_id(self, session_id: int = None) -> CurrentIndexAndNode:
+        self._ensure_node_selector()
+
+        return self._node_selector.get_node_by_session_id(session_id)
+
     def __on_failed_request_invoke(self, url: str, e: Exception):
         for event in self.__on_failed_request:
             event(FailedRequestEventArgs(self._database_name, url, e, None, None))
@@ -1229,7 +1239,7 @@ def __init__(self, index: int, response: requests.Response):
     def client_configuration_etag(self, value):
         self._client_configuration_etag = value
 
-    def __ensure_node_selector(self) -> None:
+    def _ensure_node_selector(self) -> None:
         if not self._disable_topology_updates:
             self.__wait_for_topology_update(self._first_topology_update_task)
 

ravendb/http/topology.py

@@ -96,19 +96,19 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self.__update_fastest_node_timer.__exit__()
 
     def __init__(self, topology: Topology, thread_pool: ThreadPoolExecutor):
-        self.__state = self.__NodeSelectorState(topology)
+        self._state = self.__NodeSelectorState(topology)
         self.__update_fastest_node_timer: Union[None, time] = None
         self.__thread_pool_executor = thread_pool
 
     @property
     def topology(self) -> Topology:
-        return self.__state.topology
+        return self._state.topology
 
     def node_is_available(self, index: int) -> bool:
-        return self.__state.failures[index] == 0
+        return self._state.failures[index] == 0
 
     def on_failed_request(self, node_index: int) -> None:
-        state = self.__state
+        state = self._state
         if node_index < 0 or node_index >= len(state.failures):
             return
         state.failures[node_index] += 1
@@ -117,19 +117,19 @@ def on_update_topology(self, topology: Topology, force_update: bool = False) ->
         if topology is None:
             return False
 
-        state_etag = self.__state.topology.etag if self.__state.topology.etag else 0
+        state_etag = self._state.topology.etag if self._state.topology.etag else 0
         topology_etag = topology.etag if topology.etag else 0
 
         if state_etag >= topology_etag and not force_update:
             return False
 
         state = NodeSelector.__NodeSelectorState(topology)
-        self.__state = state
+        self._state = state
 
         return True
 
     def get_requested_node(self, node_tag: str) -> CurrentIndexAndNode:
-        state = self.__state
+        state = self._state
         server_nodes = state.nodes
         for i in range(len(server_nodes)):
             if server_nodes[i].cluster_tag == node_tag:
@@ -140,7 +140,7 @@ def get_requested_node(self, node_tag: str) -> CurrentIndexAndNode:
         raise RequestedNodeUnavailableException(f"Could not find requested node {node_tag}")
 
     def get_preferred_node(self) -> CurrentIndexAndNode:
-        state = self.__state
+        state = self._state
         return self.get_preferred_node_internal(state)
 
     @classmethod
@@ -154,7 +154,7 @@ def get_preferred_node_internal(cls, state: NodeSelector.__NodeSelectorState) ->
         return cls.unlikely_everyone_faulted_choice(state)
 
     def get_preferred_node_with_topology(self) -> CurrentIndexAndNodeAndEtag:
-        state = self.__state
+        state = self._state
         preferred_node = self.get_preferred_node_internal(state)
         etag = (state.topology.etag if state.topology.etag else -2) if state.topology else -2
         return CurrentIndexAndNodeAndEtag(preferred_node.current_index, preferred_node.current_node, etag)
@@ -168,7 +168,7 @@ def unlikely_everyone_faulted_choice(state: NodeSelector.__NodeSelectorState) ->
         return state.node_when_everyone_marked_as_faulted
 
     def get_node_by_session_id(self, session_id: int) -> CurrentIndexAndNode:
-        state = self.__state
+        state = self._state
         if len(state.topology.nodes) == 0:
             raise AllTopologyNodesDownException("There are no nodes in the topology at all")
         index = abs(session_id % len(state.topology.nodes))
@@ -183,19 +183,19 @@ def get_node_by_session_id(self, session_id: int) -> CurrentIndexAndNode:
         return self.get_preferred_node()
 
     def get_fastest_node(self) -> CurrentIndexAndNode:
-        state = self.__state
+        state = self._state
         if state.failures[state.fastest] == 0 and state.nodes[state.fastest].server_role == ServerNode.Role.MEMBER:
             return CurrentIndexAndNode(state.fastest, state.nodes[state.fastest])
 
         # if the fastest node has failures, we'll immediately schedule
         # another run of finding who the fastest node is, in the mantime
         # we'll just use the server preferred node or failover as usual
 
-        self.__switch_to_speed_test_phase()
+        self._switch_to_speed_test_phase()
         return self.get_preferred_node()
 
     def restore_node_index(self, node_index: int) -> None:
-        state = self.__state
+        state = self._state
         if len(state.failures) <= node_index:
             return
 
@@ -205,24 +205,23 @@ def restore_node_index(self, node_index: int) -> None:
     def _throw_empty_topology() -> None:
         raise RuntimeError("Empty database topology, this shouldn't happen.")
 
-    def __switch_to_speed_test_phase(self) -> None:
-        state = self.__state
+    def _switch_to_speed_test_phase(self) -> None:
+        state = self._state
 
         if not state.speed_test_mode == 0:
             state.speed_test_mode = 1
             return
 
-        for i in state.fastest_records:
-            i = 0
+        state.fastest_records = len(state.fastest_records) * [0]
 
         state.speed_test_mode += 1
 
     @property
     def in_speed_test_phase(self) -> bool:
-        return self.__state.speed_test_mode > 1
+        return self._state.speed_test_mode > 1
 
     def record_fastest(self, index: int, node: ServerNode) -> None:
-        state = self.__state
+        state = self._state
         state_fastest = state.fastest_records
 
         # the following two checks are to verify that things didn't move
@@ -273,11 +272,11 @@ def __select_fastest(self, state: NodeSelector.__NodeSelectorState, index: int)
 
         else:
             self.__update_fastest_node_timer = threading.Timer(
-                datetime.timedelta(minutes=1), self.__switch_to_speed_test_phase
+                datetime.timedelta(minutes=1), self._switch_to_speed_test_phase
             )
 
     def schedule_speed_test(self) -> None:
-        self.__switch_to_speed_test_phase()
+        self._switch_to_speed_test_phase()
 
 
 class CurrentIndexAndNode: