Merge pull request #2585 from dnephin/short_signal_handlers

Fix signal handlers by moving shutdown logic out of handler

compose/cli/main.py

@@ -5,7 +5,6 @@ from __future__ import unicode_literals
 import json
 import logging
 import re
-import signal
 import sys
 from inspect import getdoc
 from operator import attrgetter
@@ -14,6 +13,7 @@ import yaml
 from docker.errors import APIError
 from requests.exceptions import ReadTimeout
 
+from . import signals
 from .. import __version__
 from ..config import config
 from ..config import ConfigurationError
@@ -682,20 +682,19 @@ def run_one_off_container(container_options, project, service, options):
         if options['--rm']:
             project.client.remove_container(container.id, force=True)
 
-    def force_shutdown(signal, frame):
+    signals.set_signal_handler_to_shutdown()
+    try:
+        try:
+            dockerpty.start(project.client, container.id, interactive=not options['-T'])
+            exit_code = container.wait()
+        except signals.ShutdownException:
+            project.client.stop(container.id)
+            exit_code = 1
+    except signals.ShutdownException:
         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)
 
@@ -710,25 +709,19 @@ 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))
+    signals.set_signal_handler_to_shutdown()
 
-    def force_shutdown(signal, frame):
+    try:
+        try:
+            log_printer.run()
+        except signals.ShutdownException:
+            print("Gracefully stopping... (press Ctrl+C again to force)")
+            project.stop(service_names=service_names, timeout=timeout)
+    except signals.ShutdownException:
         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)

compose/cli/signals.py

@@ -0,0 +1,18 @@
+import signal
+
+
+class ShutdownException(Exception):
+    pass
+
+
+def shutdown(signal, frame):
+    raise ShutdownException()
+
+
+def set_signal_handler(handler):
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+
+
+def set_signal_handler_to_shutdown():
+    set_signal_handler(shutdown)

tests/acceptance/cli_test.py

@@ -90,7 +90,8 @@ class ContainerStateCondition(object):
             return False
 
     def __str__(self):
-        return "waiting for container to have state %s" % self.expected
+        state = 'running' if self.running else 'stopped'
+        return "waiting for container to be %s" % state
 
 
 class CLITestCase(DockerClientTestCase):

tests/unit/cli/main_test.py

@@ -55,7 +55,7 @@ class CLIMainTestCase(unittest.TestCase):
         service_names = ['web', 'db']
         timeout = 12
 
-        with mock.patch('compose.cli.main.signal', autospec=True) as mock_signal:
+        with mock.patch('compose.cli.main.signals.signal', autospec=True) as mock_signal:
             attach_to_logs(project, log_printer, service_names, timeout)
 
         assert mock_signal.signal.mock_calls == [