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).
Download
Copy to clipboard1 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()
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.
Fixed, thanks.
I think there's a bug in your test_no_append() function... It has:
...which should probably be:
...since 'flag' is not defined within the scope of that function.
fixed, thanks