diff --git a/redisvl/index/index.py b/redisvl/index/index.py index 96f8aee6..002a6e55 100644 --- a/redisvl/index/index.py +++ b/redisvl/index/index.py @@ -799,10 +799,22 @@ def connect(self, redis_url: str | None = None, **kwargs): ModuleNotFoundError: If required Redis modules are not installed. """ self.invalidate_sql_schema_cache() - self.__redis_client = RedisConnectionFactory.get_redis_connection( + client = RedisConnectionFactory.get_redis_connection( redis_url=redis_url, **kwargs ) - self._register_client_finalizer(self.__redis_client) + # Release the client this index owned before, if any. Registering a + # finalizer for the new client detaches the old client's finalizer, so + # without closing it here its connections would leak. + if self._owns_redis_client: + self.disconnect() + with self._lock: + self.__redis_client = client + # This index created the client, so it owns and must close it. The + # index may have been holding a caller-provided client until now + # (constructor injection or set_client), in which case ownership + # has to be taken back or nothing would ever close this one. + self._owns_redis_client = True + self._register_client_finalizer(client) @deprecated_function("set_client", "Pass connection parameters in __init__.") def set_client(self, redis_client: SyncRedisClient, **kwargs): @@ -821,8 +833,19 @@ def set_client(self, redis_client: SyncRedisClient, **kwargs): """ RedisConnectionFactory.validate_sync_redis(redis_client) self.invalidate_sql_schema_cache() - self.__redis_client = redis_client - self._register_client_finalizer(redis_client) + # Release the client this index created for itself, if any. Skipped when + # the current client is not ours, where disconnect() would do nothing + # beyond logging that it is not disconnecting. + if self._owns_redis_client: + self.disconnect() + # Swap the client and its ownership together so no other thread can + # observe the new client while ownership still describes the old one. + with self._lock: + self.__redis_client = redis_client + # The caller owns the client they passed in, so this index must + # never close it. Matches __init__(redis_client=...) semantics. + self._owns_redis_client = False + self._detach_client_finalizer() return self def _check_svs_support(self) -> None: @@ -2023,7 +2046,8 @@ async def connect(self, redis_url: str | None = None, **kwargs): client = await RedisConnectionFactory._get_aredis_connection( redis_url=redis_url, **kwargs ) - await self.set_client(client) + # This index created the client, so it owns and must close it. + await self._swap_client(client, owns=True) @deprecated_function("set_client", "Pass connection parameters in __init__.") async def set_client(self, redis_client: AsyncRedisClient | SyncRedisClient): @@ -2031,12 +2055,37 @@ async def set_client(self, redis_client: AsyncRedisClient | SyncRedisClient): [DEPRECATED] Manually set the Redis client to use with the search index. This method is deprecated; please provide connection parameters in __init__. """ - redis_client = await self._validate_client(redis_client) + # The caller owns the client they passed in, so this index must never + # close it. This matches __init__(redis_client=...) semantics. + return await self._swap_client(redis_client, owns=False) + + async def _swap_client( + self, redis_client: AsyncRedisClient | SyncRedisClient, *, owns: bool + ): + """Replace the active client, releasing the previous one if owned. + + Args: + redis_client: The client to start using. + owns: Whether this index owns the new client and is therefore + responsible for closing it. Callers passing their own client + must use False so it is never closed on their behalf. + """ + validated_client = await self._validate_client(redis_client) self.invalidate_sql_schema_cache() - await self.disconnect() + # Release the client this index created for itself, if any. Skipped when + # the current client is not ours, where disconnect() is a no-op. + if self._owns_redis_client: + await self.disconnect() + # Swap the client and its ownership together. Setting ownership outside + # the lock would leave a window where another coroutine could see the + # new client while ownership still describes the old one, and close a + # caller-provided client on that basis. async with self._lock: - self._redis_client = redis_client - self._register_client_finalizer(redis_client) + self._redis_client = validated_client + self._owns_redis_client = owns + # No-op when owns is False; also detaches any finalizer left over + # from a previously owned client. + self._register_client_finalizer(validated_client) return self async def _get_client(self) -> AsyncRedisClient: diff --git a/tests/integration/test_async_search_index.py b/tests/integration/test_async_search_index.py index 01e30a92..a79eb1a7 100644 --- a/tests/integration/test_async_search_index.py +++ b/tests/integration/test_async_search_index.py @@ -225,9 +225,15 @@ async def test_search_index_set_client(client, redis_url, index_schema): await async_index.set_client(client) assert isinstance(async_index.client, AsyncRedis) - if async_index.client: - await async_index.disconnect() - assert async_index.client is None + # The caller supplied this client, so the index does not own it. + # disconnect() must therefore leave it in place and open, exactly + # as when a client is passed to __init__. The converted async + # wrapper shares the caller's connection pool, so closing it here + # would tear down connections the caller still relies on. + assert async_index._owns_redis_client is False + await async_index.disconnect() + assert async_index.client is not None + assert await async_index.client.ping() is True @pytest.mark.asyncio diff --git a/tests/integration/test_index_client_ownership_integration.py b/tests/integration/test_index_client_ownership_integration.py new file mode 100644 index 00000000..006ece50 --- /dev/null +++ b/tests/integration/test_index_client_ownership_integration.py @@ -0,0 +1,188 @@ +"""Client ownership semantics against a real Redis. + +Verifies with live connections that a client handed to an index via the +deprecated `set_client()` is never closed by that index, while a client the +index creates for itself (via the deprecated `connect()`) still is. + +See tests/unit/test_index_client_ownership.py for the mocked-client coverage. +""" + +import asyncio +import gc +import warnings +import weakref + +import pytest + +from redisvl.index import AsyncSearchIndex, SearchIndex + +fields = [{"name": "tag", "type": "tag"}, {"name": "num", "type": "numeric"}] + + +def collect(): + for _ in range(3): + gc.collect() + + +@pytest.fixture +def schema_dict(redis_test_name): + name = redis_test_name("ownership") + return { + "index": {"name": name, "prefix": name, "storage_type": "hash"}, + "fields": fields, + } + + +def pool_sockets_closed(sync_client) -> bool: + pool = sync_client.connection_pool + conns = list(pool._available_connections) + list(pool._in_use_connections) + return all(getattr(conn, "_sock", None) is None for conn in conns) + + +class TestSetClientLeavesCallerClientOpen: + def test_sync_caller_client_usable_after_index_collected( + self, redis_url, schema_dict, client + ): + # Index owns a client of its own first, then the caller swaps theirs in. + index = SearchIndex.from_dict(schema_dict, redis_url=redis_url) + assert index.exists() is False + index_own_client = index.client + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + index.set_client(client) + + assert index._owns_redis_client is False + # The index's own client was released when it was replaced. + assert pool_sockets_closed(index_own_client) + + ref = weakref.ref(index) + del index + collect() + + assert ref() is None, "index was not collected" + assert client.ping() is True, "caller's client was closed by the index" + + def test_sync_disconnect_leaves_caller_client_open( + self, redis_url, schema_dict, client + ): + index = SearchIndex.from_dict(schema_dict, redis_url=redis_url) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + index.set_client(client) + + index.disconnect() + + assert client.ping() is True, "disconnect() closed the caller's client" + + async def test_async_caller_client_usable_after_index_collected( + self, redis_url, schema_dict, async_client + ): + index = AsyncSearchIndex.from_dict(schema_dict, redis_url=redis_url) + assert await index.exists() is False + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + await index.set_client(async_client) + + assert index._owns_redis_client is False + + ref = weakref.ref(index) + del index + collect() + # Give any (incorrectly) scheduled close a chance to run. + await asyncio.sleep(0) + + assert ref() is None, "async index was not collected" + assert ( + await async_client.ping() is True + ), "caller's async client was closed by the index" + + async def test_async_disconnect_leaves_caller_client_open( + self, redis_url, schema_dict, async_client + ): + index = AsyncSearchIndex.from_dict(schema_dict, redis_url=redis_url) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + await index.set_client(async_client) + + await index.disconnect() + + assert ( + await async_client.ping() is True + ), "disconnect() closed the caller's async client" + + +class TestConnectStillOwnsItsClient: + """The deprecated connect() creates the client, so the index must still + close it. This guards against over-correcting the ownership fix.""" + + def test_sync_connect_created_client_closed_on_collection( + self, redis_url, schema_dict + ): + index = SearchIndex.from_dict(schema_dict) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + index.connect(redis_url=redis_url) + + assert index._owns_redis_client is True + created = index.client + assert created is not None + assert created.ping() is True + + del index + collect() + + assert pool_sockets_closed( + created + ), "client created by connect() was not closed on collection" + + def test_async_connect_created_client_closed_on_collection( + self, redis_url, schema_dict + ): + aclose_calls = [] + + async def build(): + index = AsyncSearchIndex.from_dict(schema_dict) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + await index.connect(redis_url=redis_url) + assert index._owns_redis_client is True + created = index.client + assert await created.ping() is True + + original = created.aclose + + async def recording_aclose(): + aclose_calls.append(1) + await original() + + created.aclose = recording_aclose + return index + + index = asyncio.run(build()) + ref = weakref.ref(index) + del index + collect() + + assert ref() is None + assert aclose_calls == [ + 1 + ], "client created by connect() was not closed exactly once" + + +class TestInjectedClientAtConstruction: + """Baseline: __init__ already got this right. Kept so the three entry + points (constructor, set_client, connect) are covered together.""" + + def test_sync_constructor_injected_client_survives(self, schema_dict, client): + from redisvl.schema import IndexSchema + + index = SearchIndex(IndexSchema.from_dict(schema_dict), redis_client=client) + assert index._owns_redis_client is False + assert index.exists() is False + + del index + collect() + + assert client.ping() is True diff --git a/tests/unit/test_index_client_ownership.py b/tests/unit/test_index_client_ownership.py new file mode 100644 index 00000000..6322e03d --- /dev/null +++ b/tests/unit/test_index_client_ownership.py @@ -0,0 +1,377 @@ +"""Ownership semantics for clients handed to an index after construction. + +`__init__` already treats a caller-provided client as not owned, and never +closes it. The deprecated `set_client()` did not follow that rule: an index +created with `redis_url` keeps `_owns_redis_client=True`, so after +`set_client(caller_client)` the index would close a client it never created. +Since 0.25.0 the client finalizer actually fires, which made that observable +as the caller's client being closed when the index is garbage collected. + +The deprecated `connect()` is the mirror case and must keep working: it +creates the client itself, so the index does own it and must still close it. +""" + +import asyncio +import gc +import warnings +import weakref +from unittest import mock + +from redisvl.index import AsyncSearchIndex, SearchIndex +from redisvl.schema import IndexSchema + +SCHEMA_DICT = { + "index": {"name": "ownership-probe", "prefix": "own", "storage_type": "hash"}, + "fields": [{"name": "tag", "type": "tag"}], +} + + +def collect(): + for _ in range(3): + gc.collect() + + +def sync_index_owning_client(created_client=None): + """Index built from a URL, so it owns whatever client it creates.""" + index = SearchIndex.from_dict(SCHEMA_DICT, redis_url="redis://fake:6379") + if created_client is not None: + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=created_client, + ): + assert index._redis_client is created_client + return index + + +def set_client_sync(index, client): + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.validate_sync_redis", + return_value=None, + ): + return index.set_client(client) + + +async def set_client_async(index, client): + with mock.patch.object( + AsyncSearchIndex, "_validate_client", new=mock.AsyncMock(return_value=client) + ): + return await index.set_client(client) + + +class TestCallerProvidedClientIsNotOwned: + def test_sync_set_client_marks_client_unowned(self): + index = sync_index_owning_client() + caller_client = mock.MagicMock(name="caller_client") + set_client_sync(index, caller_client) + assert index._owns_redis_client is False + + def test_sync_set_client_client_survives_gc(self): + index = sync_index_owning_client() + caller_client = mock.MagicMock(name="caller_client") + set_client_sync(index, caller_client) + + ref = weakref.ref(index) + del index + collect() + + assert ref() is None, "index was not collected" + caller_client.close.assert_not_called() + + def test_sync_set_client_client_survives_disconnect(self): + index = sync_index_owning_client() + caller_client = mock.MagicMock(name="caller_client") + set_client_sync(index, caller_client) + + index.disconnect() + + caller_client.close.assert_not_called() + + def test_async_set_client_marks_client_unowned(self): + async def run(): + index = AsyncSearchIndex.from_dict( + SCHEMA_DICT, redis_url="redis://fake:6379" + ) + caller_client = mock.MagicMock(name="caller_async_client") + caller_client.aclose = mock.AsyncMock() + await set_client_async(index, caller_client) + return index + + index = asyncio.run(run()) + assert index._owns_redis_client is False + + def test_async_set_client_client_survives_gc(self): + caller_client = mock.MagicMock(name="caller_async_client") + caller_client.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict( + SCHEMA_DICT, redis_url="redis://fake:6379" + ) + await set_client_async(index, caller_client) + return index + + index = asyncio.run(run()) + ref = weakref.ref(index) + del index + collect() + + assert ref() is None, "index was not collected" + caller_client.aclose.assert_not_awaited() + + def test_async_set_client_client_survives_disconnect(self): + caller_client = mock.MagicMock(name="caller_async_client") + caller_client.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict( + SCHEMA_DICT, redis_url="redis://fake:6379" + ) + await set_client_async(index, caller_client) + await index.disconnect() + + asyncio.run(run()) + caller_client.aclose.assert_not_awaited() + + +class TestIndexCreatedClientIsStillOwned: + """Regression guard: the deprecated connect() creates its own client, so + the index must keep ownership and still close it.""" + + def test_sync_connect_keeps_ownership_and_closes_on_gc(self): + created = mock.MagicMock(name="connect_created_client") + index = SearchIndex.from_dict(SCHEMA_DICT) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=created, + ): + index.connect(redis_url="redis://fake:6379") + + assert index._owns_redis_client is True + del index + collect() + created.close.assert_called_once() + + def test_async_connect_keeps_ownership_and_closes_on_gc(self): + created = mock.MagicMock(name="aconnect_created_client") + created.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict(SCHEMA_DICT) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory._get_aredis_connection", + new=mock.AsyncMock(return_value=created), + ): + with mock.patch.object( + AsyncSearchIndex, + "_validate_client", + new=mock.AsyncMock(return_value=created), + ): + await index.connect(redis_url="redis://fake:6379") + return index + + index = asyncio.run(run()) + assert ( + index._owns_redis_client is True + ), "connect() created the client, so the index must still own it" + del index + collect() + created.aclose.assert_awaited_once() + + +class TestConnectTakesOwnershipFromAnUnownedState: + """connect() creates the client itself, so the index must own it even when + the index was previously holding a caller-provided (unowned) client. Async + gets this right via _swap_client(owns=True); sync must match, otherwise the + client it just created is never closed.""" + + def test_sync_connect_takes_ownership_over_a_caller_client(self): + caller_client = mock.MagicMock(name="caller_client") + created = mock.MagicMock(name="connect_created_client") + schema = IndexSchema.from_dict(SCHEMA_DICT) + index = SearchIndex(schema, redis_client=caller_client) + assert index._owns_redis_client is False + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=created, + ): + index.connect(redis_url="redis://fake:6379") + + assert ( + index._owns_redis_client is True + ), "connect() created this client, so the index must own it" + del index + collect() + created.close.assert_called_once() + caller_client.close.assert_not_called() + + def test_sync_connect_takes_ownership_after_set_client(self): + caller_client = mock.MagicMock(name="caller_client") + created = mock.MagicMock(name="connect_created_client") + index = sync_index_owning_client() + set_client_sync(index, caller_client) + assert index._owns_redis_client is False + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=created, + ): + index.connect(redis_url="redis://fake:6379") + + assert index._owns_redis_client is True + del index + collect() + created.close.assert_called_once() + caller_client.close.assert_not_called() + + def test_async_connect_takes_ownership_after_set_client(self): + caller_client = mock.MagicMock(name="caller_async_client") + caller_client.aclose = mock.AsyncMock() + created = mock.MagicMock(name="aconnect_created_client") + created.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict( + SCHEMA_DICT, redis_url="redis://fake:6379" + ) + await set_client_async(index, caller_client) + assert index._owns_redis_client is False + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory._get_aredis_connection", + new=mock.AsyncMock(return_value=created), + ): + with mock.patch.object( + AsyncSearchIndex, + "_validate_client", + new=mock.AsyncMock(return_value=created), + ): + await index.connect(redis_url="redis://fake:6379") + return index + + index = asyncio.run(run()) + assert index._owns_redis_client is True + del index + collect() + created.aclose.assert_awaited_once() + caller_client.aclose.assert_not_awaited() + + +class TestConnectReleasesThePreviousOwnedClient: + """connect() replaces the active client. When the index owned the old one, + it must be closed: registering a finalizer for the new client detaches the + old client's finalizer, so nothing else would ever close it.""" + + def test_sync_repeated_connect_closes_the_first_client(self): + first = mock.MagicMock(name="first_client") + second = mock.MagicMock(name="second_client") + index = SearchIndex.from_dict(SCHEMA_DICT) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + side_effect=[first, second], + ): + index.connect(redis_url="redis://fake:6379") + index.connect(redis_url="redis://fake:6380") + + first.close.assert_called_once() + assert index.client is second + + del index + collect() + second.close.assert_called_once() + + def test_sync_connect_closes_a_lazily_created_client(self): + lazy = mock.MagicMock(name="lazy_client") + connected = mock.MagicMock(name="connect_client") + index = SearchIndex.from_dict(SCHEMA_DICT, redis_url="redis://fake:6379") + + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=lazy, + ): + assert index._redis_client is lazy + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory.get_redis_connection", + return_value=connected, + ): + index.connect(redis_url="redis://fake:6380") + + lazy.close.assert_called_once() + + def test_async_repeated_connect_closes_the_first_client(self): + first = mock.MagicMock(name="first_async_client") + first.aclose = mock.AsyncMock() + second = mock.MagicMock(name="second_async_client") + second.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict(SCHEMA_DICT) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with mock.patch( + "redisvl.index.index.RedisConnectionFactory._get_aredis_connection", + new=mock.AsyncMock(side_effect=[first, second]), + ): + with mock.patch.object( + AsyncSearchIndex, + "_validate_client", + new=mock.AsyncMock(side_effect=lambda c: c), + ): + await index.connect(redis_url="redis://fake:6379") + await index.connect(redis_url="redis://fake:6380") + return index + + index = asyncio.run(run()) + first.aclose.assert_awaited_once() + assert index.client is second + + +class TestPreviouslyOwnedClientIsReleased: + """Swapping in a caller's client must not silently abandon a client the + index created for itself.""" + + def test_sync_set_client_closes_previously_owned_client(self): + owned = mock.MagicMock(name="index_owned_client") + index = sync_index_owning_client(created_client=owned) + + caller_client = mock.MagicMock(name="caller_client") + set_client_sync(index, caller_client) + + owned.close.assert_called_once() + caller_client.close.assert_not_called() + + def test_async_set_client_closes_previously_owned_client(self): + owned = mock.MagicMock(name="index_owned_async_client") + owned.aclose = mock.AsyncMock() + caller_client = mock.MagicMock(name="caller_async_client") + caller_client.aclose = mock.AsyncMock() + + async def run(): + index = AsyncSearchIndex.from_dict( + SCHEMA_DICT, redis_url="redis://fake:6379" + ) + with mock.patch( + "redisvl.index.index.RedisConnectionFactory._get_aredis_connection", + new=mock.AsyncMock(return_value=owned), + ): + assert await index._get_client() is owned + await set_client_async(index, caller_client) + + asyncio.run(run()) + owned.aclose.assert_awaited_once() + caller_client.aclose.assert_not_awaited()