Retry iRODS objectstore reads on transient connection errors

Author: pauldg

lib/galaxy/objectstore/irods.py

@@ -2,11 +2,13 @@
 Object Store plugin for the Integrated Rule-Oriented Data System (iRODS)
 """
 
+import functools
 import logging
 import os
 import shutil
 import ssl
 import threading
+import time
 from datetime import datetime
 from pathlib import Path
 
@@ -16,11 +18,13 @@
     from irods.exception import (
         CollectionDoesNotExist,
         DataObjectDoesNotExist,
+        NetworkException,
     )
     from irods.session import iRODSSession
 except ImportError:
     irods = None
 
+
 from galaxy.util import (
     ExecutionTimer,
     string_as_bool,
@@ -35,6 +39,33 @@
 logging.getLogger("irods.connection").setLevel(logging.INFO)  # irods logging generates gigabytes of logs
 
 
+_IRODS_RETRY_ATTEMPTS = 3
+_IRODS_RETRY_BACKOFF = 0.5
+
+
+def _retry_on_connection_error(func):
+    """Retry iRODS read operations on a transient connection error."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        for attempt in range(1, _IRODS_RETRY_ATTEMPTS + 1):
+            try:
+                return func(*args, **kwargs)
+            except (NetworkException, ssl.SSLError) as exc:
+                if attempt == _IRODS_RETRY_ATTEMPTS:
+                    raise
+                log.warning(
+                    "Transient iRODS error in %s (attempt %d/%d), retrying on a fresh connection: %s",
+                    func.__name__,
+                    attempt,
+                    _IRODS_RETRY_ATTEMPTS,
+                    exc,
+                )
+                time.sleep(_IRODS_RETRY_BACKOFF * attempt)
+
+    return wrapper
+
+
 def _config_xml_error(tag):
     msg = f"No {tag} element in config XML tree"
     raise Exception(msg)
@@ -376,6 +407,7 @@ def _config_to_dict(self):
         }
 
     # rel_path is file or folder?
+    @_retry_on_connection_error
     def _get_remote_size(self, rel_path):
         ipt_timer = ExecutionTimer()
         p = Path(rel_path)
@@ -396,6 +428,7 @@ def _get_remote_size(self, rel_path):
             log.debug("irods_pt _get_remote_size: %s", ipt_timer)
 
     # rel_path is file or folder?
+    @_retry_on_connection_error
     def _exists_remotely(self, rel_path):
         ipt_timer = ExecutionTimer()
         p = Path(rel_path)
@@ -415,6 +448,7 @@ def _exists_remotely(self, rel_path):
         finally:
             log.debug("irods_pt _exists_remotely: %s", ipt_timer)
 
+    @_retry_on_connection_error
     def _download(self, rel_path):
         ipt_timer = ExecutionTimer()
         cache_path = self._get_cache_path(rel_path)

test/unit/objectstore/test_irods.py

@@ -1,8 +1,15 @@
 import os
+import ssl
 
 import pytest
 
-from galaxy.objectstore.irods import parse_config_xml
+from galaxy.objectstore import irods as irods_module
+from galaxy.objectstore.irods import (
+    _IRODS_RETRY_ATTEMPTS,
+    _retry_on_connection_error,
+    NetworkException,
+    parse_config_xml,
+)
 from galaxy.util import parse_xml
 
 SCRIPT_DIRECTORY = os.path.abspath(os.path.dirname(__file__))
@@ -80,3 +87,46 @@ def test_parse_config_xml_no_auth():
     root = tree.getroot()
     with pytest.raises(Exception, match="No auth element in config XML tree"):
         parse_config_xml(root)
+
+
+def _make_flaky(exc, fail_times):
+    calls = {"n": 0}
+
+    @_retry_on_connection_error
+    def op(_self):
+        calls["n"] += 1
+        if calls["n"] <= fail_times:
+            raise exc
+        return "ok"
+
+    return op, calls
+
+
+@pytest.fixture(autouse=True)
+def _no_sleep(monkeypatch):
+    monkeypatch.setattr(irods_module.time, "sleep", lambda *_: None)
+
+
+def test_retry_recovers_from_network_exception():
+    op, calls = _make_flaky(NetworkException("reset"), fail_times=_IRODS_RETRY_ATTEMPTS - 1)
+    assert op(object()) == "ok"
+    assert calls["n"] == _IRODS_RETRY_ATTEMPTS
+
+
+def test_retry_recovers_from_ssl_error():
+    op, calls = _make_flaky(ssl.SSLEOFError("handshake"), fail_times=1)
+    assert op(object()) == "ok"
+    assert calls["n"] == 2
+
+
+def test_retry_gives_up_after_max_attempts():
+    op, calls = _make_flaky(ssl.SSLEOFError("handshake"), fail_times=_IRODS_RETRY_ATTEMPTS)
+    with pytest.raises(ssl.SSLError):
+        op(object())
+    assert calls["n"] == _IRODS_RETRY_ATTEMPTS
+
+
+def test_no_retry_on_success():
+    op, calls = _make_flaky(NetworkException("reset"), fail_times=0)
+    assert op(object()) == "ok"
+    assert calls["n"] == 1