Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

157 строки
5.3KB

  1. from __future__ import absolute_import
  2. import six
  3. import gevent
  4. from gevent import queue
  5. from gevent.event import Event
  6. import uwsgi
  7. _websocket_available = hasattr(uwsgi, 'websocket_handshake')
  8. class Thread(gevent.Greenlet): # pragma: no cover
  9. """
  10. This wrapper class provides gevent Greenlet interface that is compatible
  11. with the standard library's Thread class.
  12. """
  13. def __init__(self, target, args=[], kwargs={}):
  14. super(Thread, self).__init__(target, *args, **kwargs)
  15. def _run(self):
  16. return self.run()
  17. class uWSGIWebSocket(object): # pragma: no cover
  18. """
  19. This wrapper class provides a uWSGI WebSocket interface that is
  20. compatible with eventlet's implementation.
  21. """
  22. def __init__(self, app):
  23. self.app = app
  24. self._sock = None
  25. def __call__(self, environ, start_response):
  26. self._sock = uwsgi.connection_fd()
  27. self.environ = environ
  28. uwsgi.websocket_handshake()
  29. self._req_ctx = None
  30. if hasattr(uwsgi, 'request_context'):
  31. # uWSGI >= 2.1.x with support for api access across-greenlets
  32. self._req_ctx = uwsgi.request_context()
  33. else:
  34. # use event and queue for sending messages
  35. from gevent.event import Event
  36. from gevent.queue import Queue
  37. from gevent.select import select
  38. self._event = Event()
  39. self._send_queue = Queue()
  40. # spawn a select greenlet
  41. def select_greenlet_runner(fd, event):
  42. """Sets event when data becomes available to read on fd."""
  43. while True:
  44. event.set()
  45. try:
  46. select([fd], [], [])[0]
  47. except ValueError:
  48. break
  49. self._select_greenlet = gevent.spawn(
  50. select_greenlet_runner,
  51. self._sock,
  52. self._event)
  53. self.app(self)
  54. def close(self):
  55. """Disconnects uWSGI from the client."""
  56. uwsgi.disconnect()
  57. if self._req_ctx is None:
  58. # better kill it here in case wait() is not called again
  59. self._select_greenlet.kill()
  60. self._event.set()
  61. def _send(self, msg):
  62. """Transmits message either in binary or UTF-8 text mode,
  63. depending on its type."""
  64. if isinstance(msg, six.binary_type):
  65. method = uwsgi.websocket_send_binary
  66. else:
  67. method = uwsgi.websocket_send
  68. if self._req_ctx is not None:
  69. method(msg, request_context=self._req_ctx)
  70. else:
  71. method(msg)
  72. def _decode_received(self, msg):
  73. """Returns either bytes or str, depending on message type."""
  74. if not isinstance(msg, six.binary_type):
  75. # already decoded - do nothing
  76. return msg
  77. # only decode from utf-8 if message is not binary data
  78. type = six.byte2int(msg[0:1])
  79. if type >= 48: # no binary
  80. return msg.decode('utf-8')
  81. # binary message, don't try to decode
  82. return msg
  83. def send(self, msg):
  84. """Queues a message for sending. Real transmission is done in
  85. wait method.
  86. Sends directly if uWSGI version is new enough."""
  87. if self._req_ctx is not None:
  88. self._send(msg)
  89. else:
  90. self._send_queue.put(msg)
  91. self._event.set()
  92. def wait(self):
  93. """Waits and returns received messages.
  94. If running in compatibility mode for older uWSGI versions,
  95. it also sends messages that have been queued by send().
  96. A return value of None means that connection was closed.
  97. This must be called repeatedly. For uWSGI < 2.1.x it must
  98. be called from the main greenlet."""
  99. while True:
  100. if self._req_ctx is not None:
  101. try:
  102. msg = uwsgi.websocket_recv(request_context=self._req_ctx)
  103. except IOError: # connection closed
  104. return None
  105. return self._decode_received(msg)
  106. else:
  107. # we wake up at least every 3 seconds to let uWSGI
  108. # do its ping/ponging
  109. event_set = self._event.wait(timeout=3)
  110. if event_set:
  111. self._event.clear()
  112. # maybe there is something to send
  113. msgs = []
  114. while True:
  115. try:
  116. msgs.append(self._send_queue.get(block=False))
  117. except gevent.queue.Empty:
  118. break
  119. for msg in msgs:
  120. self._send(msg)
  121. # maybe there is something to receive, if not, at least
  122. # ensure uWSGI does its ping/ponging
  123. try:
  124. msg = uwsgi.websocket_recv_nb()
  125. except IOError: # connection closed
  126. self._select_greenlet.kill()
  127. return None
  128. if msg: # message available
  129. return self._decode_received(msg)
  130. _async = {
  131. 'thread': Thread,
  132. 'queue': queue.JoinableQueue,
  133. 'queue_empty': queue.Empty,
  134. 'event': Event,
  135. 'websocket': uWSGIWebSocket if _websocket_available else None,
  136. 'sleep': gevent.sleep,
  137. }