Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

164 lines
6.5KB

  1. from functools import partial
  2. import uuid
  3. import json
  4. import pickle
  5. import six
  6. from .asyncio_manager import AsyncManager
  7. class AsyncPubSubManager(AsyncManager):
  8. """Manage a client list attached to a pub/sub backend under asyncio.
  9. This is a base class that enables multiple servers to share the list of
  10. clients, with the servers communicating events through a pub/sub backend.
  11. The use of a pub/sub backend also allows any client connected to the
  12. backend to emit events addressed to Socket.IO clients.
  13. The actual backends must be implemented by subclasses, this class only
  14. provides a pub/sub generic framework for asyncio applications.
  15. :param channel: The channel name on which the server sends and receives
  16. notifications.
  17. """
  18. name = 'asyncpubsub'
  19. def __init__(self, channel='socketio', write_only=False, logger=None):
  20. super().__init__()
  21. self.channel = channel
  22. self.write_only = write_only
  23. self.host_id = uuid.uuid4().hex
  24. self.logger = logger
  25. def initialize(self):
  26. super().initialize()
  27. if not self.write_only:
  28. self.thread = self.server.start_background_task(self._thread)
  29. self._get_logger().info(self.name + ' backend initialized.')
  30. async def emit(self, event, data, namespace=None, room=None, skip_sid=None,
  31. callback=None, **kwargs):
  32. """Emit a message to a single client, a room, or all the clients
  33. connected to the namespace.
  34. This method takes care or propagating the message to all the servers
  35. that are connected through the message queue.
  36. The parameters are the same as in :meth:`.Server.emit`.
  37. Note: this method is a coroutine.
  38. """
  39. if kwargs.get('ignore_queue'):
  40. return await super().emit(
  41. event, data, namespace=namespace, room=room, skip_sid=skip_sid,
  42. callback=callback)
  43. namespace = namespace or '/'
  44. if callback is not None:
  45. if self.server is None:
  46. raise RuntimeError('Callbacks can only be issued from the '
  47. 'context of a server.')
  48. if room is None:
  49. raise ValueError('Cannot use callback without a room set.')
  50. id = self._generate_ack_id(room, namespace, callback)
  51. callback = (room, namespace, id)
  52. else:
  53. callback = None
  54. await self._publish({'method': 'emit', 'event': event, 'data': data,
  55. 'namespace': namespace, 'room': room,
  56. 'skip_sid': skip_sid, 'callback': callback,
  57. 'host_id': self.host_id})
  58. async def close_room(self, room, namespace=None):
  59. await self._publish({'method': 'close_room', 'room': room,
  60. 'namespace': namespace or '/'})
  61. async def _publish(self, data):
  62. """Publish a message on the Socket.IO channel.
  63. This method needs to be implemented by the different subclasses that
  64. support pub/sub backends.
  65. """
  66. raise NotImplementedError('This method must be implemented in a '
  67. 'subclass.') # pragma: no cover
  68. async def _listen(self):
  69. """Return the next message published on the Socket.IO channel,
  70. blocking until a message is available.
  71. This method needs to be implemented by the different subclasses that
  72. support pub/sub backends.
  73. """
  74. raise NotImplementedError('This method must be implemented in a '
  75. 'subclass.') # pragma: no cover
  76. async def _handle_emit(self, message):
  77. # Events with callbacks are very tricky to handle across hosts
  78. # Here in the receiving end we set up a local callback that preserves
  79. # the callback host and id from the sender
  80. remote_callback = message.get('callback')
  81. remote_host_id = message.get('host_id')
  82. if remote_callback is not None and len(remote_callback) == 3:
  83. callback = partial(self._return_callback, remote_host_id,
  84. *remote_callback)
  85. else:
  86. callback = None
  87. await super().emit(message['event'], message['data'],
  88. namespace=message.get('namespace'),
  89. room=message.get('room'),
  90. skip_sid=message.get('skip_sid'),
  91. callback=callback)
  92. async def _handle_callback(self, message):
  93. if self.host_id == message.get('host_id'):
  94. try:
  95. sid = message['sid']
  96. namespace = message['namespace']
  97. id = message['id']
  98. args = message['args']
  99. except KeyError:
  100. return
  101. await self.trigger_callback(sid, namespace, id, args)
  102. async def _return_callback(self, host_id, sid, namespace, callback_id,
  103. *args):
  104. # When an event callback is received, the callback is returned back
  105. # the sender, which is identified by the host_id
  106. await self._publish({'method': 'callback', 'host_id': host_id,
  107. 'sid': sid, 'namespace': namespace,
  108. 'id': callback_id, 'args': args})
  109. async def _handle_close_room(self, message):
  110. await super().close_room(
  111. room=message.get('room'), namespace=message.get('namespace'))
  112. async def _thread(self):
  113. while True:
  114. try:
  115. message = await self._listen()
  116. except:
  117. import traceback
  118. traceback.print_exc()
  119. break
  120. data = None
  121. if isinstance(message, dict):
  122. data = message
  123. else:
  124. if isinstance(message, six.binary_type): # pragma: no cover
  125. try:
  126. data = pickle.loads(message)
  127. except:
  128. pass
  129. if data is None:
  130. try:
  131. data = json.loads(message)
  132. except:
  133. pass
  134. if data and 'method' in data:
  135. if data['method'] == 'emit':
  136. await self._handle_emit(data)
  137. elif data['method'] == 'callback':
  138. await self._handle_callback(data)
  139. elif data['method'] == 'close_room':
  140. await self._handle_close_room(data)