Retry ssh connection on failure

pyinfra/connectors/ssh.py

@@ -12,7 +12,9 @@
 from distutils.spawn import find_executable
 from getpass import getpass
 from os import path
+from random import uniform
 from socket import error as socket_error, gaierror
+from time import sleep
 from typing import TYPE_CHECKING, Type, Union
 
 import click
@@ -71,6 +73,10 @@ class DataKeys:
         known_hosts_file = "Custom SSH known hosts file"
         strict_host_key_checking = "Override strict host keys check setting"
 
+        connect_retries = "Number of tries to connect via ssh"
+        connect_retry_min_delay = "Lower bound for random delay between retries"
+        connect_retry_max_delay = "Upper bound for random delay between retries"
+
         paramiko_connect_kwargs = (
             "Override keyword arguments passed into paramiko's `SSHClient.connect`"
         )
@@ -226,7 +232,15 @@ def connect(state: "State", host: "Host"):
     Connect to a single host. Returns the SSH client if successful. Stateless by
     design so can be run in parallel.
     """
+    retries = host.data.get(DATA_KEYS.connect_retries, 0)
+
+    for tries_left in range(retries, -1, -1):
+        con = _connect(state, host, tries_left)
+        if con:
+            return con
+
 
+def _connect(state: "State", host: "Host", tries_left: int):
     kwargs = _make_paramiko_kwargs(state, host)
     logger.debug("Connecting to: %s (%r)", host.name, kwargs)
 
@@ -273,16 +287,26 @@ def connect(state: "State", host: "Host"):
         )
 
     except SSHException as e:
-        _raise_connect_error(host, "SSH error", e)
+        if tries_left == 0:
+            _raise_connect_error(host, "SSH error", e)
 
     except gaierror:
-        _raise_connect_error(host, "Could not resolve hostname", hostname)
+        if tries_left == 0:
+            _raise_connect_error(host, "Could not resolve hostname", hostname)
 
     except socket_error as e:
-        _raise_connect_error(host, "Could not connect", e)
+        if tries_left == 0:
+            _raise_connect_error(host, "Could not connect", e)
 
     except EOFError as e:
-        _raise_connect_error(host, "EOF error", e)
+        if tries_left == 0:
+            _raise_connect_error(host, "EOF error", e)
+
+    min_delay = host.data.get(DATA_KEYS.connect_retry_min_delay, 0.1)
+    max_delay = host.data.get(DATA_KEYS.connect_retry_max_delay, 0.5)
+    sleep(uniform(min_delay, max_delay))
+
+    return None
 
 
 def run_shell_command(

tests/test_connectors/test_ssh.py

@@ -1042,3 +1042,55 @@ def test_get_sftp_fail(self, fake_sftp_client, fake_ssh_client):
                     "not-another-file",
                     print_output=True,
                 )
+
+    @patch("pyinfra.connectors.ssh.SSHClient")
+    @patch("time.sleep")
+    def test_ssh_connect_fail_retry(self, fake_sleep, fake_ssh_client):
+        for exception_class in (
+            SSHException,
+            gaierror,
+            socket_error,
+            EOFError,
+        ):
+            inventory = make_inventory(
+                hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
+            )
+            State(inventory, Config())
+
+            unresposivehost = inventory.get_host("unresposivehost")
+            assert unresposivehost.data.ssh_connect_retries == 1
+
+            fake_ssh = MagicMock()
+            fake_ssh.connect.side_effect = exception_class()
+            fake_ssh_client.return_value = fake_ssh
+
+            with self.assertRaises(ConnectError):
+                unresposivehost.connect(show_errors=False, raise_exceptions=True)
+            assert fake_sleep.called_once()
+            assert fake_ssh_client.connect.called_twice()
+
+    @patch("pyinfra.connectors.ssh.SSHClient")
+    @patch("time.sleep")
+    def test_ssh_connect_fail_success(self, fake_sleep, fake_ssh_client):
+        for exception_class in (
+            SSHException,
+            gaierror,
+            socket_error,
+            EOFError,
+        ):
+            inventory = make_inventory(
+                hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
+            )
+            State(inventory, Config())
+
+            unresposivehost = inventory.get_host("unresposivehost")
+            assert unresposivehost.data.ssh_connect_retries == 1
+
+            connection = MagicMock()
+            fake_ssh = MagicMock()
+            fake_ssh.connect.side_effect = [exception_class(), connection]
+            fake_ssh_client.return_value = fake_ssh
+
+            unresposivehost.connect(show_errors=False, raise_exceptions=True)
+            assert fake_sleep.called_once()
+            assert fake_ssh_client.connect.called_twice()