Merge lp:~abentley/workspace-runner/robust-connections into lp:workspace-runner

[Aaron Bentley]

Make workspace-runner robust against connection failures.

The main purpose of this branch is adding retry support to workspace-runner.

This is done by creating a retry_ssh decorator, and decorating all relevant functions.

It also tries to improve separation of concerns by extracting call_with_input from run_python, and moving run_python back to SSHPrimitives.

Finally, it adds a --verbose option that causes each ssh command to be printed, and retries to be noted.

I experimented with SSH connection sharing, but determined that it does not prevent connection failures. Even after an initial successful connection, subsequent connections could fail with "Received disconnect from 54.165.231.203: 2: fork failed: Resource temporarily unavailable"

[Curtis Hovey (sinzui)]

Thank you.

review: Approve (code)

=== modified file 'workspace_runner/__init__.py'
--- workspace_runner/__init__.py	2015-06-24 20:07:17 +0000
+++ workspace_runner/__init__.py	2015-06-25 15:14:56 +0000
@@ -18,7 +18,9 @@
     parser = ArgumentParser()
     parser.add_argument('config', help='Config file to use.')
     parser.add_argument('host', help='Machine to run the command on.')
-    parser.add_argument('--private-key', '-i', help='Private SSH key to use')
+    parser.add_argument('--private-key', '-i', help='Private SSH key to use.')
+    parser.add_argument('--verbose', '-v', help='Verbose output.',
+                        action='store_true')
     return parser.parse_args(argv)
 
 
@@ -36,6 +38,25 @@
         return cls(ssh_connection, tempdir)
 
 
+def retry_ssh(func):
+    """Decorator that retries a function if it looks like SSH failed.
+
+    SSH connection failures are a CalledProcessError with returncode 255.
+    """
+    def decorated(*args, **kwargs):
+        for x in range(3):
+            if x != 0:
+                logging.info('Retrying due to ssh failure.')
+            try:
+                return func(*args, **kwargs)
+            except subprocess.CalledProcessError as e:
+                if e.returncode != 255:
+                    raise
+        else:
+            raise
+    return decorated
+
+
 class SSHConnection:
     """Class representing an SSH connection.
 
@@ -60,36 +81,34 @@
         command = ['ssh'] + self.get_ssh_options() + [self.host] + args
         return command
 
-    def run_python(self, script, argv=None):
-        """Run a python script using this connection.
-
-        :param script: The text of the script.
-        :param argv: The values to provide as sys.argv in the script.
-        """
-        if argv is None:
-            argv = []
-        python_command = ' '.join(['python', '-'] + [quote(x) for x in argv])
-        command = self.get_ssh_cmd([python_command])
+    @retry_ssh
+    def call_with_input(self, args, stdin):
+        """Run a command remotely, accepting stdin."""
+        remote_command = ' '.join(quote(x) for x in args)
+        command = self.get_ssh_cmd([remote_command])
+        logging.info('Running: {}'.format(' '.join(command)))
         proc = subprocess.Popen(command, stdin=subprocess.PIPE,
                                 stdout=subprocess.PIPE)
-        output = proc.communicate(script)
+        output = proc.communicate(stdin)
         retcode = proc.wait()
         if retcode:
             raise subprocess.CalledProcessError(
                 retcode, command, output=output)
-        return output[0].rstrip('\r\n')
+        return output[0]
 
+    @retry_ssh
     def scp_to_remote(self, sources, dest_path):
         """Copy files to the specified path on the remote host."""
         dest_loc = '{}:{}/'.format(self.host, quote(dest_path))
-        subprocess.check_call(['scp'] + self.get_ssh_options() + sources +
-                              [dest_loc])
+        command = ['scp'] + self.get_ssh_options() + sources + [dest_loc]
+        logging.info('Running: {}'.format(' '.join(command)))
+        subprocess.check_call(command)
 
 
 class SSHPrimitives(Primitives):
 
     def mkdir(self, path):
-        return self.ssh_connection.run_python(dedent("""
+        return self.run_python(self.ssh_connection, dedent("""
             import os
             import sys
 
@@ -99,24 +118,39 @@
     @classmethod
     def mktemp(cls, ssh_connection):
         """Make a temporary directory on the specified host via ssh."""
-        return ssh_connection.run_python(dedent("""
+        return cls.run_python(ssh_connection, dedent("""
             import tempfile
             print tempfile.mkdtemp()
             """))
 
     def destroy(self):
         """Destroy the workspace."""
-        return self.ssh_connection.run_python(dedent("""
+        return self.run_python(self.ssh_connection, dedent("""
             import shutil
             import sys
             shutil.rmtree(sys.argv[1])
             """), [self.workspace])
 
+    @retry_ssh
     def run(self, args, out_file):
         """Run a command in the workspace."""
-        subprocess.check_call(
-            self.ssh_connection.get_ssh_cmd(
-                ['cd', quote(self.workspace), ';'] + args), stdout=out_file)
+        command = self.ssh_connection.get_ssh_cmd(
+            ['cd', quote(self.workspace), ';'] + args)
+        logging.info('Running: {}'.format(' '.join(command)))
+        subprocess.check_call(command, stdout=out_file)
+
+    @staticmethod
+    def run_python(ssh_connection, script, argv=None):
+        """Run a python script using this connection.
+
+        :param script: The text of the script.
+        :param argv: The values to provide as sys.argv in the script.
+        """
+        if argv is None:
+            argv = []
+        python_args = ['python', '-'] + argv
+        output = ssh_connection.call_with_input(python_args, script)
+        return output.rstrip('\r\n')
 
     def install(self, sources, target):
         """Copy files into the workspace."""
@@ -149,6 +183,11 @@
 def workspace_run(argv=None, primitives_factory=SSHPrimitives):
     """Run an operation in a workspace."""
     args = parse_args(argv)
+    if args.verbose:
+        level = logging.INFO
+    else:
+        level = logging.WARNING
+    logging.basicConfig(level=level)
     with open(args.config) as config_file:
         config = safe_load(config_file)
     with workspace_context(args.host, args.private_key,

=== modified file 'workspace_runner/tests/__init__.py'
--- workspace_runner/tests/__init__.py	2015-06-24 18:59:56 +0000
+++ workspace_runner/tests/__init__.py	2015-06-25 15:14:56 +0000
@@ -1,6 +1,7 @@
 from argparse import Namespace
 from contextlib import contextmanager
 import os
+import logging
 from mock import patch
 from pipes import quote
 from shutil import (
@@ -21,6 +22,7 @@
 from workspace_runner import (
     parse_args,
     Primitives,
+    retry_ssh,
     SSHConnection,
     SSHPrimitives,
     workspace_context,
@@ -42,7 +44,7 @@
     def test_minimal(self):
         args = parse_args(['foo', 'bar'])
         self.assertEqual(args, Namespace(config='foo', host='bar',
-                         private_key=None))
+                         private_key=None, verbose=False))
 
     def test_private_key(self):
         args = parse_args(['foo', 'bar', '--private-key', 'key'])
@@ -50,6 +52,12 @@
         args = parse_args(['foo', 'bar', '-i', 'key2'])
         self.assertEqual(args.private_key, 'key2')
 
+    def test_verbose(self):
+        args = parse_args(['foo', 'bar', '--verbose'])
+        self.assertEqual(args.verbose, True)
+        args = parse_args(['foo', 'bar', '-v'])
+        self.assertEqual(args.verbose, True)
+
 
 class FakePrimitives(Primitives):
 
@@ -117,6 +125,57 @@
         return self.last_instance
 
 
+class TestRetrySSH(TestCase):
+
+    def test_retry_until_failure(self):
+        attempts = []
+
+        @retry_ssh
+        def fail_repeatedly():
+            attempts.append('attempt')
+            raise subprocess.CalledProcessError(255, [])
+
+        with self.assertRaises(subprocess.CalledProcessError):
+            fail_repeatedly()
+        self.assertEqual(len(attempts), 3)
+
+    def test_immediate_success(self):
+        attempts = []
+
+        @retry_ssh
+        def succeed():
+            attempts.append('success')
+            return 'success'
+
+        self.assertEqual(succeed(), 'success')
+        self.assertEqual(len(attempts), 1)
+
+    def test_no_retry_non_ssh(self):
+        attempts = []
+
+        @retry_ssh
+        def fail_repeatedly():
+            attempts.append('attempt')
+            raise subprocess.CalledProcessError(254, [])
+
+        with self.assertRaises(subprocess.CalledProcessError):
+            fail_repeatedly()
+        self.assertEqual(len(attempts), 1)
+
+    def test_nearly_fail(self):
+        attempts = []
+
+        @retry_ssh
+        def nearly_fail():
+            attempts.append('attempt')
+            if len(attempts) == 3:
+                return 'success'
+            raise subprocess.CalledProcessError(255, [])
+
+        self.assertEqual(nearly_fail(), 'success')
+        self.assertEqual(len(attempts), 3)
+
+
 class TestSSHConnection(TestCase):
 
     def setUp(self):
@@ -138,12 +197,12 @@
                 '-o', 'UserKnownHostsFile=/dev/null',
                 '-i', 'private-key'])
 
-    def test_run_python_error(self):
+    def test_call_with_input_error(self):
         connection = SSHConnection('host', 'key')
         with patch('subprocess.Popen') as po_mock:
             po_mock.return_value.wait.return_value = 27
             with self.assertRaises(subprocess.CalledProcessError):
-                connection.run_python('')
+                connection.call_with_input([], '')
 
 
 class TestSSHPrimitives(TestPrimitives, TestCase):
@@ -311,7 +370,10 @@
 
     def test_minimal(self):
         with self.config_file() as config_file:
-            primitives = self.run_primitives([config_file.name, 'bar'])
+            with patch('logging.root.handlers', []):
+                primitives = self.run_primitives([config_file.name, 'bar'])
+                self.assertEqual(logging.getLogger().getEffectiveLevel(),
+                                 logging.WARNING)
         self.assertEqual(primitives.run_calls, [['run', 'this']])
         self.assertFalse(os.path.exists(primitives.workspace))
         self.assertEqual(primitives.walk_results,
@@ -341,3 +403,10 @@
             (primitives.workspace, ['bin-dir'], []),
             (bin_path, [], [install_base]),
             ])
+
+    def test_verbose(self):
+        with self.config_file() as config_file:
+            with patch('logging.root.handlers', []):
+                self.run_primitives([config_file.name, 'bar', '-v'])
+                self.assertEqual(logging.getLogger().getEffectiveLevel(),
+                                 logging.INFO)