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}}). 


Discussion:

Patch:
diff --git a/docs/source/adapters/site.rst b/docs/source/adapters/site.rst
index 5ff60d0..1e965a4 100644
--- a/docs/source/adapters/site.rst
+++ b/docs/source/adapters/site.rst
@@ -241,17 +241,17 @@ Available adapter configuration options
 
 .. content-tabs:: left-col
 
-    +----------------+---------------------------------------------------------------------------------------------+-----------------+
-    | Option         | Short Description                                                                           | Requirement     |
-    +================+=============================================================================================+=================+
-    | StatusUpdate   | The result of the status call is cached for `StatusUpdate` in minutes.                      |  **Required**   |
-    +----------------+---------------------------------------------------------------------------------------------+-----------------+
-    | StartupCommand | The command executed in the batch job. (**Deprecated:** Moved to MachineTypeConfiguration!) |  **Deprecated** |
-    +----------------+---------------------------------------------------------------------------------------------+-----------------+
-    | executor       | The |executor| used to run submission and further calls to the Moab batch system.           |  **Optional**   |
-    +                +                                                                                             +                 +
-    |                | Default: ShellExecutor is used!                                                             |                 |
-    +----------------+---------------------------------------------------------------------------------------------+-----------------+
+    +----------------+-----------------------------------------------------------------------------------+-----------------+
+    | Option         | Short Description                                                                 | Requirement     |
+    +================+===================================================================================+=================+
+    | StatusUpdate   | The result of the status call is cached for `StatusUpdate` in minutes.            |  **Required**   |
+    +----------------+-----------------------------------------------------------------------------------+-----------------+
+    | StartupCommand | The command executed in the batch job.                                            |  **Required**   |
+    +----------------+-----------------------------------------------------------------------------------+-----------------+
+    | executor       | The |executor| used to run submission and further calls to the Moab batch system. |  **Optional**   |
+    +                +                                                                                   +                 +
+    |                | Default: ShellExecutor is used!                                                   |                 |
+    +----------------+-----------------------------------------------------------------------------------+-----------------+
 
     The available options in the `MachineTypeConfiguration` section are the expected `WallTime` of the placeholder jobs and
     the requested `NodeType`. For details see the Moab documentation.
@@ -273,6 +273,7 @@ Available adapter configuration options
             username: clown
             client_keys:
               - /opt/tardis/ssh/tardis
+          StartupCommand: startVM.py
           StatusUpdate: 2
           MachineTypes:
             - singularity_d2.large
@@ -281,11 +282,9 @@ Available adapter configuration options
             singularity_d2.large:
               Walltime: '02:00:00:00'
               NodeType: '1:ppn=20'
-              StartupCommand: startVM.py
             singularity_d1.large:
               Walltime: '01:00:00:00'
               NodeType: '1:ppn=20'
-              StartupCommand: startVM.py
           MachineMetaData:
             singularity_d2.large:
               Cores: 20
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index d22bc81..852eef9 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,11 +32,15 @@ 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
 -------
 
 * The SLURM adapter can now be configured to use different startup commands for each machine type.
-* The Moab adapter can now be configured to use different startup commands for each machine type.
diff --git a/docs/source/changes/139_changed_moved_startupcommand_in_machine_type_configuration_moab.yaml b/docs/source/changes/139_changed_moved_startupcommand_in_machine_type_configuration_moab.yaml
deleted file mode 100644
index 0a3e52a..0000000
--- a/docs/source/changes/139_changed_moved_startupcommand_in_machine_type_configuration_moab.yaml
+++ /dev/null
@@ -1,7 +0,0 @@
-category: changed
-summary: The Moab adapter can now be configured to use different startup commands for each machine type.
-pull requests:
-  - 139
-description: |
-  The Moab adapter can now be configured to use different startup commands for each machine type. The old behaviour of
-  providing one startup command is still supported, but will be deprecated in the next major release 0.4.0.
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/adapters/sites/moab.py b/tardis/adapters/sites/moab.py
index 3ab110a..e62038f 100644
--- a/tardis/adapters/sites/moab.py
+++ b/tardis/adapters/sites/moab.py
@@ -19,7 +19,6 @@ from datetime import datetime
 import asyncssh
 import logging
 import re
-import warnings
 from xml.dom import minidom
 
 
@@ -50,17 +49,7 @@ class MoabAdapter(SiteAdapter):
         self._configuration = getattr(Configuration(), site_name)
         self._machine_type = machine_type
         self._site_name = site_name
-
-        try:
-            self._startup_command = self.machine_type_configuration.StartupCommand
-        except AttributeError:
-            if not hasattr(self._configuration, "StartupCommand"):
-                raise
-            warnings.warn(
-                "StartupCommand has been moved to the machine_type_configuration!",
-                DeprecationWarning,
-            )
-            self._startup_command = self._configuration.StartupCommand
+        self._startup_command = self._configuration.StartupCommand
 
         self._executor = getattr(self._configuration, "executor", ShellExecutor())
 
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/adapters_t/sites_t/test_moab.py b/tests/adapters_t/sites_t/test_moab.py
index 12a580f..58a9ea2 100644
--- a/tests/adapters_t/sites_t/test_moab.py
+++ b/tests/adapters_t/sites_t/test_moab.py
@@ -9,10 +9,9 @@ from tests.utilities.utilities import mock_executor_run_command
 from tests.utilities.utilities import run_async
 
 from unittest import TestCase
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 from datetime import datetime, timedelta
-from warnings import filterwarnings
 
 import asyncio
 import asyncssh
@@ -96,16 +95,9 @@ class TestMoabAdapter(TestCase):
 
     def setUp(self):
         config = self.mock_config.return_value
-        config.TestSite = MagicMock(
-            spec=[
-                "MachineMetaData",
-                "StatusUpdate",
-                "MachineTypeConfiguration",
-                "executor",
-            ]
-        )
         self.test_site_config = config.TestSite
         self.test_site_config.MachineMetaData = self.machine_meta_data
+        self.test_site_config.StartupCommand = "startVM.py"
         self.test_site_config.StatusUpdate = 10
         self.test_site_config.MachineTypeConfiguration = self.machine_type_configuration
         self.test_site_config.executor = self.mock_executor.return_value
@@ -122,9 +114,7 @@ class TestMoabAdapter(TestCase):
     @property
     def machine_type_configuration(self):
         return AttributeDict(
-            test2large=AttributeDict(
-                NodeType="1:ppn=20", StartupCommand="startVM.py", Walltime="02:00:00:00"
-            )
+            test2large=AttributeDict(NodeType="1:ppn=20", Walltime="02:00:00:00")
         )
 
     @property
@@ -143,23 +133,6 @@ class TestMoabAdapter(TestCase):
             drone_uuid="testsite-4761849",
         )
 
-    def test_start_up_command_deprecation_warning(self):
-        # Necessary to avoid annoying message in PyCharm
-        filterwarnings(action="ignore", message="unclosed", category=ResourceWarning)
-        del self.test_site_config.MachineTypeConfiguration.test2large.StartupCommand
-
-        with self.assertRaises(AttributeError):
-            self.moab_adapter = MoabAdapter(
-                machine_type="test2large", site_name="TestSite"
-            )
-
-        self.test_site_config.StartupCommand = "startVM.py"
-
-        with self.assertWarns(DeprecationWarning):
-            self.moab_adapter = MoabAdapter(
-                machine_type="test2large", site_name="TestSite"
-            )
-
     @mock_executor_run_command(TEST_DEPLOY_RESOURCE_RESPONSE)
     def test_deploy_resource(self):
         expected_resource_attributes = self.resource_attributes
@@ -310,7 +283,7 @@ class TestMoabAdapter(TestCase):
         new_timestamp = datetime.now() - timedelta(minutes=2)
         self.moab_adapter._moab_status._last_update = new_timestamp
         with self.assertRaises(TardisResourceStatusUpdateFailed):
-            run_async(
+            response = run_async(
                 self.moab_adapter.resource_status,
                 AttributeDict(
                     resource_id=1351043,
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()