Bind Cluster sessions to single server nodes, and add naive load balancing

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)