From 95b0164eb8020ba8b2b3fd0c03b686b8a89eaaa2 Mon Sep 17 00:00:00 2001
From: MartinBelthle <102529366+MartinBelthle@users.noreply.github.com>
Date: Tue, 19 Dec 2023 16:19:01 +0100
Subject: [PATCH] feat(ssh): add retry loop around SSH Exceptions (#68)

* feat(ssh): add retry loop around ssh exceptions

* refactor(ssh): remove useless stdin

* refactor(ssh-connection): use a "retry" decorator

---------

Co-authored-by: Laurent LAPORTE <laurent.laporte.pro@gmail.com>
(cherry picked from commit 3e0d436891cd635dd9a2c4f6a4ea3df3bebb5f8f)
---
 CHANGES.md                                    |  4 ++
 .../remote_environnement/ssh_connection.py    | 42 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/CHANGES.md b/CHANGES.md
index 0c08752..bafd99a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -22,6 +22,10 @@ npx auto-changelog -l false --hide-empty-releases  -v v1.3.1 -o CHANGES.out.md
 
 - build: add a script to bump the version
 
+### Changed
+
+- feat(ssh): add retry loop around SSH Exceptions [`#68`](https://github.com/AntaresSimulatorTeam/antares-launcher/pull/68)
+
 
 ## [1.3.1] - 2023-09-26
 
diff --git a/antareslauncher/remote_environnement/ssh_connection.py b/antareslauncher/remote_environnement/ssh_connection.py
index 389151a..4099413 100644
--- a/antareslauncher/remote_environnement/ssh_connection.py
+++ b/antareslauncher/remote_environnement/ssh_connection.py
@@ -21,6 +21,48 @@
 DIRECTORY_NOT_FOUND_ERROR = "Directory not found error"
 
 
+def retry(
+    exception: t.Type[Exception],
+    *exceptions: t.Type[Exception],
+    delay_sec: float = 5,
+    max_retry: int = 5,
+    msg_fmt: str = "Retrying in {delay_sec} seconds...",
+):
+    """
+    Decorator to retry a function call if it raises an exception.
+
+    Args:
+        exception: The exception to catch.
+        exceptions: Additional exceptions to catch.
+        delay_sec: The delay (in seconds) between each retry.
+        max_retry: The maximum number of retries.
+        msg_fmt: The message to display when retrying, with the following format keys:
+            - delay_sec: The delay (in seconds) between each retry.
+            - remaining: The number of remaining retries.
+
+    Returns:
+        The decorated function.
+    """
+
+    def decorator(func):  # type: ignore
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):  # type: ignore
+            for attempt in range(max_retry):
+                try:
+                    return func(*args, **kwargs)
+                except (exception, *exceptions):
+                    logger = logging.getLogger(__name__)
+                    remaining = max_retry - attempt - 1
+                    logger.warning(msg_fmt.format(delay_sec=delay_sec, remaining=remaining))
+                    time.sleep(delay_sec)
+            # Last attempt
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def retry(
     exception: t.Type[Exception],
     *exceptions: t.Type[Exception],