diff --git a/pymilvus/client/async_grpc_handler.py b/pymilvus/client/async_grpc_handler.py index e9ff31dd8..fa44867fd 100644 --- a/pymilvus/client/async_grpc_handler.py +++ b/pymilvus/client/async_grpc_handler.py @@ -302,8 +302,24 @@ def get_server_type(self): return get_server_type(self.server_address.split(":")[0]) async def ensure_channel_ready(self, timeout: Optional[float] = None): - try: - if not self._is_channel_ready: + # Fast path: avoid the lock once a real request has already set this, + # which is the common case for every call after the first. + if self._is_channel_ready: + return + + # Without the lock, many coroutines can race in here concurrently + # while _is_channel_ready is still False: each would independently + # await _setup_identifier_interceptor_for_channel(), which appends + # its own interceptor to the shared channel's interceptor chain on + # every call. Enough concurrent callers before the first one sets + # _is_channel_ready stacks enough interceptors to blow Python's + # recursion limit on a later RPC (interceptor dispatch recurses one + # frame per interceptor). _reconnect_lock is the same lock reconnect() + # and close() already use to serialize channel-state mutation. + async with self._reconnect_lock: + if self._is_channel_ready: + return + try: wait_timeout = timeout if timeout is not None else 10 ( self._async_identifier_interceptor, @@ -317,11 +333,11 @@ async def ensure_channel_ready(self, timeout: Optional[float] = None): ) self._is_channel_ready = True - except (grpc.FutureTimeoutError, asyncio.TimeoutError, grpc.RpcError) as e: - raise MilvusException( - code=Status.CONNECT_FAILED, - message=f"Fail connecting to server on {self._address}, illegal connection params or server unavailable", - ) from e + except (grpc.FutureTimeoutError, asyncio.TimeoutError, grpc.RpcError) as e: + raise MilvusException( + code=Status.CONNECT_FAILED, + message=f"Fail connecting to server on {self._address}, illegal connection params or server unavailable", + ) from e async def _register_identifier(self, stub: Any, user: str, timeout: float = 10) -> int: host = socket.gethostname() diff --git a/tests/unit/async_grpc_handler/test_async_init.py b/tests/unit/async_grpc_handler/test_async_init.py index f90cb43bc..8451ecb7b 100644 --- a/tests/unit/async_grpc_handler/test_async_init.py +++ b/tests/unit/async_grpc_handler/test_async_init.py @@ -3,6 +3,7 @@ Coverage: Initialization, context manager, secure channel, close operations. """ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import grpc @@ -197,6 +198,39 @@ async def test_ensure_channel_ready_rpc_failure_raises_milvus_exception(self) -> with pytest.raises(MilvusException, match="Fail connecting to server on"): await handler.ensure_channel_ready() + @pytest.mark.asyncio + async def test_ensure_channel_ready_concurrent_calls_only_set_up_once(self) -> None: + """Regression test: concurrent callers racing before the channel is + ready must not each independently set up the identifier interceptor. + + Without a lock, every concurrent caller sees _is_channel_ready as + False and awaits _setup_identifier_interceptor_for_channel(), which + appends a new interceptor to the shared channel's interceptor chain + on every call. Enough concurrent callers before the first one sets + _is_channel_ready stacks enough interceptors to raise a + RecursionError on a later RPC. + """ + mock_channel = _mock_channel() + handler = AsyncGrpcHandler(channel=mock_channel) + handler._is_channel_ready = False + + call_count = 0 + + async def fake_setup(final_channel, stub, user, timeout=10): + nonlocal call_count + call_count += 1 + # Yield control so concurrent callers actually interleave here, + # matching the real await point in the unpatched code. + await asyncio.sleep(0) + return (MagicMock(), final_channel, stub) + + handler._setup_identifier_interceptor_for_channel = fake_setup + + await asyncio.gather(*(handler.ensure_channel_ready() for _ in range(20))) + + assert call_count == 1 + assert handler._is_channel_ready is True + def test_setup_authorization_interceptor_appends_header(self) -> None: """Authorization setup appends a generated interceptor to the final channel.""" mock_channel = _mock_channel()