Merge pull request #2364 from dnephin/handle_signals_properly

Handle SIGTERM/SIGINT properly for up and run

compose/cli/main.py

@@ -368,7 +368,6 @@ class TopLevelCommand(DocoptCommand):
                                   allocates a TTY.
         """
         service = project.get_service(options['SERVICE'])
-
         detach = options['-d']
 
         if IS_WINDOWS_PLATFORM and not detach:
@@ -380,22 +379,6 @@ class TopLevelCommand(DocoptCommand):
         if options['--allow-insecure-ssl']:
             log.warn(INSECURE_SSL_WARNING)
 
-        if not options['--no-deps']:
-            deps = service.get_linked_service_names()
-
-            if len(deps) > 0:
-                project.up(
-                    service_names=deps,
-                    start_deps=True,
-                    strategy=ConvergenceStrategy.never,
-                )
-            elif project.use_networking:
-                project.ensure_network_exists()
-
-        tty = True
-        if detach or options['-T'] or not sys.stdin.isatty():
-            tty = False
-
         if options['COMMAND']:
             command = [options['COMMAND']] + options['ARGS']
         else:
@@ -403,7 +386,7 @@ class TopLevelCommand(DocoptCommand):
 
         container_options = {
             'command': command,
-            'tty': tty,
+            'tty': not (detach or options['-T'] or not sys.stdin.isatty()),
             'stdin_open': not detach,
             'detach': detach,
         }
@@ -435,31 +418,7 @@ class TopLevelCommand(DocoptCommand):
         if options['--name']:
             container_options['name'] = options['--name']
 
-        try:
-            container = service.create_container(
-                quiet=True,
-                one_off=True,
-                **container_options
-            )
-        except APIError as e:
-            legacy.check_for_legacy_containers(
-                project.client,
-                project.name,
-                [service.name],
-                allow_one_off=False,
-            )
-
-            raise e
-
-        if detach:
-            container.start()
-            print(container.name)
-        else:
-            dockerpty.start(project.client, container.id, interactive=not options['-T'])
-            exit_code = container.wait()
-            if options['--rm']:
-                project.client.remove_container(container.id)
-            sys.exit(exit_code)
+        run_one_off_container(container_options, project, service, options)
 
     def scale(self, project, options):
         """
@@ -647,6 +606,58 @@ def convergence_strategy_from_opts(options):
     return ConvergenceStrategy.changed
 
 
+def run_one_off_container(container_options, project, service, options):
+    if not options['--no-deps']:
+        deps = service.get_linked_service_names()
+        if deps:
+            project.up(
+                service_names=deps,
+                start_deps=True,
+                strategy=ConvergenceStrategy.never)
+
+    if project.use_networking:
+        project.ensure_network_exists()
+
+    try:
+        container = service.create_container(
+            quiet=True,
+            one_off=True,
+            **container_options)
+    except APIError:
+        legacy.check_for_legacy_containers(
+            project.client,
+            project.name,
+            [service.name],
+            allow_one_off=False)
+        raise
+
+    if options['-d']:
+        container.start()
+        print(container.name)
+        return
+
+    def remove_container(force=False):
+        if options['--rm']:
+            project.client.remove_container(container.id, force=True)
+
+    def force_shutdown(signal, frame):
+        project.client.kill(container.id)
+        remove_container(force=True)
+        sys.exit(2)
+
+    def shutdown(signal, frame):
+        set_signal_handler(force_shutdown)
+        project.client.stop(container.id)
+        remove_container()
+        sys.exit(1)
+
+    set_signal_handler(shutdown)
+    dockerpty.start(project.client, container.id, interactive=not options['-T'])
+    exit_code = container.wait()
+    remove_container()
+    sys.exit(exit_code)
+
+
 def build_log_printer(containers, service_names, monochrome):
     if service_names:
         containers = [
@@ -657,18 +668,25 @@ def build_log_printer(containers, service_names, monochrome):
 
 
 def attach_to_logs(project, log_printer, service_names, timeout):
-    print("Attaching to", list_containers(log_printer.containers))
-    try:
-        log_printer.run()
-    finally:
-        def handler(signal, frame):
-            project.kill(service_names=service_names)
-            sys.exit(0)
-        signal.signal(signal.SIGINT, handler)
 
+    def force_shutdown(signal, frame):
+        project.kill(service_names=service_names)
+        sys.exit(2)
+
+    def shutdown(signal, frame):
+        set_signal_handler(force_shutdown)
         print("Gracefully stopping... (press Ctrl+C again to force)")
         project.stop(service_names=service_names, timeout=timeout)
 
+    print("Attaching to", list_containers(log_printer.containers))
+    set_signal_handler(shutdown)
+    log_printer.run()
+
+
+def set_signal_handler(handler):
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+
 
 def list_containers(containers):
     return ", ".join(c.name for c in containers)

tests/acceptance/cli_test.py

@@ -2,10 +2,14 @@ from __future__ import absolute_import
 
 import os
 import shlex
+import signal
 import subprocess
+import time
 from collections import namedtuple
 from operator import attrgetter
 
+from docker import errors
+
 from .. import mock
 from compose.cli.command import get_project
 from compose.cli.docker_client import docker_client
@@ -20,6 +24,64 @@ BUILD_CACHE_TEXT = 'Using cache'
 BUILD_PULL_TEXT = 'Status: Image is up to date for busybox:latest'
 
 
+def start_process(base_dir, options):
+    proc = subprocess.Popen(
+        ['docker-compose'] + options,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        cwd=base_dir)
+    print("Running process: %s" % proc.pid)
+    return proc
+
+
+def wait_on_process(proc, returncode=0):
+    stdout, stderr = proc.communicate()
+    if proc.returncode != returncode:
+        print(stderr)
+        assert proc.returncode == returncode
+    return ProcessResult(stdout.decode('utf-8'), stderr.decode('utf-8'))
+
+
+def wait_on_condition(condition, delay=0.1, timeout=5):
+    start_time = time.time()
+    while not condition():
+        if time.time() - start_time > timeout:
+            raise AssertionError("Timeout: %s" % condition)
+        time.sleep(delay)
+
+
+class ContainerCountCondition(object):
+
+    def __init__(self, project, expected):
+        self.project = project
+        self.expected = expected
+
+    def __call__(self):
+        return len(self.project.containers()) == self.expected
+
+    def __str__(self):
+        return "waiting for counter count == %s" % self.expected
+
+
+class ContainerStateCondition(object):
+
+    def __init__(self, client, name, running):
+        self.client = client
+        self.name = name
+        self.running = running
+
+    # State.Running == true
+    def __call__(self):
+        try:
+            container = self.client.inspect_container(self.name)
+            return container['State']['Running'] == self.running
+        except errors.APIError:
+            return False
+
+    def __str__(self):
+        return "waiting for container to have state %s" % self.expected
+
+
 class CLITestCase(DockerClientTestCase):
 
     def setUp(self):
@@ -42,17 +104,8 @@ class CLITestCase(DockerClientTestCase):
 
     def dispatch(self, options, project_options=None, returncode=0):
         project_options = project_options or []
-        proc = subprocess.Popen(
-            ['docker-compose'] + project_options + options,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=self.base_dir)
-        print("Running process: %s" % proc.pid)
-        stdout, stderr = proc.communicate()
-        if proc.returncode != returncode:
-            print(stderr)
-            assert proc.returncode == returncode
-        return ProcessResult(stdout.decode('utf-8'), stderr.decode('utf-8'))
+        proc = start_process(self.base_dir, project_options + options)
+        return wait_on_process(proc, returncode=returncode)
 
     def test_help(self):
         old_base_dir = self.base_dir
@@ -291,7 +344,7 @@ class CLITestCase(DockerClientTestCase):
             returncode=1)
 
     def test_up_with_timeout(self):
-        self.dispatch(['up', '-d', '-t', '1'], None)
+        self.dispatch(['up', '-d', '-t', '1'])
         service = self.project.get_service('simple')
         another = self.project.get_service('another')
         self.assertEqual(len(service.containers()), 1)
@@ -303,6 +356,20 @@ class CLITestCase(DockerClientTestCase):
         self.assertFalse(config['AttachStdout'])
         self.assertFalse(config['AttachStdin'])
 
+    def test_up_handles_sigint(self):
+        proc = start_process(self.base_dir, ['up', '-t', '2'])
+        wait_on_condition(ContainerCountCondition(self.project, 2))
+
+        os.kill(proc.pid, signal.SIGINT)
+        wait_on_condition(ContainerCountCondition(self.project, 0))
+
+    def test_up_handles_sigterm(self):
+        proc = start_process(self.base_dir, ['up', '-t', '2'])
+        wait_on_condition(ContainerCountCondition(self.project, 2))
+
+        os.kill(proc.pid, signal.SIGTERM)
+        wait_on_condition(ContainerCountCondition(self.project, 0))
+
     def test_run_service_without_links(self):
         self.base_dir = 'tests/fixtures/links-composefile'
         self.dispatch(['run', 'console', '/bin/true'])
@@ -508,6 +575,32 @@ class CLITestCase(DockerClientTestCase):
         self.assertEqual(len(networks), 1)
         self.assertEqual(container.human_readable_command, u'true')
 
+    def test_run_handles_sigint(self):
+        proc = start_process(self.base_dir, ['run', '-T', 'simple', 'top'])
+        wait_on_condition(ContainerStateCondition(
+            self.project.client,
+            'simplecomposefile_simple_run_1',
+            running=True))
+
+        os.kill(proc.pid, signal.SIGINT)
+        wait_on_condition(ContainerStateCondition(
+            self.project.client,
+            'simplecomposefile_simple_run_1',
+            running=False))
+
+    def test_run_handles_sigterm(self):
+        proc = start_process(self.base_dir, ['run', '-T', 'simple', 'top'])
+        wait_on_condition(ContainerStateCondition(
+            self.project.client,
+            'simplecomposefile_simple_run_1',
+            running=True))
+
+        os.kill(proc.pid, signal.SIGTERM)
+        wait_on_condition(ContainerStateCondition(
+            self.project.client,
+            'simplecomposefile_simple_run_1',
+            running=False))
+
     def test_rm(self):
         service = self.project.get_service('simple')
         service.create_container()

tests/unit/cli/main_test.py

@@ -57,11 +57,11 @@ class CLIMainTestCase(unittest.TestCase):
         with mock.patch('compose.cli.main.signal', autospec=True) as mock_signal:
             attach_to_logs(project, log_printer, service_names, timeout)
 
-        mock_signal.signal.assert_called_once_with(mock_signal.SIGINT, mock.ANY)
+        assert mock_signal.signal.mock_calls == [
+            mock.call(mock_signal.SIGINT, mock.ANY),
+            mock.call(mock_signal.SIGTERM, mock.ANY),
+        ]
         log_printer.run.assert_called_once_with()
-        project.stop.assert_called_once_with(
-            service_names=service_names,
-            timeout=timeout)
 
 
 class SetupConsoleHandlerTestCase(unittest.TestCase):