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.

redis.py 35KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023
  1. """
  2. kombu.transport.redis
  3. =====================
  4. Redis transport.
  5. """
  6. from __future__ import absolute_import
  7. import numbers
  8. import socket
  9. from bisect import bisect
  10. from collections import namedtuple
  11. from contextlib import contextmanager
  12. from time import time
  13. from amqp import promise
  14. from anyjson import loads, dumps
  15. from kombu.exceptions import InconsistencyError, VersionMismatch
  16. from kombu.five import Empty, values, string_t
  17. from kombu.log import get_logger
  18. from kombu.utils import cached_property, uuid
  19. from kombu.utils.eventio import poll, READ, ERR
  20. from kombu.utils.encoding import bytes_to_str
  21. from kombu.utils.url import _parse_url
  22. from . import virtual
  23. try:
  24. from billiard.util import register_after_fork
  25. except ImportError: # pragma: no cover
  26. try:
  27. from multiprocessing.util import register_after_fork # noqa
  28. except ImportError:
  29. def register_after_fork(*args, **kwargs): # noqa
  30. pass
  31. try:
  32. import redis
  33. except ImportError: # pragma: no cover
  34. redis = None # noqa
  35. logger = get_logger('kombu.transport.redis')
  36. crit, warn = logger.critical, logger.warn
  37. DEFAULT_PORT = 6379
  38. DEFAULT_DB = 0
  39. PRIORITY_STEPS = [0, 3, 6, 9]
  40. error_classes_t = namedtuple('error_classes_t', (
  41. 'connection_errors', 'channel_errors',
  42. ))
  43. NO_ROUTE_ERROR = """
  44. Cannot route message for exchange {0!r}: Table empty or key no longer exists.
  45. Probably the key ({1!r}) has been removed from the Redis database.
  46. """
  47. # This implementation may seem overly complex, but I assure you there is
  48. # a good reason for doing it this way.
  49. #
  50. # Consuming from several connections enables us to emulate channels,
  51. # which means we can have different service guarantees for individual
  52. # channels.
  53. #
  54. # So we need to consume messages from multiple connections simultaneously,
  55. # and using epoll means we don't have to do so using multiple threads.
  56. #
  57. # Also it means we can easily use PUBLISH/SUBSCRIBE to do fanout
  58. # exchanges (broadcast), as an alternative to pushing messages to fanout-bound
  59. # queues manually.
  60. def get_redis_error_classes():
  61. from redis import exceptions
  62. # This exception suddenly changed name between redis-py versions
  63. if hasattr(exceptions, 'InvalidData'):
  64. DataError = exceptions.InvalidData
  65. else:
  66. DataError = exceptions.DataError
  67. return error_classes_t(
  68. (virtual.Transport.connection_errors + tuple(filter(None, (
  69. InconsistencyError,
  70. socket.error,
  71. IOError,
  72. OSError,
  73. exceptions.ConnectionError,
  74. exceptions.AuthenticationError,
  75. getattr(exceptions, 'TimeoutError', None))))),
  76. (virtual.Transport.channel_errors + (
  77. DataError,
  78. exceptions.InvalidResponse,
  79. exceptions.ResponseError)),
  80. )
  81. def get_redis_ConnectionError():
  82. from redis import exceptions
  83. return exceptions.ConnectionError
  84. class MutexHeld(Exception):
  85. pass
  86. @contextmanager
  87. def Mutex(client, name, expire):
  88. lock_id = uuid()
  89. i_won = client.setnx(name, lock_id)
  90. try:
  91. if i_won:
  92. client.expire(name, expire)
  93. yield
  94. else:
  95. if not client.ttl(name):
  96. client.expire(name, expire)
  97. raise MutexHeld()
  98. finally:
  99. if i_won:
  100. try:
  101. with client.pipeline(True) as pipe:
  102. pipe.watch(name)
  103. if pipe.get(name) == lock_id:
  104. pipe.multi()
  105. pipe.delete(name)
  106. pipe.execute()
  107. pipe.unwatch()
  108. except redis.WatchError:
  109. pass
  110. class QoS(virtual.QoS):
  111. restore_at_shutdown = True
  112. def __init__(self, *args, **kwargs):
  113. super(QoS, self).__init__(*args, **kwargs)
  114. self._vrestore_count = 0
  115. def append(self, message, delivery_tag):
  116. delivery = message.delivery_info
  117. EX, RK = delivery['exchange'], delivery['routing_key']
  118. with self.pipe_or_acquire() as pipe:
  119. pipe.zadd(self.unacked_index_key, delivery_tag, time()) \
  120. .hset(self.unacked_key, delivery_tag,
  121. dumps([message._raw, EX, RK])) \
  122. .execute()
  123. super(QoS, self).append(message, delivery_tag)
  124. def restore_unacked(self, client=None):
  125. with self.channel.conn_or_acquire(client) as client:
  126. for tag in self._delivered:
  127. self.restore_by_tag(tag, client=client)
  128. self._delivered.clear()
  129. def ack(self, delivery_tag):
  130. self._remove_from_indices(delivery_tag).execute()
  131. super(QoS, self).ack(delivery_tag)
  132. def reject(self, delivery_tag, requeue=False):
  133. if requeue:
  134. self.restore_by_tag(delivery_tag, leftmost=True)
  135. self.ack(delivery_tag)
  136. @contextmanager
  137. def pipe_or_acquire(self, pipe=None, client=None):
  138. if pipe:
  139. yield pipe
  140. else:
  141. with self.channel.conn_or_acquire(client) as client:
  142. yield client.pipeline()
  143. def _remove_from_indices(self, delivery_tag, pipe=None):
  144. with self.pipe_or_acquire(pipe) as pipe:
  145. return pipe.zrem(self.unacked_index_key, delivery_tag) \
  146. .hdel(self.unacked_key, delivery_tag)
  147. def restore_visible(self, start=0, num=10, interval=10):
  148. self._vrestore_count += 1
  149. if (self._vrestore_count - 1) % interval:
  150. return
  151. with self.channel.conn_or_acquire() as client:
  152. ceil = time() - self.visibility_timeout
  153. try:
  154. with Mutex(client, self.unacked_mutex_key,
  155. self.unacked_mutex_expire):
  156. visible = client.zrevrangebyscore(
  157. self.unacked_index_key, ceil, 0,
  158. start=num and start, num=num, withscores=True)
  159. for tag, score in visible or []:
  160. self.restore_by_tag(tag, client)
  161. except MutexHeld:
  162. pass
  163. def restore_by_tag(self, tag, client=None, leftmost=False):
  164. with self.channel.conn_or_acquire(client) as client:
  165. with client.pipeline() as pipe:
  166. p, _, _ = self._remove_from_indices(
  167. tag, pipe.hget(self.unacked_key, tag)).execute()
  168. if p:
  169. M, EX, RK = loads(bytes_to_str(p)) # json is unicode
  170. self.channel._do_restore_message(M, EX, RK, client, leftmost)
  171. @cached_property
  172. def unacked_key(self):
  173. return self.channel.unacked_key
  174. @cached_property
  175. def unacked_index_key(self):
  176. return self.channel.unacked_index_key
  177. @cached_property
  178. def unacked_mutex_key(self):
  179. return self.channel.unacked_mutex_key
  180. @cached_property
  181. def unacked_mutex_expire(self):
  182. return self.channel.unacked_mutex_expire
  183. @cached_property
  184. def visibility_timeout(self):
  185. return self.channel.visibility_timeout
  186. class MultiChannelPoller(object):
  187. eventflags = READ | ERR
  188. #: Set by :meth:`get` while reading from the socket.
  189. _in_protected_read = False
  190. #: Set of one-shot callbacks to call after reading from socket.
  191. after_read = None
  192. def __init__(self):
  193. # active channels
  194. self._channels = set()
  195. # file descriptor -> channel map.
  196. self._fd_to_chan = {}
  197. # channel -> socket map
  198. self._chan_to_sock = {}
  199. # poll implementation (epoll/kqueue/select)
  200. self.poller = poll()
  201. # one-shot callbacks called after reading from socket.
  202. self.after_read = set()
  203. def close(self):
  204. for fd in values(self._chan_to_sock):
  205. try:
  206. self.poller.unregister(fd)
  207. except (KeyError, ValueError):
  208. pass
  209. self._channels.clear()
  210. self._fd_to_chan.clear()
  211. self._chan_to_sock.clear()
  212. def add(self, channel):
  213. self._channels.add(channel)
  214. def discard(self, channel):
  215. self._channels.discard(channel)
  216. def _on_connection_disconnect(self, connection):
  217. try:
  218. self.poller.unregister(connection._sock)
  219. except (AttributeError, TypeError):
  220. pass
  221. def _register(self, channel, client, type):
  222. if (channel, client, type) in self._chan_to_sock:
  223. self._unregister(channel, client, type)
  224. if client.connection._sock is None: # not connected yet.
  225. client.connection.connect()
  226. sock = client.connection._sock
  227. self._fd_to_chan[sock.fileno()] = (channel, type)
  228. self._chan_to_sock[(channel, client, type)] = sock
  229. self.poller.register(sock, self.eventflags)
  230. def _unregister(self, channel, client, type):
  231. self.poller.unregister(self._chan_to_sock[(channel, client, type)])
  232. def _register_BRPOP(self, channel):
  233. """enable BRPOP mode for channel."""
  234. ident = channel, channel.client, 'BRPOP'
  235. if channel.client.connection._sock is None or \
  236. ident not in self._chan_to_sock:
  237. channel._in_poll = False
  238. self._register(*ident)
  239. if not channel._in_poll: # send BRPOP
  240. channel._brpop_start()
  241. def _register_LISTEN(self, channel):
  242. """enable LISTEN mode for channel."""
  243. if channel.subclient.connection._sock is None:
  244. channel._in_listen = False
  245. self._register(channel, channel.subclient, 'LISTEN')
  246. if not channel._in_listen:
  247. channel._subscribe() # send SUBSCRIBE
  248. def on_poll_start(self):
  249. for channel in self._channels:
  250. if channel.active_queues: # BRPOP mode?
  251. if channel.qos.can_consume():
  252. self._register_BRPOP(channel)
  253. if channel.active_fanout_queues: # LISTEN mode?
  254. self._register_LISTEN(channel)
  255. def on_poll_init(self, poller):
  256. self.poller = poller
  257. for channel in self._channels:
  258. return channel.qos.restore_visible(
  259. num=channel.unacked_restore_limit,
  260. )
  261. def maybe_restore_messages(self):
  262. for channel in self._channels:
  263. if channel.active_queues:
  264. # only need to do this once, as they are not local to channel.
  265. return channel.qos.restore_visible(
  266. num=channel.unacked_restore_limit,
  267. )
  268. def on_readable(self, fileno):
  269. try:
  270. chan, type = self._fd_to_chan[fileno]
  271. except KeyError:
  272. return
  273. if chan.qos.can_consume():
  274. return chan.handlers[type]()
  275. def handle_event(self, fileno, event):
  276. if event & READ:
  277. return self.on_readable(fileno), self
  278. elif event & ERR:
  279. chan, type = self._fd_to_chan[fileno]
  280. chan._poll_error(type)
  281. def get(self, timeout=None):
  282. self._in_protected_read = True
  283. try:
  284. for channel in self._channels:
  285. if channel.active_queues: # BRPOP mode?
  286. if channel.qos.can_consume():
  287. self._register_BRPOP(channel)
  288. if channel.active_fanout_queues: # LISTEN mode?
  289. self._register_LISTEN(channel)
  290. events = self.poller.poll(timeout)
  291. for fileno, event in events or []:
  292. ret = self.handle_event(fileno, event)
  293. if ret:
  294. return ret
  295. # - no new data, so try to restore messages.
  296. # - reset active redis commands.
  297. self.maybe_restore_messages()
  298. raise Empty()
  299. finally:
  300. self._in_protected_read = False
  301. while self.after_read:
  302. try:
  303. fun = self.after_read.pop()
  304. except KeyError:
  305. break
  306. else:
  307. fun()
  308. @property
  309. def fds(self):
  310. return self._fd_to_chan
  311. class Channel(virtual.Channel):
  312. QoS = QoS
  313. _client = None
  314. _subclient = None
  315. _closing = False
  316. supports_fanout = True
  317. keyprefix_queue = '_kombu.binding.%s'
  318. keyprefix_fanout = '/{db}.'
  319. sep = '\x06\x16'
  320. _in_poll = False
  321. _in_listen = False
  322. _fanout_queues = {}
  323. ack_emulation = True
  324. unacked_key = 'unacked'
  325. unacked_index_key = 'unacked_index'
  326. unacked_mutex_key = 'unacked_mutex'
  327. unacked_mutex_expire = 300 # 5 minutes
  328. unacked_restore_limit = None
  329. visibility_timeout = 3600 # 1 hour
  330. priority_steps = PRIORITY_STEPS
  331. socket_timeout = None
  332. socket_connect_timeout = None
  333. socket_keepalive = None
  334. socket_keepalive_options = None
  335. max_connections = 10
  336. #: Transport option to enable disable fanout keyprefix.
  337. #: Should be enabled by default, but that is not
  338. #: backwards compatible. Can also be string, in which
  339. #: case it changes the default prefix ('/{db}.') into to something
  340. #: else. The prefix must include a leading slash and a trailing dot.
  341. fanout_prefix = False
  342. #: If enabled the fanout exchange will support patterns in routing
  343. #: and binding keys (like a topic exchange but using PUB/SUB).
  344. #: This will be enabled by default in a future version.
  345. fanout_patterns = False
  346. _async_pool = None
  347. _pool = None
  348. _disconnecting_pools = False
  349. from_transport_options = (
  350. virtual.Channel.from_transport_options +
  351. ('ack_emulation',
  352. 'unacked_key',
  353. 'unacked_index_key',
  354. 'unacked_mutex_key',
  355. 'unacked_mutex_expire',
  356. 'visibility_timeout',
  357. 'unacked_restore_limit',
  358. 'fanout_prefix',
  359. 'fanout_patterns',
  360. 'socket_timeout',
  361. 'socket_connect_timeout',
  362. 'socket_keepalive',
  363. 'socket_keepalive_options',
  364. 'queue_order_strategy',
  365. 'max_connections',
  366. 'priority_steps') # <-- do not add comma here!
  367. )
  368. def __init__(self, *args, **kwargs):
  369. super_ = super(Channel, self)
  370. super_.__init__(*args, **kwargs)
  371. if not self.ack_emulation: # disable visibility timeout
  372. self.QoS = virtual.QoS
  373. self._queue_cycle = []
  374. self.AsyncClient = self._get_async_client()
  375. self.Client = redis.Redis
  376. self.ResponseError = self._get_response_error()
  377. self.active_fanout_queues = set()
  378. self.auto_delete_queues = set()
  379. self._fanout_to_queue = {}
  380. self.handlers = {'BRPOP': self._brpop_read, 'LISTEN': self._receive}
  381. if self.fanout_prefix:
  382. if isinstance(self.fanout_prefix, string_t):
  383. self.keyprefix_fanout = self.fanout_prefix
  384. else:
  385. # previous versions did not set a fanout, so cannot enable
  386. # by default.
  387. self.keyprefix_fanout = ''
  388. # Evaluate connection.
  389. try:
  390. self.client.info()
  391. except Exception:
  392. self._disconnect_pools()
  393. raise
  394. self.connection.cycle.add(self) # add to channel poller.
  395. # copy errors, in case channel closed but threads still
  396. # are still waiting for data.
  397. self.connection_errors = self.connection.connection_errors
  398. register_after_fork(self, self._after_fork)
  399. def _after_fork(self):
  400. self._disconnect_pools()
  401. def _disconnect_pools(self):
  402. if not self._disconnecting_pools:
  403. self._disconnecting_pools = True
  404. try:
  405. if self._async_pool is not None:
  406. self._async_pool.disconnect()
  407. if self._pool is not None:
  408. self._pool.disconnect()
  409. self._async_pool = self._pool = None
  410. finally:
  411. self._disconnecting_pools = False
  412. def _on_connection_disconnect(self, connection):
  413. self._in_poll = False
  414. self._in_listen = False
  415. if self.connection and self.connection.cycle:
  416. self.connection.cycle._on_connection_disconnect(connection)
  417. self._disconnect_pools()
  418. if not self._closing:
  419. raise get_redis_ConnectionError()
  420. def _do_restore_message(self, payload, exchange, routing_key,
  421. client=None, leftmost=False):
  422. with self.conn_or_acquire(client) as client:
  423. try:
  424. try:
  425. payload['headers']['redelivered'] = True
  426. except KeyError:
  427. pass
  428. for queue in self._lookup(exchange, routing_key):
  429. (client.lpush if leftmost else client.rpush)(
  430. queue, dumps(payload),
  431. )
  432. except Exception:
  433. crit('Could not restore message: %r', payload, exc_info=True)
  434. def _restore(self, message, leftmost=False):
  435. if not self.ack_emulation:
  436. return super(Channel, self)._restore(message)
  437. tag = message.delivery_tag
  438. with self.conn_or_acquire() as client:
  439. with client.pipeline() as pipe:
  440. P, _ = pipe.hget(self.unacked_key, tag) \
  441. .hdel(self.unacked_key, tag) \
  442. .execute()
  443. if P:
  444. M, EX, RK = loads(bytes_to_str(P)) # json is unicode
  445. self._do_restore_message(M, EX, RK, client, leftmost)
  446. def _restore_at_beginning(self, message):
  447. return self._restore(message, leftmost=True)
  448. def basic_consume(self, queue, *args, **kwargs):
  449. if queue in self._fanout_queues:
  450. exchange, _ = self._fanout_queues[queue]
  451. self.active_fanout_queues.add(queue)
  452. self._fanout_to_queue[exchange] = queue
  453. ret = super(Channel, self).basic_consume(queue, *args, **kwargs)
  454. self._update_cycle()
  455. return ret
  456. def basic_cancel(self, consumer_tag):
  457. # If we are busy reading messages we may experience
  458. # a race condition where a message is consumed after
  459. # cancelling, so we must delay this operation until reading
  460. # is complete (Issue celery/celery#1773).
  461. connection = self.connection
  462. if connection:
  463. if connection.cycle._in_protected_read:
  464. return connection.cycle.after_read.add(
  465. promise(self._basic_cancel, (consumer_tag, )),
  466. )
  467. return self._basic_cancel(consumer_tag)
  468. def _basic_cancel(self, consumer_tag):
  469. try:
  470. queue = self._tag_to_queue[consumer_tag]
  471. except KeyError:
  472. return
  473. try:
  474. self.active_fanout_queues.remove(queue)
  475. except KeyError:
  476. pass
  477. else:
  478. self._unsubscribe_from(queue)
  479. try:
  480. exchange, _ = self._fanout_queues[queue]
  481. self._fanout_to_queue.pop(exchange)
  482. except KeyError:
  483. pass
  484. ret = super(Channel, self).basic_cancel(consumer_tag)
  485. self._update_cycle()
  486. return ret
  487. def _get_publish_topic(self, exchange, routing_key):
  488. if routing_key and self.fanout_patterns:
  489. return ''.join([self.keyprefix_fanout, exchange, '/', routing_key])
  490. return ''.join([self.keyprefix_fanout, exchange])
  491. def _get_subscribe_topic(self, queue):
  492. exchange, routing_key = self._fanout_queues[queue]
  493. return self._get_publish_topic(exchange, routing_key)
  494. def _subscribe(self):
  495. keys = [self._get_subscribe_topic(queue)
  496. for queue in self.active_fanout_queues]
  497. if not keys:
  498. return
  499. c = self.subclient
  500. if c.connection._sock is None:
  501. c.connection.connect()
  502. self._in_listen = True
  503. c.psubscribe(keys)
  504. def _unsubscribe_from(self, queue):
  505. topic = self._get_subscribe_topic(queue)
  506. c = self.subclient
  507. should_disconnect = False
  508. if c.connection._sock is None:
  509. c.connection.connect()
  510. should_disconnect = True
  511. try:
  512. c.unsubscribe([topic])
  513. finally:
  514. if should_disconnect and c.connection:
  515. c.connection.disconnect()
  516. def _handle_message(self, client, r):
  517. if bytes_to_str(r[0]) == 'unsubscribe' and r[2] == 0:
  518. client.subscribed = False
  519. elif bytes_to_str(r[0]) == 'pmessage':
  520. return {'type': r[0], 'pattern': r[1],
  521. 'channel': r[2], 'data': r[3]}
  522. else:
  523. return {'type': r[0], 'pattern': None,
  524. 'channel': r[1], 'data': r[2]}
  525. def _receive(self):
  526. c = self.subclient
  527. response = None
  528. try:
  529. response = c.parse_response()
  530. except self.connection_errors:
  531. self._in_listen = False
  532. raise Empty()
  533. if response is not None:
  534. payload = self._handle_message(c, response)
  535. if bytes_to_str(payload['type']).endswith('message'):
  536. channel = bytes_to_str(payload['channel'])
  537. if payload['data']:
  538. if channel[0] == '/':
  539. _, _, channel = channel.partition('.')
  540. try:
  541. message = loads(bytes_to_str(payload['data']))
  542. except (TypeError, ValueError):
  543. warn('Cannot process event on channel %r: %s',
  544. channel, repr(payload)[:4096], exc_info=1)
  545. raise Empty()
  546. exchange = channel.split('/', 1)[0]
  547. return message, self._fanout_to_queue[exchange]
  548. raise Empty()
  549. def _brpop_start(self, timeout=1):
  550. queues = self._consume_cycle()
  551. if not queues:
  552. return
  553. keys = [self._q_for_pri(queue, pri) for pri in PRIORITY_STEPS
  554. for queue in queues] + [timeout or 0]
  555. self._in_poll = True
  556. self.client.connection.send_command('BRPOP', *keys)
  557. def _brpop_read(self, **options):
  558. try:
  559. try:
  560. dest__item = self.client.parse_response(self.client.connection,
  561. 'BRPOP',
  562. **options)
  563. except self.connection_errors:
  564. # if there's a ConnectionError, disconnect so the next
  565. # iteration will reconnect automatically.
  566. self.client.connection.disconnect()
  567. raise Empty()
  568. if dest__item:
  569. dest, item = dest__item
  570. dest = bytes_to_str(dest).rsplit(self.sep, 1)[0]
  571. self._rotate_cycle(dest)
  572. return loads(bytes_to_str(item)), dest
  573. else:
  574. raise Empty()
  575. finally:
  576. self._in_poll = False
  577. def _poll_error(self, type, **options):
  578. if type == 'LISTEN':
  579. self.subclient.parse_response()
  580. else:
  581. self.client.parse_response(self.client.connection, type)
  582. def _get(self, queue):
  583. with self.conn_or_acquire() as client:
  584. for pri in PRIORITY_STEPS:
  585. item = client.rpop(self._q_for_pri(queue, pri))
  586. if item:
  587. return loads(bytes_to_str(item))
  588. raise Empty()
  589. def _size(self, queue):
  590. with self.conn_or_acquire() as client:
  591. with client.pipeline() as pipe:
  592. for pri in PRIORITY_STEPS:
  593. pipe = pipe.llen(self._q_for_pri(queue, pri))
  594. sizes = pipe.execute()
  595. return sum(size for size in sizes
  596. if isinstance(size, numbers.Integral))
  597. def _q_for_pri(self, queue, pri):
  598. pri = self.priority(pri)
  599. return '%s%s%s' % ((queue, self.sep, pri) if pri else (queue, '', ''))
  600. def priority(self, n):
  601. steps = self.priority_steps
  602. return steps[bisect(steps, n) - 1]
  603. def _put(self, queue, message, **kwargs):
  604. """Deliver message."""
  605. try:
  606. pri = max(min(int(
  607. message['properties']['delivery_info']['priority']), 9), 0)
  608. except (TypeError, ValueError, KeyError):
  609. pri = 0
  610. with self.conn_or_acquire() as client:
  611. client.lpush(self._q_for_pri(queue, pri), dumps(message))
  612. def _put_fanout(self, exchange, message, routing_key, **kwargs):
  613. """Deliver fanout message."""
  614. with self.conn_or_acquire() as client:
  615. client.publish(
  616. self._get_publish_topic(exchange, routing_key),
  617. dumps(message),
  618. )
  619. def _new_queue(self, queue, auto_delete=False, **kwargs):
  620. if auto_delete:
  621. self.auto_delete_queues.add(queue)
  622. def _queue_bind(self, exchange, routing_key, pattern, queue):
  623. if self.typeof(exchange).type == 'fanout':
  624. # Mark exchange as fanout.
  625. self._fanout_queues[queue] = (
  626. exchange, routing_key.replace('#', '*'),
  627. )
  628. with self.conn_or_acquire() as client:
  629. client.sadd(self.keyprefix_queue % (exchange, ),
  630. self.sep.join([routing_key or '',
  631. pattern or '',
  632. queue or '']))
  633. def _delete(self, queue, exchange, routing_key, pattern, *args):
  634. self.auto_delete_queues.discard(queue)
  635. with self.conn_or_acquire() as client:
  636. client.srem(self.keyprefix_queue % (exchange, ),
  637. self.sep.join([routing_key or '',
  638. pattern or '',
  639. queue or '']))
  640. with client.pipeline() as pipe:
  641. for pri in PRIORITY_STEPS:
  642. pipe = pipe.delete(self._q_for_pri(queue, pri))
  643. pipe.execute()
  644. def _has_queue(self, queue, **kwargs):
  645. with self.conn_or_acquire() as client:
  646. with client.pipeline() as pipe:
  647. for pri in PRIORITY_STEPS:
  648. pipe = pipe.exists(self._q_for_pri(queue, pri))
  649. return any(pipe.execute())
  650. def get_table(self, exchange):
  651. key = self.keyprefix_queue % exchange
  652. with self.conn_or_acquire() as client:
  653. values = client.smembers(key)
  654. if not values:
  655. raise InconsistencyError(NO_ROUTE_ERROR.format(exchange, key))
  656. return [tuple(bytes_to_str(val).split(self.sep)) for val in values]
  657. def _purge(self, queue):
  658. with self.conn_or_acquire() as client:
  659. with client.pipeline() as pipe:
  660. for pri in PRIORITY_STEPS:
  661. priq = self._q_for_pri(queue, pri)
  662. pipe = pipe.llen(priq).delete(priq)
  663. sizes = pipe.execute()
  664. return sum(sizes[::2])
  665. def close(self):
  666. self._closing = True
  667. self._disconnect_pools()
  668. if not self.closed:
  669. # remove from channel poller.
  670. self.connection.cycle.discard(self)
  671. # delete fanout bindings
  672. for queue in self._fanout_queues:
  673. if queue in self.auto_delete_queues:
  674. self.queue_delete(queue)
  675. self._close_clients()
  676. super(Channel, self).close()
  677. def _close_clients(self):
  678. # Close connections
  679. for attr in 'client', 'subclient':
  680. try:
  681. self.__dict__[attr].connection.disconnect()
  682. except (KeyError, AttributeError, self.ResponseError):
  683. pass
  684. def _prepare_virtual_host(self, vhost):
  685. if not isinstance(vhost, numbers.Integral):
  686. if not vhost or vhost == '/':
  687. vhost = DEFAULT_DB
  688. elif vhost.startswith('/'):
  689. vhost = vhost[1:]
  690. try:
  691. vhost = int(vhost)
  692. except ValueError:
  693. raise ValueError(
  694. 'Database is int between 0 and limit - 1, not {0}'.format(
  695. vhost,
  696. ))
  697. return vhost
  698. def _filter_tcp_connparams(self, socket_keepalive=None,
  699. socket_keepalive_options=None, **params):
  700. return params
  701. def _connparams(self, async=False, _r210_options=(
  702. 'socket_connect_timeout', 'socket_keepalive',
  703. 'socket_keepalive_options')):
  704. conninfo = self.connection.client
  705. connparams = {
  706. 'host': conninfo.hostname or '127.0.0.1',
  707. 'port': conninfo.port or DEFAULT_PORT,
  708. 'virtual_host': conninfo.virtual_host,
  709. 'password': conninfo.password,
  710. 'max_connections': self.max_connections,
  711. 'socket_timeout': self.socket_timeout,
  712. 'socket_connect_timeout': self.socket_connect_timeout,
  713. 'socket_keepalive': self.socket_keepalive,
  714. 'socket_keepalive_options': self.socket_keepalive_options,
  715. }
  716. if redis.VERSION < (2, 10):
  717. for param in _r210_options:
  718. val = connparams.pop(param, None)
  719. if val is not None:
  720. raise VersionMismatch(
  721. 'redis: {0!r} requires redis 2.10.0 or higher'.format(
  722. param))
  723. host = connparams['host']
  724. if '://' in host:
  725. scheme, _, _, _, password, path, query = _parse_url(host)
  726. if scheme == 'socket':
  727. connparams = self._filter_tcp_connparams(**connparams)
  728. connparams.update({
  729. 'connection_class': redis.UnixDomainSocketConnection,
  730. 'path': '/' + path,
  731. 'password': password}, **query)
  732. connparams.pop('socket_connect_timeout', None)
  733. connparams.pop('socket_keepalive', None)
  734. connparams.pop('socket_keepalive_options', None)
  735. connparams.pop('host', None)
  736. connparams.pop('port', None)
  737. connparams['db'] = self._prepare_virtual_host(
  738. connparams.pop('virtual_host', None))
  739. channel = self
  740. connection_cls = (
  741. connparams.get('connection_class') or
  742. redis.Connection
  743. )
  744. if async:
  745. class Connection(connection_cls):
  746. def disconnect(self):
  747. super(Connection, self).disconnect()
  748. channel._on_connection_disconnect(self)
  749. connparams['connection_class'] = Connection
  750. return connparams
  751. def _create_client(self, async=False):
  752. if async:
  753. return self.AsyncClient(connection_pool=self.async_pool)
  754. return self.Client(connection_pool=self.pool)
  755. def _get_pool(self, async=False):
  756. params = self._connparams(async=async)
  757. self.keyprefix_fanout = self.keyprefix_fanout.format(db=params['db'])
  758. return redis.ConnectionPool(**params)
  759. def _get_async_client(self):
  760. if redis.VERSION < (2, 4, 4):
  761. raise VersionMismatch(
  762. 'Redis transport requires redis-py versions 2.4.4 or later. '
  763. 'You have {0.__version__}'.format(redis))
  764. # AsyncRedis maintains a connection attribute on it's instance and
  765. # uses that when executing commands
  766. # This was added after redis-py was changed.
  767. class AsyncRedis(redis.Redis): # pragma: no cover
  768. def __init__(self, *args, **kwargs):
  769. super(AsyncRedis, self).__init__(*args, **kwargs)
  770. self.connection = self.connection_pool.get_connection('_')
  771. return AsyncRedis
  772. @contextmanager
  773. def conn_or_acquire(self, client=None):
  774. if client:
  775. yield client
  776. else:
  777. yield self._create_client()
  778. @property
  779. def pool(self):
  780. if self._pool is None:
  781. self._pool = self._get_pool()
  782. return self._pool
  783. @property
  784. def async_pool(self):
  785. if self._async_pool is None:
  786. self._async_pool = self._get_pool(async=True)
  787. return self._async_pool
  788. @cached_property
  789. def client(self):
  790. """Client used to publish messages, BRPOP etc."""
  791. return self._create_client(async=True)
  792. @cached_property
  793. def subclient(self):
  794. """Pub/Sub connection used to consume fanout queues."""
  795. client = self._create_client(async=True)
  796. pubsub = client.pubsub()
  797. pool = pubsub.connection_pool
  798. pubsub.connection = pool.get_connection('pubsub', pubsub.shard_hint)
  799. return pubsub
  800. def _update_cycle(self):
  801. """Update fair cycle between queues.
  802. We cycle between queues fairly to make sure that
  803. each queue is equally likely to be consumed from,
  804. so that a very busy queue will not block others.
  805. This works by using Redis's `BRPOP` command and
  806. by rotating the most recently used queue to the
  807. and of the list. See Kombu github issue #166 for
  808. more discussion of this method.
  809. """
  810. self._queue_cycle = list(self.active_queues)
  811. def _consume_cycle(self):
  812. """Get a fresh list of queues from the queue cycle."""
  813. active = len(self.active_queues)
  814. return self._queue_cycle[0:active]
  815. def _rotate_cycle(self, used):
  816. """Move most recently used queue to end of list."""
  817. cycle = self._queue_cycle
  818. try:
  819. cycle.append(cycle.pop(cycle.index(used)))
  820. except ValueError:
  821. pass
  822. def _get_response_error(self):
  823. from redis import exceptions
  824. return exceptions.ResponseError
  825. @property
  826. def active_queues(self):
  827. """Set of queues being consumed from (excluding fanout queues)."""
  828. return set(queue for queue in self._active_queues
  829. if queue not in self.active_fanout_queues)
  830. class Transport(virtual.Transport):
  831. Channel = Channel
  832. polling_interval = None # disable sleep between unsuccessful polls.
  833. default_port = DEFAULT_PORT
  834. supports_ev = True
  835. driver_type = 'redis'
  836. driver_name = 'redis'
  837. def __init__(self, *args, **kwargs):
  838. if redis is None:
  839. raise ImportError('Missing redis library (pip install redis)')
  840. super(Transport, self).__init__(*args, **kwargs)
  841. # Get redis-py exceptions.
  842. self.connection_errors, self.channel_errors = self._get_errors()
  843. # All channels share the same poller.
  844. self.cycle = MultiChannelPoller()
  845. def driver_version(self):
  846. return redis.__version__
  847. def register_with_event_loop(self, connection, loop):
  848. cycle = self.cycle
  849. cycle.on_poll_init(loop.poller)
  850. cycle_poll_start = cycle.on_poll_start
  851. add_reader = loop.add_reader
  852. on_readable = self.on_readable
  853. def _on_disconnect(connection):
  854. if connection._sock:
  855. loop.remove(connection._sock)
  856. cycle._on_connection_disconnect = _on_disconnect
  857. def on_poll_start():
  858. cycle_poll_start()
  859. [add_reader(fd, on_readable, fd) for fd in cycle.fds]
  860. loop.on_tick.add(on_poll_start)
  861. loop.call_repeatedly(10, cycle.maybe_restore_messages)
  862. def on_readable(self, fileno):
  863. """Handle AIO event for one of our file descriptors."""
  864. item = self.cycle.on_readable(fileno)
  865. if item:
  866. message, queue = item
  867. if not queue or queue not in self._callbacks:
  868. raise KeyError(
  869. 'Message for queue {0!r} without consumers: {1}'.format(
  870. queue, message))
  871. self._callbacks[queue](message)
  872. def _get_errors(self):
  873. """Utility to import redis-py's exceptions at runtime."""
  874. return get_redis_error_classes()