Repository: MatterMiners/tardis

Issue:
Improve exception handling and recycling of connections in SSH executors of TARDIS
This seems to be very much like the issue reported as #133 — we encountered a traceback after an SSH timeout when using {{tardis.utilities.executors.sshexecutor.SSHExecutor}} (sadly, did not catch the trace, but I guess the issue is generic). 
It would be nice if Tardis could handle that and recover :wink:. 

It might be interesting also for other users that we now use SSH multiplexing as workaround as shown here:
https://github.com/unibonn/puppet-cobald/commit/0348e098de832f6af96460b39588412dbc6d2701
This significantly reduces the latency of the SSH executor, reduces the number of connections, and seems to increase stability (you should be aware that anybody who can access the multiplex socket can hijack the connection, which is why it's usually stored in {{~/.ssh}}). 


Comments:

Patch:
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index d22bc81..fc03530 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,4 +1,4 @@
-.. Created by changelog.py at 2020-03-19, command
+.. Created by changelog.py at 2020-04-28, command
    '/Users/giffler/.cache/pre-commit/repont7o94ca/py_env-python3.7/bin/changelog docs/source/changes compile --output=docs/source/changelog.rst'
    based on the format of 'https://keepachangelog.com/'
 
@@ -32,9 +32,14 @@ Fixed
 * Fix state transitions for jobs retried by HTCondor
 * Fix state transitions and refactoring of the SLURM site adapter
 
-[Unreleased] - 2020-03-19
+[Unreleased] - 2020-04-28
 =========================
 
+Added
+-----
+
+* Add ssh connection sharing to `SSHExecutor` in order to re-use existing connection
+
 Changed
 -------
 
diff --git a/docs/source/changes/145.add_ssh_connection_sharing.yaml b/docs/source/changes/145.add_ssh_connection_sharing.yaml
new file mode 100644
index 0000000..e556827
--- /dev/null
+++ b/docs/source/changes/145.add_ssh_connection_sharing.yaml
@@ -0,0 +1,9 @@
+category: added
+summary: "Add ssh connection sharing to `SSHExecutor` in order to re-use existing connection"
+pull requests:
+  - 145
+issues:
+  - 135
+description: |
+  The `SSHExector` is now re-using existing connections. Closed connections are automatically reestablished. This will
+  avoid connection problems when bothering a remote ssh server with too many requests in too short intervals.
diff --git a/docs/source/executors/executors.rst b/docs/source/executors/executors.rst
index f0a1398..6169f8d 100644
--- a/docs/source/executors/executors.rst
+++ b/docs/source/executors/executors.rst
@@ -41,9 +41,10 @@ SSH Executor
 
 .. content-tabs:: left-col
 
-    The ssh executor is used to asynchronously execute shell commands remotely via ssh. All parameters specified in the
-    configuration are directly passed as keyword arguments to `asyncssh` `connect` call. You can find all available
-    parameters in the `asyncssh documentation`_
+    The ssh executor is used to asynchronously execute shell commands remotely via ssh. The actual ssh connection to
+    the host is preserved, recycled and automatically reestablished. All parameters specified in the configuration are
+    directly passed as keyword arguments to `asyncssh` `connect` call. You can find all available parameters in the
+    `asyncssh documentation`_
 
     .. _asyncssh documentation: https://asyncssh.readthedocs.io/en/latest/api.html#connect
 
diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py
index 7f79943..3003b0e 100644
--- a/tardis/utilities/executors/sshexecutor.py
+++ b/tardis/utilities/executors/sshexecutor.py
@@ -3,6 +3,7 @@ from ...exceptions.executorexceptions import CommandExecutionFailure
 from ...interfaces.executor import Executor
 from ..attributedict import AttributeDict
 
+import asyncio
 import asyncssh
 
 
@@ -10,13 +11,45 @@ import asyncssh
 class SSHExecutor(Executor):
     def __init__(self, **parameters):
         self._parameters = parameters
+        self._ssh_connection = None
+        self._lock = None
+
+    async def _establish_connection(self):
+        for retry in range(1, 10):
+            try:
+                return await asyncssh.connect(**self._parameters)
+            except (
+                ConnectionResetError,
+                asyncssh.DisconnectError,
+                asyncssh.ConnectionLost,
+                BrokenPipeError,
+            ):
+                await asyncio.sleep(retry * 10)
+        return await asyncssh.connect(**self._parameters)
+
+    @property
+    async def ssh_connection(self):
+        if self._ssh_connection is None:
+            async with self.lock:
+                # check that connection has not yet been initialize in a different task
+                while self._ssh_connection is None:
+                    self._ssh_connection = await self._establish_connection()
+        return self._ssh_connection
+
+    @property
+    def lock(self):
+        # Create lock once tardis event loop is running.
+        # To avoid got Future <Future pending> attached to a different loop exception
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
 
     async def run_command(self, command, stdin_input=None):
+        ssh_connection = await self.ssh_connection
         try:
-            async with asyncssh.connect(**self._parameters) as conn:
-                response = await conn.run(
-                    command, check=True, input=stdin_input and stdin_input.encode()
-                )
+            response = await ssh_connection.run(
+                command, check=True, input=stdin_input and stdin_input.encode()
+            )
         except asyncssh.ProcessError as pe:
             raise CommandExecutionFailure(
                 message=f"Run command {command} via SSHExecutor failed",
@@ -25,18 +58,15 @@ class SSHExecutor(Executor):
                 stdout=pe.stdout,
                 stderr=pe.stderr,
             ) from pe
-        except (
-            ConnectionResetError,
-            asyncssh.misc.DisconnectError,
-            asyncssh.misc.ConnectionLost,
-            BrokenPipeError,
-        ) as ce:
+        except asyncssh.ChannelOpenError as coe:
+            # Broken connection will be replaced by a new connection during next call
+            self._ssh_connection = None
             raise CommandExecutionFailure(
-                message=f"Could not run command {command} due to SSH failure: {ce}",
+                message=f"Could not run command {command} due to SSH failure: {coe}",
                 exit_code=255,
                 stdout="",
-                stderr="SSH failure",
-            ) from ce
+                stderr="SSH Broken Connection",
+            ) from coe
         else:
             return AttributeDict(
                 stdout=response.stdout,


Test patch:
diff --git a/tests/utilities_t/executors_t/test_sshexecutor.py b/tests/utilities_t/executors_t/test_sshexecutor.py
index d06fcd6..37a4476 100644
--- a/tests/utilities_t/executors_t/test_sshexecutor.py
+++ b/tests/utilities_t/executors_t/test_sshexecutor.py
@@ -1,55 +1,42 @@
-from tests.utilities.utilities import run_async
+from tests.utilities.utilities import async_return, run_async
 from tardis.utilities.attributedict import AttributeDict
 from tardis.utilities.executors.sshexecutor import SSHExecutor
 from tardis.exceptions.executorexceptions import CommandExecutionFailure
 
-from asyncssh import ProcessError
-from asyncssh.misc import ConnectionLost, DisconnectError
-
-try:
-    from contextlib import asynccontextmanager
-except ImportError:
-    from aiotools import async_ctx_manager as asynccontextmanager
+from asyncssh import ChannelOpenError, ConnectionLost, DisconnectError, ProcessError
 
 from unittest import TestCase
 from unittest.mock import patch
 
+import asyncio
 import yaml
 
 
-def generate_connect(response, exception=None):
-    @asynccontextmanager
-    async def connect(*args, **kwargs):
-        class Connection(object):
-            async def run(self, *args, input, **kwargs):
-                if exception:
-                    raise exception
-                self.stdout = input and input.decode()
-                return self
-
-            @property
-            def exit_status(self):
-                return response.exit_status
-
-            @property
-            def stderr(self):
-                return response.stderr
-
-        yield Connection()
+class MockConnection(object):
+    def __init__(self, exception=None, **kwargs):
+        self.exception = exception and exception(**kwargs)
 
-    return connect
+    async def run(self, command, input=None, **kwargs):
+        if self.exception:
+            raise self.exception
+        return AttributeDict(
+            stdout=input and input.decode(), stderr="TestError", exit_status=0
+        )
 
 
 class TestSSHExecutor(TestCase):
+    mock_asyncssh = None
+
     @classmethod
     def setUpClass(cls):
         cls.mock_asyncssh_patcher = patch(
             "tardis.utilities.executors.sshexecutor.asyncssh"
         )
         cls.mock_asyncssh = cls.mock_asyncssh_patcher.start()
+        cls.mock_asyncssh.ChannelOpenError = ChannelOpenError
+        cls.mock_asyncssh.ConnectionLost = ConnectionLost
+        cls.mock_asyncssh.DisconnectError = DisconnectError
         cls.mock_asyncssh.ProcessError = ProcessError
-        cls.mock_asyncssh.misc.ConnectionLost = ConnectionLost
-        cls.mock_asyncssh.misc.DisconnectError = DisconnectError
 
     @classmethod
     def tearDownClass(cls):
@@ -57,82 +44,105 @@ class TestSSHExecutor(TestCase):
 
     def setUp(self) -> None:
         self.response = AttributeDict(stderr="", exit_status=0)
-        self.mock_asyncssh.connect.side_effect = generate_connect(self.response)
-        self.mock_asyncssh.reset_mock()
-
-    def test_run_command(self):
-        executor = SSHExecutor(
-            host="test_host", username="test", client_keys=["TestKey"]
+        self.mock_asyncssh.connect.return_value = async_return(
+            return_value=MockConnection()
         )
-        self.assertIsNone(run_async(executor.run_command, command="Test").stdout)
-        self.mock_asyncssh.connect.assert_called_with(
+        self.test_asyncssh_params = AttributeDict(
             host="test_host", username="test", client_keys=["TestKey"]
         )
+        self.executor = SSHExecutor(**self.test_asyncssh_params)
         self.mock_asyncssh.reset_mock()
 
-        executor = SSHExecutor(
-            host="test_host", username="test", client_keys=("TestKey",)
-        )
-        self.assertIsNone(run_async(executor.run_command, command="Test").stdout)
-        self.mock_asyncssh.connect.assert_called_with(
-            host="test_host", username="test", client_keys=("TestKey",)
+    @patch("tardis.utilities.executors.sshexecutor.asyncio.sleep", async_return)
+    def test_establish_connection(self):
+        self.assertIsInstance(
+            run_async(self.executor._establish_connection), MockConnection
         )
 
-        self.mock_asyncssh.reset_mock()
+        self.mock_asyncssh.connect.assert_called_with(**self.test_asyncssh_params)
 
-        executor = SSHExecutor(
-            host="test_host", username="test", client_keys=("TestKey",)
-        )
-        self.assertEqual(
-            run_async(executor.run_command, command="Test", stdin_input="Test").stdout,
-            "Test",
-        )
+        test_exceptions = [
+            ConnectionResetError(),
+            DisconnectError(reason="test_reason", code=255),
+            ConnectionLost(reason="test_reason"),
+            BrokenPipeError(),
+        ]
+
+        for exception in test_exceptions:
+            self.mock_asyncssh.reset_mock()
+            self.mock_asyncssh.connect.side_effect = exception
+
+            with self.assertRaises(type(exception)):
+                run_async(self.executor._establish_connection)
+
+            self.assertEqual(self.mock_asyncssh.connect.call_count, 10)
+
+        self.mock_asyncssh.connect.side_effect = None
+
+    def test_connection_property(self):
+        async def helper_coroutine():
+            return await self.executor.ssh_connection
+
+        self.assertIsNone(self.executor._ssh_connection)
+        run_async(helper_coroutine)
+
+        self.assertIsInstance(self.executor._ssh_connection, MockConnection)
+
+        current_ssh_connection = self.executor._ssh_connection
+
+        run_async(helper_coroutine)
+
+        self.assertEqual(self.executor._ssh_connection, current_ssh_connection)
+
+    def test_lock(self):
+        self.assertIsInstance(self.executor.lock, asyncio.Lock)
+
+    def test_run_command(self):
+        self.assertIsNone(run_async(self.executor.run_command, command="Test").stdout)
         self.mock_asyncssh.connect.assert_called_with(
-            host="test_host", username="test", client_keys=("TestKey",)
+            host="test_host", username="test", client_keys=["TestKey"]
         )
+        self.mock_asyncssh.reset_mock()
 
-    def test_run_raises_process_error(self):
-        test_exception = ProcessError(
-            env="Test",
-            command="Test",
-            subsystem="Test",
-            exit_status=1,
-            exit_signal=None,
-            returncode=1,
-            stdout="TestError",
-            stderr="TestError",
+        response = run_async(
+            self.executor.run_command, command="Test", stdin_input="Test"
         )
 
-        self.mock_asyncssh.connect.side_effect = generate_connect(
-            self.response, exception=test_exception
-        )
+        self.assertEqual(response.stdout, "Test")
 
-        executor = SSHExecutor(
-            host="test_host", username="test", client_keys=("TestKey",)
-        )
+        self.assertEqual(response.stderr, "TestError")
 
-        with self.assertRaises(CommandExecutionFailure):
-            run_async(executor.run_command, command="Test", stdin_input="Test")
+        self.assertEqual(response.exit_code, 0)
 
-    def test_run_raises_ssh_errors(self):
-        test_exceptions = [
-            ConnectionResetError,
-            DisconnectError(reason="test_reason", code=255),
-            ConnectionLost(reason="test_reason"),
-            BrokenPipeError,
-        ]
+        raising_executor = SSHExecutor(**self.test_asyncssh_params)
 
-        for test_exception in test_exceptions:
-            self.mock_asyncssh.connect.side_effect = generate_connect(
-                self.response, exception=test_exception
+        self.mock_asyncssh.connect.return_value = async_return(
+            return_value=MockConnection(
+                exception=ProcessError,
+                env="Test",
+                command="Test",
+                subsystem="Test",
+                exit_status=1,
+                exit_signal=None,
+                returncode=1,
+                stdout="TestError",
+                stderr="TestError",
             )
+        )
 
-            executor = SSHExecutor(
-                host="test_host", username="test", client_keys=("TestKey",)
+        with self.assertRaises(CommandExecutionFailure):
+            run_async(raising_executor.run_command, command="Test", stdin_input="Test")
+
+        raising_executor = SSHExecutor(**self.test_asyncssh_params)
+
+        self.mock_asyncssh.connect.return_value = async_return(
+            return_value=MockConnection(
+                exception=ChannelOpenError, reason="test_reason", code=255
             )
+        )
 
-            with self.assertRaises(CommandExecutionFailure):
-                run_async(executor.run_command, command="Test", stdin_input="Test")
+        with self.assertRaises(CommandExecutionFailure):
+            run_async(raising_executor.run_command, command="Test", stdin_input="Test")
 
     def test_construction_by_yaml(self):
         executor = yaml.safe_load(
@@ -144,9 +154,7 @@ class TestSSHExecutor(TestCase):
                     - TestKey
                    """
         )
-        response = AttributeDict(stderr="", exit_status=0)
 
-        self.mock_asyncssh.connect.side_effect = generate_connect(response)
         self.assertEqual(
             run_async(executor.run_command, command="Test", stdin_input="Test").stdout,
             "Test",
@@ -154,5 +162,3 @@ class TestSSHExecutor(TestCase):
         self.mock_asyncssh.connect.assert_called_with(
             host="test_host", username="test", client_keys=["TestKey"]
         )
-        self.mock_asyncssh.connect.side_effect = None
-        self.mock_asyncssh.reset_mock()