From 9fb0376bf888bafb392813869ba0d37790f8e061 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Fri, 26 Apr 2024 13:52:58 +0800 Subject: [PATCH 01/44] fix scan iter command issued to different replicas --- redis/asyncio/sentinel.py | 121 +++++++++++++++++++++++++++++++++++++- redis/commands/core.py | 27 +++++++-- 2 files changed, 142 insertions(+), 6 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 6fd233adc8..0c642e1b1a 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -1,7 +1,8 @@ import asyncio import random import weakref -from typing import AsyncIterator, Iterable, Mapping, Optional, Sequence, Tuple, Type +import uuid +from typing import AsyncIterator, Iterable, Mapping, Optional, Sequence, Tuple, Type, Any from redis.asyncio.client import Redis from redis.asyncio.connection import ( @@ -65,6 +66,22 @@ async def connect(self): self._connect_retry, lambda error: asyncio.sleep(0), ) + + async def _connect_to_address_retry(self, host: str, port: int) -> None: + if self._reader: + return # already connected + try: + return await self.connect_to((host, port)) + except ConnectionError as exc: + raise SlaveNotFoundError + + async def connect_to_address(self, host: str, port: int) -> None: + # Connect to the specified host and port + # instead of connecting to the master / rotated slaves + return await self.retry.call_with_retry( + lambda: self._connect_to_address_retry(host, port), + lambda error: asyncio.sleep(0), + ) async def read_response( self, @@ -122,6 +139,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs): self.sentinel_manager = sentinel_manager self.master_address = None self.slave_rr_counter = None + self._request_id_to_replica_address = {} def __repr__(self): return ( @@ -152,6 +170,11 @@ async def get_master_address(self): async def rotate_slaves(self) -> AsyncIterator: """Round-robin slave balancer""" + ( + server_host, + server_port, + ) = self._request_id_to_replica_address.get(iter_req_id, (None, None)) + slaves = await self.sentinel_manager.discover_slaves(self.service_name) if slaves: if self.slave_rr_counter is None: @@ -167,6 +190,102 @@ async def rotate_slaves(self) -> AsyncIterator: pass raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") + async def get_connection( + self, command_name: str, *keys: Any, **options: Any + ) -> SentinelManagedConnection: + """ + Get a connection from the pool. + `xxx_scan_iter` commands needs to be handled specially. + If the client is created using a connection pool, in replica mode, + all `scan` command-equivalent of the `xxx_scan_iter` commands needs + to be issued to the same Redis replica. + + The way each server positions each key is different with one another, + and the cursor acts as the 'offset' of the scan. + Hence, all scans coming from a single xxx_scan_iter_channel command + should go to the same replica. + """ + # If not an iter command or in master mode, call super() + # No custom logic for master, because there's only 1 master. + # The bug is only when Redis has the possibility to connect to multiple replicas + if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: + return await super().get_connection(command_name, *keys, **options) # type: ignore[no-any-return] + + # Check if this iter request has already been directed to a particular server + # Check if this iter request has already been directed to a particular server + ( + server_host, + server_port, + ) = self._request_id_to_replica_address.get(iter_req_id, (None, None)) + connection = None + # If this is the first scan request of the iter command, + # get a connection from the pool + if server_host is None or server_port is None: + try: + connection = self._available_connections.pop() # type: ignore [assignment] + except IndexError: + connection = self.make_connection() + # If this is not the first scan request of the iter command + else: + # Check from the available connections, if any of the connection + # is connected to the host and port that we want + # If yes, use that connection + for available_connection in self._available_connections.copy(): + if ( + available_connection.host == server_host + and available_connection.port == server_port + ): + self._available_connections.remove(available_connection) + connection = available_connection # type: ignore[assignment] + # If not, make a new dummy connection object, and set its host and port + # to the one that we want later in the call to ``connect_to_address`` + if not connection: + connection = self.make_connection() + assert connection + self._in_use_connections.add(connection) + try: + # ensure this connection is connected to Redis + # If this is the first scan request, + # just call the SentinelManagedConnection.connect() + # This will call rotate_slaves + # and connect to a random replica + if server_port is None or server_port is None: + await connection.connect() + # If this is not the first scan request, + # connect to the particular address and port + else: + # This will connect to the host and port that we've specified above + await connection.connect_to_address(server_host, server_port) # type: ignore[arg-type] + # connections that the pool provides should be ready to send + # a command. if not, the connection was either returned to the + # pool before all data has been read or the socket has been + # closed. either way, reconnect and verify everything is good. + try: + # type ignore below: + # attr Not defined in redis stubs and + # we don't need to create a subclass to help with this single attr + if await connection.can_read_destructive(): # type: ignore[attr-defined] + raise ConnectionError("Connection has data") from None + except (ConnectionError, OSError): + await connection.disconnect() + await connection.connect() + # type ignore below: similar to above + if await connection.can_read_destructive(): # type: ignore[attr-defined] + raise ConnectionError("Connection not ready") from None + except BaseException: + # release the connection back to the pool so that we don't + # leak it + await self.release(connection) + raise + # Store the connection to the dictionary + self._request_id_to_replica_address[iter_req_id] = ( + connection.host, + connection.port, + ) + + return connection + + class Sentinel(AsyncSentinelCommands): """ diff --git a/redis/commands/core.py b/redis/commands/core.py index b356d101ee..cc996b9974 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -1,5 +1,6 @@ # from __future__ import annotations +import uuid import datetime import hashlib import warnings @@ -3050,6 +3051,7 @@ def sscan( cursor: int = 0, match: Union[PatternT, None] = None, count: Union[int, None] = None, + **kwargs, ) -> ResponseT: """ Incrementally return lists of elements in a set. Also return a cursor @@ -3066,7 +3068,7 @@ def sscan( pieces.extend([b"MATCH", match]) if count is not None: pieces.extend([b"COUNT", count]) - return self.execute_command("SSCAN", *pieces) + return self.execute_command("SSCAN", *pieces,) def sscan_iter( self, @@ -3094,6 +3096,7 @@ def hscan( match: Union[PatternT, None] = None, count: Union[int, None] = None, no_values: Union[bool, None] = None, + **kwargs, ) -> ResponseT: """ Incrementally return key/value slices in a hash. Also return a cursor @@ -3114,7 +3117,7 @@ def hscan( pieces.extend([b"COUNT", count]) if no_values is not None: pieces.extend([b"NOVALUES"]) - return self.execute_command("HSCAN", *pieces, no_values=no_values) + return self.execute_command("HSCAN", *pieces, no_values=no_values, **kwargs) def hscan_iter( self, @@ -3150,6 +3153,7 @@ def zscan( match: Union[PatternT, None] = None, count: Union[int, None] = None, score_cast_func: Union[type, Callable] = float, + **kwargs, ) -> ResponseT: """ Incrementally return lists of elements in a sorted set. Also return a @@ -3169,7 +3173,7 @@ def zscan( if count is not None: pieces.extend([b"COUNT", count]) options = {"score_cast_func": score_cast_func} - return self.execute_command("ZSCAN", *pieces, **options) + return self.execute_command("ZSCAN", *pieces, **options, **kwargs) def zscan_iter( self, @@ -3222,10 +3226,12 @@ async def scan_iter( HASH, LIST, SET, STREAM, STRING, ZSET Additionally, Redis modules can expose other types as well. """ + # DO NOT inline this statement to the scan call + iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica cursor = "0" while cursor != 0: cursor, data = await self.scan( - cursor=cursor, match=match, count=count, _type=_type, **kwargs + cursor=cursor, match=match, count=count, _type=_type, _iter_req_id=iter_req_id, **kwargs ) for d in data: yield d @@ -3244,10 +3250,12 @@ async def sscan_iter( ``count`` allows for hint the minimum number of returns """ + # DO NOT inline this statement to the scan call + iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica cursor = "0" while cursor != 0: cursor, data = await self.sscan( - name, cursor=cursor, match=match, count=count + name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id ) for d in data: yield d @@ -3269,10 +3277,16 @@ async def hscan_iter( ``no_values`` indicates to return only the keys, without values """ + # DO NOT inline this statement to the scan call + iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica cursor = "0" while cursor != 0: cursor, data = await self.hscan( +<<<<<<< HEAD name, cursor=cursor, match=match, count=count, no_values=no_values +======= + name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id +>>>>>>> 39aaec6 (fix scan iter command issued to different replicas) ) if no_values: for it in data: @@ -3298,6 +3312,8 @@ async def zscan_iter( ``score_cast_func`` a callable used to cast the score return value """ + # DO NOT inline this statement to the scan call + iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica cursor = "0" while cursor != 0: cursor, data = await self.zscan( @@ -3306,6 +3322,7 @@ async def zscan_iter( match=match, count=count, score_cast_func=score_cast_func, + _iter_req_id=iter_req_id ) for d in data: yield d From 48034de70ac8396ead2232a5450ab3f85df88c4a Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 27 Apr 2024 00:13:13 +0800 Subject: [PATCH 02/44] add tests --- .../test_sentinel_managed_connection.py | 169 +++++++++++++++++- 1 file changed, 168 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index cae4b9581f..2c6741fdc0 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,9 +1,14 @@ import socket import pytest +import pytest_asyncio from redis.asyncio.retry import Retry -from redis.asyncio.sentinel import SentinelManagedConnection +from redis.asyncio.sentinel import SentinelManagedConnection, SentinelConnectionPool, Sentinel from redis.backoff import NoBackoff +from unittest.mock import AsyncMock +from typing import AsyncIterator + +import pytest from .compat import mock @@ -35,3 +40,165 @@ async def mock_connect(): await conn.connect() assert conn._connect.call_count == 3 await conn.disconnect() + + +def same_address( + connection_1: SentinelManagedConnection, + connection_2: SentinelManagedConnection, +) -> bool: + return bool( + connection_1.host == connection_2.host and connection_1.port == connection_2.port + ) + +class SentinelManagedConnectionMock(SentinelManagedConnection): + async def connect_to_address(self, host: str, port: int) -> None: + self.host = host + self.port = port + + async def can_read_destructive(self) -> bool: + # Mock this function to always return False. + # Trrue means there's still data to be read and hence we can't reconnect + # to this connection yet + return False + + +class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock): + async def connect(self) -> None: + """ + This simulates the behavior of connect when the ``redis.asyncio.sentinel.SentinelConnectionPool`` + is in replica mode. + It'll call rotate_slaves and connect to the next replica + """ + import random + import time + + self.host = f"host-{random.randint(0, 10)}" + self.port = time.time() + + +class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): + async def connect_to(self, address: tuple[str, int]) -> None: + """ + This simulates the behavior of connect_to when the ``redis.asyncio.sentinel.SentinelConnectionPool`` + is in master mode. + It'll try to connect to master but we should not create any connection in test, + so just set the host and port without actually connecting. + """ + self.host, self.port = address + + +@pytest_asyncio.fixture() +async def connection_pool_replica_mock() -> SentinelConnectionPool: + sentinel_manager = Sentinel([["master", 400]]) + # Give a random slave + sentinel_manager.discover_slaves = AsyncMock(return_value=["replica", 5000]) # type: ignore[method-assign] + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=False, + connection_class=SentinelManagedConnectionMockForReplicaMode, + ) + return connection_pool + + +@pytest_asyncio.fixture() +async def connection_pool_master_mock() -> SentinelConnectionPool: + sentinel_manager = Sentinel([["master", 400]]) + # Give a random slave + sentinel_manager.discover_master = AsyncMock(return_value=["replica", 5000]) # type: ignore[method-assign] + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=True, + connection_class=SentinelManagedConnectionMockForMasterMode, + ) + return connection_pool + + +async def test_connection_pool_connects_to_same_address_if_same_iter_req_id_in_replica_mode( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + # Assert that the connection address is the same if the _iter_req_id is the same + connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + assert same_address( + await connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ), + connection_for_req_1, + ) + + +async def test_connection_pool_returns_same_conn_object_if_same_iter_req_id_and_released_in_replica_mode( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + # Assert that the connection object is the same if the _iter_req_id is the same + connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + await connection_pool_replica_mock.release(connection_for_req_1) + assert ( + await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + == connection_for_req_1 + ) + + +async def test_connection_pool_connects_to_diff_address_if_no_iter_req_id_in_replica_mode( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + # Assert that the connection object is different if no _iter_req_id is supplied + # In reality, they can be the same, but in this case, we're not releasing the connection + # to the pool so they should always be different. + connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + connection_for_random_req = await connection_pool_replica_mock.get_connection("ANYY") + assert not same_address(connection_for_random_req, connection_for_req_1) + assert not same_address( + await connection_pool_replica_mock.get_connection("ANY_COMMAND"), + connection_for_random_req, + ) + assert not same_address( + await connection_pool_replica_mock.get_connection( + "ANY_COMMAND" + ), + connection_for_req_1, + ) + + +async def test_connection_pool_connects_to_same_address_if_same_iter_req_id_in_master_mode( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + # Assert that the connection address is the same if the _iter_req_id is the same + connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + assert same_address( + await connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ), + connection_for_req_1, + ) + + +async def test_connection_pool_returns_same_conn_object_if_same_iter_req_id_and_released_in_master_mode( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + # Assert that the connection address is the same if the _iter_req_id is the same + connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + assert same_address( + await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + connection_for_req_1, + ) + +async def test_connection_pool_connects_to_same_address_if_no_iter_req_id_in_master_mode( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + # Assert that connection address is always the same regardless if it's an iter command or not + connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + connection_for_random_req = await connection_pool_master_mock.get_connection("ANYY") + assert same_address(connection_for_random_req, connection_for_req_1) + assert same_address( + await connection_pool_master_mock.get_connection("ANY_COMMAND"), + connection_for_random_req + ) + + assert same_address( + await connection_pool_master_mock.get_connection("ANY_COMMAND"), + connection_for_req_1, + ) From 824eada260a257c3f62781ca2a7dac2840aecc6f Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 27 Apr 2024 10:40:43 +0800 Subject: [PATCH 03/44] reorder --- .../test_sentinel_managed_connection.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 2c6741fdc0..d4740ebf03 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -41,15 +41,6 @@ async def mock_connect(): assert conn._connect.call_count == 3 await conn.disconnect() - -def same_address( - connection_1: SentinelManagedConnection, - connection_2: SentinelManagedConnection, -) -> bool: - return bool( - connection_1.host == connection_2.host and connection_1.port == connection_2.port - ) - class SentinelManagedConnectionMock(SentinelManagedConnection): async def connect_to_address(self, host: str, port: int) -> None: self.host = host @@ -104,7 +95,7 @@ async def connection_pool_replica_mock() -> SentinelConnectionPool: @pytest_asyncio.fixture() async def connection_pool_master_mock() -> SentinelConnectionPool: - sentinel_manager = Sentinel([["master", 400]]) + sentinel_manager = Sentinel([["master", 400]]) # Give a random slave sentinel_manager.discover_master = AsyncMock(return_value=["replica", 5000]) # type: ignore[method-assign] # Create connection pool with our mock connection object @@ -117,6 +108,14 @@ async def connection_pool_master_mock() -> SentinelConnectionPool: return connection_pool +def same_address( + connection_1: SentinelManagedConnection, + connection_2: SentinelManagedConnection, +) -> bool: + return bool( + connection_1.host == connection_2.host and connection_1.port == connection_2.port + ) + async def test_connection_pool_connects_to_same_address_if_same_iter_req_id_in_replica_mode( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: @@ -197,7 +196,7 @@ async def test_connection_pool_connects_to_same_address_if_no_iter_req_id_in_mas await connection_pool_master_mock.get_connection("ANY_COMMAND"), connection_for_random_req ) - + assert same_address( await connection_pool_master_mock.get_connection("ANY_COMMAND"), connection_for_req_1, From 8f2ec5f18d147dcfb5dd76abee1232563e89b068 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 27 Apr 2024 10:46:10 +0800 Subject: [PATCH 04/44] remove ignore --- tests/test_asyncio/test_sentinel_managed_connection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index d4740ebf03..17ff5d7802 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -82,7 +82,7 @@ async def connect_to(self, address: tuple[str, int]) -> None: async def connection_pool_replica_mock() -> SentinelConnectionPool: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave - sentinel_manager.discover_slaves = AsyncMock(return_value=["replica", 5000]) # type: ignore[method-assign] + sentinel_manager.discover_slaves = AsyncMock(return_value=["replica", 5000]) # Create connection pool with our mock connection object connection_pool = SentinelConnectionPool( "usasm", @@ -97,7 +97,7 @@ async def connection_pool_replica_mock() -> SentinelConnectionPool: async def connection_pool_master_mock() -> SentinelConnectionPool: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave - sentinel_manager.discover_master = AsyncMock(return_value=["replica", 5000]) # type: ignore[method-assign] + sentinel_manager.discover_master = AsyncMock(return_value=["replica", 5000]) # Create connection pool with our mock connection object connection_pool = SentinelConnectionPool( "usasm", From ca9bd63eb3f16123fb8ea3861408a86109d3ffaf Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 27 Apr 2024 11:59:31 +0800 Subject: [PATCH 05/44] lint and format --- redis/asyncio/sentinel.py | 42 +++--- redis/commands/core.py | 35 +++-- .../test_sentinel_managed_connection.py | 126 ++++++++++++------ 3 files changed, 128 insertions(+), 75 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 0c642e1b1a..c6e4e926c8 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -1,8 +1,16 @@ import asyncio import random import weakref -import uuid -from typing import AsyncIterator, Iterable, Mapping, Optional, Sequence, Tuple, Type, Any +from typing import ( + Any, + AsyncIterator, + Iterable, + Mapping, + Optional, + Sequence, + Tuple, + Type, +) from redis.asyncio.client import Redis from redis.asyncio.connection import ( @@ -66,13 +74,13 @@ async def connect(self): self._connect_retry, lambda error: asyncio.sleep(0), ) - + async def _connect_to_address_retry(self, host: str, port: int) -> None: if self._reader: return # already connected try: return await self.connect_to((host, port)) - except ConnectionError as exc: + except ConnectionError: raise SlaveNotFoundError async def connect_to_address(self, host: str, port: int) -> None: @@ -170,11 +178,6 @@ async def get_master_address(self): async def rotate_slaves(self) -> AsyncIterator: """Round-robin slave balancer""" - ( - server_host, - server_port, - ) = self._request_id_to_replica_address.get(iter_req_id, (None, None)) - slaves = await self.sentinel_manager.discover_slaves(self.service_name) if slaves: if self.slave_rr_counter is None: @@ -201,15 +204,15 @@ async def get_connection( to be issued to the same Redis replica. The way each server positions each key is different with one another, - and the cursor acts as the 'offset' of the scan. - Hence, all scans coming from a single xxx_scan_iter_channel command + and the cursor acts as the 'offset' of the scan. + Hence, all scans coming from a single xxx_scan_iter_channel command should go to the same replica. """ # If not an iter command or in master mode, call super() # No custom logic for master, because there's only 1 master. # The bug is only when Redis has the possibility to connect to multiple replicas if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: - return await super().get_connection(command_name, *keys, **options) # type: ignore[no-any-return] + return await super().get_connection(command_name, *keys, **options) # Check if this iter request has already been directed to a particular server # Check if this iter request has already been directed to a particular server @@ -222,7 +225,7 @@ async def get_connection( # get a connection from the pool if server_host is None or server_port is None: try: - connection = self._available_connections.pop() # type: ignore [assignment] + connection = self._available_connections.pop() except IndexError: connection = self.make_connection() # If this is not the first scan request of the iter command @@ -236,7 +239,7 @@ async def get_connection( and available_connection.port == server_port ): self._available_connections.remove(available_connection) - connection = available_connection # type: ignore[assignment] + connection = available_connection # If not, make a new dummy connection object, and set its host and port # to the one that we want later in the call to ``connect_to_address`` if not connection: @@ -255,22 +258,18 @@ async def get_connection( # connect to the particular address and port else: # This will connect to the host and port that we've specified above - await connection.connect_to_address(server_host, server_port) # type: ignore[arg-type] + await connection.connect_to_address(server_host, server_port) # connections that the pool provides should be ready to send # a command. if not, the connection was either returned to the # pool before all data has been read or the socket has been # closed. either way, reconnect and verify everything is good. try: - # type ignore below: - # attr Not defined in redis stubs and - # we don't need to create a subclass to help with this single attr - if await connection.can_read_destructive(): # type: ignore[attr-defined] + if await connection.can_read_destructive(): raise ConnectionError("Connection has data") from None except (ConnectionError, OSError): await connection.disconnect() await connection.connect() - # type ignore below: similar to above - if await connection.can_read_destructive(): # type: ignore[attr-defined] + if await connection.can_read_destructive(): raise ConnectionError("Connection not ready") from None except BaseException: # release the connection back to the pool so that we don't @@ -286,7 +285,6 @@ async def get_connection( return connection - class Sentinel(AsyncSentinelCommands): """ Redis Sentinel cluster client diff --git a/redis/commands/core.py b/redis/commands/core.py index cc996b9974..d267f20f26 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -1,8 +1,8 @@ # from __future__ import annotations -import uuid import datetime import hashlib +import uuid import warnings from typing import ( TYPE_CHECKING, @@ -3226,12 +3226,19 @@ async def scan_iter( HASH, LIST, SET, STREAM, STRING, ZSET Additionally, Redis modules can expose other types as well. """ - # DO NOT inline this statement to the scan call - iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica + # DO NOT inline this statement to the scan call + # Each iter command should have an ID to maintain + # connection to the same replica + iter_req_id = uuid.uuid4() cursor = "0" while cursor != 0: cursor, data = await self.scan( - cursor=cursor, match=match, count=count, _type=_type, _iter_req_id=iter_req_id, **kwargs + cursor=cursor, + match=match, + count=count, + _type=_type, + _iter_req_id=iter_req_id, + **kwargs, ) for d in data: yield d @@ -3250,8 +3257,10 @@ async def sscan_iter( ``count`` allows for hint the minimum number of returns """ - # DO NOT inline this statement to the scan call - iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica + # DO NOT inline this statement to the scan call + # Each iter command should have an ID to maintain + # connection to the same replica + iter_req_id = uuid.uuid4() cursor = "0" while cursor != 0: cursor, data = await self.sscan( @@ -3277,8 +3286,10 @@ async def hscan_iter( ``no_values`` indicates to return only the keys, without values """ - # DO NOT inline this statement to the scan call - iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica + # DO NOT inline this statement to the scan call + # Each iter command should have an ID to maintain + # connection to the same replica + iter_req_id = uuid.uuid4() cursor = "0" while cursor != 0: cursor, data = await self.hscan( @@ -3312,8 +3323,10 @@ async def zscan_iter( ``score_cast_func`` a callable used to cast the score return value """ - # DO NOT inline this statement to the scan call - iter_req_id = uuid.uuid4() # each iter command should have an ID to maintain connection to the same replica + # DO NOT inline this statement to the scan call + # Each iter command should have an ID to maintain + # connection to the same replica + iter_req_id = uuid.uuid4() cursor = "0" while cursor != 0: cursor, data = await self.zscan( @@ -3322,7 +3335,7 @@ async def zscan_iter( match=match, count=count, score_cast_func=score_cast_func, - _iter_req_id=iter_req_id + _iter_req_id=iter_req_id, ) for d in data: yield d diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 17ff5d7802..222d121b3e 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,14 +1,15 @@ import socket +from unittest.mock import AsyncMock import pytest import pytest_asyncio from redis.asyncio.retry import Retry -from redis.asyncio.sentinel import SentinelManagedConnection, SentinelConnectionPool, Sentinel +from redis.asyncio.sentinel import ( + Sentinel, + SentinelConnectionPool, + SentinelManagedConnection, +) from redis.backoff import NoBackoff -from unittest.mock import AsyncMock -from typing import AsyncIterator - -import pytest from .compat import mock @@ -41,6 +42,7 @@ async def mock_connect(): assert conn._connect.call_count == 3 await conn.disconnect() + class SentinelManagedConnectionMock(SentinelManagedConnection): async def connect_to_address(self, host: str, port: int) -> None: self.host = host @@ -56,9 +58,10 @@ async def can_read_destructive(self) -> bool: class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock): async def connect(self) -> None: """ - This simulates the behavior of connect when the ``redis.asyncio.sentinel.SentinelConnectionPool`` + This simulates the behavior of connect when + :py:class:`~redis.SentinelConnectionPool` is in replica mode. - It'll call rotate_slaves and connect to the next replica + It'll call rotate_slaves and connect to the next replica. """ import random import time @@ -70,10 +73,11 @@ async def connect(self) -> None: class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): async def connect_to(self, address: tuple[str, int]) -> None: """ - This simulates the behavior of connect_to when the ``redis.asyncio.sentinel.SentinelConnectionPool`` + This simulates the behavior of connect_to when + :py:class:`~redis.SentinelConnectionPool` is in master mode. - It'll try to connect to master but we should not create any connection in test, - so just set the host and port without actually connecting. + It'll try to connect to master. In this mock class, + it'll just set the host and port without actually connecting. """ self.host, self.port = address @@ -113,27 +117,41 @@ def same_address( connection_2: SentinelManagedConnection, ) -> bool: return bool( - connection_1.host == connection_2.host and connection_1.port == connection_2.port + connection_1.host == connection_2.host + and connection_1.port == connection_2.port ) -async def test_connection_pool_connects_to_same_address_if_same_iter_req_id_in_replica_mode( + +async def test_connects_to_same_address_if_same_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: - # Assert that the connection address is the same if the _iter_req_id is the same - connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in replica mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + """ + connection_for_req_1 = await connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) assert same_address( - await connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 - ), + await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1), connection_for_req_1, ) -async def test_connection_pool_returns_same_conn_object_if_same_iter_req_id_and_released_in_replica_mode( +async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: - # Assert that the connection object is the same if the _iter_req_id is the same - connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + """ + Assert that the connection object is the same if the ``_iter_req_id`` is the same + when we are in replica mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + and if we release the connection back to the connection pool before + trying to connect again + """ + connection_for_req_1 = await connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) await connection_pool_replica_mock.release(connection_for_req_1) assert ( await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) @@ -141,60 +159,84 @@ async def test_connection_pool_returns_same_conn_object_if_same_iter_req_id_and_ ) -async def test_connection_pool_connects_to_diff_address_if_no_iter_req_id_in_replica_mode( +async def test_connects_to_diff_address_if_no_iter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: - # Assert that the connection object is different if no _iter_req_id is supplied - # In reality, they can be the same, but in this case, we're not releasing the connection - # to the pool so they should always be different. - connection_for_req_1 = await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) - connection_for_random_req = await connection_pool_replica_mock.get_connection("ANYY") + """ + Assert that the connection object is different if no _iter_req_id is supplied + In reality, they can be the same, but in this case, we're not + releasing the connection to the pool so they should always be different. + """ + connection_for_req_1 = await connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) + connection_for_random_req = await connection_pool_replica_mock.get_connection( + "ANYY" + ) assert not same_address(connection_for_random_req, connection_for_req_1) assert not same_address( await connection_pool_replica_mock.get_connection("ANY_COMMAND"), connection_for_random_req, ) assert not same_address( - await connection_pool_replica_mock.get_connection( - "ANY_COMMAND" - ), + await connection_pool_replica_mock.get_connection("ANY_COMMAND"), connection_for_req_1, ) -async def test_connection_pool_connects_to_same_address_if_same_iter_req_id_in_master_mode( +async def test_connects_to_same_address_if_same_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: - # Assert that the connection address is the same if the _iter_req_id is the same - connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + """ + connection_for_req_1 = await connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) assert same_address( - await connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 - ), + await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), connection_for_req_1, ) -async def test_connection_pool_returns_same_conn_object_if_same_iter_req_id_and_released_in_master_mode( +async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: - # Assert that the connection address is the same if the _iter_req_id is the same - connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + and if we release the connection back to the connection pool before + trying to connect again + """ + connection_for_req_1 = await connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) assert same_address( await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), connection_for_req_1, ) -async def test_connection_pool_connects_to_same_address_if_no_iter_req_id_in_master_mode( + +async def test_connects_to_same_address_if_no_iter_req_id_in_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: - # Assert that connection address is always the same regardless if it's an iter command or not - connection_for_req_1 = await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1) + """ + Assert that connection address is always the same regardless if + there's an ``iter_req_id`` or not + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + """ + connection_for_req_1 = await connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) connection_for_random_req = await connection_pool_master_mock.get_connection("ANYY") assert same_address(connection_for_random_req, connection_for_req_1) assert same_address( await connection_pool_master_mock.get_connection("ANY_COMMAND"), - connection_for_random_req + connection_for_random_req, ) assert same_address( From 076e5924471fcf4ccd64b4720110615a5b0230c9 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Mon, 29 Apr 2024 21:36:13 +0800 Subject: [PATCH 06/44] better inline --- redis/asyncio/sentinel.py | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index c6e4e926c8..cf45d8b3a7 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -198,23 +198,21 @@ async def get_connection( ) -> SentinelManagedConnection: """ Get a connection from the pool. - `xxx_scan_iter` commands needs to be handled specially. + 'xxxscan_iter' ('scan_iter', 'hscan_iter', 'sscan_iter', 'zscan_iter') + commands needs to be handled specially. If the client is created using a connection pool, in replica mode, - all `scan` command-equivalent of the `xxx_scan_iter` commands needs + all 'scan' command-equivalent of the 'xxx_scan_iter' commands needs to be issued to the same Redis replica. The way each server positions each key is different with one another, - and the cursor acts as the 'offset' of the scan. - Hence, all scans coming from a single xxx_scan_iter_channel command + and the cursor acts as the offset of the scan. + Hence, all scans coming from a single 'xxx_scan_iter_channel' command should go to the same replica. """ - # If not an iter command or in master mode, call super() - # No custom logic for master, because there's only 1 master. - # The bug is only when Redis has the possibility to connect to multiple replicas + # If not an iter command or in master mode, call superclass' implementation if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: return await super().get_connection(command_name, *keys, **options) - # Check if this iter request has already been directed to a particular server # Check if this iter request has already been directed to a particular server ( server_host, @@ -232,8 +230,8 @@ async def get_connection( else: # Check from the available connections, if any of the connection # is connected to the host and port that we want - # If yes, use that connection for available_connection in self._available_connections.copy(): + # if yes, use that connection if ( available_connection.host == server_host and available_connection.port == server_port @@ -247,22 +245,20 @@ async def get_connection( assert connection self._in_use_connections.add(connection) try: - # ensure this connection is connected to Redis - # If this is the first scan request, - # just call the SentinelManagedConnection.connect() - # This will call rotate_slaves - # and connect to a random replica + # Ensure this connection is connected to Redis + # If this is the first scan request, it will + # call rotate_slaves and connect to a random replica if server_port is None or server_port is None: await connection.connect() # If this is not the first scan request, - # connect to the particular address and port + # connect to the previous replica. + # This will connect to the host and port of the replica else: - # This will connect to the host and port that we've specified above await connection.connect_to_address(server_host, server_port) - # connections that the pool provides should be ready to send - # a command. if not, the connection was either returned to the + # Connections that the pool provides should be ready to send + # a command. If not, the connection was either returned to the # pool before all data has been read or the socket has been - # closed. either way, reconnect and verify everything is good. + # closed. Either way, reconnect and verify everything is good. try: if await connection.can_read_destructive(): raise ConnectionError("Connection has data") from None @@ -272,7 +268,7 @@ async def get_connection( if await connection.can_read_destructive(): raise ConnectionError("Connection not ready") from None except BaseException: - # release the connection back to the pool so that we don't + # Release the connection back to the pool so that we don't # leak it await self.release(connection) raise @@ -281,7 +277,6 @@ async def get_connection( connection.host, connection.port, ) - return connection From 224f2536754bca499d5c09fa7f8433c4d77ecd2f Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Mon, 29 Apr 2024 22:22:07 +0800 Subject: [PATCH 07/44] backward compatible typing --- tests/test_asyncio/test_sentinel_managed_connection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 222d121b3e..369097220b 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,4 +1,5 @@ import socket +from typing import Tuple from unittest.mock import AsyncMock import pytest @@ -71,7 +72,7 @@ async def connect(self) -> None: class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): - async def connect_to(self, address: tuple[str, int]) -> None: + async def connect_to(self, address: Tuple[str, int]) -> None: """ This simulates the behavior of connect_to when :py:class:`~redis.SentinelConnectionPool` From 8081e03062df10da2e460a5e0a0caefc73bee20b Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Fri, 10 May 2024 22:32:40 +0800 Subject: [PATCH 08/44] test inline docs --- .../test_sentinel_managed_connection.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 369097220b..1707fb6b6f 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -129,7 +129,7 @@ async def test_connects_to_same_address_if_same_id_replica( """ Assert that the connection address is the same if the ``_iter_req_id`` is the same when we are in replica mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_replica_mock.get_connection( "ANY", _iter_req_id=1 @@ -148,7 +148,7 @@ async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica when we are in replica mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before - trying to connect again + trying to connect again. """ connection_for_req_1 = await connection_pool_replica_mock.get_connection( "ANY", _iter_req_id=1 @@ -164,7 +164,7 @@ async def test_connects_to_diff_address_if_no_iter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection object is different if no _iter_req_id is supplied + Assert that the connection object is different if no _iter_req_id is supplied. In reality, they can be the same, but in this case, we're not releasing the connection to the pool so they should always be different. """ @@ -191,7 +191,7 @@ async def test_connects_to_same_address_if_same_iter_req_id_master( """ Assert that the connection address is the same if the ``_iter_req_id`` is the same when we are in master mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( "ANY", _iter_req_id=1 @@ -210,7 +210,7 @@ async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_mas when we are in master mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before - trying to connect again + trying to connect again. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( "ANY", _iter_req_id=1 @@ -221,14 +221,14 @@ async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_mas ) -async def test_connects_to_same_address_if_no_iter_req_id_in_master( +async def test_connects_to_same_address_if_no_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ Assert that connection address is always the same regardless if there's an ``iter_req_id`` or not when we are in master mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( "ANY", _iter_req_id=1 From ed2b539f50cbf6e0ff077ec9963296991302d93d Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Fri, 10 May 2024 23:13:41 +0800 Subject: [PATCH 09/44] add tests for all scan iter family --- redis/commands/core.py | 8 ++---- .../test_sentinel_managed_connection.py | 26 ++++++++++++++++++- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/redis/commands/core.py b/redis/commands/core.py index d267f20f26..d750be721d 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3068,7 +3068,7 @@ def sscan( pieces.extend([b"MATCH", match]) if count is not None: pieces.extend([b"COUNT", count]) - return self.execute_command("SSCAN", *pieces,) + return self.execute_command("SSCAN", *pieces, **kwargs) def sscan_iter( self, @@ -3293,11 +3293,7 @@ async def hscan_iter( cursor = "0" while cursor != 0: cursor, data = await self.hscan( -<<<<<<< HEAD - name, cursor=cursor, match=match, count=count, no_values=no_values -======= - name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id ->>>>>>> 39aaec6 (fix scan iter command issued to different replicas) + name, cursor=cursor, match=match, count=count, no_values=no_values, _iter_req_id=iter_req_id ) if no_values: for it in data: diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 1707fb6b6f..adf741f611 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,10 +1,12 @@ import socket from typing import Tuple -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, ANY import pytest import pytest_asyncio +import uuid from redis.asyncio.retry import Retry +from redis.commands.core import AsyncScanCommands from redis.asyncio.sentinel import ( Sentinel, SentinelConnectionPool, @@ -244,3 +246,25 @@ async def test_connects_to_same_address_if_no_iter_req_id_master( await connection_pool_master_mock.get_connection("ANY_COMMAND"), connection_for_req_1, ) + + +async def test_scan_iter_family_executes_commands_with_same_iter_req_id(): + """Assert that all calls to execute_command receives the _iter_req_id kwarg""" + with mock.patch.object( + AsyncScanCommands, + "execute_command", + AsyncMock(return_value=(0, [])) + ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): + [a async for a in AsyncScanCommands().scan_iter()] + mock_execute_command.assert_called_with('SCAN', '0', _iter_req_id="uuid") + [a async for a in AsyncScanCommands().sscan_iter("")] + mock_execute_command.assert_called_with('SSCAN', '', '0', _iter_req_id="uuid") + with mock.patch.object( + AsyncScanCommands, + "execute_command", + AsyncMock(return_value=(0, {})) + ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): + [a async for a in AsyncScanCommands().hscan_iter("")] + mock_execute_command.assert_called_with('HSCAN', '', '0', no_values=None, _iter_req_id="uuid") + [a async for a in AsyncScanCommands().zscan_iter("")] + mock_execute_command.assert_called_with('ZSCAN', '', '0', score_cast_func=ANY, _iter_req_id="uuid") From 5d7800e7a8475c12ef2bc5118b904ab20a8dbf51 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 11 May 2024 00:06:40 +0800 Subject: [PATCH 10/44] lint --- redis/commands/core.py | 7 ++++- .../test_sentinel_managed_connection.py | 26 +++++++++---------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/redis/commands/core.py b/redis/commands/core.py index d750be721d..9ace469c01 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3293,7 +3293,12 @@ async def hscan_iter( cursor = "0" while cursor != 0: cursor, data = await self.hscan( - name, cursor=cursor, match=match, count=count, no_values=no_values, _iter_req_id=iter_req_id + name, + cursor=cursor, + match=match, + count=count, + no_values=no_values, + _iter_req_id=iter_req_id, ) if no_values: for it in data: diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index adf741f611..18195cf53e 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,18 +1,18 @@ import socket +import uuid from typing import Tuple -from unittest.mock import AsyncMock, ANY +from unittest.mock import ANY, AsyncMock import pytest import pytest_asyncio -import uuid from redis.asyncio.retry import Retry -from redis.commands.core import AsyncScanCommands from redis.asyncio.sentinel import ( Sentinel, SentinelConnectionPool, SentinelManagedConnection, ) from redis.backoff import NoBackoff +from redis.commands.core import AsyncScanCommands from .compat import mock @@ -251,20 +251,20 @@ async def test_connects_to_same_address_if_no_iter_req_id_master( async def test_scan_iter_family_executes_commands_with_same_iter_req_id(): """Assert that all calls to execute_command receives the _iter_req_id kwarg""" with mock.patch.object( - AsyncScanCommands, - "execute_command", - AsyncMock(return_value=(0, [])) + AsyncScanCommands, "execute_command", AsyncMock(return_value=(0, [])) ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): [a async for a in AsyncScanCommands().scan_iter()] - mock_execute_command.assert_called_with('SCAN', '0', _iter_req_id="uuid") + mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") [a async for a in AsyncScanCommands().sscan_iter("")] - mock_execute_command.assert_called_with('SSCAN', '', '0', _iter_req_id="uuid") + mock_execute_command.assert_called_with("SSCAN", "", "0", _iter_req_id="uuid") with mock.patch.object( - AsyncScanCommands, - "execute_command", - AsyncMock(return_value=(0, {})) + AsyncScanCommands, "execute_command", AsyncMock(return_value=(0, {})) ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): [a async for a in AsyncScanCommands().hscan_iter("")] - mock_execute_command.assert_called_with('HSCAN', '', '0', no_values=None, _iter_req_id="uuid") + mock_execute_command.assert_called_with( + "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" + ) [a async for a in AsyncScanCommands().zscan_iter("")] - mock_execute_command.assert_called_with('ZSCAN', '', '0', score_cast_func=ANY, _iter_req_id="uuid") + mock_execute_command.assert_called_with( + "ZSCAN", "", "0", score_cast_func=ANY, _iter_req_id="uuid" + ) From b964b9b2cd08db0249290349b6a65b35a507b6be Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 9 Jul 2024 22:15:13 +0800 Subject: [PATCH 11/44] implement in sync client --- redis/asyncio/client.py | 2 + redis/asyncio/connection.py | 14 +- redis/asyncio/sentinel.py | 45 ++-- redis/commands/core.py | 23 +- redis/connection.py | 94 ++++--- redis/sentinel.py | 80 +++++- tests/test_asyncio/test_commands.py | 33 +++ .../test_sentinel_managed_connection.py | 46 ++-- tests/test_commands.py | 31 +++ tests/test_connection.py | 16 +- tests/test_sentinel_managed_connection.py | 253 ++++++++++++++++++ 11 files changed, 540 insertions(+), 97 deletions(-) create mode 100644 tests/test_sentinel_managed_connection.py diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index 1845b7252f..fc3f2e7569 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -651,6 +651,8 @@ async def execute_command(self, *args, **options): finally: if not self.connection: await pool.release(conn) + if "ITER" in command_name.upper(): + pool.cleanup_scan(iter_req_id=options.get("_iter_req_id", None)) async def parse_response( self, connection: Connection, command_name: Union[str, bytes], **options diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 2ac6637986..0a907b40b7 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -38,7 +38,7 @@ from redis.asyncio.retry import Retry from redis.backoff import NoBackoff -from redis.connection import DEFAULT_RESP_VERSION +from redis.connection import DEFAULT_RESP_VERSION, ConnectionsIndexer from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider from redis.exceptions import ( AuthenticationError, @@ -1057,6 +1057,13 @@ class ConnectionPool: ``connection_class``. """ + @abstractmethod + def cleanup_scan(self, **options): + """ + Additional cleanup operations that the connection pool might + need to do after a SCAN ITER family command is executed + """ + @classmethod def from_url(cls: Type[_CP], url: str, **kwargs) -> _CP: """ @@ -1118,7 +1125,7 @@ def __init__( self.connection_kwargs = connection_kwargs self.max_connections = max_connections - self._available_connections: List[AbstractConnection] = [] + self._available_connections: ConnectionsIndexer = ConnectionsIndexer() self._in_use_connections: Set[AbstractConnection] = set() self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) @@ -1324,3 +1331,6 @@ async def release(self, connection: AbstractConnection): async with self._condition: await super().release(connection) self._condition.notify() + + def cleanup_scan(self, **options): + pass diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index cf45d8b3a7..5c210845ca 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -147,7 +147,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs): self.sentinel_manager = sentinel_manager self.master_address = None self.slave_rr_counter = None - self._request_id_to_replica_address = {} + self._iter_req_id_to_replica_address = {} def __repr__(self): return ( @@ -193,6 +193,14 @@ async def rotate_slaves(self) -> AsyncIterator: pass raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") + def cleanup_scan(self, **options): + """ + Remove the SCAN ITER family command's request id from the dictionary + """ + self._iter_req_id_to_replica_address.pop( + options.get("_iter_req_id", None), None + ) + async def get_connection( self, command_name: str, *keys: Any, **options: Any ) -> SentinelManagedConnection: @@ -217,7 +225,7 @@ async def get_connection( ( server_host, server_port, - ) = self._request_id_to_replica_address.get(iter_req_id, (None, None)) + ) = self._iter_req_id_to_replica_address.get(iter_req_id, (None, None)) connection = None # If this is the first scan request of the iter command, # get a connection from the pool @@ -228,18 +236,12 @@ async def get_connection( connection = self.make_connection() # If this is not the first scan request of the iter command else: - # Check from the available connections, if any of the connection - # is connected to the host and port that we want - for available_connection in self._available_connections.copy(): - # if yes, use that connection - if ( - available_connection.host == server_host - and available_connection.port == server_port - ): - self._available_connections.remove(available_connection) - connection = available_connection - # If not, make a new dummy connection object, and set its host and port - # to the one that we want later in the call to ``connect_to_address`` + # Get the connection that has the same host and port + connection = self._available_connections.get_connection( + host=server_host, port=server_port + ) + # If not, make a new dummy connection object, and set its host and + # port to the one that we want later in the call to ``connect_to_address`` if not connection: connection = self.make_connection() assert connection @@ -255,25 +257,14 @@ async def get_connection( # This will connect to the host and port of the replica else: await connection.connect_to_address(server_host, server_port) - # Connections that the pool provides should be ready to send - # a command. If not, the connection was either returned to the - # pool before all data has been read or the socket has been - # closed. Either way, reconnect and verify everything is good. - try: - if await connection.can_read_destructive(): - raise ConnectionError("Connection has data") from None - except (ConnectionError, OSError): - await connection.disconnect() - await connection.connect() - if await connection.can_read_destructive(): - raise ConnectionError("Connection not ready") from None + self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it await self.release(connection) raise # Store the connection to the dictionary - self._request_id_to_replica_address[iter_req_id] = ( + self._iter_req_id_to_replica_address[iter_req_id] = ( connection.host, connection.port, ) diff --git a/redis/commands/core.py b/redis/commands/core.py index 9ace469c01..999a14622a 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3039,9 +3039,15 @@ def scan_iter( Additionally, Redis modules can expose other types as well. """ cursor = "0" + iter_req_id = uuid.uuid4() while cursor != 0: cursor, data = self.scan( - cursor=cursor, match=match, count=count, _type=_type, **kwargs + cursor=cursor, + match=match, + count=count, + _type=_type, + _iter_req_id=iter_req_id, + **kwargs, ) yield from data @@ -3085,8 +3091,11 @@ def sscan_iter( ``count`` allows for hint the minimum number of returns """ cursor = "0" + iter_req_id = uuid.uuid4() while cursor != 0: - cursor, data = self.sscan(name, cursor=cursor, match=match, count=count) + cursor, data = self.sscan( + name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id + ) yield from data def hscan( @@ -3137,9 +3146,15 @@ def hscan_iter( ``no_values`` indicates to return only the keys, without values """ cursor = "0" + iter_req_id = uuid.uuid4() while cursor != 0: cursor, data = self.hscan( - name, cursor=cursor, match=match, count=count, no_values=no_values + name, + cursor=cursor, + match=match, + count=count, + no_values=no_values, + _iter_req_id=iter_req_id, ) if no_values: yield from data @@ -3193,6 +3208,7 @@ def zscan_iter( ``score_cast_func`` a callable used to cast the score return value """ cursor = "0" + iter_req_id = uuid.uuid4() while cursor != 0: cursor, data = self.zscan( name, @@ -3200,6 +3216,7 @@ def zscan_iter( match=match, count=count, score_cast_func=score_cast_func, + _iter_req_id=iter_req_id, ) yield from data diff --git a/redis/connection.py b/redis/connection.py index 1f862d0371..3ac0b23cdc 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -6,10 +6,11 @@ import threading import weakref from abc import abstractmethod +from collections import defaultdict from itertools import chain from queue import Empty, Full, LifoQueue from time import time -from typing import Any, Callable, List, Optional, Sequence, Type, Union +from typing import Any, Callable, Iterable, List, Optional, Sequence, Type, Union from urllib.parse import parse_qs, unquote, urlparse from ._cache import ( @@ -734,6 +735,43 @@ def _host_error(self): return f"{self.host}:{self.port}" +class ConnectionsIndexer(Iterable): + """ + Data structure that manages a list of available connections which + is also indexed based on the address (ip and port) of the connection + """ + + def __init__(self): + self._connections = [] + # Map the address to a dictionary of connections + # The inner dictionary is a map between the object id to the object itself + # This is to support O(1) operations on all of the class' methods + self._address_to_connections = defaultdict(dict) + + def pop(self): + connection = self._connections.pop() + del self._address_to_connections[(connection.host, connection.port)][ + id(connection) + ] + return connection + + def append(self, connection: Connection): + self._connections.append(connection) + self._address_to_connections[(connection.host, connection.port)][ + id(connection) + ] = connection + + def get_connection(self, host: str, port: int): + try: + connection = self._address_to_connections[(host, port)].popitem() + except KeyError: + return None + return connection + + def __iter__(self): + return iter(self._connections) + + class SSLConnection(Connection): """Manages SSL connections to and from the Redis server(s). This class extends the Connection class, adding SSL functionality, and making @@ -1113,7 +1151,7 @@ def __repr__(self) -> (str, str): def reset(self) -> None: self._lock = threading.Lock() self._created_connections = 0 - self._available_connections = [] + self._available_connections = ConnectionsIndexer() self._in_use_connections = set() # this must be the last operation in this method. while reset() is @@ -1174,6 +1212,25 @@ def _checkpid(self) -> None: finally: self._fork_lock.release() + def ensure_connection(self, connection: AbstractConnection): + # ensure this connection is connected to Redis + connection.connect() + # if client caching is not enabled connections that the pool + # provides should be ready to send a command. + # if not, the connection was either returned to the + # pool before all data has been read or the socket has been + # closed. either way, reconnect and verify everything is good. + # (if caching enabled the connection will not always be ready + # to send a command because it may contain invalidation messages) + try: + if connection.can_read() and connection.client_cache is None: + raise ConnectionError("Connection has data") + except (ConnectionError, OSError): + connection.disconnect() + connection.connect() + if connection.can_read(): + raise ConnectionError("Connection not ready") + def get_connection(self, command_name: str, *keys, **options) -> "Connection": "Get a connection from the pool" self._checkpid() @@ -1185,23 +1242,7 @@ def get_connection(self, command_name: str, *keys, **options) -> "Connection": self._in_use_connections.add(connection) try: - # ensure this connection is connected to Redis - connection.connect() - # if client caching is not enabled connections that the pool - # provides should be ready to send a command. - # if not, the connection was either returned to the - # pool before all data has been read or the socket has been - # closed. either way, reconnect and verify everything is good. - # (if caching enabled the connection will not always be ready - # to send a command because it may contain invalidation messages) - try: - if connection.can_read() and connection.client_cache is None: - raise ConnectionError("Connection has data") - except (ConnectionError, OSError): - connection.disconnect() - connection.connect() - if connection.can_read(): - raise ConnectionError("Connection not ready") + self.ensure_connection(connection) except BaseException: # release the connection back to the pool so that we don't # leak it @@ -1414,20 +1455,7 @@ def get_connection(self, command_name, *keys, **options): connection = self.make_connection() try: - # ensure this connection is connected to Redis - connection.connect() - # connections that the pool provides should be ready to send - # a command. if not, the connection was either returned to the - # pool before all data has been read or the socket has been - # closed. either way, reconnect and verify everything is good. - try: - if connection.can_read(): - raise ConnectionError("Connection has data") - except (ConnectionError, OSError): - connection.disconnect() - connection.connect() - if connection.can_read(): - raise ConnectionError("Connection not ready") + self.ensure_connection(connection) except BaseException: # release the connection back to the pool so that we don't leak it self.release(connection) diff --git a/redis/sentinel.py b/redis/sentinel.py index 72b5bef548..fefd3f4e21 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -1,6 +1,6 @@ import random import weakref -from typing import Optional +from typing import Any, Optional from redis.client import Redis from redis.commands import SentinelCommands @@ -164,6 +164,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs): self.connection_kwargs["connection_pool"] = self.proxy self.service_name = service_name self.sentinel_manager = sentinel_manager + self._iter_req_id_to_replica_address = {} def __repr__(self): role = "master" if self.is_master else "slave" @@ -194,6 +195,83 @@ def rotate_slaves(self): "Round-robin slave balancer" return self.proxy.rotate_slaves() + def cleanup_scan(self, **options): + """ + Remove the SCAN ITER family command's request id from the dictionary + """ + self._iter_req_id_to_replica_address.pop( + options.get("_iter_req_id", None), None + ) + + def get_connection( + self, command_name: str, *keys: Any, **options: Any + ) -> SentinelManagedConnection: + """ + Get a connection from the pool. + 'xxxscan_iter' ('scan_iter', 'hscan_iter', 'sscan_iter', 'zscan_iter') + commands needs to be handled specially. + If the client is created using a connection pool, in replica mode, + all 'scan' command-equivalent of the 'xxx_scan_iter' commands needs + to be issued to the same Redis replica. + + The way each server positions each key is different with one another, + and the cursor acts as the offset of the scan. + Hence, all scans coming from a single 'xxx_scan_iter_channel' command + should go to the same replica. + """ + # If not an iter command or in master mode, call superclass' implementation + if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: + return super().get_connection(command_name, *keys, **options) + + # Check if this iter request has already been directed to a particular server + ( + server_host, + server_port, + ) = self._iter_req_id_to_replica_address.get(iter_req_id, (None, None)) + connection = None + # If this is the first scan request of the iter command, + # get a connection from the pool + if server_host is None or server_port is None: + try: + connection = self._available_connections.pop() + except IndexError: + connection = self.make_connection() + # If this is not the first scan request of the iter command + else: + # Get the connection that has the same host and port + connection = self._available_connections.get_connection( + host=server_host, port=server_port + ) + # If not, make a new dummy connection object, and set its host and port + # to the one that we want later in the call to ``connect_to_address`` + if not connection: + connection = self.make_connection() + assert connection + self._in_use_connections.add(connection) + try: + # Ensure this connection is connected to Redis + # If this is the first scan request, it will + # call rotate_slaves and connect to a random replica + if server_port is None or server_port is None: + connection.connect() + # If this is not the first scan request, + # connect to the previous replica. + # This will connect to the host and port of the replica + else: + connection.connect_to_address(server_host, server_port) + self.ensure_connection(connection) + except BaseException: + # Release the connection back to the pool so that we don't + # leak it + self.release(connection) + raise + # Store the connection to the dictionary + self._iter_req_id_to_replica_address[iter_req_id] = ( + connection.host, + connection.port, + ) + return connection + class Sentinel(SentinelCommands): """ diff --git a/tests/test_asyncio/test_commands.py b/tests/test_asyncio/test_commands.py index 61c00541cb..4b2b2c1743 100644 --- a/tests/test_asyncio/test_commands.py +++ b/tests/test_asyncio/test_commands.py @@ -1402,6 +1402,39 @@ async def test_zscan_iter(self, r: redis.Redis): pairs = [k async for k in r.zscan_iter("a", match="a")] assert set(pairs) == {(b"a", 1)} + async def test_scan_iter_family_executes_commands_with_same_iter_req_id(): + """Assert that all calls to execute_command receives the _iter_req_id kwarg""" + import uuid + + from redis.commands.core import AsyncScanCommands + + from .compat import mock + + with mock.patch.object( + AsyncScanCommands, "execute_command", mock.AsyncMock(return_value=(0, [])) + ) as mock_execute_command, mock.patch.object( + uuid, "uuid4", return_value="uuid" + ): + [a async for a in AsyncScanCommands().scan_iter()] + mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") + [a async for a in AsyncScanCommands().sscan_iter("")] + mock_execute_command.assert_called_with( + "SSCAN", "", "0", _iter_req_id="uuid" + ) + with mock.patch.object( + AsyncScanCommands, "execute_command", mock.AsyncMock(return_value=(0, {})) + ) as mock_execute_command, mock.patch.object( + uuid, "uuid4", return_value="uuid" + ): + [a async for a in AsyncScanCommands().hscan_iter("")] + mock_execute_command.assert_called_with( + "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" + ) + [a async for a in AsyncScanCommands().zscan_iter("")] + mock_execute_command.assert_called_with( + "ZSCAN", "", "0", score_cast_func=mock.ANY, _iter_req_id="uuid" + ) + # SET COMMANDS async def test_sadd(self, r: redis.Redis): members = {b"1", b"2", b"3"} diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 18195cf53e..6ccd5b5077 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,10 +1,7 @@ import socket -import uuid from typing import Tuple -from unittest.mock import ANY, AsyncMock import pytest -import pytest_asyncio from redis.asyncio.retry import Retry from redis.asyncio.sentinel import ( Sentinel, @@ -12,7 +9,6 @@ SentinelManagedConnection, ) from redis.backoff import NoBackoff -from redis.commands.core import AsyncScanCommands from .compat import mock @@ -85,11 +81,11 @@ async def connect_to(self, address: Tuple[str, int]) -> None: self.host, self.port = address -@pytest_asyncio.fixture() -async def connection_pool_replica_mock() -> SentinelConnectionPool: +@pytest.fixture() +def connection_pool_replica_mock() -> SentinelConnectionPool: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave - sentinel_manager.discover_slaves = AsyncMock(return_value=["replica", 5000]) + sentinel_manager.discover_slaves = mock.AsyncMock(return_value=["replica", 5000]) # Create connection pool with our mock connection object connection_pool = SentinelConnectionPool( "usasm", @@ -100,11 +96,11 @@ async def connection_pool_replica_mock() -> SentinelConnectionPool: return connection_pool -@pytest_asyncio.fixture() -async def connection_pool_master_mock() -> SentinelConnectionPool: +@pytest.fixture() +def connection_pool_master_mock() -> SentinelConnectionPool: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave - sentinel_manager.discover_master = AsyncMock(return_value=["replica", 5000]) + sentinel_manager.discover_master = mock.AsyncMock(return_value=["replica", 5000]) # Create connection pool with our mock connection object connection_pool = SentinelConnectionPool( "usasm", @@ -248,23 +244,13 @@ async def test_connects_to_same_address_if_no_iter_req_id_master( ) -async def test_scan_iter_family_executes_commands_with_same_iter_req_id(): - """Assert that all calls to execute_command receives the _iter_req_id kwarg""" - with mock.patch.object( - AsyncScanCommands, "execute_command", AsyncMock(return_value=(0, [])) - ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): - [a async for a in AsyncScanCommands().scan_iter()] - mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") - [a async for a in AsyncScanCommands().sscan_iter("")] - mock_execute_command.assert_called_with("SSCAN", "", "0", _iter_req_id="uuid") - with mock.patch.object( - AsyncScanCommands, "execute_command", AsyncMock(return_value=(0, {})) - ) as mock_execute_command, mock.patch.object(uuid, "uuid4", return_value="uuid"): - [a async for a in AsyncScanCommands().hscan_iter("")] - mock_execute_command.assert_called_with( - "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" - ) - [a async for a in AsyncScanCommands().zscan_iter("")] - mock_execute_command.assert_called_with( - "ZSCAN", "", "0", score_cast_func=ANY, _iter_req_id="uuid" - ) +async def test_scan_iter_family_cleans_up( + connection_pool_replica_mock: SentinelConnectionPool, +): + """Test that connection pool is correctly cleaned up""" + from redis.asyncio import Redis + + r = Redis(connection_pool=connection_pool_replica_mock) + + [k async for k in r.scan_iter("a")] + assert not connection_pool_replica_mock._iter_req_id_to_replica_address diff --git a/tests/test_commands.py b/tests/test_commands.py index 42376b50d8..1a79a45a4c 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2235,6 +2235,37 @@ def test_zscan_iter(self, r): pairs = list(r.zscan_iter("a", match="a")) assert set(pairs) == {(b"a", 1)} + def test_scan_iter_family_executes_commands_with_same_iter_req_id(): + """Assert that all calls to execute_command receives the _iter_req_id kwarg""" + import uuid + + from redis.commands.core import ScanCommands + + with mock.patch.object( + ScanCommands, "execute_command", mock.Mock(return_value=(0, [])) + ) as mock_execute_command, mock.patch.object( + uuid, "uuid4", return_value="uuid" + ): + [a for a in ScanCommands().scan_iter()] + mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") + [a for a in ScanCommands().sscan_iter("")] + mock_execute_command.assert_called_with( + "SSCAN", "", "0", _iter_req_id="uuid" + ) + with mock.patch.object( + ScanCommands, "execute_command", mock.Mock(return_value=(0, {})) + ) as mock_execute_command, mock.patch.object( + uuid, "uuid4", return_value="uuid" + ): + [a for a in ScanCommands().hscan_iter("")] + mock_execute_command.assert_called_with( + "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" + ) + [a for a in ScanCommands().zscan_iter("")] + mock_execute_command.assert_called_with( + "ZSCAN", "", "0", score_cast_func=mock.ANY, _iter_req_id="uuid" + ) + # SET COMMANDS def test_sadd(self, r): members = {b"1", b"2", b"3"} diff --git a/tests/test_connection.py b/tests/test_connection.py index 69275d58c0..e5bbb597f5 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -10,6 +10,7 @@ from redis.backoff import NoBackoff from redis.connection import ( Connection, + ConnectionsIndexer, SSLConnection, UnixDomainSocketConnection, parse_url, @@ -297,7 +298,6 @@ def mock_disconnect(_): assert called == 1 pool.disconnect() - @pytest.mark.parametrize( "conn, error, expected_message", [ @@ -346,3 +346,17 @@ def test_unix_socket_connection_failure(): str(e.value) == "Error 2 connecting to unix:///tmp/a.sock. No such file or directory." ) + +def test_connections_indexer_operations(): + ci = ConnectionsIndexer() + c1 = Connection(host="1", port=2) + ci.append(c1) + assert list(ci) == [c1] + assert ci.pop() == c1 + + c2 = Connection(host="3", port=4) + ci.append(c1) + ci.append(c2) + assert ci.get_connection("3", 4) == c2 + assert not ci.get_connection("5", 6) + assert list(ci) == [c1] diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py new file mode 100644 index 0000000000..b4c927daa7 --- /dev/null +++ b/tests/test_sentinel_managed_connection.py @@ -0,0 +1,253 @@ +import socket +from typing import Tuple +from unittest import mock + +import pytest +from redis.asyncio.retry import Retry +from redis.asyncio.sentinel import ( + Sentinel, + SentinelConnectionPool, + SentinelManagedConnection, +) +from redis.backoff import NoBackoff + +pytestmark = pytest.mark.asyncio + + +def test_connect_retry_on_timeout_error(connect_args): + """Test that the _connect function is retried in case of a timeout""" + connection_pool = mock.AsyncMock() + connection_pool.get_master_address = mock.AsyncMock( + return_value=(connect_args["host"], connect_args["port"]) + ) + conn = SentinelManagedConnection( + retry_on_timeout=True, + retry=Retry(NoBackoff(), 3), + connection_pool=connection_pool, + ) + origin_connect = conn._connect + conn._connect = mock.AsyncMock() + + async def mock_connect(): + # connect only on the last retry + if conn._connect.call_count <= 2: + raise socket.timeout + else: + return await origin_connect() + + conn._connect.side_effect = mock_connect + conn.connect() + assert conn._connect.call_count == 3 + conn.disconnect() + + +class SentinelManagedConnectionMock(SentinelManagedConnection): + def connect_to_address(self, host: str, port: int) -> None: + self.host = host + self.port = port + + def can_read_destructive(self) -> bool: + # Mock this function to always return False. + # Trrue means there's still data to be read and hence we can't reconnect + # to this connection yet + return False + + +class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock): + def connect(self) -> None: + """ + This simulates the behavior of connect when + :py:class:`~redis.SentinelConnectionPool` + is in replica mode. + It'll call rotate_slaves and connect to the next replica. + """ + import random + import time + + self.host = f"host-{random.randint(0, 10)}" + self.port = time.time() + + +class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): + def connect_to(self, address: Tuple[str, int]) -> None: + """ + This simulates the behavior of connect_to when + :py:class:`~redis.SentinelConnectionPool` + is in master mode. + It'll try to connect to master. In this mock class, + it'll just set the host and port without actually connecting. + """ + self.host, self.port = address + + +@pytest.fixture() +def connection_pool_replica_mock() -> SentinelConnectionPool: + sentinel_manager = Sentinel([["master", 400]]) + # Give a random slave + sentinel_manager.discover_slaves = mock.Mock(return_value=["replica", 5000]) + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=False, + connection_class=SentinelManagedConnectionMockForReplicaMode, + ) + return connection_pool + + +@pytest.fixture() +def connection_pool_master_mock() -> SentinelConnectionPool: + sentinel_manager = Sentinel([["master", 400]]) + # Give a random slave + sentinel_manager.discover_master = mock.Mock(return_value=["replica", 5000]) + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=True, + connection_class=SentinelManagedConnectionMockForMasterMode, + ) + return connection_pool + + +def same_address( + connection_1: SentinelManagedConnection, + connection_2: SentinelManagedConnection, +) -> bool: + return bool( + connection_1.host == connection_2.host + and connection_1.port == connection_2.port + ) + + +def test_connects_to_same_address_if_same_id_replica( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in replica mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + """ + connection_for_req_1 = connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) + assert same_address( + connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1), + connection_for_req_1, + ) + + +def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + """ + Assert that the connection object is the same if the ``_iter_req_id`` is the same + when we are in replica mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + and if we release the connection back to the connection pool before + trying to connect again. + """ + connection_for_req_1 = connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) + connection_pool_replica_mock.release(connection_for_req_1) + assert ( + connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + == connection_for_req_1 + ) + + +def test_connects_to_diff_address_if_no_iter_req_id_replica( + connection_pool_replica_mock: SentinelConnectionPool, +) -> None: + """ + Assert that the connection object is different if no _iter_req_id is supplied. + In reality, they can be the same, but in this case, we're not + releasing the connection to the pool so they should always be different. + """ + connection_for_req_1 = connection_pool_replica_mock.get_connection( + "ANY", _iter_req_id=1 + ) + connection_for_random_req = connection_pool_replica_mock.get_connection("ANYY") + assert not same_address(connection_for_random_req, connection_for_req_1) + assert not same_address( + connection_pool_replica_mock.get_connection("ANY_COMMAND"), + connection_for_random_req, + ) + assert not same_address( + connection_pool_replica_mock.get_connection("ANY_COMMAND"), + connection_for_req_1, + ) + + +def test_connects_to_same_address_if_same_iter_req_id_master( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + """ + connection_for_req_1 = connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) + assert same_address( + connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + connection_for_req_1, + ) + + +def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + """ + Assert that the connection address is the same if the ``_iter_req_id`` is the same + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + and if we release the connection back to the connection pool before + trying to connect again. + """ + connection_for_req_1 = connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) + assert same_address( + connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + connection_for_req_1, + ) + + +def test_connects_to_same_address_if_no_iter_req_id_master( + connection_pool_master_mock: SentinelConnectionPool, +) -> None: + """ + Assert that connection address is always the same regardless if + there's an ``iter_req_id`` or not + when we are in master mode using a + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + """ + connection_for_req_1 = connection_pool_master_mock.get_connection( + "ANY", _iter_req_id=1 + ) + connection_for_random_req = connection_pool_master_mock.get_connection("ANYY") + assert same_address(connection_for_random_req, connection_for_req_1) + assert same_address( + connection_pool_master_mock.get_connection("ANY_COMMAND"), + connection_for_random_req, + ) + + assert same_address( + connection_pool_master_mock.get_connection("ANY_COMMAND"), + connection_for_req_1, + ) + + +def test_scan_iter_family_cleans_up( + connection_pool_replica_mock: SentinelConnectionPool, +): + """Test that connection pool is correctly cleaned up""" + from redis.asyncio import Redis + + r = Redis(connection_pool=connection_pool_replica_mock) + + [k for k in r.scan_iter("a")] + assert not connection_pool_replica_mock._iter_req_id_to_replica_address From e5c74ace6d36d8997e62a91782809b49f25a6251 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 12:40:46 +0000 Subject: [PATCH 12/44] more features for ConnectionsINdexer --- redis/asyncio/sentinel.py | 2 +- redis/connection.py | 36 +++++++++++++++++++++---------- redis/sentinel.py | 45 +++++++++++++++++++++++++++++++++++---- 3 files changed, 67 insertions(+), 16 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 5c210845ca..ac24ff0507 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -257,7 +257,7 @@ async def get_connection( # This will connect to the host and port of the replica else: await connection.connect_to_address(server_host, server_port) - self.ensure_connection(connection) + await self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/connection.py b/redis/connection.py index 3ac0b23cdc..87f41bb366 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -737,39 +737,53 @@ def _host_error(self): class ConnectionsIndexer(Iterable): """ - Data structure that manages a list of available connections which - is also indexed based on the address (ip and port) of the connection + Data structure that simulates a list of available connections. + Instead of list, we keep 2 additional DS to support O(1) operations + on all of the class' methods. + The first DS is indexed on the connection object's ID. + The second DS is indexed on the address (ip and port) of the connection. """ def __init__(self): - self._connections = [] + # Map the id to the connection object + self._id_to_connection = {} # Map the address to a dictionary of connections # The inner dictionary is a map between the object id to the object itself - # This is to support O(1) operations on all of the class' methods + # Both of these DS support O(1) operations on all of the class' methods self._address_to_connections = defaultdict(dict) def pop(self): - connection = self._connections.pop() - del self._address_to_connections[(connection.host, connection.port)][ - id(connection) - ] + try: + _, connection = self._id_to_connection.popitem() + del self._address_to_connections[(connection.host, connection.port)][ + id(connection) + ] + except KeyError: + # We are simulating a list, hence we raise IndexError + # when there's no item in the dictionary + raise IndexError() return connection def append(self, connection: Connection): - self._connections.append(connection) + self._id_to_connection[id(connection)] = connection self._address_to_connections[(connection.host, connection.port)][ id(connection) ] = connection def get_connection(self, host: str, port: int): try: - connection = self._address_to_connections[(host, port)].popitem() + _, connection = self._address_to_connections[(host, port)].popitem() + del self._id_to_connection[id(connection)] except KeyError: return None return connection def __iter__(self): - return iter(self._connections) + # This is an O(1) operation in python3.7 and later + return iter(self._id_to_connection.values()) + + def __len__(self): + return len(self._id_to_connection) class SSLConnection(Connection): diff --git a/redis/sentinel.py b/redis/sentinel.py index fefd3f4e21..e53f5537d8 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -41,9 +41,16 @@ def connect_to(self, address): if str_if_bytes(self.read_response()) != "PONG": raise ConnectionError("PING failed") - def _connect_retry(self): + def _connect_retry(self, same_server: bool = False): if self._sock: return # already connected + # If same_server is True, it means that the connection + # is not rotating to the next slave (if the connection pool is not master) + if same_server: + self.connect_to(self.host, self.port) + return + # If same_server is False, connnect to master in master mode + # and rotate to the next slave in slave mode if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) else: @@ -54,8 +61,11 @@ def _connect_retry(self): continue raise SlaveNotFoundError # Never be here - def connect(self): - return self.retry.call_with_retry(self._connect_retry, lambda error: None) + def connect(self, same_server: bool = False): + return self.retry.call_with_retry( + lambda: self._connect_retry(same_server), + lambda error: None + ) def read_response( self, @@ -195,6 +205,29 @@ def rotate_slaves(self): "Round-robin slave balancer" return self.proxy.rotate_slaves() + def ensure_connection_connected_to_address(self, connection: SentinelManagedConnection): + """ + Ensure the connection is already connected to the server that this connection + object wants to connect to + + Similar to self.ensure_connection, but calling connection.connect() + in SentinelManagedConnection (replica mode) will cause the + connection object to connect to the next replica in rotation, + and we don't wnat behavior. Look at get_connection inline docs for details. + + Here, we just try to make sure that the connection is already connected + to the replica we wanted it to. + """ + connection.connect(same_address=True) + try: + if connection.can_read(same_address=True) and connection.client_cache is None: + raise ConnectionError("Connection has data") + except (ConnectionError, OSError): + connection.disconnect() + connection.connect(same_address=True) + if connection.can_read(): + raise ConnectionError("Connection has data") + def cleanup_scan(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary @@ -248,6 +281,7 @@ def get_connection( connection = self.make_connection() assert connection self._in_use_connections.add(connection) + breakpoint() try: # Ensure this connection is connected to Redis # If this is the first scan request, it will @@ -259,7 +293,9 @@ def get_connection( # This will connect to the host and port of the replica else: connection.connect_to_address(server_host, server_port) - self.ensure_connection(connection) + breakpoint() + self.ensure_connection_connected_to_address(connection) + breakpoint() except BaseException: # Release the connection back to the pool so that we don't # leak it @@ -270,6 +306,7 @@ def get_connection( connection.host, connection.port, ) + breakpoint() return connection From 83dc599b99eb4f9c3394ea38745e62b0ef09b22e Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 21:42:16 +0800 Subject: [PATCH 13/44] add _same_addres methods for sentinels --- redis/asyncio/client.py | 2 +- redis/asyncio/connection.py | 4 ++-- redis/asyncio/sentinel.py | 37 ++++++++++++++++++++++++++++-- redis/sentinel.py | 45 +++++++++++++++++++++++++------------ 4 files changed, 69 insertions(+), 19 deletions(-) diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index fc3f2e7569..a74ff82542 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -652,7 +652,7 @@ async def execute_command(self, *args, **options): if not self.connection: await pool.release(conn) if "ITER" in command_name.upper(): - pool.cleanup_scan(iter_req_id=options.get("_iter_req_id", None)) + pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) async def parse_response( self, connection: Connection, command_name: Union[str, bytes], **options diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 0a907b40b7..b0d511f176 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1058,7 +1058,7 @@ class ConnectionPool: """ @abstractmethod - def cleanup_scan(self, **options): + def cleanup(self, **options): """ Additional cleanup operations that the connection pool might need to do after a SCAN ITER family command is executed @@ -1332,5 +1332,5 @@ async def release(self, connection: AbstractConnection): await super().release(connection) self._condition.notify() - def cleanup_scan(self, **options): + def cleanup(self, **options): pass diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index ac24ff0507..8a73178547 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -91,6 +91,16 @@ async def connect_to_address(self, host: str, port: int) -> None: lambda error: asyncio.sleep(0), ) + async def connect_to_same_address(self): + """ + Similar to connect, but instead of rotating to the next slave (if not in master mode), + it just connects to the same address of the connection object. + """ + return await self.retry.call_with_retry( + lambda: self._connect_retry(same_address=True), + lambda error: asyncio.sleep(0), + ) + async def read_response( self, disable_decoding: bool = False, @@ -193,7 +203,30 @@ async def rotate_slaves(self) -> AsyncIterator: pass raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") - def cleanup_scan(self, **options): + async def ensure_connection_connected_to_address(self, connection: SentinelManagedConnection): + """ + Ensure the connection is already connected to the server that this connection + object wants to connect to + + Similar to self.ensure_connection, but calling connection.connect() + in SentinelManagedConnection (replica mode) will cause the + connection object to connect to the next replica in rotation, + and we don't wnat behavior. Look at get_connection inline docs for details. + + Here, we just try to make sure that the connection is already connected + to the replica we wanted it to. + """ + await connection.connect_to_same_address() + try: + if await connection.can_read_destructive() and connection.client_cache is None: + raise ConnectionError("Connection has data") + except (ConnectionError, OSError): + await connection.disconnect() + await connection.connect_to_same_address() + if await connection.can_read_destructive(): + raise ConnectionError("Connection has data") + + def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ @@ -257,7 +290,7 @@ async def get_connection( # This will connect to the host and port of the replica else: await connection.connect_to_address(server_host, server_port) - await self.ensure_connection(connection) + self.ensure_connection_connected_to_address(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/sentinel.py b/redis/sentinel.py index e53f5537d8..53d1b33fc9 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -41,15 +41,15 @@ def connect_to(self, address): if str_if_bytes(self.read_response()) != "PONG": raise ConnectionError("PING failed") - def _connect_retry(self, same_server: bool = False): + def _connect_retry(self, same_address: bool = False): if self._sock: return # already connected # If same_server is True, it means that the connection # is not rotating to the next slave (if the connection pool is not master) - if same_server: + if same_address: self.connect_to(self.host, self.port) return - # If same_server is False, connnect to master in master mode + # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) @@ -61,12 +61,33 @@ def _connect_retry(self, same_server: bool = False): continue raise SlaveNotFoundError # Never be here - def connect(self, same_server: bool = False): + def connect(self): + return self.retry.call_with_retry(lambda: self._connect_retry(), lambda error: None) + + def connect_to_same_address(self): + """ + Similar to connect, but instead of rotating to the next slave (if not in master mode), + it just connects to the same address of the connection object. + """ return self.retry.call_with_retry( - lambda: self._connect_retry(same_server), - lambda error: None + lambda: self._connect_retry(same_address=True), + lambda error: None ) + def can_read_same_address(self, timeout=0): + """Similar to can_read_same_address, but calls connect_to_same_address instead of connect""" + sock = self._sock + if not sock: + self.connect_to_same_address() + + host_error = self._host_error() + + try: + return self._parser.can_read(timeout) + except OSError as e: + self.disconnect() + raise ConnectionError(f"Error while reading from {host_error}: {e.args}") + def read_response( self, disable_decoding=False, @@ -218,17 +239,17 @@ def ensure_connection_connected_to_address(self, connection: SentinelManagedConn Here, we just try to make sure that the connection is already connected to the replica we wanted it to. """ - connection.connect(same_address=True) + connection.connect_to_same_address() try: - if connection.can_read(same_address=True) and connection.client_cache is None: + if connection.can_read_same_address() and connection.client_cache is None: raise ConnectionError("Connection has data") except (ConnectionError, OSError): connection.disconnect() - connection.connect(same_address=True) + connection.connect_to_same_address() if connection.can_read(): raise ConnectionError("Connection has data") - def cleanup_scan(self, **options): + def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ @@ -281,7 +302,6 @@ def get_connection( connection = self.make_connection() assert connection self._in_use_connections.add(connection) - breakpoint() try: # Ensure this connection is connected to Redis # If this is the first scan request, it will @@ -293,9 +313,7 @@ def get_connection( # This will connect to the host and port of the replica else: connection.connect_to_address(server_host, server_port) - breakpoint() self.ensure_connection_connected_to_address(connection) - breakpoint() except BaseException: # Release the connection back to the pool so that we don't # leak it @@ -306,7 +324,6 @@ def get_connection( connection.host, connection.port, ) - breakpoint() return connection From 546b443990d22b1fbc4ffe86356f4e9dee5a155c Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 21:51:32 +0800 Subject: [PATCH 14/44] fix connect_to args --- redis/asyncio/sentinel.py | 29 +++++++------------ redis/sentinel.py | 6 ++-- .../test_sentinel_managed_connection.py | 5 ++-- tests/test_sentinel_managed_connection.py | 5 ++-- 4 files changed, 17 insertions(+), 28 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 8a73178547..9f84437e5e 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -56,9 +56,16 @@ async def connect_to(self, address): if str_if_bytes(await self.read_response()) != "PONG": raise ConnectionError("PING failed") - async def _connect_retry(self): + async def _connect_retry(self, same_address: bool = False): if self._reader: return # already connected + # If same_server is True, it means that the connection + # is not rotating to the next slave (if the connection pool is not master) + if same_address: + self.connect_to(self.host, self.port) + return + # If same_server is False, connnect to master in master mode + # and rotate to the next slave in slave mode if self.connection_pool.is_master: await self.connect_to(await self.connection_pool.get_master_address()) else: @@ -75,22 +82,6 @@ async def connect(self): lambda error: asyncio.sleep(0), ) - async def _connect_to_address_retry(self, host: str, port: int) -> None: - if self._reader: - return # already connected - try: - return await self.connect_to((host, port)) - except ConnectionError: - raise SlaveNotFoundError - - async def connect_to_address(self, host: str, port: int) -> None: - # Connect to the specified host and port - # instead of connecting to the master / rotated slaves - return await self.retry.call_with_retry( - lambda: self._connect_to_address_retry(host, port), - lambda error: asyncio.sleep(0), - ) - async def connect_to_same_address(self): """ Similar to connect, but instead of rotating to the next slave (if not in master mode), @@ -274,7 +265,7 @@ async def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and - # port to the one that we want later in the call to ``connect_to_address`` + # port to the one that we want later in the call to ``connect_to_same_address`` if not connection: connection = self.make_connection() assert connection @@ -289,7 +280,7 @@ async def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - await connection.connect_to_address(server_host, server_port) + await connection.connect_to_same_address() self.ensure_connection_connected_to_address(connection) except BaseException: # Release the connection back to the pool so that we don't diff --git a/redis/sentinel.py b/redis/sentinel.py index 53d1b33fc9..b18177e042 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -47,7 +47,7 @@ def _connect_retry(self, same_address: bool = False): # If same_server is True, it means that the connection # is not rotating to the next slave (if the connection pool is not master) if same_address: - self.connect_to(self.host, self.port) + self.connect_to((self.host, self.port)) return # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode @@ -297,7 +297,7 @@ def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and port - # to the one that we want later in the call to ``connect_to_address`` + # to the one that we want later in the call to ``connect_to_same_address`` if not connection: connection = self.make_connection() assert connection @@ -312,7 +312,7 @@ def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - connection.connect_to_address(server_host, server_port) + connection.connect_to_same_address() self.ensure_connection_connected_to_address(connection) except BaseException: # Release the connection back to the pool so that we don't diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 6ccd5b5077..372d85ab4b 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -43,9 +43,8 @@ async def mock_connect(): class SentinelManagedConnectionMock(SentinelManagedConnection): - async def connect_to_address(self, host: str, port: int) -> None: - self.host = host - self.port = port + async def connect_to_same_address(self) -> None: + pass async def can_read_destructive(self) -> bool: # Mock this function to always return False. diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index b4c927daa7..aa244faa9d 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -42,9 +42,8 @@ async def mock_connect(): class SentinelManagedConnectionMock(SentinelManagedConnection): - def connect_to_address(self, host: str, port: int) -> None: - self.host = host - self.port = port + def connect_to_same_address(self) -> None: + pass def can_read_destructive(self) -> bool: # Mock this function to always return False. From 803b29118a7c8c3d536e401b16d5e4bc0e8fab09 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:14:08 +0800 Subject: [PATCH 15/44] fix tests --- tests/test_asyncio/test_connection_pool.py | 3 +- tests/test_connection_pool.py | 2 +- tests/test_sentinel_managed_connection.py | 47 ++++------------------ 3 files changed, 10 insertions(+), 42 deletions(-) diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index 2f5bbfb621..1547b712bd 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -1,5 +1,6 @@ import asyncio import re +from itertools import chain import pytest import pytest_asyncio @@ -35,7 +36,7 @@ async def create_two_conn(r: redis.Redis): def has_no_connected_connections(pool: redis.ConnectionPool): return not any( x.is_connected - for x in pool._available_connections + list(pool._in_use_connections) + for x in chain(pool._available_connections, pool._in_use_connections) ) async def test_auto_disconnect_redis_created_pool(self, r: redis.Redis): diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index dee7c554d3..c15f7dd424 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -502,7 +502,7 @@ def test_on_connect_error(self): bad_connection.info() pool = bad_connection.connection_pool assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._sock + assert not list(pool._available_connections)[0]._sock @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index aa244faa9d..a5abcf2bc2 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -1,45 +1,12 @@ -import socket from typing import Tuple from unittest import mock import pytest -from redis.asyncio.retry import Retry -from redis.asyncio.sentinel import ( +from redis.sentinel import ( Sentinel, SentinelConnectionPool, SentinelManagedConnection, ) -from redis.backoff import NoBackoff - -pytestmark = pytest.mark.asyncio - - -def test_connect_retry_on_timeout_error(connect_args): - """Test that the _connect function is retried in case of a timeout""" - connection_pool = mock.AsyncMock() - connection_pool.get_master_address = mock.AsyncMock( - return_value=(connect_args["host"], connect_args["port"]) - ) - conn = SentinelManagedConnection( - retry_on_timeout=True, - retry=Retry(NoBackoff(), 3), - connection_pool=connection_pool, - ) - origin_connect = conn._connect - conn._connect = mock.AsyncMock() - - async def mock_connect(): - # connect only on the last retry - if conn._connect.call_count <= 2: - raise socket.timeout - else: - return await origin_connect() - - conn._connect.side_effect = mock_connect - conn.connect() - assert conn._connect.call_count == 3 - conn.disconnect() - class SentinelManagedConnectionMock(SentinelManagedConnection): def connect_to_same_address(self) -> None: @@ -125,7 +92,7 @@ def test_connects_to_same_address_if_same_id_replica( """ Assert that the connection address is the same if the ``_iter_req_id`` is the same when we are in replica mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + :py:class:`~redis.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = connection_pool_replica_mock.get_connection( "ANY", _iter_req_id=1 @@ -142,7 +109,7 @@ def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica( """ Assert that the connection object is the same if the ``_iter_req_id`` is the same when we are in replica mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + :py:class:`~redis.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before trying to connect again. """ @@ -185,7 +152,7 @@ def test_connects_to_same_address_if_same_iter_req_id_master( """ Assert that the connection address is the same if the ``_iter_req_id`` is the same when we are in master mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + :py:class:`~redis.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = connection_pool_master_mock.get_connection( "ANY", _iter_req_id=1 @@ -202,7 +169,7 @@ def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( """ Assert that the connection address is the same if the ``_iter_req_id`` is the same when we are in master mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` + :py:class:`~redis.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before trying to connect again. """ @@ -222,7 +189,7 @@ def test_connects_to_same_address_if_no_iter_req_id_master( Assert that connection address is always the same regardless if there's an ``iter_req_id`` or not when we are in master mode using a - :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + :py:class:`~redis.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = connection_pool_master_mock.get_connection( "ANY", _iter_req_id=1 @@ -244,7 +211,7 @@ def test_scan_iter_family_cleans_up( connection_pool_replica_mock: SentinelConnectionPool, ): """Test that connection pool is correctly cleaned up""" - from redis.asyncio import Redis + from redis import Redis r = Redis(connection_pool=connection_pool_replica_mock) From c7708047487b2b39a48f89608fed530fd6e9cd7d Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:16:13 +0800 Subject: [PATCH 16/44] add self --- tests/test_asyncio/test_commands.py | 2 +- tests/test_commands.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncio/test_commands.py b/tests/test_asyncio/test_commands.py index 4b2b2c1743..50d6f58063 100644 --- a/tests/test_asyncio/test_commands.py +++ b/tests/test_asyncio/test_commands.py @@ -1402,7 +1402,7 @@ async def test_zscan_iter(self, r: redis.Redis): pairs = [k async for k in r.zscan_iter("a", match="a")] assert set(pairs) == {(b"a", 1)} - async def test_scan_iter_family_executes_commands_with_same_iter_req_id(): + async def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): """Assert that all calls to execute_command receives the _iter_req_id kwarg""" import uuid diff --git a/tests/test_commands.py b/tests/test_commands.py index 1a79a45a4c..e41c590b0a 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2235,7 +2235,7 @@ def test_zscan_iter(self, r): pairs = list(r.zscan_iter("a", match="a")) assert set(pairs) == {(b"a", 1)} - def test_scan_iter_family_executes_commands_with_same_iter_req_id(): + def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): """Assert that all calls to execute_command receives the _iter_req_id kwarg""" import uuid From 6add0ef08abdc9be12d0ef3ee072419af64f8324 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:22:02 +0800 Subject: [PATCH 17/44] convert ConnectionsIndexer to list before indexing --- tests/test_asyncio/test_connection_pool.py | 10 +++++----- tests/test_connection_pool.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index 1547b712bd..1a27b53731 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -57,7 +57,7 @@ async def test_do_not_auto_disconnect_redis_created_pool(self, r2: redis.Redis): assert r2.connection_pool._in_use_connections == {new_conn} assert new_conn.is_connected assert len(r2.connection_pool._available_connections) == 1 - assert r2.connection_pool._available_connections[0].is_connected + assert r2.connection_pool._list(available_connections)[0].is_connected async def test_auto_release_override_true_manual_created_pool(self, r: redis.Redis): assert r.auto_close_connection_pool is True, "This is from the class fixture" @@ -85,7 +85,7 @@ async def test_negate_auto_close_client_pool( await r.aclose(close_connection_pool=False) assert not self.has_no_connected_connections(r.connection_pool) assert r.connection_pool._in_use_connections == {new_conn} - assert r.connection_pool._available_connections[0].is_connected + assert r.connection_pool._list(available_connections)[0].is_connected assert self.get_total_connected_connections(r.connection_pool) == 2 @@ -579,7 +579,7 @@ async def test_on_connect_error(self): await bad_connection.info() pool = bad_connection.connection_pool assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._reader + assert not pool._list(available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -610,7 +610,7 @@ async def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._reader + assert not pool._list(available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -627,7 +627,7 @@ async def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._reader + assert not pool._list(available_connections)[0]._reader @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index c15f7dd424..41aa951bc2 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -530,7 +530,7 @@ def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._sock + assert not pool._list(available_connections)[0]._sock @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -547,7 +547,7 @@ def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._available_connections[0]._sock + assert not pool._list(available_connections)[0]._sock @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() From 7d6c2a43402deeaaa9fb50c77b86fa9dec8b7fa6 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:27:18 +0800 Subject: [PATCH 18/44] convert ConnectionsIndexer to list before indexing --- tests/test_asyncio/test_connection_pool.py | 10 +++++----- tests/test_connection_pool.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index 1a27b53731..799bd2ddeb 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -57,7 +57,7 @@ async def test_do_not_auto_disconnect_redis_created_pool(self, r2: redis.Redis): assert r2.connection_pool._in_use_connections == {new_conn} assert new_conn.is_connected assert len(r2.connection_pool._available_connections) == 1 - assert r2.connection_pool._list(available_connections)[0].is_connected + assert r2.connection_pool.list(available_connections)[0].is_connected async def test_auto_release_override_true_manual_created_pool(self, r: redis.Redis): assert r.auto_close_connection_pool is True, "This is from the class fixture" @@ -85,7 +85,7 @@ async def test_negate_auto_close_client_pool( await r.aclose(close_connection_pool=False) assert not self.has_no_connected_connections(r.connection_pool) assert r.connection_pool._in_use_connections == {new_conn} - assert r.connection_pool._list(available_connections)[0].is_connected + assert r.connection_pool.list(available_connections)[0].is_connected assert self.get_total_connected_connections(r.connection_pool) == 2 @@ -579,7 +579,7 @@ async def test_on_connect_error(self): await bad_connection.info() pool = bad_connection.connection_pool assert len(pool._available_connections) == 1 - assert not pool._list(available_connections)[0]._reader + assert not pool.list(available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -610,7 +610,7 @@ async def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._list(available_connections)[0]._reader + assert not pool.list(available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -627,7 +627,7 @@ async def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._list(available_connections)[0]._reader + assert not pool.list(available_connections)[0]._reader @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index 41aa951bc2..19dcc41f62 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -530,7 +530,7 @@ def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._list(available_connections)[0]._sock + assert not pool.list(available_connections)[0]._sock @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -547,7 +547,7 @@ def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool._list(available_connections)[0]._sock + assert not pool.list(available_connections)[0]._sock @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() From 9b52903890d55ab400be8287c17c4d9525bd67bf Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:50:02 +0800 Subject: [PATCH 19/44] fix typo --- tests/test_asyncio/test_connection_pool.py | 10 +++++----- tests/test_connection_pool.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index 799bd2ddeb..a01437ed93 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -57,7 +57,7 @@ async def test_do_not_auto_disconnect_redis_created_pool(self, r2: redis.Redis): assert r2.connection_pool._in_use_connections == {new_conn} assert new_conn.is_connected assert len(r2.connection_pool._available_connections) == 1 - assert r2.connection_pool.list(available_connections)[0].is_connected + assert r2.connection_list(pool._available_connections)[0].is_connected async def test_auto_release_override_true_manual_created_pool(self, r: redis.Redis): assert r.auto_close_connection_pool is True, "This is from the class fixture" @@ -85,7 +85,7 @@ async def test_negate_auto_close_client_pool( await r.aclose(close_connection_pool=False) assert not self.has_no_connected_connections(r.connection_pool) assert r.connection_pool._in_use_connections == {new_conn} - assert r.connection_pool.list(available_connections)[0].is_connected + assert r.connection_list(pool._available_connections)[0].is_connected assert self.get_total_connected_connections(r.connection_pool) == 2 @@ -579,7 +579,7 @@ async def test_on_connect_error(self): await bad_connection.info() pool = bad_connection.connection_pool assert len(pool._available_connections) == 1 - assert not pool.list(available_connections)[0]._reader + assert not list(pool._available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -610,7 +610,7 @@ async def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool.list(available_connections)[0]._reader + assert not list(pool._available_connections)[0]._reader @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -627,7 +627,7 @@ async def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool.list(available_connections)[0]._reader + assert not list(pool._available_connections)[0]._reader @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index 19dcc41f62..e2b6c173fb 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -530,7 +530,7 @@ def test_busy_loading_from_pipeline_immediate_command(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool.list(available_connections)[0]._sock + assert not list(pool._available_connections)[0]._sock @pytest.mark.onlynoncluster @skip_if_server_version_lt("2.8.8") @@ -547,7 +547,7 @@ def test_busy_loading_from_pipeline(self, r): pool = r.connection_pool assert not pipe.connection assert len(pool._available_connections) == 1 - assert not pool.list(available_connections)[0]._sock + assert not list(pool._available_connections)[0]._sock @skip_if_server_version_lt("2.8.8") @skip_if_redis_enterprise() From e6375b2c721be34be497a8b09aa8740d84002b43 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 22:55:49 +0800 Subject: [PATCH 20/44] fix --- tests/test_asyncio/test_connection_pool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index a01437ed93..beb1edacb3 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -57,7 +57,7 @@ async def test_do_not_auto_disconnect_redis_created_pool(self, r2: redis.Redis): assert r2.connection_pool._in_use_connections == {new_conn} assert new_conn.is_connected assert len(r2.connection_pool._available_connections) == 1 - assert r2.connection_list(pool._available_connections)[0].is_connected + assert list(r2.connection_pool._available_connections)[0].is_connected async def test_auto_release_override_true_manual_created_pool(self, r: redis.Redis): assert r.auto_close_connection_pool is True, "This is from the class fixture" @@ -85,7 +85,7 @@ async def test_negate_auto_close_client_pool( await r.aclose(close_connection_pool=False) assert not self.has_no_connected_connections(r.connection_pool) assert r.connection_pool._in_use_connections == {new_conn} - assert r.connection_list(pool._available_connections)[0].is_connected + assert list(r.connection_pool._available_connections)[0].is_connected assert self.get_total_connected_connections(r.connection_pool) == 2 From d7f8e901386db775c4a739ce57151da7005a04b8 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sun, 14 Jul 2024 15:43:40 +0000 Subject: [PATCH 21/44] fix connect_to_address --- redis/asyncio/sentinel.py | 12 ++++++++++-- redis/sentinel.py | 10 +++++++++- tests/test_asyncio/test_connection_pool.py | 2 ++ tests/test_connection_pool.py | 2 ++ 4 files changed, 23 insertions(+), 3 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 9f84437e5e..4fb236bf8c 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -92,6 +92,14 @@ async def connect_to_same_address(self): lambda error: asyncio.sleep(0), ) + async def connect_to_address(self, address): + """ + Similar to connect, but instead of rotating to the next slave (if not in master mode), + it just connects to the same address of the connection object. + """ + self.host, self.port = address + return await self.connect_to_same_address() + async def read_response( self, disable_decoding: bool = False, @@ -280,8 +288,8 @@ async def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - await connection.connect_to_same_address() - self.ensure_connection_connected_to_address(connection) + await connection.connect_to_address((server_host, server_port)) + await self.ensure_connection_connected_to_address(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/sentinel.py b/redis/sentinel.py index b18177e042..b76a097f0d 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -74,6 +74,14 @@ def connect_to_same_address(self): lambda error: None ) + def connect_to_address(self, address): + """ + Similar to connect, but instead of rotating to the next slave (if not in master mode), + it just connects to the address supplied. + """ + self.host, self.port = address + return self.connect_to_same_address() + def can_read_same_address(self, timeout=0): """Similar to can_read_same_address, but calls connect_to_same_address instead of connect""" sock = self._sock @@ -312,7 +320,7 @@ def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - connection.connect_to_same_address() + connection.connect_to_address((server_host, server_port)) self.ensure_connection_connected_to_address(connection) except BaseException: # Release the connection back to the pool so that we don't diff --git a/tests/test_asyncio/test_connection_pool.py b/tests/test_asyncio/test_connection_pool.py index beb1edacb3..f027e7c289 100644 --- a/tests/test_asyncio/test_connection_pool.py +++ b/tests/test_asyncio/test_connection_pool.py @@ -94,6 +94,8 @@ class DummyConnection(Connection): def __init__(self, **kwargs): self.kwargs = kwargs + self.host = kwargs.get("host", None) + self.port = kwargs.get("port", None) def repr_pieces(self): return [("id", id(self)), ("kwargs", self.kwargs)] diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index e2b6c173fb..313dd96287 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -20,6 +20,8 @@ class DummyConnection: def __init__(self, **kwargs): self.kwargs = kwargs self.pid = os.getpid() + self.host = kwargs.get("host", None) + self.port = kwargs.get("port", None) def connect(self): pass From d1c25b9618ae018f992a085ad475fea1fb829853 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 12:59:42 +0000 Subject: [PATCH 22/44] cleanup in sync client --- redis/client.py | 3 +++ redis/sentinel.py | 1 + tests/test_sentinel_managed_connection.py | 7 +++++-- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/redis/client.py b/redis/client.py index b7a1f88d92..ceb0ab5947 100755 --- a/redis/client.py +++ b/redis/client.py @@ -581,6 +581,9 @@ def execute_command(self, *args, **options): finally: if not self.connection: pool.release(conn) + if "SCAN" in command_name.upper(): + breakpoint() + pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) def parse_response(self, connection, command_name, **options): """Parses a response from the Redis server""" diff --git a/redis/sentinel.py b/redis/sentinel.py index b76a097f0d..d09c1a4208 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -261,6 +261,7 @@ def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ + breakpoint() self._iter_req_id_to_replica_address.pop( options.get("_iter_req_id", None), None ) diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index a5abcf2bc2..561c0a9bd6 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -213,7 +213,10 @@ def test_scan_iter_family_cleans_up( """Test that connection pool is correctly cleaned up""" from redis import Redis - r = Redis(connection_pool=connection_pool_replica_mock) + from redis.commands.core import ScanCommands - [k for k in r.scan_iter("a")] + r = Redis(connection_pool=connection_pool_replica_mock) + with mock.patch.object(r, "_send_command_parse_response", return_value=(0, [])): + [k for k in r.scan_iter("a")] assert not connection_pool_replica_mock._iter_req_id_to_replica_address + From b8aeb71e142813e08e5b1ec91efc883b1254ee5a Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 21:03:37 +0800 Subject: [PATCH 23/44] rename kwargs to no underscore for consistency --- redis/asyncio/client.py | 2 +- redis/asyncio/sentinel.py | 4 +- redis/commands/core.py | 16 ++++---- redis/sentinel.py | 4 +- tests/test_asyncio/test_commands.py | 12 +++--- .../test_sentinel_managed_connection.py | 40 +++++++++---------- tests/test_commands.py | 12 +++--- 7 files changed, 45 insertions(+), 45 deletions(-) diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index a74ff82542..bb5a7d2a65 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -652,7 +652,7 @@ async def execute_command(self, *args, **options): if not self.connection: await pool.release(conn) if "ITER" in command_name.upper(): - pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) + pool.cleanup(iter_req_id=options.get("iter_req_id", None)) async def parse_response( self, connection: Connection, command_name: Union[str, bytes], **options diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 4fb236bf8c..9ad520fdac 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -230,7 +230,7 @@ def cleanup(self, **options): Remove the SCAN ITER family command's request id from the dictionary """ self._iter_req_id_to_replica_address.pop( - options.get("_iter_req_id", None), None + options.get("iter_req_id", None), None ) async def get_connection( @@ -250,7 +250,7 @@ async def get_connection( should go to the same replica. """ # If not an iter command or in master mode, call superclass' implementation - if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: + if not (iter_req_id := options.get("iter_req_id", None)) or self.is_master: return await super().get_connection(command_name, *keys, **options) # Check if this iter request has already been directed to a particular server diff --git a/redis/commands/core.py b/redis/commands/core.py index 999a14622a..2f6f93ced7 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3046,7 +3046,7 @@ def scan_iter( match=match, count=count, _type=_type, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, **kwargs, ) yield from data @@ -3094,7 +3094,7 @@ def sscan_iter( iter_req_id = uuid.uuid4() while cursor != 0: cursor, data = self.sscan( - name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id + name, cursor=cursor, match=match, count=count, iter_req_id=iter_req_id ) yield from data @@ -3154,7 +3154,7 @@ def hscan_iter( match=match, count=count, no_values=no_values, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, ) if no_values: yield from data @@ -3216,7 +3216,7 @@ def zscan_iter( match=match, count=count, score_cast_func=score_cast_func, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, ) yield from data @@ -3254,7 +3254,7 @@ async def scan_iter( match=match, count=count, _type=_type, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, **kwargs, ) for d in data: @@ -3281,7 +3281,7 @@ async def sscan_iter( cursor = "0" while cursor != 0: cursor, data = await self.sscan( - name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id + name, cursor=cursor, match=match, count=count, iter_req_id=iter_req_id ) for d in data: yield d @@ -3315,7 +3315,7 @@ async def hscan_iter( match=match, count=count, no_values=no_values, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, ) if no_values: for it in data: @@ -3353,7 +3353,7 @@ async def zscan_iter( match=match, count=count, score_cast_func=score_cast_func, - _iter_req_id=iter_req_id, + iter_req_id=iter_req_id, ) for d in data: yield d diff --git a/redis/sentinel.py b/redis/sentinel.py index d09c1a4208..cb2dd1ead4 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -263,7 +263,7 @@ def cleanup(self, **options): """ breakpoint() self._iter_req_id_to_replica_address.pop( - options.get("_iter_req_id", None), None + options.get("iter_req_id", None), None ) def get_connection( @@ -283,7 +283,7 @@ def get_connection( should go to the same replica. """ # If not an iter command or in master mode, call superclass' implementation - if not (iter_req_id := options.get("_iter_req_id", None)) or self.is_master: + if not (iter_req_id := options.get("iter_req_id", None)) or self.is_master: return super().get_connection(command_name, *keys, **options) # Check if this iter request has already been directed to a particular server diff --git a/tests/test_asyncio/test_commands.py b/tests/test_asyncio/test_commands.py index 50d6f58063..78ca92bf49 100644 --- a/tests/test_asyncio/test_commands.py +++ b/tests/test_asyncio/test_commands.py @@ -1402,8 +1402,8 @@ async def test_zscan_iter(self, r: redis.Redis): pairs = [k async for k in r.zscan_iter("a", match="a")] assert set(pairs) == {(b"a", 1)} - async def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): - """Assert that all calls to execute_command receives the _iter_req_id kwarg""" + async def test_scan_iter_family_executes_commands_with_sameiter_req_id(self): + """Assert that all calls to execute_command receives the iter_req_id kwarg""" import uuid from redis.commands.core import AsyncScanCommands @@ -1416,10 +1416,10 @@ async def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): uuid, "uuid4", return_value="uuid" ): [a async for a in AsyncScanCommands().scan_iter()] - mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") + mock_execute_command.assert_called_with("SCAN", "0", iter_req_id="uuid") [a async for a in AsyncScanCommands().sscan_iter("")] mock_execute_command.assert_called_with( - "SSCAN", "", "0", _iter_req_id="uuid" + "SSCAN", "", "0", iter_req_id="uuid" ) with mock.patch.object( AsyncScanCommands, "execute_command", mock.AsyncMock(return_value=(0, {})) @@ -1428,11 +1428,11 @@ async def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): ): [a async for a in AsyncScanCommands().hscan_iter("")] mock_execute_command.assert_called_with( - "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" + "HSCAN", "", "0", no_values=None, iter_req_id="uuid" ) [a async for a in AsyncScanCommands().zscan_iter("")] mock_execute_command.assert_called_with( - "ZSCAN", "", "0", score_cast_func=mock.ANY, _iter_req_id="uuid" + "ZSCAN", "", "0", score_cast_func=mock.ANY, iter_req_id="uuid" ) # SET COMMANDS diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 372d85ab4b..2c79247c2a 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -124,15 +124,15 @@ async def test_connects_to_same_address_if_same_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection address is the same if the ``_iter_req_id`` is the same + Assert that the connection address is the same if the ``iter_req_id`` is the same when we are in replica mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) assert same_address( - await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1), + await connection_pool_replica_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) @@ -141,32 +141,32 @@ async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection object is the same if the ``_iter_req_id`` is the same + Assert that the connection object is the same if the ``iter_req_id`` is the same when we are in replica mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before trying to connect again. """ connection_for_req_1 = await connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) await connection_pool_replica_mock.release(connection_for_req_1) assert ( - await connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + await connection_pool_replica_mock.get_connection("ANY", iter_req_id=1) == connection_for_req_1 ) -async def test_connects_to_diff_address_if_no_iter_req_id_replica( +async def test_connects_to_diff_address_if_noiter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection object is different if no _iter_req_id is supplied. + Assert that the connection object is different if no iter_req_id is supplied. In reality, they can be the same, but in this case, we're not releasing the connection to the pool so they should always be different. """ connection_for_req_1 = await connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) connection_for_random_req = await connection_pool_replica_mock.get_connection( "ANYY" @@ -182,43 +182,43 @@ async def test_connects_to_diff_address_if_no_iter_req_id_replica( ) -async def test_connects_to_same_address_if_same_iter_req_id_master( +async def test_connects_to_same_address_if_sameiter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection address is the same if the ``_iter_req_id`` is the same + Assert that the connection address is the same if the ``iter_req_id`` is the same when we are in master mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) assert same_address( - await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + await connection_pool_master_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) -async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( +async def test_connects_to_same_conn_object_if_sameiter_req_id_and_released_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection address is the same if the ``_iter_req_id`` is the same + Assert that the connection address is the same if the ``iter_req_id`` is the same when we are in master mode using a :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before trying to connect again. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) assert same_address( - await connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + await connection_pool_master_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) -async def test_connects_to_same_address_if_no_iter_req_id_master( +async def test_connects_to_same_address_if_noiter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ @@ -228,7 +228,7 @@ async def test_connects_to_same_address_if_no_iter_req_id_master( :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = await connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) connection_for_random_req = await connection_pool_master_mock.get_connection("ANYY") assert same_address(connection_for_random_req, connection_for_req_1) @@ -252,4 +252,4 @@ async def test_scan_iter_family_cleans_up( r = Redis(connection_pool=connection_pool_replica_mock) [k async for k in r.scan_iter("a")] - assert not connection_pool_replica_mock._iter_req_id_to_replica_address + assert not connection_pool_replica_mock.iter_req_id_to_replica_address diff --git a/tests/test_commands.py b/tests/test_commands.py index e41c590b0a..a2dc98a7c6 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2235,8 +2235,8 @@ def test_zscan_iter(self, r): pairs = list(r.zscan_iter("a", match="a")) assert set(pairs) == {(b"a", 1)} - def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): - """Assert that all calls to execute_command receives the _iter_req_id kwarg""" + def test_scan_iter_family_executes_commands_with_sameiter_req_id(self): + """Assert that all calls to execute_command receives the iter_req_id kwarg""" import uuid from redis.commands.core import ScanCommands @@ -2247,10 +2247,10 @@ def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): uuid, "uuid4", return_value="uuid" ): [a for a in ScanCommands().scan_iter()] - mock_execute_command.assert_called_with("SCAN", "0", _iter_req_id="uuid") + mock_execute_command.assert_called_with("SCAN", "0", iter_req_id="uuid") [a for a in ScanCommands().sscan_iter("")] mock_execute_command.assert_called_with( - "SSCAN", "", "0", _iter_req_id="uuid" + "SSCAN", "", "0", iter_req_id="uuid" ) with mock.patch.object( ScanCommands, "execute_command", mock.Mock(return_value=(0, {})) @@ -2259,11 +2259,11 @@ def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): ): [a for a in ScanCommands().hscan_iter("")] mock_execute_command.assert_called_with( - "HSCAN", "", "0", no_values=None, _iter_req_id="uuid" + "HSCAN", "", "0", no_values=None, iter_req_id="uuid" ) [a for a in ScanCommands().zscan_iter("")] mock_execute_command.assert_called_with( - "ZSCAN", "", "0", score_cast_func=mock.ANY, _iter_req_id="uuid" + "ZSCAN", "", "0", score_cast_func=mock.ANY, iter_req_id="uuid" ) # SET COMMANDS From ca7d0ff257f344a8a6a8cc8ab3bde5da14d73576 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 21:27:46 +0800 Subject: [PATCH 24/44] add cleanup tests for pipeline --- redis/client.py | 2 +- tests/test_sentinel_managed_connection.py | 24 ++++++++++++++++++++++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/redis/client.py b/redis/client.py index ceb0ab5947..7d9fe343ca 100755 --- a/redis/client.py +++ b/redis/client.py @@ -582,7 +582,6 @@ def execute_command(self, *args, **options): if not self.connection: pool.release(conn) if "SCAN" in command_name.upper(): - breakpoint() pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) def parse_response(self, connection, command_name, **options): @@ -764,6 +763,7 @@ def reset(self) -> None: self.patterns = {} self.pending_unsubscribe_patterns = set() self.subscribed_event.clear() + self.connection_pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) def close(self) -> None: self.reset() diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 561c0a9bd6..86d47a5594 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -207,7 +207,22 @@ def test_connects_to_same_address_if_no_iter_req_id_master( ) -def test_scan_iter_family_cleans_up( +def test_scan_iter_in_redis_cleans_up( + connection_pool_replica_mock: SentinelConnectionPool, +): + """Test that connection pool is correctly cleaned up""" + from redis import Redis + + r = Redis(connection_pool=connection_pool_replica_mock) + # Patch the actual sending and parsing response from the Connection object + # but still let the connection pool does all the necessary work + with mock.patch.object(r, "_send_command_parse_response", return_value=(0, [])): + [k for k in r.scan_iter("a")] + # Test that the iter_req_id for the scan command is cleared at the + # end of the SCAN ITER command + assert not connection_pool_replica_mock._iter_req_id_to_replica_address + +def test_scan_iter_in_pipeline_cleans_up( connection_pool_replica_mock: SentinelConnectionPool, ): """Test that connection pool is correctly cleaned up""" @@ -216,7 +231,14 @@ def test_scan_iter_family_cleans_up( from redis.commands.core import ScanCommands r = Redis(connection_pool=connection_pool_replica_mock) + r.pipeline() + r.scan_iter("a") + # Patch the actual sending and parsing response from the Connection object + # but still let the connection pool does all the necessary work with mock.patch.object(r, "_send_command_parse_response", return_value=(0, [])): + r.execute() [k for k in r.scan_iter("a")] + # Test that the iter_req_id for the scan command is cleared at the + # end of the SCAN ITER command assert not connection_pool_replica_mock._iter_req_id_to_replica_address From b21f1ea1d5c99669c9351fc8fe16b2dec62dbc99 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 13:51:47 +0000 Subject: [PATCH 25/44] remove test for pipeline --- redis/sentinel.py | 1 - tests/test_sentinel_managed_connection.py | 20 -------------------- 2 files changed, 21 deletions(-) diff --git a/redis/sentinel.py b/redis/sentinel.py index cb2dd1ead4..9983a98f4b 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -261,7 +261,6 @@ def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ - breakpoint() self._iter_req_id_to_replica_address.pop( options.get("iter_req_id", None), None ) diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 86d47a5594..39f748e9f7 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -222,23 +222,3 @@ def test_scan_iter_in_redis_cleans_up( # end of the SCAN ITER command assert not connection_pool_replica_mock._iter_req_id_to_replica_address -def test_scan_iter_in_pipeline_cleans_up( - connection_pool_replica_mock: SentinelConnectionPool, -): - """Test that connection pool is correctly cleaned up""" - from redis import Redis - - from redis.commands.core import ScanCommands - - r = Redis(connection_pool=connection_pool_replica_mock) - r.pipeline() - r.scan_iter("a") - # Patch the actual sending and parsing response from the Connection object - # but still let the connection pool does all the necessary work - with mock.patch.object(r, "_send_command_parse_response", return_value=(0, [])): - r.execute() - [k for k in r.scan_iter("a")] - # Test that the iter_req_id for the scan command is cleared at the - # end of the SCAN ITER command - assert not connection_pool_replica_mock._iter_req_id_to_replica_address - From 24b30722c4386278c6c609c7e97dc8bc833fa37e Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 14:08:36 +0000 Subject: [PATCH 26/44] lints --- redis/asyncio/sentinel.py | 21 ++++++++++++-------- redis/client.py | 1 - redis/connection.py | 4 ++-- redis/sentinel.py | 24 +++++++++++++++-------- tests/test_sentinel_managed_connection.py | 11 +++-------- 5 files changed, 34 insertions(+), 27 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 9ad520fdac..9eb422d204 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -60,11 +60,11 @@ async def _connect_retry(self, same_address: bool = False): if self._reader: return # already connected # If same_server is True, it means that the connection - # is not rotating to the next slave (if the connection pool is not master) + # is not rotating to the next slave (if the connection pool is in replica mode) if same_address: self.connect_to(self.host, self.port) return - # If same_server is False, connnect to master in master mode + # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode if self.connection_pool.is_master: await self.connect_to(await self.connection_pool.get_master_address()) @@ -84,7 +84,7 @@ async def connect(self): async def connect_to_same_address(self): """ - Similar to connect, but instead of rotating to the next slave (if not in master mode), + Similar to connect, but instead of rotating to the next slave (in replica mode), it just connects to the same address of the connection object. """ return await self.retry.call_with_retry( @@ -94,7 +94,7 @@ async def connect_to_same_address(self): async def connect_to_address(self, address): """ - Similar to connect, but instead of rotating to the next slave (if not in master mode), + Similar to connect, but instead of rotating to the next slave (in replica mode), it just connects to the same address of the connection object. """ self.host, self.port = address @@ -202,13 +202,15 @@ async def rotate_slaves(self) -> AsyncIterator: pass raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") - async def ensure_connection_connected_to_address(self, connection: SentinelManagedConnection): + async def ensure_connection_connected_to_address( + self, connection: SentinelManagedConnection + ): """ Ensure the connection is already connected to the server that this connection object wants to connect to Similar to self.ensure_connection, but calling connection.connect() - in SentinelManagedConnection (replica mode) will cause the + in SentinelManagedConnection (replica mode) will cause the connection object to connect to the next replica in rotation, and we don't wnat behavior. Look at get_connection inline docs for details. @@ -217,7 +219,10 @@ async def ensure_connection_connected_to_address(self, connection: SentinelManag """ await connection.connect_to_same_address() try: - if await connection.can_read_destructive() and connection.client_cache is None: + if ( + await connection.can_read_destructive() + and connection.client_cache is None + ): raise ConnectionError("Connection has data") except (ConnectionError, OSError): await connection.disconnect() @@ -273,7 +278,7 @@ async def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and - # port to the one that we want later in the call to ``connect_to_same_address`` + # port to the one that we want later in the call to ``connect_to_address`` if not connection: connection = self.make_connection() assert connection diff --git a/redis/client.py b/redis/client.py index 7d9fe343ca..e39e88edbd 100755 --- a/redis/client.py +++ b/redis/client.py @@ -763,7 +763,6 @@ def reset(self) -> None: self.patterns = {} self.pending_unsubscribe_patterns = set() self.subscribed_event.clear() - self.connection_pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) def close(self) -> None: self.reset() diff --git a/redis/connection.py b/redis/connection.py index 87f41bb366..39aa8470ab 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -737,10 +737,10 @@ def _host_error(self): class ConnectionsIndexer(Iterable): """ - Data structure that simulates a list of available connections. + Data structure that simulates a list of available connections. Instead of list, we keep 2 additional DS to support O(1) operations on all of the class' methods. - The first DS is indexed on the connection object's ID. + The first DS is indexed on the connection object's ID. The second DS is indexed on the address (ip and port) of the connection. """ diff --git a/redis/sentinel.py b/redis/sentinel.py index 9983a98f4b..fa44c7a79a 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -49,7 +49,7 @@ def _connect_retry(self, same_address: bool = False): if same_address: self.connect_to((self.host, self.port)) return - # If same_server is False, connnect to master in master mode + # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) @@ -62,11 +62,14 @@ def _connect_retry(self, same_address: bool = False): raise SlaveNotFoundError # Never be here def connect(self): - return self.retry.call_with_retry(lambda: self._connect_retry(), lambda error: None) + return self.retry.call_with_retry( + lambda: self._connect_retry(), + lambda error: None + ) def connect_to_same_address(self): """ - Similar to connect, but instead of rotating to the next slave (if not in master mode), + Similar to connect, but instead of rotating to the next slave (in replica mode), it just connects to the same address of the connection object. """ return self.retry.call_with_retry( @@ -76,14 +79,17 @@ def connect_to_same_address(self): def connect_to_address(self, address): """ - Similar to connect, but instead of rotating to the next slave (if not in master mode), + Similar to connect, but instead of rotating to the next slave (in replica mode), it just connects to the address supplied. """ self.host, self.port = address return self.connect_to_same_address() def can_read_same_address(self, timeout=0): - """Similar to can_read_same_address, but calls connect_to_same_address instead of connect""" + """ + Similar to can_read_same_address, but calls + connect_to_same_address instead of connect + """ sock = self._sock if not sock: self.connect_to_same_address() @@ -234,13 +240,15 @@ def rotate_slaves(self): "Round-robin slave balancer" return self.proxy.rotate_slaves() - def ensure_connection_connected_to_address(self, connection: SentinelManagedConnection): + def ensure_connection_connected_to_address( + self, connection: SentinelManagedConnection + ): """ Ensure the connection is already connected to the server that this connection - object wants to connect to + object wants to connect to. Similar to self.ensure_connection, but calling connection.connect() - in SentinelManagedConnection (replica mode) will cause the + in SentinelManagedConnection (replica mode) will cause the connection object to connect to the next replica in rotation, and we don't wnat behavior. Look at get_connection inline docs for details. diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 39f748e9f7..5d2f0d1b8b 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -2,11 +2,9 @@ from unittest import mock import pytest -from redis.sentinel import ( - Sentinel, - SentinelConnectionPool, - SentinelManagedConnection, -) +from redis import Redis +from redis.sentinel import Sentinel, SentinelConnectionPool, SentinelManagedConnection + class SentinelManagedConnectionMock(SentinelManagedConnection): def connect_to_same_address(self) -> None: @@ -211,8 +209,6 @@ def test_scan_iter_in_redis_cleans_up( connection_pool_replica_mock: SentinelConnectionPool, ): """Test that connection pool is correctly cleaned up""" - from redis import Redis - r = Redis(connection_pool=connection_pool_replica_mock) # Patch the actual sending and parsing response from the Connection object # but still let the connection pool does all the necessary work @@ -221,4 +217,3 @@ def test_scan_iter_in_redis_cleans_up( # Test that the iter_req_id for the scan command is cleared at the # end of the SCAN ITER command assert not connection_pool_replica_mock._iter_req_id_to_replica_address - From 1d29914d7fe0c892292e2ee703f1eb017461ce29 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 14:13:28 +0000 Subject: [PATCH 27/44] reformat --- redis/asyncio/sentinel.py | 4 +--- redis/sentinel.py | 10 +++------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 9eb422d204..96dbfcead1 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -234,9 +234,7 @@ def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ - self._iter_req_id_to_replica_address.pop( - options.get("iter_req_id", None), None - ) + self._iter_req_id_to_replica_address.pop(options.get("iter_req_id", None), None) async def get_connection( self, command_name: str, *keys: Any, **options: Any diff --git a/redis/sentinel.py b/redis/sentinel.py index fa44c7a79a..84004ca369 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -63,8 +63,7 @@ def _connect_retry(self, same_address: bool = False): def connect(self): return self.retry.call_with_retry( - lambda: self._connect_retry(), - lambda error: None + lambda: self._connect_retry(), lambda error: None ) def connect_to_same_address(self): @@ -73,8 +72,7 @@ def connect_to_same_address(self): it just connects to the same address of the connection object. """ return self.retry.call_with_retry( - lambda: self._connect_retry(same_address=True), - lambda error: None + lambda: self._connect_retry(same_address=True), lambda error: None ) def connect_to_address(self, address): @@ -269,9 +267,7 @@ def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary """ - self._iter_req_id_to_replica_address.pop( - options.get("iter_req_id", None), None - ) + self._iter_req_id_to_replica_address.pop(options.get("iter_req_id", None), None) def get_connection( self, command_name: str, *keys: Any, **options: Any From a4e20ac956236ab4cc01d3f99db4e020f5c87c36 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 22:17:19 +0800 Subject: [PATCH 28/44] def cleanup in base class --- redis/asyncio/connection.py | 7 +++++-- redis/connection.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index b0d511f176..2cca66bc65 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1057,12 +1057,15 @@ class ConnectionPool: ``connection_class``. """ - @abstractmethod def cleanup(self, **options): """ Additional cleanup operations that the connection pool might - need to do after a SCAN ITER family command is executed + need to do after a SCAN ITER family command is executed. + + See SentinelManagedConnection for an example cleanup operation that + might need to be done. """ + pass @classmethod def from_url(cls: Type[_CP], url: str, **kwargs) -> _CP: diff --git a/redis/connection.py b/redis/connection.py index 39aa8470ab..750e1ec702 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -1162,6 +1162,16 @@ def __repr__(self) -> (str, str): f"({repr(self.connection_class(**self.connection_kwargs))})>" ) + def cleanup(self, **options): + """ + Additional cleanup operations that the connection pool might + need to do after a SCAN ITER family command is executed. + + See SentinelManagedConnection for an example cleanup operation that + might need to be done. + """ + pass + def reset(self) -> None: self._lock = threading.Lock() self._created_connections = 0 From 12e6495dff05c57291ea59f207bb236d3bcb8da9 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 14:50:11 +0000 Subject: [PATCH 29/44] fix some tests --- redis/asyncio/client.py | 4 +++- redis/asyncio/sentinel.py | 3 +++ redis/client.py | 4 +++- redis/sentinel.py | 1 + tests/test_asyncio/test_sentinel_managed_connection.py | 5 +++-- 5 files changed, 13 insertions(+), 4 deletions(-) diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index bb5a7d2a65..abc141e8c9 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -651,7 +651,9 @@ async def execute_command(self, *args, **options): finally: if not self.connection: await pool.release(conn) - if "ITER" in command_name.upper(): + # Do additional cleanup if this is part of a SCAN ITER family command. + # It's possible that this is just a pure SCAN family command though. + if "SCAN" in command_name.upper(): pool.cleanup(iter_req_id=options.get("iter_req_id", None)) async def parse_response( diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 96dbfcead1..2491cdab33 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -253,6 +253,7 @@ async def get_connection( should go to the same replica. """ # If not an iter command or in master mode, call superclass' implementation + breakpoint() if not (iter_req_id := options.get("iter_req_id", None)) or self.is_master: return await super().get_connection(command_name, *keys, **options) @@ -292,7 +293,9 @@ async def get_connection( # This will connect to the host and port of the replica else: await connection.connect_to_address((server_host, server_port)) + breakpoint() await self.ensure_connection_connected_to_address(connection) + breakpoint() except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/client.py b/redis/client.py index e39e88edbd..c29a1d226c 100755 --- a/redis/client.py +++ b/redis/client.py @@ -581,8 +581,10 @@ def execute_command(self, *args, **options): finally: if not self.connection: pool.release(conn) + # Do additional cleanup if this is part of a SCAN ITER family command. + # It's possible that this is just a pure SCAN family command though. if "SCAN" in command_name.upper(): - pool.cleanup(iter_req_id=options.get("_iter_req_id", None)) + pool.cleanup(iter_req_id=options.get("iter_req_id", None)) def parse_response(self, connection, command_name, **options): """Parses a response from the Redis server""" diff --git a/redis/sentinel.py b/redis/sentinel.py index 84004ca369..03915481f7 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -54,6 +54,7 @@ def _connect_retry(self, same_address: bool = False): if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) else: + breakpoint() for slave in self.connection_pool.rotate_slaves(): try: return self.connect_to(slave) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 2c79247c2a..5c7af46c17 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -251,5 +251,6 @@ async def test_scan_iter_family_cleans_up( r = Redis(connection_pool=connection_pool_replica_mock) - [k async for k in r.scan_iter("a")] - assert not connection_pool_replica_mock.iter_req_id_to_replica_address + with mock.patch.object(r, "_send_command_parse_response", return_value=(0, [])): + [k async for k in r.scan_iter("a")] + assert not connection_pool_replica_mock._iter_req_id_to_replica_address From a7ffda89091a02cbdfc8edd02da2ba5f89dc0329 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 22:53:26 +0800 Subject: [PATCH 30/44] rename iter_req_id properly --- .../test_sentinel_managed_connection.py | 8 +++---- tests/test_sentinel_managed_connection.py | 24 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 5c7af46c17..8993dde416 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -157,7 +157,7 @@ async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica ) -async def test_connects_to_diff_address_if_noiter_req_id_replica( +async def test_connects_to_diff_address_if_no_iter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ @@ -182,7 +182,7 @@ async def test_connects_to_diff_address_if_noiter_req_id_replica( ) -async def test_connects_to_same_address_if_sameiter_req_id_master( +async def test_connects_to_same_address_if_same_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ @@ -199,7 +199,7 @@ async def test_connects_to_same_address_if_sameiter_req_id_master( ) -async def test_connects_to_same_conn_object_if_sameiter_req_id_and_released_master( +async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ @@ -218,7 +218,7 @@ async def test_connects_to_same_conn_object_if_sameiter_req_id_and_released_mast ) -async def test_connects_to_same_address_if_noiter_req_id_master( +async def test_connects_to_same_address_if_no_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: """ diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 5d2f0d1b8b..d639f96540 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -88,15 +88,15 @@ def test_connects_to_same_address_if_same_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection address is the same if the ``_iter_req_id`` is the same + Assert that the connection address is the same if the ``iter_req_id`` is the same when we are in replica mode using a :py:class:`~redis.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) assert same_address( - connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1), + connection_pool_replica_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) @@ -105,18 +105,18 @@ def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection object is the same if the ``_iter_req_id`` is the same + Assert that the connection object is the same if the ``iter_req_id`` is the same when we are in replica mode using a :py:class:`~redis.sentinel.SentinelConnectionPool` and if we release the connection back to the connection pool before trying to connect again. """ connection_for_req_1 = connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) connection_pool_replica_mock.release(connection_for_req_1) assert ( - connection_pool_replica_mock.get_connection("ANY", _iter_req_id=1) + connection_pool_replica_mock.get_connection("ANY", iter_req_id=1) == connection_for_req_1 ) @@ -125,12 +125,12 @@ def test_connects_to_diff_address_if_no_iter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: """ - Assert that the connection object is different if no _iter_req_id is supplied. + Assert that the connection object is different if no iter_req_id is supplied. In reality, they can be the same, but in this case, we're not releasing the connection to the pool so they should always be different. """ connection_for_req_1 = connection_pool_replica_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) connection_for_random_req = connection_pool_replica_mock.get_connection("ANYY") assert not same_address(connection_for_random_req, connection_for_req_1) @@ -156,7 +156,7 @@ def test_connects_to_same_address_if_same_iter_req_id_master( "ANY", _iter_req_id=1 ) assert same_address( - connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + connection_pool_master_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) @@ -172,10 +172,10 @@ def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( trying to connect again. """ connection_for_req_1 = connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) assert same_address( - connection_pool_master_mock.get_connection("ANY", _iter_req_id=1), + connection_pool_master_mock.get_connection("ANY", iter_req_id=1), connection_for_req_1, ) @@ -190,7 +190,7 @@ def test_connects_to_same_address_if_no_iter_req_id_master( :py:class:`~redis.sentinel.SentinelConnectionPool`. """ connection_for_req_1 = connection_pool_master_mock.get_connection( - "ANY", _iter_req_id=1 + "ANY", iter_req_id=1 ) connection_for_random_req = connection_pool_master_mock.get_connection("ANYY") assert same_address(connection_for_random_req, connection_for_req_1) From 8a1bc0660966c170aefbb4cc2d41e7b100982cea Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Tue, 16 Jul 2024 15:23:03 +0000 Subject: [PATCH 31/44] fix tests --- redis/asyncio/sentinel.py | 3 --- redis/sentinel.py | 1 - 2 files changed, 4 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 2491cdab33..96dbfcead1 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -253,7 +253,6 @@ async def get_connection( should go to the same replica. """ # If not an iter command or in master mode, call superclass' implementation - breakpoint() if not (iter_req_id := options.get("iter_req_id", None)) or self.is_master: return await super().get_connection(command_name, *keys, **options) @@ -293,9 +292,7 @@ async def get_connection( # This will connect to the host and port of the replica else: await connection.connect_to_address((server_host, server_port)) - breakpoint() await self.ensure_connection_connected_to_address(connection) - breakpoint() except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/sentinel.py b/redis/sentinel.py index 03915481f7..84004ca369 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -54,7 +54,6 @@ def _connect_retry(self, same_address: bool = False): if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) else: - breakpoint() for slave in self.connection_pool.rotate_slaves(): try: return self.connect_to(slave) From f504741baf3c901b146d23847fce2e0a9567b014 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 00:56:00 +0800 Subject: [PATCH 32/44] set fix address as a property of SentinelManagedConnection --- redis/asyncio/connection.py | 6 ++- redis/asyncio/sentinel.py | 65 +++++++---------------------- redis/connection.py | 6 ++- redis/sentinel.py | 82 ++++++++----------------------------- tests/test_connection.py | 2 + 5 files changed, 42 insertions(+), 119 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 2cca66bc65..7e88e8b1e5 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1118,6 +1118,7 @@ def __init__( self, connection_class: Type[AbstractConnection] = Connection, max_connections: Optional[int] = None, + index_available_connections: bool = False, **connection_kwargs, ): max_connections = max_connections or 2**31 @@ -1130,6 +1131,7 @@ def __init__( self._available_connections: ConnectionsIndexer = ConnectionsIndexer() self._in_use_connections: Set[AbstractConnection] = set() + self._index_available_connections= index_available_connections self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) def __repr__(self): @@ -1139,7 +1141,9 @@ def __repr__(self): ) def reset(self): - self._available_connections = [] + self._available_connections: ConnectionsIndexer | list = ( + ConnectionsIndexer() if self._index_available_connections else [] + ) self._in_use_connections = weakref.WeakSet() def can_get_connection(self) -> bool: diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 96dbfcead1..98a8466e55 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -35,6 +35,11 @@ class SlaveNotFoundError(ConnectionError): class SentinelManagedConnection(Connection): def __init__(self, **kwargs): self.connection_pool = kwargs.pop("connection_pool") + # To be set to True if we want to prevent + # the sentinel managed connection to connect + # to the most relevant sentinel in the pool and just + # connect to the current self.host and self.port + self._is_address_fixed = False super().__init__(**kwargs) def __repr__(self): @@ -48,6 +53,10 @@ def __repr__(self): s += host_info return s + ")>" + def fix_address(self, address): + self.host, self.port = address + self._is_address_fixed = True + async def connect_to(self, address): self.host, self.port = address await super().connect() @@ -56,12 +65,12 @@ async def connect_to(self, address): if str_if_bytes(await self.read_response()) != "PONG": raise ConnectionError("PING failed") - async def _connect_retry(self, same_address: bool = False): + async def _connect_retry(self): if self._reader: return # already connected - # If same_server is True, it means that the connection + # If address is fixed, it means that the connection # is not rotating to the next slave (if the connection pool is in replica mode) - if same_address: + if self._is_address_fixed: self.connect_to(self.host, self.port) return # If same_server is False, connnect to master in master mode @@ -82,24 +91,6 @@ async def connect(self): lambda error: asyncio.sleep(0), ) - async def connect_to_same_address(self): - """ - Similar to connect, but instead of rotating to the next slave (in replica mode), - it just connects to the same address of the connection object. - """ - return await self.retry.call_with_retry( - lambda: self._connect_retry(same_address=True), - lambda error: asyncio.sleep(0), - ) - - async def connect_to_address(self, address): - """ - Similar to connect, but instead of rotating to the next slave (in replica mode), - it just connects to the same address of the connection object. - """ - self.host, self.port = address - return await self.connect_to_same_address() - async def read_response( self, disable_decoding: bool = False, @@ -202,34 +193,6 @@ async def rotate_slaves(self) -> AsyncIterator: pass raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") - async def ensure_connection_connected_to_address( - self, connection: SentinelManagedConnection - ): - """ - Ensure the connection is already connected to the server that this connection - object wants to connect to - - Similar to self.ensure_connection, but calling connection.connect() - in SentinelManagedConnection (replica mode) will cause the - connection object to connect to the next replica in rotation, - and we don't wnat behavior. Look at get_connection inline docs for details. - - Here, we just try to make sure that the connection is already connected - to the replica we wanted it to. - """ - await connection.connect_to_same_address() - try: - if ( - await connection.can_read_destructive() - and connection.client_cache is None - ): - raise ConnectionError("Connection has data") - except (ConnectionError, OSError): - await connection.disconnect() - await connection.connect_to_same_address() - if await connection.can_read_destructive(): - raise ConnectionError("Connection has data") - def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary @@ -291,8 +254,8 @@ async def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - await connection.connect_to_address((server_host, server_port)) - await self.ensure_connection_connected_to_address(connection) + connection.fix_address((server_host, server_port)) + await self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/redis/connection.py b/redis/connection.py index 750e1ec702..7ce73dcb35 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -1135,6 +1135,7 @@ def __init__( self, connection_class=Connection, max_connections: Optional[int] = None, + index_available_connections: bool = False, **connection_kwargs, ): max_connections = max_connections or 2**31 @@ -1154,6 +1155,7 @@ def __init__( # will notice the first thread already did the work and simply # release the lock. self._fork_lock = threading.Lock() + self._index_available_connections= index_available_connections self.reset() def __repr__(self) -> (str, str): @@ -1175,7 +1177,9 @@ def cleanup(self, **options): def reset(self) -> None: self._lock = threading.Lock() self._created_connections = 0 - self._available_connections = ConnectionsIndexer() + self._available_connections: ConnectionsIndexer | list = ( + ConnectionsIndexer() if self._index_available_connections else [] + ) self._in_use_connections = set() # this must be the last operation in this method. while reset() is diff --git a/redis/sentinel.py b/redis/sentinel.py index 84004ca369..65612ce611 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -20,6 +20,11 @@ class SlaveNotFoundError(ConnectionError): class SentinelManagedConnection(Connection): def __init__(self, **kwargs): self.connection_pool = kwargs.pop("connection_pool") + # To be set to True if we want to prevent + # the sentinel managed connection to connect + # to the most relevant sentinel in the pool and just + # connect to the current self.host and self.port + self._is_address_fixed = False super().__init__(**kwargs) def __repr__(self): @@ -33,6 +38,10 @@ def __repr__(self): s = s % host_info return s + def fix_address(self, address): + self.host, self.port = address + self._is_address_fixed = True + def connect_to(self, address): self.host, self.port = address super().connect() @@ -41,12 +50,12 @@ def connect_to(self, address): if str_if_bytes(self.read_response()) != "PONG": raise ConnectionError("PING failed") - def _connect_retry(self, same_address: bool = False): + def _connect_retry(self): if self._sock: return # already connected - # If same_server is True, it means that the connection + # If address is fixed, it means that the connection # is not rotating to the next slave (if the connection pool is not master) - if same_address: + if self._is_address_fixed: self.connect_to((self.host, self.port)) return # If same_server is False, connnect to master in master mode @@ -66,40 +75,6 @@ def connect(self): lambda: self._connect_retry(), lambda error: None ) - def connect_to_same_address(self): - """ - Similar to connect, but instead of rotating to the next slave (in replica mode), - it just connects to the same address of the connection object. - """ - return self.retry.call_with_retry( - lambda: self._connect_retry(same_address=True), lambda error: None - ) - - def connect_to_address(self, address): - """ - Similar to connect, but instead of rotating to the next slave (in replica mode), - it just connects to the address supplied. - """ - self.host, self.port = address - return self.connect_to_same_address() - - def can_read_same_address(self, timeout=0): - """ - Similar to can_read_same_address, but calls - connect_to_same_address instead of connect - """ - sock = self._sock - if not sock: - self.connect_to_same_address() - - host_error = self._host_error() - - try: - return self._parser.can_read(timeout) - except OSError as e: - self.disconnect() - raise ConnectionError(f"Error while reading from {host_error}: {e.args}") - def read_response( self, disable_decoding=False, @@ -203,7 +178,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs): service_name=service_name, sentinel_manager=sentinel_manager, ) - super().__init__(**kwargs) + super().__init__(index_available_connections=True, **kwargs) self.connection_kwargs["connection_pool"] = self.proxy self.service_name = service_name self.sentinel_manager = sentinel_manager @@ -238,31 +213,6 @@ def rotate_slaves(self): "Round-robin slave balancer" return self.proxy.rotate_slaves() - def ensure_connection_connected_to_address( - self, connection: SentinelManagedConnection - ): - """ - Ensure the connection is already connected to the server that this connection - object wants to connect to. - - Similar to self.ensure_connection, but calling connection.connect() - in SentinelManagedConnection (replica mode) will cause the - connection object to connect to the next replica in rotation, - and we don't wnat behavior. Look at get_connection inline docs for details. - - Here, we just try to make sure that the connection is already connected - to the replica we wanted it to. - """ - connection.connect_to_same_address() - try: - if connection.can_read_same_address() and connection.client_cache is None: - raise ConnectionError("Connection has data") - except (ConnectionError, OSError): - connection.disconnect() - connection.connect_to_same_address() - if connection.can_read(): - raise ConnectionError("Connection has data") - def cleanup(self, **options): """ Remove the SCAN ITER family command's request id from the dictionary @@ -309,7 +259,7 @@ def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and port - # to the one that we want later in the call to ``connect_to_same_address`` + # to the one that we want later in the call to ``fix_address`` if not connection: connection = self.make_connection() assert connection @@ -324,8 +274,8 @@ def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - connection.connect_to_address((server_host, server_port)) - self.ensure_connection_connected_to_address(connection) + connection.fix_address((server_host, server_port)) + self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't # leak it diff --git a/tests/test_connection.py b/tests/test_connection.py index e5bbb597f5..910f4b1586 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -298,6 +298,7 @@ def mock_disconnect(_): assert called == 1 pool.disconnect() + @pytest.mark.parametrize( "conn, error, expected_message", [ @@ -347,6 +348,7 @@ def test_unix_socket_connection_failure(): == "Error 2 connecting to unix:///tmp/a.sock. No such file or directory." ) + def test_connections_indexer_operations(): ci = ConnectionsIndexer() c1 = Connection(host="1", port=2) From c75db8d0a7e8af881fad4c7966d9735f4dfe0bb5 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 04:50:35 +0000 Subject: [PATCH 33/44] lint --- redis/asyncio/connection.py | 2 +- redis/connection.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 7e88e8b1e5..7c044347fe 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1131,7 +1131,7 @@ def __init__( self._available_connections: ConnectionsIndexer = ConnectionsIndexer() self._in_use_connections: Set[AbstractConnection] = set() - self._index_available_connections= index_available_connections + self._index_available_connections = index_available_connections self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) def __repr__(self): diff --git a/redis/connection.py b/redis/connection.py index 7ce73dcb35..4e73ba7fc6 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -1155,7 +1155,7 @@ def __init__( # will notice the first thread already did the work and simply # release the lock. self._fork_lock = threading.Lock() - self._index_available_connections= index_available_connections + self._index_available_connections = index_available_connections self.reset() def __repr__(self) -> (str, str): From 0f5e07966bf730826612a8e63df46d036850f8a1 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 12:56:30 +0800 Subject: [PATCH 34/44] make mock class have same behavior as actual class --- redis/sentinel.py | 3 ++ .../test_sentinel_managed_connection.py | 42 +++++-------------- tests/test_sentinel_managed_connection.py | 33 ++++++--------- 3 files changed, 26 insertions(+), 52 deletions(-) diff --git a/redis/sentinel.py b/redis/sentinel.py index 65612ce611..9cb1a8ef97 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -58,6 +58,9 @@ def _connect_retry(self): if self._is_address_fixed: self.connect_to((self.host, self.port)) return + self._connect_to_sentinel() + + def _connect_to_sentinel(self): # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode if self.connection_pool.is_master: diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 8993dde416..732220f456 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -43,41 +43,21 @@ async def mock_connect(): class SentinelManagedConnectionMock(SentinelManagedConnection): - async def connect_to_same_address(self) -> None: - pass - - async def can_read_destructive(self) -> bool: - # Mock this function to always return False. - # Trrue means there's still data to be read and hence we can't reconnect - # to this connection yet - return False - - -class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock): - async def connect(self) -> None: + async def _connect_to_sentinel(self) -> None: """ - This simulates the behavior of connect when + This simulates the behavior of _connect_to_sentinel when :py:class:`~redis.SentinelConnectionPool` is in replica mode. It'll call rotate_slaves and connect to the next replica. """ - import random - import time - - self.host = f"host-{random.randint(0, 10)}" - self.port = time.time() - + if self.connection_pool.is_master: + self.host, self.port = ("master", 1) + else: + import random + import time -class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): - async def connect_to(self, address: Tuple[str, int]) -> None: - """ - This simulates the behavior of connect_to when - :py:class:`~redis.SentinelConnectionPool` - is in master mode. - It'll try to connect to master. In this mock class, - it'll just set the host and port without actually connecting. - """ - self.host, self.port = address + self.host = f"host-{random.randint(0, 10)}" + self.port = time.time() @pytest.fixture() @@ -90,7 +70,7 @@ def connection_pool_replica_mock() -> SentinelConnectionPool: "usasm", sentinel_manager, is_master=False, - connection_class=SentinelManagedConnectionMockForReplicaMode, + connection_class=SentinelManagedConnectionMock, ) return connection_pool @@ -105,7 +85,7 @@ def connection_pool_master_mock() -> SentinelConnectionPool: "usasm", sentinel_manager, is_master=True, - connection_class=SentinelManagedConnectionMockForMasterMode, + connection_class=SentinelManagedConnectionMock, ) return connection_pool diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index d639f96540..b07cfbaa63 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -7,32 +7,24 @@ class SentinelManagedConnectionMock(SentinelManagedConnection): - def connect_to_same_address(self) -> None: - pass - - def can_read_destructive(self) -> bool: - # Mock this function to always return False. - # Trrue means there's still data to be read and hence we can't reconnect - # to this connection yet - return False - - -class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock): - def connect(self) -> None: + def _connect_to_sentinel(self) -> None: """ - This simulates the behavior of connect when + This simulates the behavior of _connect_to_sentinel when :py:class:`~redis.SentinelConnectionPool` is in replica mode. It'll call rotate_slaves and connect to the next replica. """ - import random - import time + if self.connection_pool.is_master: + self.host, self.port = ("master", 1) + else: + import random + import time - self.host = f"host-{random.randint(0, 10)}" - self.port = time.time() + self.host = f"host-{random.randint(0, 10)}" + self.port = time.time() -class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock): +class SentinelManagedConnectionMock(SentinelManagedConnectionMock): def connect_to(self, address: Tuple[str, int]) -> None: """ This simulates the behavior of connect_to when @@ -41,7 +33,6 @@ def connect_to(self, address: Tuple[str, int]) -> None: It'll try to connect to master. In this mock class, it'll just set the host and port without actually connecting. """ - self.host, self.port = address @pytest.fixture() @@ -54,7 +45,7 @@ def connection_pool_replica_mock() -> SentinelConnectionPool: "usasm", sentinel_manager, is_master=False, - connection_class=SentinelManagedConnectionMockForReplicaMode, + connection_class=SentinelManagedConnectionMock, ) return connection_pool @@ -69,7 +60,7 @@ def connection_pool_master_mock() -> SentinelConnectionPool: "usasm", sentinel_manager, is_master=True, - connection_class=SentinelManagedConnectionMockForMasterMode, + connection_class=SentinelManagedConnectionMock, ) return connection_pool From ef0d969450f8ae75c7b0f87ff4550aa0f35450eb Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 12:58:10 +0800 Subject: [PATCH 35/44] define _connect_to_sentinel in async server --- redis/asyncio/sentinel.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 98a8466e55..b111b341fb 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -73,6 +73,9 @@ async def _connect_retry(self): if self._is_address_fixed: self.connect_to(self.host, self.port) return + await self._connect_to_sentinel() + + async def _connect_to_sentinel(self): # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode if self.connection_pool.is_master: From bb69c855101caf64eac264a9b85c9780ba25f78f Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 05:52:50 +0000 Subject: [PATCH 36/44] mock can_read_destructive for parser --- redis/asyncio/sentinel.py | 4 +- redis/sentinel.py | 2 +- .../test_sentinel_managed_connection.py | 49 ++++++++++++------- tests/test_sentinel_managed_connection.py | 9 +--- 4 files changed, 35 insertions(+), 29 deletions(-) diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index b111b341fb..9bac55167a 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -71,10 +71,10 @@ async def _connect_retry(self): # If address is fixed, it means that the connection # is not rotating to the next slave (if the connection pool is in replica mode) if self._is_address_fixed: - self.connect_to(self.host, self.port) + await self.connect_to((self.host, self.port)) return await self._connect_to_sentinel() - + async def _connect_to_sentinel(self): # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode diff --git a/redis/sentinel.py b/redis/sentinel.py index 9cb1a8ef97..5fd6c0151f 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -59,7 +59,7 @@ def _connect_retry(self): self.connect_to((self.host, self.port)) return self._connect_to_sentinel() - + def _connect_to_sentinel(self): # If same_server is False, connnect to master in master mode # and rotate to the next slave in slave mode diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 732220f456..ee7ba09c7f 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -1,5 +1,5 @@ import socket -from typing import Tuple +from typing import Iterator, Tuple import pytest from redis.asyncio.retry import Retry @@ -59,35 +59,46 @@ async def _connect_to_sentinel(self) -> None: self.host = f"host-{random.randint(0, 10)}" self.port = time.time() + async def connect_to(self, address: Tuple[str, int]) -> None: + """ + Do nothing, just mock. + """ + @pytest.fixture() -def connection_pool_replica_mock() -> SentinelConnectionPool: +def connection_pool_replica_mock() -> Iterator[SentinelConnectionPool]: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave sentinel_manager.discover_slaves = mock.AsyncMock(return_value=["replica", 5000]) - # Create connection pool with our mock connection object - connection_pool = SentinelConnectionPool( - "usasm", - sentinel_manager, - is_master=False, - connection_class=SentinelManagedConnectionMock, - ) - return connection_pool + with mock.patch( + "redis._parsers._AsyncRESP2Parser.can_read_destructive", return_value=False + ): + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=False, + connection_class=SentinelManagedConnectionMock, + ) + yield connection_pool @pytest.fixture() -def connection_pool_master_mock() -> SentinelConnectionPool: +def connection_pool_master_mock() -> Iterator[SentinelConnectionPool]: sentinel_manager = Sentinel([["master", 400]]) # Give a random slave sentinel_manager.discover_master = mock.AsyncMock(return_value=["replica", 5000]) - # Create connection pool with our mock connection object - connection_pool = SentinelConnectionPool( - "usasm", - sentinel_manager, - is_master=True, - connection_class=SentinelManagedConnectionMock, - ) - return connection_pool + with mock.patch( + "redis._parsers._AsyncRESP2Parser.can_read_destructive", return_value=False + ): + # Create connection pool with our mock connection object + connection_pool = SentinelConnectionPool( + "usasm", + sentinel_manager, + is_master=True, + connection_class=SentinelManagedConnectionMock, + ) + yield connection_pool def same_address( diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index b07cfbaa63..22019d9691 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -23,16 +23,11 @@ def _connect_to_sentinel(self) -> None: self.host = f"host-{random.randint(0, 10)}" self.port = time.time() - -class SentinelManagedConnectionMock(SentinelManagedConnectionMock): def connect_to(self, address: Tuple[str, int]) -> None: """ - This simulates the behavior of connect_to when - :py:class:`~redis.SentinelConnectionPool` - is in master mode. - It'll try to connect to master. In this mock class, - it'll just set the host and port without actually connecting. + Do nothing, this is just to mock. """ + pass @pytest.fixture() From a4601a74c923afb64d38000921da01ca6452fa12 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 08:35:16 +0000 Subject: [PATCH 37/44] skip test sentinel managed connection if hirediswq --- redis/asyncio/connection.py | 4 +++- tests/test_asyncio/test_sentinel_managed_connection.py | 7 +++++++ tests/test_sentinel_managed_connection.py | 2 ++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 7c044347fe..5f10d0e85e 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1129,7 +1129,9 @@ def __init__( self.connection_kwargs = connection_kwargs self.max_connections = max_connections - self._available_connections: ConnectionsIndexer = ConnectionsIndexer() + self._available_connections: ConnectionsIndexer = ( + ConnectionsIndexer() if index_available_connections else [] + ) self._in_use_connections: Set[AbstractConnection] = set() self._index_available_connections = index_available_connections self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index ee7ba09c7f..5ff83f7eaa 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -111,6 +111,7 @@ def same_address( ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_same_address_if_same_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: @@ -128,6 +129,7 @@ async def test_connects_to_same_address_if_same_id_replica( ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: @@ -148,6 +150,7 @@ async def test_connects_to_same_conn_object_if_same_id_and_conn_released_replica ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_diff_address_if_no_iter_req_id_replica( connection_pool_replica_mock: SentinelConnectionPool, ) -> None: @@ -173,6 +176,7 @@ async def test_connects_to_diff_address_if_no_iter_req_id_replica( ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_same_address_if_same_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: @@ -190,6 +194,7 @@ async def test_connects_to_same_address_if_same_iter_req_id_master( ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: @@ -209,6 +214,7 @@ async def test_connects_to_same_conn_object_if_same_iter_req_id_and_released_mas ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_connects_to_same_address_if_no_iter_req_id_master( connection_pool_master_mock: SentinelConnectionPool, ) -> None: @@ -234,6 +240,7 @@ async def test_connects_to_same_address_if_no_iter_req_id_master( ) +@pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") async def test_scan_iter_family_cleans_up( connection_pool_replica_mock: SentinelConnectionPool, ): diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 22019d9691..dbd83e2eb6 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -5,6 +5,8 @@ from redis import Redis from redis.sentinel import Sentinel, SentinelConnectionPool, SentinelManagedConnection +pytestmark = pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") + class SentinelManagedConnectionMock(SentinelManagedConnection): def _connect_to_sentinel(self) -> None: From 88501a0f6c0194949bf169491fefecf4ddbc4631 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 08:44:02 +0000 Subject: [PATCH 38/44] undo ensure_connection deduplication in BlockingConnectionPool --- redis/connection.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/redis/connection.py b/redis/connection.py index 4e73ba7fc6..26a3b5b3dd 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -1483,7 +1483,20 @@ def get_connection(self, command_name, *keys, **options): connection = self.make_connection() try: - self.ensure_connection(connection) + # ensure this connection is connected to Redis + connection.connect() + # connections that the pool provides should be ready to send + # a command. if not, the connection was either returned to the + # pool before all data has been read or the socket has been + # closed. either way, reconnect and verify everything is good. + try: + if connection.can_read(): + raise ConnectionError("Connection has data") + except (ConnectionError, OSError): + connection.disconnect() + connection.connect() + if connection.can_read(): + raise ConnectionError("Connection not ready") except BaseException: # release the connection back to the pool so that we don't leak it self.release(connection) From c75262b30e3808e67e5205addc21cf67fcca7792 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 08:45:33 +0000 Subject: [PATCH 39/44] import HIREDIS --- tests/test_asyncio/test_sentinel_managed_connection.py | 1 + tests/test_sentinel_managed_connection.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 5ff83f7eaa..2ee1247203 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -9,6 +9,7 @@ SentinelManagedConnection, ) from redis.backoff import NoBackoff +from redis.utils import HIREDIS_AVAILABLE from .compat import mock diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index dbd83e2eb6..382e6a53fd 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -4,6 +4,7 @@ import pytest from redis import Redis from redis.sentinel import Sentinel, SentinelConnectionPool, SentinelManagedConnection +from redis.utils import HIREDIS_AVAILABLE pytestmark = pytest.mark.skipif(HIREDIS_AVAILABLE, reason="PythonParser only") From cdbe9578227c0abf4d8bc6f9c382d7140af1eeee Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 17:10:13 +0800 Subject: [PATCH 40/44] polymorphism for reset available connections instead --- redis/asyncio/connection.py | 15 +++++---- redis/asyncio/sentinel.py | 4 +++ redis/connection.py | 63 +++---------------------------------- redis/sentinel.py | 59 ++++++++++++++++++++++++++++++++-- tests/test_connection.py | 15 --------- tests/test_sentinel.py | 16 ++++++++++ 6 files changed, 91 insertions(+), 81 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 5f10d0e85e..a8357cb4e5 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -38,7 +38,7 @@ from redis.asyncio.retry import Retry from redis.backoff import NoBackoff -from redis.connection import DEFAULT_RESP_VERSION, ConnectionsIndexer +from redis.connection import DEFAULT_RESP_VERSION from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider from redis.exceptions import ( AuthenticationError, @@ -1118,7 +1118,6 @@ def __init__( self, connection_class: Type[AbstractConnection] = Connection, max_connections: Optional[int] = None, - index_available_connections: bool = False, **connection_kwargs, ): max_connections = max_connections or 2**31 @@ -1129,11 +1128,14 @@ def __init__( self.connection_kwargs = connection_kwargs self.max_connections = max_connections +<<<<<<< HEAD self._available_connections: ConnectionsIndexer = ( ConnectionsIndexer() if index_available_connections else [] ) +======= + self._available_connections = self.reset_available_connections() +>>>>>>> bafbc03 (polymorphism for reset available connections instead) self._in_use_connections: Set[AbstractConnection] = set() - self._index_available_connections = index_available_connections self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) def __repr__(self): @@ -1143,11 +1145,12 @@ def __repr__(self): ) def reset(self): - self._available_connections: ConnectionsIndexer | list = ( - ConnectionsIndexer() if self._index_available_connections else [] - ) + self._available_connections = self.reset_available_connections() self._in_use_connections = weakref.WeakSet() + def reset_available_connections(self): + return [] + def can_get_connection(self) -> bool: """Return True if a connection can be retrieved from the pool.""" return ( diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 9bac55167a..2e2a0165bc 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -21,6 +21,7 @@ ) from redis.commands import AsyncSentinelCommands from redis.exceptions import ConnectionError, ReadOnlyError, ResponseError, TimeoutError +from redis.sentinel import ConnectionsIndexer from redis.utils import str_if_bytes @@ -163,6 +164,9 @@ def reset(self): self.master_address = None self.slave_rr_counter = None + def reset_available_connections(self): + return ConnectionsIndexer() + def owns_connection(self, connection: Connection): check = not self.is_master or ( self.is_master and self.master_address == (connection.host, connection.port) diff --git a/redis/connection.py b/redis/connection.py index 26a3b5b3dd..4d6c2e9e4b 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -6,11 +6,10 @@ import threading import weakref from abc import abstractmethod -from collections import defaultdict from itertools import chain from queue import Empty, Full, LifoQueue from time import time -from typing import Any, Callable, Iterable, List, Optional, Sequence, Type, Union +from typing import Any, Callable, List, Optional, Sequence, Type, Union from urllib.parse import parse_qs, unquote, urlparse from ._cache import ( @@ -735,57 +734,6 @@ def _host_error(self): return f"{self.host}:{self.port}" -class ConnectionsIndexer(Iterable): - """ - Data structure that simulates a list of available connections. - Instead of list, we keep 2 additional DS to support O(1) operations - on all of the class' methods. - The first DS is indexed on the connection object's ID. - The second DS is indexed on the address (ip and port) of the connection. - """ - - def __init__(self): - # Map the id to the connection object - self._id_to_connection = {} - # Map the address to a dictionary of connections - # The inner dictionary is a map between the object id to the object itself - # Both of these DS support O(1) operations on all of the class' methods - self._address_to_connections = defaultdict(dict) - - def pop(self): - try: - _, connection = self._id_to_connection.popitem() - del self._address_to_connections[(connection.host, connection.port)][ - id(connection) - ] - except KeyError: - # We are simulating a list, hence we raise IndexError - # when there's no item in the dictionary - raise IndexError() - return connection - - def append(self, connection: Connection): - self._id_to_connection[id(connection)] = connection - self._address_to_connections[(connection.host, connection.port)][ - id(connection) - ] = connection - - def get_connection(self, host: str, port: int): - try: - _, connection = self._address_to_connections[(host, port)].popitem() - del self._id_to_connection[id(connection)] - except KeyError: - return None - return connection - - def __iter__(self): - # This is an O(1) operation in python3.7 and later - return iter(self._id_to_connection.values()) - - def __len__(self): - return len(self._id_to_connection) - - class SSLConnection(Connection): """Manages SSL connections to and from the Redis server(s). This class extends the Connection class, adding SSL functionality, and making @@ -1135,7 +1083,6 @@ def __init__( self, connection_class=Connection, max_connections: Optional[int] = None, - index_available_connections: bool = False, **connection_kwargs, ): max_connections = max_connections or 2**31 @@ -1155,7 +1102,6 @@ def __init__( # will notice the first thread already did the work and simply # release the lock. self._fork_lock = threading.Lock() - self._index_available_connections = index_available_connections self.reset() def __repr__(self) -> (str, str): @@ -1177,9 +1123,7 @@ def cleanup(self, **options): def reset(self) -> None: self._lock = threading.Lock() self._created_connections = 0 - self._available_connections: ConnectionsIndexer | list = ( - ConnectionsIndexer() if self._index_available_connections else [] - ) + self._available_connections = self.reset_available_connections() self._in_use_connections = set() # this must be the last operation in this method. while reset() is @@ -1193,6 +1137,9 @@ def reset(self) -> None: # reset() and they will immediately release _fork_lock and continue on. self.pid = os.getpid() + def reset_available_connections(self): + return [] + def _checkpid(self) -> None: # _checkpid() attempts to keep ConnectionPool fork-safe on modern # systems. this is called by all ConnectionPool methods that diff --git a/redis/sentinel.py b/redis/sentinel.py index 5fd6c0151f..e9b66a0286 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -1,6 +1,7 @@ import random import weakref -from typing import Any, Optional +from typing import Any, Iterable, Optional +from collections import defaultdict from redis.client import Redis from redis.commands import SentinelCommands @@ -155,6 +156,57 @@ def rotate_slaves(self): raise SlaveNotFoundError(f"No slave found for {self.service_name!r}") +class ConnectionsIndexer(Iterable): + """ + Data structure that stores available connections in a pool. + Instead of list, we keep 2 additional DS to support O(1) operations + on all of the class' methods. + The first DS is indexed on the connection object's ID. + The second DS is indexed on the address (ip and port) of the connection. + """ + + def __init__(self): + # Map the id to the connection object + self._id_to_connection = {} + # Map the address to a dictionary of connections + # The inner dictionary is a map between the object id to the object itself + # Both of these DS support O(1) operations on all of the class' methods + self._address_to_connections = defaultdict(dict) + + def pop(self): + try: + _, connection = self._id_to_connection.popitem() + del self._address_to_connections[(connection.host, connection.port)][ + id(connection) + ] + except KeyError: + # We are simulating a list, hence we raise IndexError + # when there's no item in the dictionary + raise IndexError() + return connection + + def append(self, connection: Connection): + self._id_to_connection[id(connection)] = connection + self._address_to_connections[(connection.host, connection.port)][ + id(connection) + ] = connection + + def get_connection(self, host: str, port: int): + try: + _, connection = self._address_to_connections[(host, port)].popitem() + del self._id_to_connection[id(connection)] + except KeyError: + return None + return connection + + def __iter__(self): + # This is an O(1) operation in python3.7 and later + return iter(self._id_to_connection.values()) + + def __len__(self): + return len(self._id_to_connection) + + class SentinelConnectionPool(ConnectionPool): """ Sentinel backed connection pool. @@ -181,7 +233,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs): service_name=service_name, sentinel_manager=sentinel_manager, ) - super().__init__(index_available_connections=True, **kwargs) + super().__init__(**kwargs) self.connection_kwargs["connection_pool"] = self.proxy self.service_name = service_name self.sentinel_manager = sentinel_manager @@ -198,6 +250,9 @@ def reset(self): super().reset() self.proxy.reset() + def reset_available_connections(self): + return ConnectionsIndexer() + @property def master_address(self): return self.proxy.master_address diff --git a/tests/test_connection.py b/tests/test_connection.py index 910f4b1586..0a48a60744 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -10,7 +10,6 @@ from redis.backoff import NoBackoff from redis.connection import ( Connection, - ConnectionsIndexer, SSLConnection, UnixDomainSocketConnection, parse_url, @@ -348,17 +347,3 @@ def test_unix_socket_connection_failure(): == "Error 2 connecting to unix:///tmp/a.sock. No such file or directory." ) - -def test_connections_indexer_operations(): - ci = ConnectionsIndexer() - c1 = Connection(host="1", port=2) - ci.append(c1) - assert list(ci) == [c1] - assert ci.pop() == c1 - - c2 = Connection(host="3", port=4) - ci.append(c1) - ci.append(c2) - assert ci.get_connection("3", 4) == c2 - assert not ci.get_connection("5", 6) - assert list(ci) == [c1] diff --git a/tests/test_sentinel.py b/tests/test_sentinel.py index 54b9647098..609f98781b 100644 --- a/tests/test_sentinel.py +++ b/tests/test_sentinel.py @@ -5,6 +5,7 @@ import redis.sentinel from redis import exceptions from redis.sentinel import ( + ConnectionsIndexer, MasterNotFoundError, Sentinel, SentinelConnectionPool, @@ -266,3 +267,18 @@ def mock_disconnect(): assert calls == 1 pool.disconnect() + + +def test_connections_indexer_operations(): + ci = ConnectionsIndexer() + c1 = Connection(host="1", port=2) + ci.append(c1) + assert list(ci) == [c1] + assert ci.pop() == c1 + + c2 = Connection(host="3", port=4) + ci.append(c1) + ci.append(c2) + assert ci.get_connection("3", 4) == c2 + assert not ci.get_connection("5", 6) + assert list(ci) == [c1] From ac20752de179c0b4de57d79be3806b8466e92faa Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 09:12:03 +0000 Subject: [PATCH 41/44] merge --- redis/asyncio/connection.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index a8357cb4e5..51d9d73cb6 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1128,13 +1128,7 @@ def __init__( self.connection_kwargs = connection_kwargs self.max_connections = max_connections -<<<<<<< HEAD - self._available_connections: ConnectionsIndexer = ( - ConnectionsIndexer() if index_available_connections else [] - ) -======= self._available_connections = self.reset_available_connections() ->>>>>>> bafbc03 (polymorphism for reset available connections instead) self._in_use_connections: Set[AbstractConnection] = set() self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder) From 9405695152775c06117cf81805b106184676f23c Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 09:17:38 +0000 Subject: [PATCH 42/44] lint --- redis/sentinel.py | 2 +- tests/test_connection.py | 1 - tests/test_sentinel.py | 1 + 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/redis/sentinel.py b/redis/sentinel.py index e9b66a0286..f46db7bf38 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -1,7 +1,7 @@ import random import weakref -from typing import Any, Iterable, Optional from collections import defaultdict +from typing import Any, Iterable, Optional from redis.client import Redis from redis.commands import SentinelCommands diff --git a/tests/test_connection.py b/tests/test_connection.py index 0a48a60744..69275d58c0 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -346,4 +346,3 @@ def test_unix_socket_connection_failure(): str(e.value) == "Error 2 connecting to unix:///tmp/a.sock. No such file or directory." ) - diff --git a/tests/test_sentinel.py b/tests/test_sentinel.py index 609f98781b..9b5bce0fea 100644 --- a/tests/test_sentinel.py +++ b/tests/test_sentinel.py @@ -4,6 +4,7 @@ import pytest import redis.sentinel from redis import exceptions +from redis.connection import Connection from redis.sentinel import ( ConnectionsIndexer, MasterNotFoundError, From a1ece925371b78180960dfc1ff344a0a37f40ab5 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 18:05:48 +0800 Subject: [PATCH 43/44] fix inline comments + rename --- redis/asyncio/connection.py | 4 +-- redis/asyncio/sentinel.py | 25 ++++++++++--------- redis/connection.py | 4 +-- redis/sentinel.py | 22 ++++++++-------- tests/test_asyncio/test_commands.py | 2 +- .../test_sentinel_managed_connection.py | 9 ++++--- tests/test_commands.py | 2 +- tests/test_sentinel_managed_connection.py | 9 ++++--- 8 files changed, 39 insertions(+), 38 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 51d9d73cb6..35ef22272a 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1059,9 +1059,7 @@ class ConnectionPool: def cleanup(self, **options): """ - Additional cleanup operations that the connection pool might - need to do after a SCAN ITER family command is executed. - + Additional cleanup operations that the connection pool might need to do. See SentinelManagedConnection for an example cleanup operation that might need to be done. """ diff --git a/redis/asyncio/sentinel.py b/redis/asyncio/sentinel.py index 2e2a0165bc..dea7b5f7fe 100644 --- a/redis/asyncio/sentinel.py +++ b/redis/asyncio/sentinel.py @@ -37,10 +37,9 @@ class SentinelManagedConnection(Connection): def __init__(self, **kwargs): self.connection_pool = kwargs.pop("connection_pool") # To be set to True if we want to prevent - # the sentinel managed connection to connect - # to the most relevant sentinel in the pool and just - # connect to the current self.host and self.port - self._is_address_fixed = False + # the connection to connect to the most relevant sentinel + # in the pool and just connect to the current host and port + self._is_address_set = False super().__init__(**kwargs) def __repr__(self): @@ -54,9 +53,13 @@ def __repr__(self): s += host_info return s + ")>" - def fix_address(self, address): + def set_address(self, address): + """ + By setting the address, the connection will just connect + to the current host and port the next time connect is called. + """ self.host, self.port = address - self._is_address_fixed = True + self._is_address_set = True async def connect_to(self, address): self.host, self.port = address @@ -70,15 +73,13 @@ async def _connect_retry(self): if self._reader: return # already connected # If address is fixed, it means that the connection - # is not rotating to the next slave (if the connection pool is in replica mode) - if self._is_address_fixed: + # just connect to the current host and port + if self._is_address_set: await self.connect_to((self.host, self.port)) return await self._connect_to_sentinel() async def _connect_to_sentinel(self): - # If same_server is False, connnect to master in master mode - # and rotate to the next slave in slave mode if self.connection_pool.is_master: await self.connect_to(await self.connection_pool.get_master_address()) else: @@ -246,7 +247,7 @@ async def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and - # port to the one that we want later in the call to ``connect_to_address`` + # port to the one that we want later in the call to ``set_address`` if not connection: connection = self.make_connection() assert connection @@ -261,7 +262,7 @@ async def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - connection.fix_address((server_host, server_port)) + connection.set_address((server_host, server_port)) await self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't diff --git a/redis/connection.py b/redis/connection.py index 4d6c2e9e4b..09f472da79 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -1112,9 +1112,7 @@ def __repr__(self) -> (str, str): def cleanup(self, **options): """ - Additional cleanup operations that the connection pool might - need to do after a SCAN ITER family command is executed. - + Additional cleanup operations that the connection pool might need to do. See SentinelManagedConnection for an example cleanup operation that might need to be done. """ diff --git a/redis/sentinel.py b/redis/sentinel.py index f46db7bf38..768c7c245a 100644 --- a/redis/sentinel.py +++ b/redis/sentinel.py @@ -25,7 +25,7 @@ def __init__(self, **kwargs): # the sentinel managed connection to connect # to the most relevant sentinel in the pool and just # connect to the current self.host and self.port - self._is_address_fixed = False + self._is_address_set = False super().__init__(**kwargs) def __repr__(self): @@ -39,9 +39,13 @@ def __repr__(self): s = s % host_info return s - def fix_address(self, address): + def set_address(self, address): + """ + By setting the address, the connection will just connect + to the current host and port the next time connect is called. + """ self.host, self.port = address - self._is_address_fixed = True + self._is_address_set = True def connect_to(self, address): self.host, self.port = address @@ -54,16 +58,14 @@ def connect_to(self, address): def _connect_retry(self): if self._sock: return # already connected - # If address is fixed, it means that the connection - # is not rotating to the next slave (if the connection pool is not master) - if self._is_address_fixed: + # If address is set, it means that the connection + # will just connect to the current host and port. + if self._is_address_set: self.connect_to((self.host, self.port)) return self._connect_to_sentinel() def _connect_to_sentinel(self): - # If same_server is False, connnect to master in master mode - # and rotate to the next slave in slave mode if self.connection_pool.is_master: self.connect_to(self.connection_pool.get_master_address()) else: @@ -317,7 +319,7 @@ def get_connection( host=server_host, port=server_port ) # If not, make a new dummy connection object, and set its host and port - # to the one that we want later in the call to ``fix_address`` + # to the one that we want later in the call to ``set_address`` if not connection: connection = self.make_connection() assert connection @@ -332,7 +334,7 @@ def get_connection( # connect to the previous replica. # This will connect to the host and port of the replica else: - connection.fix_address((server_host, server_port)) + connection.set_address((server_host, server_port)) self.ensure_connection(connection) except BaseException: # Release the connection back to the pool so that we don't diff --git a/tests/test_asyncio/test_commands.py b/tests/test_asyncio/test_commands.py index 78ca92bf49..52967e26e0 100644 --- a/tests/test_asyncio/test_commands.py +++ b/tests/test_asyncio/test_commands.py @@ -1402,7 +1402,7 @@ async def test_zscan_iter(self, r: redis.Redis): pairs = [k async for k in r.zscan_iter("a", match="a")] assert set(pairs) == {(b"a", 1)} - async def test_scan_iter_family_executes_commands_with_sameiter_req_id(self): + async def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): """Assert that all calls to execute_command receives the iter_req_id kwarg""" import uuid diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index 2ee1247203..b6b4e68054 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -47,9 +47,9 @@ class SentinelManagedConnectionMock(SentinelManagedConnection): async def _connect_to_sentinel(self) -> None: """ This simulates the behavior of _connect_to_sentinel when - :py:class:`~redis.SentinelConnectionPool` - is in replica mode. - It'll call rotate_slaves and connect to the next replica. + :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. + In master mode, it'll connect to the master. + In non-master mode, it'll call rotate_slaves and connect to the next replica. """ if self.connection_pool.is_master: self.host, self.port = ("master", 1) @@ -62,7 +62,8 @@ async def _connect_to_sentinel(self) -> None: async def connect_to(self, address: Tuple[str, int]) -> None: """ - Do nothing, just mock. + Do nothing, just mock so it won't try to make a connection to the + dummy address. """ diff --git a/tests/test_commands.py b/tests/test_commands.py index a2dc98a7c6..c26734bf42 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2235,7 +2235,7 @@ def test_zscan_iter(self, r): pairs = list(r.zscan_iter("a", match="a")) assert set(pairs) == {(b"a", 1)} - def test_scan_iter_family_executes_commands_with_sameiter_req_id(self): + def test_scan_iter_family_executes_commands_with_same_iter_req_id(self): """Assert that all calls to execute_command receives the iter_req_id kwarg""" import uuid diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 382e6a53fd..36807216c2 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -13,9 +13,9 @@ class SentinelManagedConnectionMock(SentinelManagedConnection): def _connect_to_sentinel(self) -> None: """ This simulates the behavior of _connect_to_sentinel when - :py:class:`~redis.SentinelConnectionPool` - is in replica mode. - It'll call rotate_slaves and connect to the next replica. + :py:class:`~redis.SentinelConnectionPool`. + In master mode, it'll connect to the master. + In non-master mode, it'll call rotate_slaves and connect to the next replica. """ if self.connection_pool.is_master: self.host, self.port = ("master", 1) @@ -28,7 +28,8 @@ def _connect_to_sentinel(self) -> None: def connect_to(self, address: Tuple[str, int]) -> None: """ - Do nothing, this is just to mock. + Do nothing, just mock so it won't try to make a connection to the + dummy address. """ pass From eede649a1fa24e89b1090696e9130909101c1b58 Mon Sep 17 00:00:00 2001 From: Agnes Natasya Date: Sat, 20 Jul 2024 18:06:22 +0800 Subject: [PATCH 44/44] lint --- tests/test_asyncio/test_sentinel_managed_connection.py | 2 +- tests/test_sentinel_managed_connection.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncio/test_sentinel_managed_connection.py b/tests/test_asyncio/test_sentinel_managed_connection.py index b6b4e68054..822756b81e 100644 --- a/tests/test_asyncio/test_sentinel_managed_connection.py +++ b/tests/test_asyncio/test_sentinel_managed_connection.py @@ -48,7 +48,7 @@ async def _connect_to_sentinel(self) -> None: """ This simulates the behavior of _connect_to_sentinel when :py:class:`~redis.asyncio.sentinel.SentinelConnectionPool`. - In master mode, it'll connect to the master. + In master mode, it'll connect to the master. In non-master mode, it'll call rotate_slaves and connect to the next replica. """ if self.connection_pool.is_master: diff --git a/tests/test_sentinel_managed_connection.py b/tests/test_sentinel_managed_connection.py index 36807216c2..f59ff3f3c2 100644 --- a/tests/test_sentinel_managed_connection.py +++ b/tests/test_sentinel_managed_connection.py @@ -14,7 +14,7 @@ def _connect_to_sentinel(self) -> None: """ This simulates the behavior of _connect_to_sentinel when :py:class:`~redis.SentinelConnectionPool`. - In master mode, it'll connect to the master. + In master mode, it'll connect to the master. In non-master mode, it'll call rotate_slaves and connect to the next replica. """ if self.connection_pool.is_master: