When closing a client, also close all its sessions (#191)

## What is the goal of this PR?

The `close` method of `GraknClient` now automatically closes all `Session` objects that were opened from that client.

## What are the changes implemented in this PR?

The `close` method of `GraknClient` now automatically closes all `Session` objects that were opened from that client.

Author: Alex Walker

dependencies/graknlabs/artifacts.bzl

@@ -37,5 +37,5 @@ def graknlabs_grakn_cluster_artifacts():
         artifact_name = "grakn-cluster-all-{platform}-{version}.{ext}",
         tag_source = deployment_private["artifact.release"],
         commit_source = deployment_private["artifact.snapshot"],
-        commit = "cc3382fe6be239973f1e7d8437a71683e3dff809",
+        commit = "a5cf161f63b33010009c63e2b955d2fa81e4e736",
     )

grakn/client.py

@@ -91,12 +91,15 @@ def __init__(self, address: str):
         self._address = address
         self._channel = grpc.insecure_channel(self._address)
         self._databases = _DatabaseManagerRPC(self._channel)
+        self._sessions: Dict[bytes, _SessionRPC] = {}
         self._is_open = True
 
     def session(self, database: str, session_type: SessionType, options=None) -> Session:
         if not options:
             options = GraknOptions.core()
-        return _SessionRPC(self, database, session_type, options)
+        session = _SessionRPC(self, database, session_type, options)
+        self._sessions[session.session_id()] = session
+        return session
 
     def databases(self):
         return self._databases
@@ -105,6 +108,8 @@ def is_open(self):
         return self._is_open
 
     def close(self):
+        for session_id in self._sessions:
+            self._sessions[session_id].close()
         self._channel.close()
         self._is_open = False
 
@@ -121,6 +126,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         else:
             return False
 
+    def remove_session(self, session: _SessionRPC):
+        del self._sessions[session.session_id()]
+
     def channel(self):
         return self._channel
 

grakn/rpc/session.py

@@ -20,15 +20,13 @@
 import sched
 import time
 from abc import ABC, abstractmethod
-from threading import Thread
+from threading import Thread, Lock
 
 import grakn_protocol.protobuf.session_pb2 as session_proto
 import grpc
 from grakn_protocol.protobuf.grakn_pb2_grpc import GraknStub
-from grpc import RpcError
 
 from grakn import grakn_proto_builder
-from grakn.common.exception import GraknClientException
 from grakn.options import GraknOptions
 from grakn.rpc.database import Database, _DatabaseRPC
 from grakn.rpc.transaction import Transaction, TransactionType
@@ -83,19 +81,21 @@ class _SessionRPC(Session):
     def __init__(self, client, database: str, session_type: SessionType, options: GraknOptions = None):
         if not options:
             options = GraknOptions.core()
+        self._client = client
         self._address = client._address
         self._channel = grpc.insecure_channel(client._address)
         self._scheduler = sched.scheduler(time.time, time.sleep)
         self._database = _DatabaseRPC(database_manager=client.databases(), name=database)
         self._session_type = session_type
         self._grpc_stub = GraknStub(self._channel)
+        self._lock = Lock()
 
         open_req = session_proto.Session.Open.Req()
         open_req.database = database
         open_req.type = _session_type_proto(session_type)
         open_req.options.CopyFrom(grakn_proto_builder.options(options))
 
-        self._session_id = self._grpc_stub.session_open(open_req).session_id
+        self._session_id: bytes = self._grpc_stub.session_open(open_req).session_id
         self._is_open = True
         self._pulse = self._scheduler.enter(delay=self._PULSE_FREQUENCY_SECONDS, priority=1, action=self._transmit_pulse, argument=())
         Thread(target=self._scheduler.run, name="session_pulse_{}".format(self._session_id.hex()), daemon=True).start()
@@ -112,8 +112,11 @@ def is_open(self) -> bool:
         return self._is_open
 
     def close(self) -> None:
+        self._lock.acquire(blocking=True)
         if self._is_open:
             self._is_open = False
+            self._client.remove_session(self)
+            self._lock.release()
             self._scheduler.cancel(self._pulse)
             self._scheduler.empty()
             req = session_proto.Session.Close.Req()
@@ -122,10 +125,15 @@ def close(self) -> None:
                 self._grpc_stub.session_close(req)
             finally:
                 self._channel.close()
+        else:
+            self._lock.release()
 
     def database(self) -> Database:
         return self._database
 
+    def session_id(self) -> bytes:
+        return self._session_id
+
     def _transmit_pulse(self) -> None:
         if not self._is_open:
             return

grakn/rpc/transaction.py

@@ -45,7 +45,7 @@ class TransactionType(enum.Enum):
 
 class Transaction:
 
-    def __init__(self, address: str, session_id: str, transaction_type: TransactionType, options: GraknOptions = None):
+    def __init__(self, address: str, session_id: bytes, transaction_type: TransactionType, options: GraknOptions = None):
         if not options:
             options = GraknOptions.core()
         self._transaction_type = transaction_type

tools/behave_rule.bzl

@@ -110,7 +110,7 @@ def _rule_implementation(ctx):
     cmd += " && rm -rf " + steps_out_dir
     cmd += " && mkdir " + steps_out_dir + " && "
     cmd += " && ".join(["cp %s %s" % (step_file.path, steps_out_dir) for step_file in ctx.files.steps])
-    cmd += " && behave %s -D port=$PORT && export RESULT=0 || export RESULT=1" % feats_dir
+    cmd += " && behave %s --no-capture -D port=$PORT && export RESULT=0 || export RESULT=1" % feats_dir
     cmd += """
            echo Tests concluded with exit value $RESULT
            echo Stopping server.