| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381 |
- import asyncio
- import random
- import weakref
- from typing import AsyncIterator, Iterable, Mapping, Optional, Sequence, Tuple, Type
- from redis.asyncio.client import Redis
- from redis.asyncio.connection import (
- Connection,
- ConnectionPool,
- EncodableT,
- SSLConnection,
- )
- from redis.commands import AsyncSentinelCommands
- from redis.exceptions import ConnectionError, ReadOnlyError, ResponseError, TimeoutError
- from redis.utils import str_if_bytes
- class MasterNotFoundError(ConnectionError):
- pass
- class SlaveNotFoundError(ConnectionError):
- pass
- class SentinelManagedConnection(Connection):
- def __init__(self, **kwargs):
- self.connection_pool = kwargs.pop("connection_pool")
- super().__init__(**kwargs)
- def __repr__(self):
- pool = self.connection_pool
- s = (
- f"<{self.__class__.__module__}.{self.__class__.__name__}"
- f"(service={pool.service_name}"
- )
- if self.host:
- host_info = f",host={self.host},port={self.port}"
- s += host_info
- return s + ")>"
- async def connect_to(self, address):
- self.host, self.port = address
- await super().connect()
- if self.connection_pool.check_connection:
- await self.send_command("PING")
- if str_if_bytes(await self.read_response()) != "PONG":
- raise ConnectionError("PING failed")
- async def _connect_retry(self):
- if self._reader:
- return # already connected
- if self.connection_pool.is_master:
- await self.connect_to(await self.connection_pool.get_master_address())
- else:
- async for slave in self.connection_pool.rotate_slaves():
- try:
- return await self.connect_to(slave)
- except ConnectionError:
- continue
- raise SlaveNotFoundError # Never be here
- async def connect(self):
- return await self.retry.call_with_retry(
- self._connect_retry,
- lambda error: asyncio.sleep(0),
- )
- async def read_response(
- self,
- disable_decoding: bool = False,
- timeout: Optional[float] = None,
- *,
- disconnect_on_error: Optional[float] = True,
- push_request: Optional[bool] = False,
- ):
- try:
- return await super().read_response(
- disable_decoding=disable_decoding,
- timeout=timeout,
- disconnect_on_error=disconnect_on_error,
- push_request=push_request,
- )
- except ReadOnlyError:
- if self.connection_pool.is_master:
- # When talking to a master, a ReadOnlyError when likely
- # indicates that the previous master that we're still connected
- # to has been demoted to a slave and there's a new master.
- # calling disconnect will force the connection to re-query
- # sentinel during the next connect() attempt.
- await self.disconnect()
- raise ConnectionError("The previous master is now a slave")
- raise
- class SentinelManagedSSLConnection(SentinelManagedConnection, SSLConnection):
- pass
- class SentinelConnectionPool(ConnectionPool):
- """
- Sentinel backed connection pool.
- If ``check_connection`` flag is set to True, SentinelManagedConnection
- sends a PING command right after establishing the connection.
- """
- def __init__(self, service_name, sentinel_manager, **kwargs):
- kwargs["connection_class"] = kwargs.get(
- "connection_class",
- SentinelManagedSSLConnection
- if kwargs.pop("ssl", False)
- else SentinelManagedConnection,
- )
- self.is_master = kwargs.pop("is_master", True)
- self.check_connection = kwargs.pop("check_connection", False)
- super().__init__(**kwargs)
- self.connection_kwargs["connection_pool"] = weakref.proxy(self)
- self.service_name = service_name
- self.sentinel_manager = sentinel_manager
- self.master_address = None
- self.slave_rr_counter = None
- def __repr__(self):
- return (
- f"<{self.__class__.__module__}.{self.__class__.__name__}"
- f"(service={self.service_name}({self.is_master and 'master' or 'slave'}))>"
- )
- def reset(self):
- super().reset()
- self.master_address = None
- self.slave_rr_counter = None
- def owns_connection(self, connection: Connection):
- check = not self.is_master or (
- self.is_master and self.master_address == (connection.host, connection.port)
- )
- return check and super().owns_connection(connection)
- async def get_master_address(self):
- master_address = await self.sentinel_manager.discover_master(self.service_name)
- if self.is_master:
- if self.master_address != master_address:
- self.master_address = master_address
- # disconnect any idle connections so that they reconnect
- # to the new master the next time that they are used.
- await self.disconnect(inuse_connections=False)
- return master_address
- async def rotate_slaves(self) -> AsyncIterator:
- """Round-robin slave balancer"""
- slaves = await self.sentinel_manager.discover_slaves(self.service_name)
- if slaves:
- if self.slave_rr_counter is None:
- self.slave_rr_counter = random.randint(0, len(slaves) - 1)
- for _ in range(len(slaves)):
- self.slave_rr_counter = (self.slave_rr_counter + 1) % len(slaves)
- slave = slaves[self.slave_rr_counter]
- yield slave
- # Fallback to the master connection
- try:
- yield await self.get_master_address()
- except MasterNotFoundError:
- pass
- raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
- class Sentinel(AsyncSentinelCommands):
- """
- Redis Sentinel cluster client
- >>> from redis.sentinel import Sentinel
- >>> sentinel = Sentinel([('localhost', 26379)], socket_timeout=0.1)
- >>> master = sentinel.master_for('mymaster', socket_timeout=0.1)
- >>> await master.set('foo', 'bar')
- >>> slave = sentinel.slave_for('mymaster', socket_timeout=0.1)
- >>> await slave.get('foo')
- b'bar'
- ``sentinels`` is a list of sentinel nodes. Each node is represented by
- a pair (hostname, port).
- ``min_other_sentinels`` defined a minimum number of peers for a sentinel.
- When querying a sentinel, if it doesn't meet this threshold, responses
- from that sentinel won't be considered valid.
- ``sentinel_kwargs`` is a dictionary of connection arguments used when
- connecting to sentinel instances. Any argument that can be passed to
- a normal Redis connection can be specified here. If ``sentinel_kwargs`` is
- not specified, any socket_timeout and socket_keepalive options specified
- in ``connection_kwargs`` will be used.
- ``connection_kwargs`` are keyword arguments that will be used when
- establishing a connection to a Redis server.
- """
- def __init__(
- self,
- sentinels,
- min_other_sentinels=0,
- sentinel_kwargs=None,
- **connection_kwargs,
- ):
- # if sentinel_kwargs isn't defined, use the socket_* options from
- # connection_kwargs
- if sentinel_kwargs is None:
- sentinel_kwargs = {
- k: v for k, v in connection_kwargs.items() if k.startswith("socket_")
- }
- self.sentinel_kwargs = sentinel_kwargs
- self.sentinels = [
- Redis(host=hostname, port=port, **self.sentinel_kwargs)
- for hostname, port in sentinels
- ]
- self.min_other_sentinels = min_other_sentinels
- self.connection_kwargs = connection_kwargs
- async def execute_command(self, *args, **kwargs):
- """
- Execute Sentinel command in sentinel nodes.
- once - If set to True, then execute the resulting command on a single
- node at random, rather than across the entire sentinel cluster.
- """
- once = bool(kwargs.get("once", False))
- if "once" in kwargs.keys():
- kwargs.pop("once")
- if once:
- await random.choice(self.sentinels).execute_command(*args, **kwargs)
- else:
- tasks = [
- asyncio.Task(sentinel.execute_command(*args, **kwargs))
- for sentinel in self.sentinels
- ]
- await asyncio.gather(*tasks)
- return True
- def __repr__(self):
- sentinel_addresses = []
- for sentinel in self.sentinels:
- sentinel_addresses.append(
- f"{sentinel.connection_pool.connection_kwargs['host']}:"
- f"{sentinel.connection_pool.connection_kwargs['port']}"
- )
- return (
- f"<{self.__class__}.{self.__class__.__name__}"
- f"(sentinels=[{','.join(sentinel_addresses)}])>"
- )
- def check_master_state(self, state: dict, service_name: str) -> bool:
- if not state["is_master"] or state["is_sdown"] or state["is_odown"]:
- return False
- # Check if our sentinel doesn't see other nodes
- if state["num-other-sentinels"] < self.min_other_sentinels:
- return False
- return True
- async def discover_master(self, service_name: str):
- """
- Asks sentinel servers for the Redis master's address corresponding
- to the service labeled ``service_name``.
- Returns a pair (address, port) or raises MasterNotFoundError if no
- master is found.
- """
- collected_errors = list()
- for sentinel_no, sentinel in enumerate(self.sentinels):
- try:
- masters = await sentinel.sentinel_masters()
- except (ConnectionError, TimeoutError) as e:
- collected_errors.append(f"{sentinel} - {e!r}")
- continue
- state = masters.get(service_name)
- if state and self.check_master_state(state, service_name):
- # Put this sentinel at the top of the list
- self.sentinels[0], self.sentinels[sentinel_no] = (
- sentinel,
- self.sentinels[0],
- )
- return state["ip"], state["port"]
- error_info = ""
- if len(collected_errors) > 0:
- error_info = f" : {', '.join(collected_errors)}"
- raise MasterNotFoundError(f"No master found for {service_name!r}{error_info}")
- def filter_slaves(
- self, slaves: Iterable[Mapping]
- ) -> Sequence[Tuple[EncodableT, EncodableT]]:
- """Remove slaves that are in an ODOWN or SDOWN state"""
- slaves_alive = []
- for slave in slaves:
- if slave["is_odown"] or slave["is_sdown"]:
- continue
- slaves_alive.append((slave["ip"], slave["port"]))
- return slaves_alive
- async def discover_slaves(
- self, service_name: str
- ) -> Sequence[Tuple[EncodableT, EncodableT]]:
- """Returns a list of alive slaves for service ``service_name``"""
- for sentinel in self.sentinels:
- try:
- slaves = await sentinel.sentinel_slaves(service_name)
- except (ConnectionError, ResponseError, TimeoutError):
- continue
- slaves = self.filter_slaves(slaves)
- if slaves:
- return slaves
- return []
- def master_for(
- self,
- service_name: str,
- redis_class: Type[Redis] = Redis,
- connection_pool_class: Type[SentinelConnectionPool] = SentinelConnectionPool,
- **kwargs,
- ):
- """
- Returns a redis client instance for the ``service_name`` master.
- A :py:class:`~redis.sentinel.SentinelConnectionPool` class is
- used to retrieve the master's address before establishing a new
- connection.
- NOTE: If the master's address has changed, any cached connections to
- the old master are closed.
- By default clients will be a :py:class:`~redis.Redis` instance.
- Specify a different class to the ``redis_class`` argument if you
- desire something different.
- The ``connection_pool_class`` specifies the connection pool to
- use. The :py:class:`~redis.sentinel.SentinelConnectionPool`
- will be used by default.
- All other keyword arguments are merged with any connection_kwargs
- passed to this class and passed to the connection pool as keyword
- arguments to be used to initialize Redis connections.
- """
- kwargs["is_master"] = True
- connection_kwargs = dict(self.connection_kwargs)
- connection_kwargs.update(kwargs)
- connection_pool = connection_pool_class(service_name, self, **connection_kwargs)
- # The Redis object "owns" the pool
- return redis_class.from_pool(connection_pool)
- def slave_for(
- self,
- service_name: str,
- redis_class: Type[Redis] = Redis,
- connection_pool_class: Type[SentinelConnectionPool] = SentinelConnectionPool,
- **kwargs,
- ):
- """
- Returns redis client instance for the ``service_name`` slave(s).
- A SentinelConnectionPool class is used to retrieve the slave's
- address before establishing a new connection.
- By default clients will be a :py:class:`~redis.Redis` instance.
- Specify a different class to the ``redis_class`` argument if you
- desire something different.
- The ``connection_pool_class`` specifies the connection pool to use.
- The SentinelConnectionPool will be used by default.
- All other keyword arguments are merged with any connection_kwargs
- passed to this class and passed to the connection pool as keyword
- arguments to be used to initialize Redis connections.
- """
- kwargs["is_master"] = False
- connection_kwargs = dict(self.connection_kwargs)
- connection_kwargs.update(kwargs)
- connection_pool = connection_pool_class(service_name, self, **connection_kwargs)
- # The Redis object "owns" the pool
- return redis_class.from_pool(connection_pool)
|