common.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. # -*- coding: utf-8 -*-
  2. """
  3. This module contains utilities added by billiard, to keep
  4. "non-core" functionality out of ``.util``."""
  5. from __future__ import absolute_import
  6. import os
  7. import signal
  8. import sys
  9. import pickle as pypickle
  10. try:
  11. import cPickle as cpickle
  12. except ImportError: # pragma: no cover
  13. cpickle = None # noqa
  14. from .exceptions import RestartFreqExceeded
  15. from .five import monotonic
  16. if sys.version_info < (2, 6): # pragma: no cover
  17. # cPickle does not use absolute_imports
  18. pickle = pypickle
  19. pickle_load = pypickle.load
  20. pickle_loads = pypickle.loads
  21. else:
  22. pickle = cpickle or pypickle
  23. pickle_load = pickle.load
  24. pickle_loads = pickle.loads
  25. # cPickle.loads does not support buffer() objects,
  26. # but we can just create a StringIO and use load.
  27. if sys.version_info[0] == 3:
  28. from io import BytesIO
  29. else:
  30. try:
  31. from cStringIO import StringIO as BytesIO # noqa
  32. except ImportError:
  33. from StringIO import StringIO as BytesIO # noqa
  34. SIGMAP = dict(
  35. (getattr(signal, n), n) for n in dir(signal) if n.startswith('SIG')
  36. )
  37. for _alias_sig in ('SIGHUP', 'SIGABRT'):
  38. try:
  39. # Alias for deprecated signal overwrites the name we want
  40. SIGMAP[getattr(signal, _alias_sig)] = _alias_sig
  41. except AttributeError:
  42. pass
  43. TERM_SIGNAL, TERM_SIGNAME = signal.SIGTERM, 'SIGTERM'
  44. REMAP_SIGTERM = os.environ.get('REMAP_SIGTERM')
  45. if REMAP_SIGTERM:
  46. TERM_SIGNAL, TERM_SIGNAME = (
  47. getattr(signal, REMAP_SIGTERM), REMAP_SIGTERM)
  48. TERMSIGS_IGNORE = {'SIGTERM'} if REMAP_SIGTERM else set()
  49. TERMSIGS_FORCE = {'SIGQUIT'} if REMAP_SIGTERM else set()
  50. EX_SOFTWARE = 70
  51. TERMSIGS_DEFAULT = {
  52. 'SIGHUP',
  53. 'SIGQUIT',
  54. TERM_SIGNAME,
  55. 'SIGUSR1',
  56. 'SIGUSR2'
  57. }
  58. TERMSIGS_FULL = {
  59. 'SIGHUP',
  60. 'SIGQUIT',
  61. 'SIGTRAP',
  62. 'SIGABRT',
  63. 'SIGEMT',
  64. 'SIGSYS',
  65. 'SIGPIPE',
  66. 'SIGALRM',
  67. TERM_SIGNAME,
  68. 'SIGXCPU',
  69. 'SIGXFSZ',
  70. 'SIGVTALRM',
  71. 'SIGPROF',
  72. 'SIGUSR1',
  73. 'SIGUSR2',
  74. }
  75. #: set by signal handlers just before calling exit.
  76. #: if this is true after the sighandler returns it means that something
  77. #: went wrong while terminating the process, and :func:`os._exit`
  78. #: must be called ASAP.
  79. _should_have_exited = [False]
  80. def human_status(status):
  81. if (status or 0) < 0:
  82. try:
  83. return 'signal {0} ({1})'.format(-status, SIGMAP[-status])
  84. except KeyError:
  85. return 'signal {0}'.format(-status)
  86. return 'exitcode {0}'.format(status)
  87. def pickle_loads(s, load=pickle_load):
  88. # used to support buffer objects
  89. return load(BytesIO(s))
  90. def maybe_setsignal(signum, handler):
  91. try:
  92. signal.signal(signum, handler)
  93. except (OSError, AttributeError, ValueError, RuntimeError):
  94. pass
  95. def _shutdown_cleanup(signum, frame):
  96. # we will exit here so if the signal is received a second time
  97. # we can be sure that something is very wrong and we may be in
  98. # a crashing loop.
  99. if _should_have_exited[0]:
  100. os._exit(EX_SOFTWARE)
  101. maybe_setsignal(signum, signal.SIG_DFL)
  102. _should_have_exited[0] = True
  103. sys.exit(-(256 - signum))
  104. def signum(sig):
  105. return getattr(signal, sig, None)
  106. def _should_override_term_signal(sig, current):
  107. return (
  108. sig in TERMSIGS_FORCE or
  109. (current is not None and current != signal.SIG_IGN)
  110. )
  111. def reset_signals(handler=_shutdown_cleanup, full=False):
  112. for sig in TERMSIGS_FULL if full else TERMSIGS_DEFAULT:
  113. num = signum(sig)
  114. if num:
  115. if _should_override_term_signal(sig, signal.getsignal(num)):
  116. maybe_setsignal(num, handler)
  117. for sig in TERMSIGS_IGNORE:
  118. num = signum(sig)
  119. if num:
  120. maybe_setsignal(num, signal.SIG_IGN)
  121. class restart_state(object):
  122. RestartFreqExceeded = RestartFreqExceeded
  123. def __init__(self, maxR, maxT):
  124. self.maxR, self.maxT = maxR, maxT
  125. self.R, self.T = 0, None
  126. def step(self, now=None):
  127. now = monotonic() if now is None else now
  128. R = self.R
  129. if self.T and now - self.T >= self.maxT:
  130. # maxT passed, reset counter and time passed.
  131. self.T, self.R = now, 0
  132. elif self.maxR and self.R >= self.maxR:
  133. # verify that R has a value as the result handler
  134. # resets this when a job is accepted. If a job is accepted
  135. # the startup probably went fine (startup restart burst
  136. # protection)
  137. if self.R: # pragma: no cover
  138. self.R = 0 # reset in case someone catches the error
  139. raise self.RestartFreqExceeded("%r in %rs" % (R, self.maxT))
  140. # first run sets T
  141. if self.T is None:
  142. self.T = now
  143. self.R += 1