Welcome, guest | Sign In | My Account | Store | Cart

A context manager which properly handles SIGTERM (SystemExit) and SIGINT (KeyboardInterrupt) signals, registering a function which is always guaranteed to be called on interpreter exit. Also, it makes sure to execute previously registered functions as well (if any).

Python, 157 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# Author: Giampaolo Rodola' <g.rodola [AT] gmail [DOT] com>
# License: MIT

from __future__ import with_statement
import contextlib
import signal
import sys


def _sigterm_handler(signum, frame):
    sys.exit(0)
_sigterm_handler.__enter_ctx__ = False


@contextlib.contextmanager
def handle_exit(callback=None, append=False):
    """A context manager which properly handles SIGTERM and SIGINT
    (KeyboardInterrupt) signals, registering a function which is
    guaranteed to be called after signals are received.
    Also, it makes sure to execute previously registered signal
    handlers as well (if any).

    >>> app = App()
    >>> with handle_exit(app.stop):
    ...     app.start()
    ...
    >>>

    If append == False raise RuntimeError if there's already a handler
    registered for SIGTERM, otherwise both new and old handlers are
    executed in this order.
    """
    old_handler = signal.signal(signal.SIGTERM, _sigterm_handler)
    if (old_handler != signal.SIG_DFL) and (old_handler != _sigterm_handler):
        if not append:
            raise RuntimeError("there is already a handler registered for "
                               "SIGTERM: %r" % old_handler)

        def handler(signum, frame):
            try:
                _sigterm_handler(signum, frame)
            finally:
                old_handler(signum, frame)
        signal.signal(signal.SIGTERM, handler)

    if _sigterm_handler.__enter_ctx__:
        raise RuntimeError("can't use nested contexts")
    _sigterm_handler.__enter_ctx__ = True

    try:
        yield
    except KeyboardInterrupt:
        pass
    except SystemExit, err:
        # code != 0 refers to an application error (e.g. explicit
        # sys.exit('some error') call).
        # We don't want that to pass silently.
        # Nevertheless, the 'finally' clause below will always
        # be executed.
        if err.code != 0:
            raise
    finally:
        _sigterm_handler.__enter_ctx__ = False
        if callback is not None:
            callback()


if __name__ == '__main__':
    # ===============================================================
    # --- test suite
    # ===============================================================

    import unittest
    import os

    class TestOnExit(unittest.TestCase):

        def setUp(self):
            # reset signal handlers
            signal.signal(signal.SIGTERM, signal.SIG_DFL)
            self.flag = None

        def tearDown(self):
            # make sure we exited the ctx manager
            self.assertTrue(self.flag is not None)

        def test_base(self):
            with handle_exit():
                pass
            self.flag = True

        def test_callback(self):
            callback = []
            with handle_exit(lambda: callback.append(None)):
                pass
            self.flag = True
            self.assertEqual(callback, [None])

        def test_kinterrupt(self):
            with handle_exit():
                raise KeyboardInterrupt
            self.flag = True

        def test_sigterm(self):
            with handle_exit():
                os.kill(os.getpid(), signal.SIGTERM)
            self.flag = True

        def test_sigint(self):
            with handle_exit():
                os.kill(os.getpid(), signal.SIGINT)
            self.flag = True

        def test_sigterm_old(self):
            # make sure the old handler gets executed
            queue = []
            signal.signal(signal.SIGTERM, lambda s, f: queue.append('old'))
            with handle_exit(lambda: queue.append('new'), append=True):
                os.kill(os.getpid(), signal.SIGTERM)
            self.flag = True
            self.assertEqual(queue, ['old', 'new'])

        def test_sigint_old(self):
            # make sure the old handler gets executed
            queue = []
            signal.signal(signal.SIGINT, lambda s, f: queue.append('old'))
            with handle_exit(lambda: queue.append('new'), append=True):
                os.kill(os.getpid(), signal.SIGINT)
            self.flag = True
            self.assertEqual(queue, ['old', 'new'])

        def test_no_append(self):
            # make sure we can't use the context manager if there's
            # already a handler registered for SIGTERM
            signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
            try:
                with handle_exit(lambda: self.flag.append(None)):
                    pass
            except RuntimeError:
                pass
            else:
                self.fail("exception not raised")
            finally:
                self.flag = True

        def test_nested_context(self):
            self.flag = True
            try:
                with handle_exit():
                    with handle_exit():
                        pass
            except RuntimeError:
                pass
            else:
                self.fail("exception not raised")

    unittest.main()

The problem

One trap I often fall into is using atexit module to register an exit function and then discover it does not handle SIGTERM signal by default:

import atexit
import time
import os
import signal

@atexit.register
def cleanup():
    # ==== XXX ====
    # this never gets called
    print "exiting"

def main():
   print "starting"
   time.sleep(1)
   os.kill(os.getpid(), signal.SIGTERM)

if __name__ == '__main__':
    main()

This is documented behavior. From http://docs.python.org/library/atexit.html:

the functions registered via this module are not called when the program is killed by a signal not handled by Python

I understand the rationale behind that, but it's something I usually forget. So, what one would usually do next is a try/finally statement, like this:

import atexit
import time
import os
import signal

@atexit.register
def cleanup():
    # ==== XXX ====
    # this never gets called
    print "exiting"

def main():
    print "starting"
    time.sleep(1)
    os.kill(os.getpid(), signal.SIGTERM)

if __name__ == '__main__':
    try:
        main()
    except (KeyboardInterrupt, SystemExit):
        pass
    finally:
        cleanup()

This code has a bug though: in case of SIGTERM the finally clause is never executed and cleanup() never gets called.

So basically, what's actually needed in order to gracefully shutdown an application and call an exit function is registering your own SIGTERM/SIGINT handler via the signal module.

This approach has yet one downside though: care should be taken in case a third-party module has already registered a handler for SIGTERM/SIGINT signals, because your new handler will overwrite the old one. The proof:

import os
import signal

def old(*args):
    # ==== XXX ====
    # this never gets called
    print "old"

def new(*args):
    print "new"

signal.signal(signal.SIGTERM, old)
signal.signal(signal.SIGTERM, new)
os.kill(os.getpid(), signal.SIGTERM)

The solution

Ok, now you know executing an exit function is not as easy as you imagined. The behavior I would expect when facing such a problem is the following:

  • I want my exit function to always be called, no matter what
    • I'm aware that it won't be called in case of SIGKILL, SIGSTOP or os._exit() though
  • I do not want to override previously registered signal handlers (if any)
    • Instead I want my exit function be called first and the old handler be called last

If you've read this far, you now know what this recipe is about: correctly handling SIGINT/SIGTERM signals and making sure that all registered functions/handlers are executed on exit.

Example usages

Silly example, just to prove it works:

def cleanup():
    print "exiting"

def main():
    import time, os
    print "starting"
    time.sleep(1)
    os.kill(os.getpid(), signal.SIGTERM)

if __name__ == '__main__':
    with handle_exit(cleanup):
        main()

More real world example, using SimpleHTTPServer:

import SimpleHTTPServer
import SocketServer

handler = SimpleHTTPServer.SimpleHTTPRequestHandler
server = SocketServer.TCPServer(("", 8080), handler)
with handle_exit(server.shutdown):
    server.serve_forever()

Starting from python 3.2, thanks to http://docs.python.org/dev/whatsnew/3.2.html#contextlib, this can also be used as a decorator:

def cleanup():
    print("exiting")

@handle_exit(cleanup)
def main():
    import time
    while 1:
        time.sleep(1)

if __name__ == '__main__':
    main()

4 comments

Garron Moore 9 years, 11 months ago  # | flag

I think there is a bug in your append logic. The "handler" function defined in-line only calls the old handler and doesn't call _sigterm_handler.

Giampaolo Rodolà (author) 9 years, 11 months ago  # | flag

Fixed, thanks.

Dan McDougall 8 years, 2 months ago  # | flag

I think there's a bug in your test_no_append() function... It has:

with handle_exit(lambda: flag.append(None)):

...which should probably be:

with handle_exit(lambda: self.flag.append(None)):

...since 'flag' is not defined within the scope of that function.

Giampaolo Rodolà (author) 8 years, 2 months ago  # | flag

fixed, thanks