Use connection manager for keepalives (#20)

Author: DavidStirling

README.md

@@ -166,6 +166,24 @@ with omero2pandas.OMEROConnection(server='my.server', port=4064,
 
 The context manager will handle session creation and cleanup automatically.
 
+### Connection Management
+
+omero2pandas keeps track of any active connector objects and shuts them down 
+safely when Python exits. Deleting all references to a connector will also 
+handle closing the connection to OMERO gracefully. You can also call 
+`connector.shutdown()` to close a connection manually.
+
+By default omero2pandas also keeps active connections alive by pinging the 
+server once per minute (otherwise the session may timeout and require 
+reconnecting). This can be disabled as follows 
+
+```python
+omero2pandas.connect_to_omero(keep_alive=False)
+```
+
+N.b. omero2pandas uses a different system from the native OMERO API's
+`client.enableKeepAlive` function, using both is unnecessary.
+
 ### Querying tables
 
 You can also supply [PyTables condition syntax](https://www.pytables.org/usersguide/condition_syntax.html) to the `read_table` and `download_table` functions.

omero2pandas/__init__.py

@@ -377,7 +377,7 @@ def _validate_requested_object(file_id, annotation_id):
 
 def connect_to_omero(client=None, server=None, port=4064,
                      username=None, password=None, session_key=None,
-                     allow_token=True, interactive=True):
+                     allow_token=True, interactive=True, keep_alive=True):
     """
     Connect to OMERO and return an OMEROConnection object.
     :param client: An existing omero.client object to be used instead of
@@ -391,12 +391,13 @@ def connect_to_omero(client=None, server=None, port=4064,
     :param allow_token: True/False Search for omero_user_token before trying
     to use credentials. Default True.
     :param interactive: Prompt user for missing login details. Default True.
+    :param keep_alive: Periodically ping the server to prevent session timeout.
     :return: OMEROConnection object wrapping a client and Blitz Gateway object,
     with automatic session management and cleanup.
     """
     connector = OMEROConnection(server=server, port=port,
                                 session_key=session_key, username=username,
                                 password=password, client=client,
                                 allow_token=allow_token)
-    connector.connect(interactive=interactive)
+    connector.connect(interactive=interactive, keep_alive=keep_alive)
     return connector

omero2pandas/connect.py

@@ -10,14 +10,17 @@
 import getpass
 import importlib.util
 import logging
+import threading
+import time
 import weakref
 
+import Ice
 import omero
 from omero.gateway import BlitzGateway
 
 LOGGER = logging.getLogger(__name__)
-
 ACTIVE_CONNECTORS = weakref.WeakSet()
+KEEPALIVE_THREAD = None
 
 
 class OMEROConnection:
@@ -83,6 +86,7 @@ def shutdown(self):
             LOGGER.debug("Closing OMERO session")
             self.client.closeSession()
             self.client = None
+            self.session = None
 
     def __del__(self):
         # Make sure we close sessions on deletion.
@@ -104,7 +108,7 @@ def need_connection_details(self):
             return True
         return False
 
-    def connect(self, interactive=True):
+    def connect(self, interactive=True, keep_alive=True):
         if self.connected:
             return True
         # Attempt to establish a connection.
@@ -122,25 +126,33 @@ def connect(self, interactive=True):
         if self.session_key is not None:
             try:
                 self.client.joinSession(self.session_key)
+                self.session = self.client.getSession()
+                self.session.detachOnDestroy()
             except Exception as e:
                 print(f"Failed to join session, token may have expired: {e}")
                 self.client = None
+                self.session = None
                 return False
         elif self.username is not None:
             try:
                 self.session = self.client.createSession(
                     username=self.username, password=self.password)
-                self.client.enableKeepAlive(60)
                 self.session.detachOnDestroy()
             except Exception as e:
                 print(f"Failed to create session: {e}")
                 self.client = None
+                self.session = None
                 return False
         else:
             self.client = None
+            self.session = None
             raise Exception(
                 "Not enough details to create a server connection.")
         print(f"Connected to {self.server}")
+        if keep_alive:
+            # Use o2p keep alive instead of omero-py
+            self.client.stopKeepAlive()
+            start_keep_alive()
         return True
 
     def connect_widget(self):
@@ -268,6 +280,15 @@ def get_client(self):
             LOGGER.warning("Client connection not initialised")
         return self.client
 
+    def keep_alive(self):
+        if self.client is not None and self.session is not None:
+            try:
+                self.session.keepAlive(None)
+            except Ice.CommunicatorDestroyedException:
+                self.session = None  # Was shut down
+            except Exception as e:
+                LOGGER.warning(f"Failed to keep alive: {e}")
+
 
 def detect_jupyter():
     # Determine whether we're running in a Jupyter Notebook.
@@ -293,4 +314,21 @@ def cleanup_sessions():
         connector.shutdown()
 
 
+def keep_sessions_alive():
+    while ACTIVE_CONNECTORS:
+        time.sleep(60)
+        for connector in ACTIVE_CONNECTORS:
+            connector.keep_alive()
+        connector = None  # Don't keep a reference (would prevent shutdown!)
+
+
+def start_keep_alive():
+    global KEEPALIVE_THREAD
+    if KEEPALIVE_THREAD is None or not KEEPALIVE_THREAD.is_alive():
+        KEEPALIVE_THREAD = threading.Thread(target=keep_sessions_alive,
+                                            name="omero2pandas_keepalive",
+                                            daemon=True)
+        KEEPALIVE_THREAD.start()
+
+
 atexit.register(cleanup_sessions)