Bind Cluster sessions to single server nodes, and add naive load balancing (#183)

## What is the goal of this PR?

Previously, each Cluster session would lazily create Core sessions on-demand as Transactions were opened. This delegated too much responsibility to the Transaction. Now, each Cluster session is bound to a single server node, which may be a primary or a secondary node depending on the Options passed to the Session.

We also select secondary nodes in a fairer, more balanced way; the server indicates which node is its first preference for secondary sessions, and the client tries to obey the server's decision, with fallback to backup nodes only in the event of a failure.

## What are the changes implemented in this PR?

- On creating a Cluster session, a Core session is immediately created to a server node of the appropriate type
- Node selection is based on the `read_any_replica` Option passed to the Session
- When `read_any_replica` is `true`, the client will first select nodes that the server has indicated are "preferred secondary", with fallback to backup nodes only in the event of a failure

Author: Alex Walker

dependencies/graknlabs/artifacts.bzl

@@ -27,7 +27,7 @@ def graknlabs_grakn_core_artifacts():
         artifact_name = "grakn-core-server-{platform}-{version}.{ext}",
         tag_source = deployment["artifact.release"],
         commit_source = deployment["artifact.snapshot"],
-        commit = "3227b9e07c9c2317c0e7eab29259204f92433d76",
+        commit = "4d449aa198fd5cceca54cb3889114ab5aa1b8e5e",
     )
 
 def graknlabs_grakn_cluster_artifacts():
@@ -37,5 +37,5 @@ def graknlabs_grakn_cluster_artifacts():
         artifact_name = "grakn-cluster-server-{platform}-{version}.{ext}",
         tag_source = deployment_private["artifact.release"],
         commit_source = deployment_private["artifact.snapshot"],
-        commit = "93cbc149d7b9ce588e52d8132c8a0b2265658a3c",
+        commit = "55bb7a5bb99fecf6990b0e8ed3220b264f8de452",
     )

grakn/client.py

@@ -27,10 +27,12 @@
 # Repackaging these symbols allows them to be imported from "grakn.client"
 from grakn.common.exception import GraknClientException  # noqa # pylint: disable=unused-import
 from grakn.concept.type.value_type import ValueType  # noqa # pylint: disable=unused-import
-from grakn.options import GraknOptions
+from grakn.options import GraknOptions, GraknClusterOptions
+from grakn.rpc.cluster.failsafe_task import FailsafeTask
+from grakn.rpc.cluster.replica_info import ReplicaInfo
 from grakn.rpc.cluster.server_address import ServerAddress
 from grakn.rpc.cluster.database_manager import _DatabaseManagerClusterRPC
-from grakn.rpc.cluster.session import _SessionClusterRPC
+from grakn.rpc.cluster.session import SessionClusterRPC
 from grakn.rpc.database_manager import DatabaseManager, _DatabaseManagerRPC
 from grakn.rpc.session import Session, SessionType, _SessionRPC
 from grakn.rpc.transaction import TransactionType  # noqa # pylint: disable=unused-import
@@ -109,22 +111,29 @@ def channel(self):
         return self._channel
 
 
-# _RPCGraknClientCluster must live in this package because of circular ref with GraknClient
+# _ClientClusterRPC must live in this package because of circular ref with GraknClient
 class _ClientClusterRPC(GraknClient):
 
     def __init__(self, addresses: List[str]):
-        self._core_clients: Dict[ServerAddress, _ClientRPC] = {addr: _ClientRPC(addr.client()) for addr in self._discover_cluster(addresses)}
+        self._core_clients: Dict[ServerAddress, _ClientRPC] = {addr: _ClientRPC(addr.client()) for addr in self._fetch_cluster_servers(addresses)}
         self._grakn_cluster_grpc_stubs = {addr: GraknClusterStub(client.channel()) for (addr, client) in self._core_clients.items()}
-        self._databases = _DatabaseManagerClusterRPC({addr: client.databases() for (addr, client) in self._core_clients.items()})
+        self._database_managers = _DatabaseManagerClusterRPC({addr: client.databases() for (addr, client) in self._core_clients.items()})
+        self._replica_info_map: Dict[str, ReplicaInfo] = {}
         self._is_open = True
 
     def session(self, database: str, session_type: SessionType, options=None) -> Session:
         if not options:
             options = GraknOptions.cluster()
-        return _SessionClusterRPC(self, database, session_type, options)
+        return self._session_any_replica(database, session_type, options) if options.read_any_replica else self._session_primary_replica(database, session_type, options)
+
+    def _session_primary_replica(self, database: str, session_type: SessionType, options=None) -> SessionClusterRPC:
+        return _OpenSessionFailsafeTask(database, session_type, options, self).run_primary_replica()
+
+    def _session_any_replica(self, database: str, session_type: SessionType, options=None) -> SessionClusterRPC:
+        return _OpenSessionFailsafeTask(database, session_type, options, self).run_any_replica()
 
     def databases(self) -> DatabaseManager:
-        return self._databases
+        return self._database_managers
 
     def is_open(self) -> bool:
         return self._is_open
@@ -140,6 +149,9 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
+    def replica_info_map(self) -> Dict[str, ReplicaInfo]:
+        return self._replica_info_map
+
     def cluster_members(self) -> Set[ServerAddress]:
         return set(self._core_clients.keys())
 
@@ -149,16 +161,27 @@ def core_client(self, address: ServerAddress) -> _ClientRPC:
     def grakn_cluster_grpc_stub(self, address: ServerAddress) -> GraknClusterStub:
         return self._grakn_cluster_grpc_stubs.get(address)
 
-    def _discover_cluster(self, addresses: List[str]) -> Set[ServerAddress]:
+    def _fetch_cluster_servers(self, addresses: List[str]) -> Set[ServerAddress]:
         for address in addresses:
             try:
                 with _ClientRPC(address) as client:
                     print("Performing cluster discovery to %s..." % address)
                     grakn_cluster_stub = GraknClusterStub(client.channel())
-                    res = grakn_cluster_stub.cluster_discover(cluster_proto.Cluster.Discover.Req())
+                    res = grakn_cluster_stub.cluster_servers(cluster_proto.Cluster.Servers.Req())
                     members = set([ServerAddress.parse(srv) for srv in res.servers])
                     print("Discovered %s" % [str(member) for member in members])
                     return members
-            except RpcError:
-                print("Cluster discovery to %s failed." % address)
+            except RpcError as e:
+                print("Cluster discovery to %s failed. %s" % (address, str(e)))
         raise GraknClientException("Unable to connect to Grakn Cluster. Attempted connecting to the cluster members, but none are available: %s" % str(addresses))
+
+
+class _OpenSessionFailsafeTask(FailsafeTask):
+
+    def __init__(self, database: str, session_type: SessionType, options: GraknClusterOptions, client: "_ClientClusterRPC"):
+        super().__init__(client, database)
+        self.session_type = session_type
+        self.options = options
+
+    def run(self, replica: ReplicaInfo.Replica):
+        return SessionClusterRPC(self.client, replica.address(), self.database, self.session_type, self.options)

grakn/rpc/cluster/failsafe_task.py

@@ -0,0 +1,123 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+#
+import time
+from abc import ABC, abstractmethod
+
+import grakn_protocol.protobuf.cluster.database_pb2 as database_proto
+from grpc import RpcError, StatusCode
+
+from grakn.common.exception import GraknClientException
+from grakn.rpc.cluster.replica_info import ReplicaInfo
+
+
+class FailsafeTask(ABC):
+
+    PRIMARY_REPLICA_TASK_MAX_RETRIES = 10
+    FETCH_REPLICAS_MAX_RETRIES = 10
+    WAIT_FOR_PRIMARY_REPLICA_SELECTION_SECONDS: float = 2
+
+    def __init__(self, client, database: str):
+        self.client = client
+        self.database = database
+
+    @abstractmethod
+    def run(self, replica: ReplicaInfo.Replica):
+        pass
+
+    def rerun(self, replica: ReplicaInfo.Replica):
+        return self.run(replica)
+
+    def run_primary_replica(self):
+        if self.database not in self.client.replica_info_map() or not self.client.replica_info_map()[self.database].primary_replica():
+            self._seek_primary_replica()
+        replica = self.client.replica_info_map()[self.database].primary_replica()
+        retries = 0
+        while True:
+            try:
+                return self.run(replica) if retries == 0 else self.rerun(replica)
+            except GraknClientException as e:
+                # TODO: propagate exception from the server in a less brittle way
+                if "[RPL01]" in str(e):  # The contacted replica reported that it was not the primary replica
+                    print("Unable to open a session or transaction, retrying in 2s... %s" % str(e))
+                    time.sleep(self.WAIT_FOR_PRIMARY_REPLICA_SELECTION_SECONDS)
+                    replica = self._seek_primary_replica()
+                else:
+                    raise e
+                # TODO: introduce a special type that extends RpcError and Call
+            except RpcError as e:
+                # TODO: this logic should be extracted into GraknClientException
+                # TODO: error message should be checked in a less brittle way
+                if e.code() == StatusCode.UNAVAILABLE or "[INT07]" in str(e) or "Received RST_STREAM" in str(e):
+                    print("Unable to open a session or transaction, retrying in 2s... %s" % str(e))
+                    time.sleep(self.WAIT_FOR_PRIMARY_REPLICA_SELECTION_SECONDS)
+                    replica = self._seek_primary_replica()
+                else:
+                    raise e
+            retries += 1
+            if retries > self.PRIMARY_REPLICA_TASK_MAX_RETRIES:
+                raise self._cluster_not_available_exception()
+
+    def run_any_replica(self):
+        if self.database in self.client.replica_info_map():
+            replica_info = self.client.replica_info_map()[self.database]
+        else:
+            replica_info = self._fetch_database_replicas()
+
+        replicas = [replica_info.preferred_secondary_replica()] + [replica for replica in replica_info.replicas() if not replica.is_preferred_secondary()]
+        retries = 0
+        for replica in replicas:
+            try:
+                return self.run(replica) if retries == 0 else self.rerun(replica)
+            except RpcError as e:
+                if e.code() == StatusCode.UNAVAILABLE or "[INT07]" in str(e) or "Received RST_STREAM" in str(e):
+                    print("Unable to open a session or transaction to %s. Attempting next replica. %s" % (str(replica.replica_id()), str(e)))
+                else:
+                    raise e
+            retries += 1
+        raise self._cluster_not_available_exception()
+
+    def _seek_primary_replica(self) -> ReplicaInfo.Replica:
+        retries = 0
+        while retries < self.FETCH_REPLICAS_MAX_RETRIES:
+            replica_info = self._fetch_database_replicas()
+            if replica_info.primary_replica():
+                return replica_info.primary_replica()
+            else:
+                time.sleep(self.WAIT_FOR_PRIMARY_REPLICA_SELECTION_SECONDS)
+                retries += 1
+        raise self._cluster_not_available_exception()
+
+    def _fetch_database_replicas(self) -> ReplicaInfo:
+        for server_address in self.client.cluster_members():
+            try:
+                print("Fetching replica info from %s" % server_address)
+                db_replicas_req = database_proto.Database.Replicas.Req()
+                db_replicas_req.database = self.database
+                res = self.client.grakn_cluster_grpc_stub(server_address).database_replicas(db_replicas_req)
+                replica_info = ReplicaInfo.of_proto(res)
+                print("Requested database discovery from peer %s, and got response: %s" % (str(server_address), str([str(replica) for replica in replica_info.replicas()])))
+                self.client.replica_info_map()[self.database] = replica_info
+                return replica_info
+            except RpcError as e:
+                print("Unable to perform database discovery to %s. Attempting next address. %s" % (str(server_address), str(e)))
+        raise self._cluster_not_available_exception()
+
+    def _cluster_not_available_exception(self) -> GraknClientException:
+        addresses = str([str(addr) for addr in self.client.cluster_members()])
+        return GraknClientException("Unable to connect to Grakn Cluster. Attempted connecting to the cluster members, but none are available: '%s'" % addresses)

grakn/rpc/cluster/replica_info.py

@@ -30,7 +30,7 @@ def __init__(self, replicas: Dict["ReplicaInfo.Replica.Id", "ReplicaInfo.Replica
         self._replicas = replicas
 
     @staticmethod
-    def of_proto(res: database_proto.Database.Discover.Res) -> "ReplicaInfo":
+    def of_proto(res: database_proto.Database.Replicas.Res) -> "ReplicaInfo":
         replica_map: Dict["ReplicaInfo.Replica.Id", "ReplicaInfo.Replica"] = {}
         for replica_proto in res.replicas:
             replica_id = ReplicaInfo.Replica.Id(ServerAddress.parse(replica_proto.address), replica_proto.database)
@@ -41,22 +41,30 @@ def primary_replica(self) -> Optional["ReplicaInfo.Replica"]:
         primaries = [replica for replica in self._replicas.values() if replica.is_primary()]
         return max(primaries, key=lambda r: r.term) if primaries else None
 
+    def preferred_secondary_replica(self) -> "ReplicaInfo.Replica":
+        return next(iter([replica for replica in self._replicas.values() if replica.is_preferred_secondary()]), next(iter(self._replicas.values())))
+
     def replicas(self):
         return self._replicas.values()
 
+    def __str__(self):
+        return str([str(replica) for replica in self._replicas.values()])
+
     class Replica:
 
-        def __init__(self, replica_id: "ReplicaInfo.Replica.Id", term: int, is_primary: bool):
+        def __init__(self, replica_id: "ReplicaInfo.Replica.Id", term: int, is_primary: bool, is_preferred_secondary: bool):
             self._replica_id = replica_id
             self._term = term
             self._is_primary = is_primary
+            self._is_preferred_secondary = is_preferred_secondary
 
         @staticmethod
-        def of_proto(replica_proto: database_proto.Database.Discover.Res.Replica) -> "ReplicaInfo.Replica":
+        def of_proto(replica_proto: database_proto.Database.Replica) -> "ReplicaInfo.Replica":
             return ReplicaInfo.Replica(
                 replica_id=ReplicaInfo.Replica.Id(ServerAddress.parse(replica_proto.address), replica_proto.database),
                 term=replica_proto.term,
-                is_primary=replica_proto.is_primary
+                is_primary=replica_proto.primary,
+                is_preferred_secondary=replica_proto.preferred_secondary
             )
 
         def replica_id(self) -> "ReplicaInfo.Replica.Id":
@@ -68,15 +76,21 @@ def term(self) -> int:
         def is_primary(self) -> bool:
             return self._is_primary
 
+        def is_preferred_secondary(self) -> bool:
+            return self._is_preferred_secondary
+
+        def address(self) -> ServerAddress:
+            return self._replica_id.address()
+
         def __eq__(self, other):
             if self is other:
                 return True
             if not other or type(self) != type(other):
                 return False
-            return self._term == other.term() and self._is_primary == other.is_primary()
+            return self._term == other.term() and self._is_primary == other.is_primary() and self._is_preferred_secondary == other.is_preferred_secondary()
 
         def __hash__(self):
-            return hash((self._is_primary, self._term))
+            return hash((self._is_primary, self._is_preferred_secondary, self._term))
 
         def __str__(self):
             return "%s:%s:%d" % (str(self._replica_id), "P" if self._is_primary else "S", self._term)