Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

rdb.py 4.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.contrib.rdb
  4. ==================
  5. Remote debugger for Celery tasks running in multiprocessing pool workers.
  6. Inspired by http://snippets.dzone.com/posts/show/7248
  7. **Usage**
  8. .. code-block:: python
  9. from celery.contrib import rdb
  10. from celery import task
  11. @task()
  12. def add(x, y):
  13. result = x + y
  14. rdb.set_trace()
  15. return result
  16. **Environment Variables**
  17. .. envvar:: CELERY_RDB_HOST
  18. Hostname to bind to. Default is '127.0.01', which means the socket
  19. will only be accessible from the local host.
  20. .. envvar:: CELERY_RDB_PORT
  21. Base port to bind to. Default is 6899.
  22. The debugger will try to find an available port starting from the
  23. base port. The selected port will be logged by the worker.
  24. """
  25. from __future__ import absolute_import, print_function
  26. import errno
  27. import os
  28. import socket
  29. import sys
  30. from pdb import Pdb
  31. from billiard import current_process
  32. from celery.five import range
  33. __all__ = ['CELERY_RDB_HOST', 'CELERY_RDB_PORT', 'default_port',
  34. 'Rdb', 'debugger', 'set_trace']
  35. default_port = 6899
  36. CELERY_RDB_HOST = os.environ.get('CELERY_RDB_HOST') or '127.0.0.1'
  37. CELERY_RDB_PORT = int(os.environ.get('CELERY_RDB_PORT') or default_port)
  38. #: Holds the currently active debugger.
  39. _current = [None]
  40. _frame = getattr(sys, '_getframe')
  41. NO_AVAILABLE_PORT = """\
  42. {self.ident}: Couldn't find an available port.
  43. Please specify one using the CELERY_RDB_PORT environment variable.
  44. """
  45. BANNER = """\
  46. {self.ident}: Please telnet into {self.host} {self.port}.
  47. Type `exit` in session to continue.
  48. {self.ident}: Waiting for client...
  49. """
  50. SESSION_STARTED = '{self.ident}: Now in session with {self.remote_addr}.'
  51. SESSION_ENDED = '{self.ident}: Session with {self.remote_addr} ended.'
  52. class Rdb(Pdb):
  53. me = 'Remote Debugger'
  54. _prev_outs = None
  55. _sock = None
  56. def __init__(self, host=CELERY_RDB_HOST, port=CELERY_RDB_PORT,
  57. port_search_limit=100, port_skew=+0, out=sys.stdout):
  58. self.active = True
  59. self.out = out
  60. self._prev_handles = sys.stdin, sys.stdout
  61. self._sock, this_port = self.get_avail_port(
  62. host, port, port_search_limit, port_skew,
  63. )
  64. self._sock.setblocking(1)
  65. self._sock.listen(1)
  66. self.ident = '{0}:{1}'.format(self.me, this_port)
  67. self.host = host
  68. self.port = this_port
  69. self.say(BANNER.format(self=self))
  70. self._client, address = self._sock.accept()
  71. self._client.setblocking(1)
  72. self.remote_addr = ':'.join(str(v) for v in address)
  73. self.say(SESSION_STARTED.format(self=self))
  74. self._handle = sys.stdin = sys.stdout = self._client.makefile('rw')
  75. Pdb.__init__(self, completekey='tab',
  76. stdin=self._handle, stdout=self._handle)
  77. def get_avail_port(self, host, port, search_limit=100, skew=+0):
  78. try:
  79. _, skew = current_process().name.split('-')
  80. skew = int(skew)
  81. except ValueError:
  82. pass
  83. this_port = None
  84. for i in range(search_limit):
  85. _sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  86. this_port = port + skew + i
  87. try:
  88. _sock.bind((host, this_port))
  89. except socket.error as exc:
  90. if exc.errno in [errno.EADDRINUSE, errno.EINVAL]:
  91. continue
  92. raise
  93. else:
  94. return _sock, this_port
  95. else:
  96. raise Exception(NO_AVAILABLE_PORT.format(self=self))
  97. def say(self, m):
  98. print(m, file=self.out)
  99. def __enter__(self):
  100. return self
  101. def __exit__(self, *exc_info):
  102. self._close_session()
  103. def _close_session(self):
  104. self.stdin, self.stdout = sys.stdin, sys.stdout = self._prev_handles
  105. if self.active:
  106. if self._handle is not None:
  107. self._handle.close()
  108. if self._client is not None:
  109. self._client.close()
  110. if self._sock is not None:
  111. self._sock.close()
  112. self.active = False
  113. self.say(SESSION_ENDED.format(self=self))
  114. def do_continue(self, arg):
  115. self._close_session()
  116. self.set_continue()
  117. return 1
  118. do_c = do_cont = do_continue
  119. def do_quit(self, arg):
  120. self._close_session()
  121. self.set_quit()
  122. return 1
  123. do_q = do_exit = do_quit
  124. def set_quit(self):
  125. # this raises a BdbQuit exception that we are unable to catch.
  126. sys.settrace(None)
  127. def debugger():
  128. """Return the current debugger instance (if any),
  129. or creates a new one."""
  130. rdb = _current[0]
  131. if rdb is None or not rdb.active:
  132. rdb = _current[0] = Rdb()
  133. return rdb
  134. def set_trace(frame=None):
  135. """Set breakpoint at current location, or a specified frame"""
  136. if frame is None:
  137. frame = _frame().f_back
  138. return debugger().set_trace(frame)