From cb63103830a98c15426f27268da9bea6a31b1e58 Mon Sep 17 00:00:00 2001 From: David Stirling Date: Mon, 26 Aug 2024 16:05:59 +0100 Subject: [PATCH] Use connection manager for keepalives (#20) --- README.md | 18 ++++++++++++++++ omero2pandas/__init__.py | 5 +++-- omero2pandas/connect.py | 44 +++++++++++++++++++++++++++++++++++++--- 3 files changed, 62 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 77db672..7ee4c6f 100644 --- a/README.md +++ b/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. diff --git a/omero2pandas/__init__.py b/omero2pandas/__init__.py index ecfb52b..dd8e687 100644 --- a/omero2pandas/__init__.py +++ b/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,6 +391,7 @@ 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. """ @@ -398,5 +399,5 @@ def connect_to_omero(client=None, server=None, port=4064, 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 diff --git a/omero2pandas/connect.py b/omero2pandas/connect.py index f1ccc40..06cfc8b 100644 --- a/omero2pandas/connect.py +++ b/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)