diff --git a/pymilvus/__init__.py b/pymilvus/__init__.py index 9e3284e99..a8a9ba44c 100644 --- a/pymilvus/__init__.py +++ b/pymilvus/__init__.py @@ -30,6 +30,7 @@ TopHits, ) from .client.search_result import Hit, Hits, SearchResult +from .client.telemetry import TelemetryConfig, new_client_request_id from .client.types import ( BulkInsertState, DataType, @@ -142,6 +143,7 @@ "Shard", "Status", "StructFieldSchema", + "TelemetryConfig", "TopHits", "WeightedRanker", "__version__", @@ -167,6 +169,7 @@ "mkts_from_datetime", "mkts_from_hybridts", "mkts_from_unixtime", + "new_client_request_id", "reset_password", "transfer_node", "transfer_replica", diff --git a/pymilvus/client/async_grpc_handler.py b/pymilvus/client/async_grpc_handler.py index e9ff31dd8..beffc42f5 100644 --- a/pymilvus/client/async_grpc_handler.py +++ b/pymilvus/client/async_grpc_handler.py @@ -2,7 +2,9 @@ import base64 import logging import socket +import threading import time +import weakref from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union from urllib import parse @@ -46,6 +48,11 @@ from .embedding_list import EmbeddingList from .prepare import Prepare from .search_result import SearchResult +from .telemetry import ( + AsyncClientTelemetryManager, + AsyncTelemetryUnaryUnaryInterceptor, + telemetry_operation, +) from .types import ( AnalyzeResult, CompactionState, @@ -89,6 +96,10 @@ def __init__( ) -> None: self._async_stub = None self._async_channel = channel + self._client_telemetry_lock = threading.RLock() + self._client_telemetry_bindings = weakref.WeakKeyDictionary() + self._client_telemetry_generation = 0 + self._client_owned_telemetry = bool(kwargs.pop("_client_owned_telemetry", False)) addr = kwargs.get("address") self._address = addr if addr is not None else self.__get_address(uri, host, port) @@ -96,6 +107,25 @@ def __init__( self._user = kwargs.get("user") self._connect_reserved = kwargs.get("option", {}) self._grpc_options = kwargs.get("grpc_options", {}) + self._current_db_name = kwargs.get("db_name", "") + self._telemetry_stub = None + self._telemetry = None + self._telemetry_interceptor = None + if not self._client_owned_telemetry: + self._telemetry = AsyncClientTelemetryManager( + lambda: self._telemetry_stub, + kwargs.get("telemetry_config"), + user=self._user, + database_provider=lambda: self._current_db_name, + config_provider=lambda: { + "address": self._address, + "username": self._user or "", + "db_name": self._current_db_name, + "secure": self._secure, + }, + runtime_client_id=kwargs.get("_telemetry_client_id", ""), + ) + self._telemetry_interceptor = AsyncTelemetryUnaryUnaryInterceptor(self._telemetry) self._set_authorization(**kwargs) self._reconnect_lock = asyncio.Lock() self._setup_grpc_channel(**kwargs) @@ -104,6 +134,104 @@ def __init__( self.callbacks = [] # Do nothing self._server_info_cache = None + def _rebind_telemetry_stub(self, stub: Any) -> None: + client_lock = getattr(self, "_client_telemetry_lock", None) + if client_lock is None: + self._telemetry_stub = stub + telemetry = getattr(self, "_telemetry", None) + if telemetry is not None: + telemetry.rebind_stub(stub) + return + with client_lock: + self._telemetry_stub = stub + self._client_telemetry_generation += 1 + generation = self._client_telemetry_generation + telemetry = getattr(self, "_telemetry", None) + bindings = list(self._client_telemetry_bindings.items()) + + if telemetry is not None: + self._stabilize_legacy_telemetry(telemetry, stub, generation) + for manager, (binding_token, _database) in bindings: + try: + self._stabilize_client_telemetry(manager, binding_token, stub, generation) + except BaseException: + logger.warning("failed to rebind logical async client telemetry", exc_info=True) + + def _stabilize_legacy_telemetry( + self, telemetry: AsyncClientTelemetryManager, stub: Any, generation: int + ) -> None: + while True: + telemetry.rebind_stub(stub) + with self._client_telemetry_lock: + if telemetry is not self._telemetry: + return + if generation == self._client_telemetry_generation: + return + stub = self._telemetry_stub + generation = self._client_telemetry_generation + + def _stabilize_client_telemetry( + self, + manager: AsyncClientTelemetryManager, + binding_token: Any, + stub: Any, + generation: int, + ) -> None: + while True: + manager.rebind_transport(stub, binding_token) + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + if current is None or current[0] is not binding_token: + return + if generation == self._client_telemetry_generation: + return + stub = self._telemetry_stub + generation = self._client_telemetry_generation + + def register_client_telemetry( + self, manager: AsyncClientTelemetryManager, database: str = "" + ) -> None: + """Attach one logical client's manager to this authenticated transport.""" + + with self._client_telemetry_lock: + existing = self._client_telemetry_bindings.get(manager) + if existing is not None and existing[1] == database: + return + binding_token = existing[0] if existing is not None else object() + self._client_telemetry_bindings[manager] = (binding_token, database) + stub = self._telemetry_stub + generation = self._client_telemetry_generation + try: + manager.bind_transport(stub, database, binding_token) + except BaseException: + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + if current is not None and current[0] is binding_token: + self._client_telemetry_bindings.pop(manager, None) + manager.unbind_transport(binding_token) + raise + + while True: + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + detached = current is None or current[0] is not binding_token + stable = generation == self._client_telemetry_generation + if not stable: + stub = self._telemetry_stub + generation = self._client_telemetry_generation + if detached: + manager.unbind_transport(binding_token) + return + if stable: + return + manager.rebind_transport(stub, binding_token) + + def unregister_client_telemetry(self, manager: AsyncClientTelemetryManager) -> None: + with self._client_telemetry_lock: + binding = self._client_telemetry_bindings.pop(manager, None) + if binding is not None: + manager.unbind_transport(binding[0]) + def __get_address(self, uri: str, host: str, port: str) -> str: if host != "" and port != "" and is_legal_host(host) and is_legal_port(port): return f"{host}:{port}" @@ -135,6 +263,9 @@ def __exit__(self: object, exc_type: object, exc_val: object, exc_tb: object): async def close(self): async with self._reconnect_lock: + if self._telemetry is not None: + await self._telemetry.stop_async() + self._rebind_telemetry_stub(None) if self._async_channel: await self._async_channel.close() self._async_channel = None @@ -184,6 +315,9 @@ async def reconnect(self, address: Optional[str] = None, timeout: float = 10): self._final_channel = new_final_channel self._async_stub = new_stub self._async_identifier_interceptor = new_identifier_interceptor + self._rebind_telemetry_stub( + milvus_pb2_grpc.ClientTelemetryServiceStub(new_final_channel) + ) self._is_channel_ready = True if old_channel: @@ -285,6 +419,8 @@ def _build_stub(self, channel: Any, **kwargs: Any) -> Tuple[Any, Any]: ) final_channel._unary_unary_interceptors.append(async_log_level_interceptor) self._log_level = None + if self._telemetry_interceptor is not None: + final_channel._unary_unary_interceptors.append(self._telemetry_interceptor) return final_channel, milvus_pb2_grpc.MilvusServiceStub(final_channel) def _setup_grpc_channel(self, **kwargs): @@ -316,7 +452,12 @@ async def ensure_channel_ready(self, timeout: Optional[float] = None): timeout=wait_timeout, ) + self._rebind_telemetry_stub( + milvus_pb2_grpc.ClientTelemetryServiceStub(self._final_channel) + ) self._is_channel_ready = True + if self._telemetry is not None: + self._telemetry.start() except (grpc.FutureTimeoutError, asyncio.TimeoutError, grpc.RpcError) as e: raise MilvusException( code=Status.CONNECT_FAILED, @@ -344,6 +485,18 @@ async def _setup_identifier_interceptor_for_channel( stub = milvus_pb2_grpc.MilvusServiceStub(final_channel) return async_identifier_interceptor, final_channel, stub + @property + def telemetry(self) -> Optional[AsyncClientTelemetryManager]: + """Return the legacy direct-handler telemetry manager, if enabled.""" + + return self._telemetry + + @property + def telemetry_stub(self) -> Any: + """Return the current authenticated heartbeat transport.""" + + return self._telemetry_stub + @retry_on_rpc_failure() async def create_collection( self, @@ -733,6 +886,7 @@ async def release_collection( ) check_status(response) + @telemetry_operation("Insert") @retry_on_rpc_failure() @retry_on_schema_mismatch() async def insert_rows( @@ -810,6 +964,7 @@ async def get_persistent_segment_infos( check_status(response.status) return response.infos + @telemetry_operation("Delete") async def delete( self, collection_name: str, @@ -880,6 +1035,7 @@ async def _prepare_batch_upsert_request( ) ) + @telemetry_operation("Upsert") @retry_on_rpc_failure() async def upsert( self, @@ -942,6 +1098,7 @@ async def _prepare_row_upsert_request( field_ops=field_ops, ) + @telemetry_operation("Upsert") @retry_on_rpc_failure() @retry_on_schema_mismatch() async def upsert_rows( @@ -1004,6 +1161,7 @@ async def _execute_hybrid_search( round_decimal = kwargs.get("round_decimal", -1) return SearchResult(response.results, round_decimal, status=response.status) + @telemetry_operation("Search") @retry_on_rpc_failure() async def search( self, @@ -1084,6 +1242,7 @@ async def search( request, timeout, context=context, round_decimal=round_decimal, **kwargs ) + @telemetry_operation("HybridSearch") @retry_on_rpc_failure() async def hybrid_search( self, @@ -1433,6 +1592,7 @@ async def list_partitions( check_status(response.status) return list(response.partition_names) + @telemetry_operation("Query") @retry_on_rpc_failure() async def query( self, @@ -1952,9 +2112,10 @@ async def list_aliases( def reset_db_name(self, db_name: str): """Deprecated: db_name is now passed per-request via kwargs. - This method is kept for backward compatibility but does nothing. + This method is kept for backward compatibility and updates telemetry identity. Use AsyncMilvusClient.use_database() instead. """ + self._current_db_name = db_name @retry_on_rpc_failure() async def create_database( @@ -2665,6 +2826,7 @@ async def get_compaction_state( response.completedPlanNo, ) + @telemetry_operation("RunAnalyzer") @retry_on_rpc_failure() async def run_analyzer( self, diff --git a/pymilvus/client/asynch.py b/pymilvus/client/asynch.py index 485c88fa0..3827dded4 100644 --- a/pymilvus/client/asynch.py +++ b/pymilvus/client/asynch.py @@ -1,6 +1,7 @@ import abc import inspect import threading +from contextlib import suppress from typing import Any, Callable, Optional import grpc @@ -94,11 +95,54 @@ def __init__( self._results = None self._exception = pre_exception self._callback_called = False # callback function should be called only once + self._processed_callback_list = [] + self._processed = False + self._processed_error = None self._kwargs = kwargs def add_callback(self, func: Callable): self._done_cb_list.append(func) + def _add_processed_callback(self, func: Callable[[Optional[BaseException]], None]) -> None: + """Run an internal callback after response parsing finishes. + + Unlike the underlying gRPC future's done callback, this completion point includes + ``on_response()`` and the wrapper's result processing. It is intentionally separate + from user callbacks, whose signature is based on the parsed result. + """ + with self._condition: + if not self._processed: + self._processed_callback_list.append(func) + return + error = self._processed_error + func(error) + + def _mark_processed_locked( + self, error: Optional[BaseException] + ) -> list[Callable[[Optional[BaseException]], None]]: + """Atomically fix the first response-processing outcome. + + The caller must hold ``_condition`` through parsing and user callbacks, so a + concurrent ``result()`` or ``done()`` cannot publish success between a failing + processing step and this state transition. + """ + if self._processed: + return [] + self._processed = True + self._processed_error = error + callbacks = self._processed_callback_list + self._processed_callback_list = [] + return callbacks + + @staticmethod + def _dispatch_processed_callbacks( + callbacks: list[Callable[[Optional[BaseException]], None]], + error: Optional[BaseException], + ) -> None: + for callback in callbacks: + with suppress(BaseException): + callback(error) + def __del__(self) -> None: self._future = None @@ -123,36 +167,48 @@ def _callback(self): self._callback_called = True def result(self, **kwargs): - self.exception() - with self._condition: - # future not finished. wait callback being called. - to = kwargs.get("timeout") - if to is None: - to = self._kwargs.get("timeout", None) - - if self._future and self._results is None: + processed_callbacks = [] + processed_error = None + try: + with self._condition: try: - self._response = self._future.result(timeout=to) - except Exception as e: - raise MilvusException(message=str(e)) from e - if self._response is None: - raise _build_none_response_exception(self._future) - self._results = self.on_response(self._response) - - self._callback() + self.exception() + # future not finished. wait callback being called. + to = kwargs.get("timeout") + if to is None: + to = self._kwargs.get("timeout", None) + + if self._future and self._results is None: + try: + self._response = self._future.result(timeout=to) + except Exception as e: + raise MilvusException(message=str(e)) from e + if self._response is None: + raise _build_none_response_exception(self._future) + self._results = self.on_response(self._response) - self._done = True + self._callback() - self._condition.notify_all() + self._done = True + self._condition.notify_all() - self.exception() - if kwargs.get("raw", False) is True: - # just return response object received from gRPC - return self._response - - if self._results is not None: - return self._results - return self.on_response(self._response) + self.exception() + if kwargs.get("raw", False) is True: + # just return response object received from gRPC + result = self._response + elif self._results is not None: + result = self._results + else: + result = self.on_response(self._response) + except BaseException as exc: + processed_error = exc + processed_callbacks = self._mark_processed_locked(exc) + raise + else: + processed_callbacks = self._mark_processed_locked(None) + return result + finally: + self._dispatch_processed_callbacks(processed_callbacks, processed_error) def cancel(self): with self._condition: @@ -164,21 +220,36 @@ def is_done(self): return self._done def done(self): - with self._condition: - if self._future and self._results is None: + processed_callbacks = [] + processed_error = None + try: + with self._condition: + if self._processed: + self._done = True + self._condition.notify_all() + return try: - self._response = self._future.result() - if self._response is None: - self._exception = _build_none_response_exception(self._future) - else: - self._results = self.on_response(self._response) - self._callback() # https://github.com/milvus-io/milvus/issues/6160 - except Exception as e: - self._exception = e - - self._done = True - - self._condition.notify_all() + if self._future and self._results is None: + try: + self._response = self._future.result() + if self._response is None: + self._exception = _build_none_response_exception(self._future) + else: + self._results = self.on_response(self._response) + self._callback() # https://github.com/milvus-io/milvus/issues/6160 + except Exception as e: + self._exception = e + + self._done = True + self._condition.notify_all() + processed_error = self._exception + processed_callbacks = self._mark_processed_locked(processed_error) + except BaseException as exc: + processed_error = exc + processed_callbacks = self._mark_processed_locked(exc) + raise + finally: + self._dispatch_processed_callbacks(processed_callbacks, processed_error) def exception(self): if self._exception: diff --git a/pymilvus/client/call_context.py b/pymilvus/client/call_context.py index b6e1574c3..c8eb49c0c 100644 --- a/pymilvus/client/call_context.py +++ b/pymilvus/client/call_context.py @@ -1,8 +1,16 @@ -from typing import Optional +from typing import Any, Optional from pymilvus.client.utils import current_time_ms +def is_valid_client_request_id(value: Any) -> bool: + """Return whether value is a non-zero lowercase OpenTelemetry TraceID.""" + + if not isinstance(value, str) or len(value) != 32 or value == "0" * 32: + return False + return all(character in "0123456789abcdef" for character in value) + + def _api_level_md(context: Optional["CallContext"]) -> Optional[list]: if context is None: return None @@ -15,11 +23,16 @@ def __init__(self, db_name: str = "", client_request_id: str = ""): self._client_request_id = client_request_id def to_grpc_metadata(self): - return [ + metadata = [ ("dbname", self._db_name), - ("client-request-id", self._client_request_id), ("client-request-unixmsec", current_time_ms()), ] + # Preserve the legacy access-log contract: callers may attach any nonempty ID to + # the wire. Telemetry correlation validates the stricter OTel trace-ID shape at + # the recording boundary instead of silently dropping an existing header here. + if self._client_request_id: + metadata.append(("client-request-id", self._client_request_id)) + return metadata def get_db_name(self): return self._db_name diff --git a/pymilvus/client/connection_manager.py b/pymilvus/client/connection_manager.py index 46f994612..67df42ad9 100644 --- a/pymilvus/client/connection_manager.py +++ b/pymilvus/client/connection_manager.py @@ -36,6 +36,28 @@ DEFAULT_PORT = 19530 +def _telemetry_connection_key(value: Any) -> str: + """Validate logical-client telemetry without splitting the transport pool. + + Telemetry identity/configuration belongs to MilvusClient, not the pooled handler. + Different client IDs and sampling policies can therefore share one authenticated + transport safely. Retain this helper so invalid configs still fail before a handler + is acquired, and preserve the historical four-part structural key with an empty tail. + """ + from pymilvus.client.telemetry import TelemetryConfig # noqa: PLC0415 + + TelemetryConfig.from_value(value) + return "" + + +def _pooled_handler_kwargs(config: "ConnectionConfig") -> Dict[str, Any]: + """Mark handlers as transport-only for logical MilvusClient pooling.""" + + kwargs = config.get_handler_kwargs() + kwargs["_client_owned_telemetry"] = True + return kwargs + + @dataclass class ConnectionConfig: """Configuration for a Milvus connection. @@ -52,15 +74,18 @@ class ConnectionConfig: token: str = "" db_name: str = "" handler_kwargs: Tuple = () + telemetry_key: str = "" def get_handler_kwargs(self) -> Dict[str, Any]: """Return handler_kwargs as a dict.""" return dict(self.handler_kwargs) @property - def key(self) -> str: - """Return deduplication key: address|token.""" - return f"{self.address}|{self.token}" + def key(self) -> Tuple[str, str, str, str]: + """Return deduplication key for connection-owned settings.""" + # Keep field boundaries structural. Tokens and database names are + # caller-controlled strings and can contain any delimiter we choose. + return (self.address, self.token, self.db_name, self.telemetry_key) @property def is_global(self) -> bool: @@ -120,6 +145,7 @@ def from_uri( token=token or "", db_name=db_name or "", handler_kwargs=tuple(kwargs.items()), + telemetry_key=_telemetry_connection_key(kwargs.get("telemetry_config")), ) # --- Normal URI parsing --- @@ -157,6 +183,11 @@ def from_uri( # For token and db_name, empty string means "use URI value" final_token = token if token else uri_token final_db_name = db_name if db_name else uri_db_name + # A username-only URI is the supported token shorthand + # (https://@host). Do not expose that credential as the telemetry + # user identity. Only user:password credentials contain a distinct user. + if parsed.username and parsed.password is not None and not token and "user" not in kwargs: + kwargs["user"] = parsed.username # Auto-detect secure from https:// scheme if parsed.scheme == "https" and "secure" not in kwargs: @@ -168,6 +199,7 @@ def from_uri( token=final_token, db_name=final_db_name, handler_kwargs=tuple(kwargs.items()), + telemetry_key=_telemetry_connection_key(kwargs.get("telemetry_config")), ) @@ -295,7 +327,7 @@ def create_handler(self, config: ConnectionConfig) -> "GrpcHandler": address=config.address, token=config.token, db_name=config.db_name, - **config.get_handler_kwargs(), + **_pooled_handler_kwargs(config), ) def on_unavailable(self, managed: ManagedConnection) -> bool: @@ -413,7 +445,7 @@ def create_handler(self, config: ConnectionConfig) -> "GrpcHandler": uri=primary.endpoint, token=config.token, db_name=config.db_name, - **config.get_handler_kwargs(), + **_pooled_handler_kwargs(config), ) def close(self, managed: ManagedConnection) -> None: @@ -442,7 +474,7 @@ class ConnectionManager: def __init__(self): self._lock = threading.RLock() - self._registry: Dict[str, ManagedConnection] = {} # key -> ManagedConnection + self._registry: Dict[Tuple[str, str, str, str], ManagedConnection] = {} self._dedicated: Dict[int, ManagedConnection] = {} # handler_id -> ManagedConnection @classmethod @@ -488,13 +520,12 @@ def get_or_create( key = config.key if key in self._registry: managed = self._registry[key] - if client: - managed.add_client(client) - # Health check if idle too long (before touch so idle_time is accurate) if managed.idle_time > IDLE_THRESHOLD_SECONDS and not self._check_health(managed): self._recover(managed) + if client: + managed.add_client(client) managed.touch() return managed.handler @@ -521,16 +552,21 @@ def _create_shared( strategy = self._get_strategy(config) handler = strategy.create_handler(config) self._register_error_callback(handler) - - # Wait for channel ready - handler._wait_for_channel_ready(timeout=timeout) - managed = ManagedConnection( handler=handler, config=config, strategy=strategy, connect_timeout=timeout, ) + try: + # Wait for channel ready before publishing ownership. + handler._wait_for_channel_ready(timeout=timeout) + except Exception: + try: + strategy.close(managed) + except Exception: + logger.warning("Failed to close rejected shared connection", exc_info=True) + raise if client: managed.add_client(client) @@ -547,16 +583,21 @@ def _create_dedicated( strategy = self._get_strategy(config) handler = strategy.create_handler(config) self._register_error_callback(handler) - - # Wait for channel ready - handler._wait_for_channel_ready(timeout=timeout) - managed = ManagedConnection( handler=handler, config=config, strategy=strategy, connect_timeout=timeout, ) + try: + # Wait for channel ready before publishing ownership. + handler._wait_for_channel_ready(timeout=timeout) + except Exception: + try: + strategy.close(managed) + except Exception: + logger.warning("Failed to close rejected dedicated connection", exc_info=True) + raise if client: managed.add_client(client) @@ -610,6 +651,19 @@ def release( managed = self._get_managed(handler) if managed and client: managed.remove_client(client) + if not managed.has_clients: + key = next( + (key for key, value in self._registry.items() if value is managed), + None, + ) + if key is not None: + self._registry.pop(key, None) + try: + managed.strategy.close(managed) + except Exception: + logger.warning( + "Failed to close unused shared connection", exc_info=True + ) def close_all(self) -> None: """Close all connections.""" @@ -719,6 +773,13 @@ def handle_error( if should_recover: with self._lock: + # The last logical client may have released and closed this + # connection while on_unavailable() was doing network I/O + # outside the manager lock. recovery_gen only fences another + # recovery; it does not prove that the manager still owns the + # handler. Never resurrect an already-unpublished transport. + if self._get_managed(handler) is not managed: + return False # Re-verify no other thread has already recovered this # connection while we were outside the lock. if managed.recovery_gen != saved_gen: @@ -736,7 +797,7 @@ def get_stats(self) -> Dict[str, Any]: with self._lock: shared = [ { - "key": key, + "key": str(id(m.handler)), "address": m.config.address, "idle_time": m.idle_time, "client_count": len(m.clients), @@ -779,7 +840,7 @@ def create_handler(self, config: ConnectionConfig) -> "AsyncGrpcHandler": address=config.address, token=config.token, db_name=config.db_name, - **config.get_handler_kwargs(), + **_pooled_handler_kwargs(config), ) def on_unavailable(self, managed: ManagedConnection) -> bool: @@ -826,7 +887,7 @@ def create_handler(self, config: ConnectionConfig) -> "AsyncGrpcHandler": uri=primary.endpoint, token=config.token, db_name=config.db_name, - **config.get_handler_kwargs(), + **_pooled_handler_kwargs(config), ) async def close_async(self, managed: ManagedConnection) -> None: @@ -860,7 +921,7 @@ def __init__(self): # the current thread. Deferring creation to the first async call avoids # RuntimeError when get_instance() is called from sync code. self._lock: Optional[asyncio.Lock] = None - self._registry: Dict[str, ManagedConnection] = {} + self._registry: Dict[Tuple[str, str, str, str], ManagedConnection] = {} self._dedicated: Dict[int, ManagedConnection] = {} def _get_lock(self) -> asyncio.Lock: @@ -918,15 +979,14 @@ async def get_or_create( key = config.key if key in self._registry: managed = self._registry[key] - if client: - managed.add_client(client) - # Health check if idle too long (before touch so idle_time is accurate) if managed.idle_time > IDLE_THRESHOLD_SECONDS and not await self._check_health( managed ): await self._recover(managed) + if client: + managed.add_client(client) managed.touch() return managed.handler @@ -962,15 +1022,23 @@ async def _create_shared( strategy = self._get_strategy(config) handler = strategy.create_handler(config) self._register_error_callback(handler) - - await handler.ensure_channel_ready(timeout=timeout) - managed = ManagedConnection( handler=handler, config=config, strategy=strategy, connect_timeout=timeout, ) + try: + await handler.ensure_channel_ready(timeout=timeout) + except BaseException: + try: + await strategy.close_async(managed) + except Exception: + logger.warning( + "Failed to close rejected async shared connection", + exc_info=True, + ) + raise if client: managed.add_client(client) @@ -987,15 +1055,23 @@ async def _create_dedicated( strategy = self._get_strategy(config) handler = strategy.create_handler(config) self._register_error_callback(handler) - - await handler.ensure_channel_ready(timeout=timeout) - managed = ManagedConnection( handler=handler, config=config, strategy=strategy, connect_timeout=timeout, ) + try: + await handler.ensure_channel_ready(timeout=timeout) + except BaseException: + try: + await strategy.close_async(managed) + except Exception: + logger.warning( + "Failed to close rejected async dedicated connection", + exc_info=True, + ) + raise if client: managed.add_client(client) @@ -1065,6 +1141,20 @@ async def release( managed = self._get_managed(handler) if managed and client: managed.remove_client(client) + if not managed.has_clients: + key = next( + (key for key, value in self._registry.items() if value is managed), + None, + ) + if key is not None: + self._registry.pop(key, None) + try: + await managed.strategy.close_async(managed) + except Exception: + logger.warning( + "Failed to close unused async shared connection", + exc_info=True, + ) async def handle_error( self, diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py index a28910b33..a3e952e96 100644 --- a/pymilvus/client/grpc_handler.py +++ b/pymilvus/client/grpc_handler.py @@ -3,6 +3,7 @@ import socket import threading import time +import weakref from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union from urllib import parse @@ -57,6 +58,11 @@ from .embedding_list import EmbeddingList from .prepare import Prepare from .search_result import SearchResult +from .telemetry import ( + ClientTelemetryManager, + TelemetryUnaryUnaryInterceptor, + telemetry_operation, +) from .types import ( AnalyzeResult, BulkInsertState, @@ -121,23 +127,20 @@ def check_state_and_reconnect_later(self): with self.reconnect_lock: logger.info("reconnect on idle state") self.is_idle_state = False - try: - logger.debug("try disconnecting old connection...") - self.conns.disconnect(self.connection_name) - except Exception: - logger.warning("disconnect failed: {e}") - finally: - reconnected = False - while not reconnected: - try: - logger.debug("try reconnecting...") - self.conns.connect(self.connection_name, **self._kwargs) - reconnected = True - except Exception as e: - logger.warning( - f"reconnect failed: {e}, try again after {check_after_seconds} seconds" - ) - time.sleep(check_after_seconds) + timeout = self._kwargs.get("timeout") + timeout = timeout if isinstance(timeout, (int, float)) else Config.MILVUS_CONN_TIMEOUT + reconnected = False + while not reconnected: + try: + logger.debug("try reconnecting existing handler...") + old_handler = self.conns._fetch_handler(self.connection_name) + old_handler.reconnect(timeout=timeout) + reconnected = True + except Exception as e: + logger.warning( + f"reconnect failed: {e}, try again after {check_after_seconds} seconds" + ) + time.sleep(check_after_seconds) logger.info("reconnected") def reconnect_on_idle(self, state: object): @@ -164,6 +167,10 @@ def __init__( self._stub = None self._channel = channel self._channel_swap_lock = threading.Lock() + self._client_telemetry_lock = threading.RLock() + self._client_telemetry_bindings = weakref.WeakKeyDictionary() + self._client_telemetry_generation = 0 + self._client_owned_telemetry = bool(kwargs.pop("_client_owned_telemetry", False)) addr = kwargs.get("address") self._address = addr if addr is not None else self.__get_address(uri, host, port) @@ -172,11 +179,132 @@ def __init__( self._connect_reserved = kwargs.get("option", {}) self._server_info_cache = None self._grpc_options = kwargs.get("grpc_options", {}) + self._current_db_name = kwargs.get("db_name", "") + self._telemetry_stub = None + self._telemetry = None + self._telemetry_interceptor = None + if not self._client_owned_telemetry: + self._telemetry = ClientTelemetryManager( + lambda: self._telemetry_stub, + kwargs.get("telemetry_config"), + user=self._user, + database_provider=lambda: self._current_db_name, + config_provider=lambda: { + "address": self._address, + "username": self._user or "", + "db_name": self._current_db_name, + "secure": self._secure, + }, + runtime_client_id=kwargs.get("_telemetry_client_id", ""), + ) + self._telemetry_interceptor = TelemetryUnaryUnaryInterceptor(self._telemetry) self._set_authorization(**kwargs) self._setup_grpc_channel() self.callbacks = [] self._reconnect_handler = None + def _rebind_telemetry_stub(self, stub: Any) -> None: + client_lock = getattr(self, "_client_telemetry_lock", None) + if client_lock is None: + self._telemetry_stub = stub + telemetry = getattr(self, "_telemetry", None) + if telemetry is not None: + telemetry.rebind_stub(stub) + return + with client_lock: + self._telemetry_stub = stub + self._client_telemetry_generation += 1 + generation = self._client_telemetry_generation + telemetry = getattr(self, "_telemetry", None) + bindings = list(self._client_telemetry_bindings.items()) + + # Manager callbacks take the manager endpoint lock. Keep them outside the + # handler lock because a custom telemetry command may call client lifecycle + # APIs while holding that endpoint lock. The generation loops make the + # two-phase publication converge without introducing the reverse lock order. + if telemetry is not None: + self._stabilize_legacy_telemetry(telemetry, stub, generation) + for manager, (binding_token, _database) in bindings: + try: + self._stabilize_client_telemetry(manager, binding_token, stub, generation) + except BaseException: + logger.warning("failed to rebind logical client telemetry", exc_info=True) + + def _stabilize_legacy_telemetry( + self, telemetry: ClientTelemetryManager, stub: Any, generation: int + ) -> None: + while True: + telemetry.rebind_stub(stub) + with self._client_telemetry_lock: + if telemetry is not self._telemetry: + return + if generation == self._client_telemetry_generation: + return + stub = self._telemetry_stub + generation = self._client_telemetry_generation + + def _stabilize_client_telemetry( + self, + manager: ClientTelemetryManager, + binding_token: Any, + stub: Any, + generation: int, + ) -> None: + while True: + manager.rebind_transport(stub, binding_token) + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + if current is None or current[0] is not binding_token: + return + if generation == self._client_telemetry_generation: + return + stub = self._telemetry_stub + generation = self._client_telemetry_generation + + def register_client_telemetry( + self, manager: ClientTelemetryManager, database: str = "" + ) -> None: + """Attach one logical client's manager to this authenticated transport.""" + + with self._client_telemetry_lock: + existing = self._client_telemetry_bindings.get(manager) + if existing is not None and existing[1] == database: + return + binding_token = existing[0] if existing is not None else object() + self._client_telemetry_bindings[manager] = (binding_token, database) + stub = self._telemetry_stub + generation = self._client_telemetry_generation + try: + manager.bind_transport(stub, database, binding_token) + except BaseException: + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + if current is not None and current[0] is binding_token: + self._client_telemetry_bindings.pop(manager, None) + manager.unbind_transport(binding_token) + raise + + while True: + with self._client_telemetry_lock: + current = self._client_telemetry_bindings.get(manager) + detached = current is None or current[0] is not binding_token + stable = generation == self._client_telemetry_generation + if not stable: + stub = self._telemetry_stub + generation = self._client_telemetry_generation + if detached: + manager.unbind_transport(binding_token) + return + if stable: + return + manager.rebind_transport(stub, binding_token) + + def unregister_client_telemetry(self, manager: ClientTelemetryManager) -> None: + with self._client_telemetry_lock: + binding = self._client_telemetry_bindings.pop(manager, None) + if binding is not None: + manager.unbind_transport(binding[0]) + def register_reconnect_handler(self, handler: ReconnectHandler): if handler is not None: self._reconnect_handler = handler @@ -270,10 +398,17 @@ def _wait_for_channel_ready( self.close() raise else: + if update_self and self._telemetry is not None: + self._telemetry.start() return target_final_channel, target_stub def close(self): + # Command handlers can re-enter APIs that take _channel_swap_lock. Stop and + # join the telemetry worker before acquiring it to avoid lock inversion. + if self._telemetry is not None: + self._telemetry.stop() with self._channel_swap_lock: + self._rebind_telemetry_stub(None) self.deregister_state_change_callbacks() if self._channel: self._channel.close() @@ -331,15 +466,19 @@ def reconnect(self, address: Optional[str] = None, timeout: float = 10): self._channel = new_channel self._final_channel = new_final_channel self._stub = new_stub + self._rebind_telemetry_stub( + milvus_pb2_grpc.ClientTelemetryServiceStub(new_final_channel) + ) self._address = target_address self._move_state_change_callbacks(old_channel, new_channel) def reset_db_name(self, db_name: str): """Deprecated: db_name is now passed per-request via kwargs. - This method is kept for backward compatibility but does nothing. + This method is kept for backward compatibility and updates telemetry identity. Use MilvusClient.use_database() instead. """ + self._current_db_name = db_name def _setup_authorization_interceptor(self, user: str, password: str, token: str): keys = [] @@ -423,10 +562,14 @@ def _setup_grpc_channel( ) final_channel = grpc.intercept_channel(final_channel, log_level_interceptor) self._log_level = None + if self._telemetry_interceptor is not None: + final_channel = grpc.intercept_channel(final_channel, self._telemetry_interceptor) stub = milvus_pb2_grpc.MilvusServiceStub(final_channel) + telemetry_stub = milvus_pb2_grpc.ClientTelemetryServiceStub(final_channel) if update_self: self._final_channel = final_channel self._stub = stub + self._rebind_telemetry_stub(telemetry_stub) return final_channel, stub def set_onetime_loglevel(self, log_level: str): @@ -459,8 +602,23 @@ def _setup_identifier_interceptor( self._identifier_interceptor = identifier_interceptor self._final_channel = target_final_channel self._stub = target_stub + self._rebind_telemetry_stub( + milvus_pb2_grpc.ClientTelemetryServiceStub(target_final_channel) + ) return target_final_channel, target_stub + @property + def telemetry(self) -> Optional[ClientTelemetryManager]: + """Return the legacy direct-handler telemetry manager, if enabled.""" + + return self._telemetry + + @property + def telemetry_stub(self) -> Any: + """Return the current authenticated heartbeat transport.""" + + return self._telemetry_stub + @property def server_address(self): return self._address @@ -955,6 +1113,7 @@ def get_partition_stats( check_status(status) return response.stats + @telemetry_operation("Insert") @retry_on_rpc_failure() @retry_on_schema_mismatch() def insert_rows( @@ -1081,6 +1240,7 @@ def _prepare_batch_insert_request( else Prepare.batch_insert_param(collection_name, entities, partition_name, fields_info) ) + @telemetry_operation("Insert") @retry_on_rpc_failure() def batch_insert( self, @@ -1127,6 +1287,7 @@ def batch_insert( else: return m + @telemetry_operation("Delete") @retry_on_rpc_failure() def delete( self, @@ -1215,6 +1376,7 @@ def _prepare_batch_upsert_request( ) ) + @telemetry_operation("Upsert") @retry_on_rpc_failure() def upsert( self, @@ -1295,6 +1457,7 @@ def _prepare_row_upsert_request( field_ops=field_ops, ) + @telemetry_operation("Upsert") @retry_on_rpc_failure() @retry_on_schema_mismatch() def upsert_rows( @@ -1387,6 +1550,7 @@ def _execute_hybrid_search( return SearchFuture(None, None, e) raise + @telemetry_operation("Search") @retry_on_rpc_failure() def search( self, @@ -1462,6 +1626,7 @@ def search( request, timeout, round_decimal=round_decimal, context=context, **kwargs ) + @telemetry_operation("HybridSearch") @retry_on_rpc_failure() def hybrid_search( self, @@ -2324,6 +2489,7 @@ def dummy( request = Prepare.dummy_request(request_type) return self._stub.Dummy(request, timeout=timeout, metadata=_api_level_md(context)) + @telemetry_operation("Query") @retry_on_rpc_failure() def query( self, @@ -3377,6 +3543,7 @@ def remove_privileges_from_group( ) check_status(resp) + @telemetry_operation("RunAnalyzer") @retry_on_rpc_failure() def run_analyzer( self, diff --git a/pymilvus/client/iterator/query_iterator.py b/pymilvus/client/iterator/query_iterator.py index c34ff7380..9ec8ce57b 100644 --- a/pymilvus/client/iterator/query_iterator.py +++ b/pymilvus/client/iterator/query_iterator.py @@ -25,6 +25,7 @@ REDUCE_STOP_FOR_BEST, UNLIMITED, ) +from pymilvus.client.telemetry import suppress_telemetry from pymilvus.client.types import DataType from pymilvus.client.utils import mkts_from_datetime from pymilvus.exceptions import MilvusException, ParamError @@ -152,15 +153,16 @@ def seek_offset_by_batch(batch: int, expr: str) -> int: if self._has_element_cursor(): query_params[QUERY_ITER_LAST_PK] = self._next_id query_params[QUERY_ITER_LAST_ELEMENT_OFFSET] = self._next_element_offset - res = self._handler.query( - collection_name=self._collection_name, - expr=expr, - output_fields=[], - partition_names=self._partition_names, - timeout=self._timeout, - context=self._context, - **query_params, - ) + with suppress_telemetry(): + res = self._handler.query( + collection_name=self._collection_name, + expr=expr, + output_fields=[], + partition_names=self._partition_names, + timeout=self._timeout, + context=self._context, + **query_params, + ) self.__update_cursor(res) return len(res) @@ -286,15 +288,16 @@ def __setup_ts_by_request(self): init_ts_kwargs[OFFSET] = 0 init_ts_kwargs[MILVUS_LIMIT] = 1 # just to set up mvccTs for iterator, no need correct limit - res = self._handler.query( - collection_name=self._collection_name, - expr=self._expr, - output_fields=[], - partition_names=[], - timeout=self._timeout, - context=self._context, - **init_ts_kwargs, - ) + with suppress_telemetry(): + res = self._handler.query( + collection_name=self._collection_name, + expr=self._expr, + output_fields=[], + partition_names=[], + timeout=self._timeout, + context=self._context, + **init_ts_kwargs, + ) if res is None: raise MilvusException( message="failed to connect to milvus for setting up " @@ -380,15 +383,16 @@ def next(self): current_expr = self.__setup_next_expr() log.debug(f"query_iterator_next_expr:{current_expr}") query_params = self.__setup_query_params() - res = self._handler.query( - collection_name=self._collection_name, - expr=current_expr, - output_fields=self._output_fields, - partition_names=self._partition_names, - timeout=self._timeout, - context=self._context, - **query_params, - ) + with suppress_telemetry(): + res = self._handler.query( + collection_name=self._collection_name, + expr=current_expr, + output_fields=self._output_fields, + partition_names=self._partition_names, + timeout=self._timeout, + context=self._context, + **query_params, + ) self.__maybe_cache(res) ret = res[0 : min(self._query_options[BATCH_SIZE], len(res))] diff --git a/pymilvus/client/iterator/search_iterator.py b/pymilvus/client/iterator/search_iterator.py index ef5b87622..cdd4e9234 100644 --- a/pymilvus/client/iterator/search_iterator.py +++ b/pymilvus/client/iterator/search_iterator.py @@ -36,6 +36,7 @@ RANGE_FILTER, UNLIMITED, ) +from pymilvus.client.telemetry import suppress_telemetry from pymilvus.client.types import DataType from pymilvus.exceptions import ( ExceptionsMessage, @@ -420,23 +421,24 @@ def __execute_next_search( self, next_params: dict, next_expr: str, to_extend_batch: bool ) -> SearchPage: log.debug(f"search_iterator_next_expr:{next_expr}, next_params:{next_params}") - res = self._handler.search( - collection_name=self._iterator_params["collection_name"], - anns_field=self._iterator_params["ann_field"], - param=next_params, - limit=extend_batch_size( - self._iterator_params[BATCH_SIZE], next_params, to_extend_batch - ), - data=self._iterator_params["data"], - expression=next_expr, - partition_names=self._iterator_params["partition_names"], - output_fields=self._iterator_params["output_fields"], - round_decimal=self._iterator_params["round_decimal"], - timeout=self._iterator_params["timeout"], - schema=self._schema, - context=self._context, - **self._search_options, - ) + with suppress_telemetry(): + res = self._handler.search( + collection_name=self._iterator_params["collection_name"], + anns_field=self._iterator_params["ann_field"], + param=next_params, + limit=extend_batch_size( + self._iterator_params[BATCH_SIZE], next_params, to_extend_batch + ), + data=self._iterator_params["data"], + expression=next_expr, + partition_names=self._iterator_params["partition_names"], + output_fields=self._iterator_params["output_fields"], + round_decimal=self._iterator_params["round_decimal"], + timeout=self._iterator_params["timeout"], + schema=self._schema, + context=self._context, + **self._search_options, + ) return SearchPage(res[0], res.get_session_ts()) # at present, the range_filter parameter means 'larger/less and equal', @@ -566,7 +568,8 @@ def _probe_for_compability(self, params: dict): dummy_batch_size = 1 dummy_params["limit"] = dummy_batch_size dummy_params[ITER_SEARCH_BATCH_SIZE_KEY] = dummy_batch_size - probe_result = self._handler.search(context=self._context, **dummy_params) + with suppress_telemetry(): + probe_result = self._handler.search(context=self._context, **dummy_params) iter_info = probe_result.get_search_iterator_v2_results_info() self._check_token_exists(iter_info.token) # Pin GUARANTEE_TIMESTAMP from probe call's session_ts so that all subsequent @@ -585,7 +588,8 @@ def _probe_for_compability(self, params: dict): # internal next function, do not use this outside of this class def _next(self): - res = self._handler.search(context=self._context, **self._params) + with suppress_telemetry(): + res = self._handler.search(context=self._context, **self._params) iter_info = res.get_search_iterator_v2_results_info() self._check_token_exists(iter_info.token) self._params[ITER_SEARCH_LAST_BOUND_KEY] = iter_info.last_bound diff --git a/pymilvus/client/telemetry.py b/pymilvus/client/telemetry.py new file mode 100644 index 000000000..04a68e496 --- /dev/null +++ b/pymilvus/client/telemetry.py @@ -0,0 +1,1641 @@ +"""Client-side telemetry, heartbeat, and server command support.""" + +from __future__ import annotations + +import asyncio +import contextvars +import functools +import hashlib +import heapq +import inspect +import json +import math +import os +import re +import socket +import threading +import time +import uuid +from array import array +from collections import deque +from contextlib import contextmanager, suppress +from dataclasses import dataclass, field +from typing import Any, Callable, Iterable, Mapping + +import grpc + +from pymilvus.client import __version__ +from pymilvus.client.call_context import is_valid_client_request_id +from pymilvus.grpc_gen import common_pb2, milvus_pb2 + +# Seconds between heartbeats, and therefore the metrics window: each heartbeat carries the +# operations since the last one. The coordinator answers a telemetry query from the window +# before the newest, so what a caller reads is between one and two intervals old. +_DEFAULT_HEARTBEAT_INTERVAL = 10.0 +_HEARTBEAT_RPC_TIMEOUT = 10.0 +# Give the synchronous RPC its full deadline plus a small scheduling margin. If +# custom handler code outlives that bound, retain the live thread handle so a later +# stop can finish joining it instead of losing ownership of the worker. +_HEARTBEAT_STOP_JOIN_TIMEOUT = _HEARTBEAT_RPC_TIMEOUT + 1.0 +_MAX_UNSUPPORTED_BACKOFF = 30 * 60.0 +_MAX_REPLY_PAYLOAD_SIZE = 1024 * 1024 +_MAX_INT64 = 2**63 - 1 +_MAX_HEARTBEAT_INTERVAL = _MAX_INT64 / 1000.0 +# Keep timer conversions within a conservative signed 32-bit millisecond range. +# Longer protocol-valid intervals are represented by cancellable chunks. +_MAX_WAIT_CHUNK_SECONDS = (2**31 - 1) / 1000.0 +# Fixed-point unit for accumulating a fractional sampling rate. A rate becomes an integer +# step of this many units, so the smallest rate that still samples is 1e-9 -- far below +# anything an operator would set, which is the point: a configured rate must never round +# down to "off". +_SAMPLING_SCALE = 1_000_000_000 +_LATENCY_SAMPLE_SIZE = 1000 +_HISTORY_LATENCY_SAMPLE_SIZE = 128 +_SNAPSHOT_HISTORY_TTL_MS = 60 * 60 * 1000 +_SNAPSHOT_HARD_LIMIT = 4096 +_RFC3339_PATTERN = re.compile( + r"^(?P\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})" + r"(?P\.\d+)?(?PZ|[+-]\d{2}:\d{2})$" +) + +# A logical-operation wrapper stays active across validation, schema/cache work, retries, +# result parsing, and (for asyncio) awaits. The gRPC interceptor consults this depth and +# leaves recording to the outer wrapper, preventing one logical call from being counted +# once per transport attempt. +_LOGICAL_TELEMETRY_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( + "pymilvus_logical_telemetry_depth", default=0 +) + +_OPERATION_NAMES = { + "Insert": "Insert", + "Delete": "Delete", + "Upsert": "Upsert", + "Search": "Search", + "HybridSearch": "HybridSearch", + "Query": "Query", + "RunAnalyzer": "RunAnalyzer", +} + +_PUSH_CONFIG_KEYS = { + "enabled", + "heartbeat_interval_ms", + "sampling_rate", +} +_UNSET = object() + + +@contextmanager +def suppress_telemetry(): + """Exclude an internal operation and its transport attempts from public metrics.""" + + depth = _LOGICAL_TELEMETRY_DEPTH.get() + token = _LOGICAL_TELEMETRY_DEPTH.set(depth + 1) + try: + yield + finally: + _LOGICAL_TELEMETRY_DEPTH.reset(token) + + +def _reject_json_constant(value: str) -> None: + msg = f"invalid JSON constant: {value}" + raise ValueError(msg) + + +def _json_object(payload: bytes) -> dict[str, Any]: + if not payload: + return {} + value = json.loads(payload.decode(), parse_constant=_reject_json_constant) + if not isinstance(value, dict): + msg = "command payload must be a JSON object" + raise TypeError(msg) + return value + + +def _optional_bool(payload: Mapping[str, Any], key: str) -> bool | None: + value = payload.get(key) + if value is None: + return None + if not isinstance(value, bool): + msg = f"{key} must be a boolean" + raise TypeError(msg) + return value + + +def _optional_int(payload: Mapping[str, Any], key: str) -> int | None: + value = payload.get(key) + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + msg = f"{key} must be an integer" + raise TypeError(msg) + if value < -(2**63) or value > 2**63 - 1: + msg = f"{key} must fit in a signed 64-bit integer" + raise ValueError(msg) + return value + + +def _optional_number(payload: Mapping[str, Any], key: str) -> float | None: + value = payload.get(key) + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + msg = f"{key} must be a number" + raise TypeError(msg) + result = float(value) + if not math.isfinite(result): + msg = f"{key} must be finite" + raise ValueError(msg) + return result + + +def _optional_string_list(payload: Mapping[str, Any], key: str) -> list[str] | None: + value = payload.get(key) + if value is None: + return None + if not isinstance(value, list) or any(not isinstance(item, str) for item in value): + msg = f"{key} must be an array of strings" + raise TypeError(msg) + return list(value) + + +def _safe_exception_message(exc: BaseException) -> str: + try: + return str(exc) + except BaseException: + try: + return f"{type(exc).__name__} (failed to format exception)" + except BaseException: + return "command handler raised an unprintable exception" + + +def _logical_call_info( + signature: inspect.Signature, args: tuple[Any, ...], kwargs: dict[str, Any] +) -> tuple[str, str]: + try: + arguments = signature.bind_partial(*args, **kwargs).arguments + except TypeError: + arguments = {} + extra = arguments.get("kwargs") + if not isinstance(extra, Mapping): + extra = {} + collection = arguments.get("collection_name") or extra.get("collection_name") or "" + if not collection: + positional = arguments.get("args") + if isinstance(positional, tuple) and positional: + collection = positional[0] + context = arguments.get("context") + request_id = getattr(context, "_client_request_id", "") if context is not None else "" + if not request_id: + request_id = extra.get("client_request_id") or extra.get("client-request-id", "") + return str(collection), request_id if is_valid_client_request_id(request_id) else "" + + +def _record_logical_operation( + owner: Any, + operation: str, + collection: str, + started_at: float, + error: BaseException | None, + request_id: str, +) -> None: + manager = getattr(owner, "_telemetry", None) + if manager is None: + manager = getattr(getattr(owner, "_handler", None), "_telemetry", None) + if manager is None: + parent = getattr(owner, "_parent", None) + manager = getattr(parent, "_telemetry", None) + if manager is None: + manager = getattr(getattr(parent, "_handler", None), "_telemetry", None) + if manager is None: + return + try: + manager.record_operation(operation, collection, started_at, error, request_id) + except BaseException: + # Telemetry is best-effort and must never replace the operation's own result. + return + + +def _defer_sync_future_recording( + result: Any, + owner: Any, + operation: str, + collection: str, + started_at: float, + request_id: str, +) -> bool: + pre_exception = getattr(result, "_exception", None) + if pre_exception is not None: + _record_logical_operation( + owner, operation, collection, started_at, pre_exception, request_id + ) + return True + + add_processed_callback = getattr(result, "_add_processed_callback", None) + if callable(add_processed_callback): + + def processed(error: BaseException | None) -> None: + _record_logical_operation(owner, operation, collection, started_at, error, request_id) + + registered = False + with suppress(BaseException): + add_processed_callback(processed) + registered = True + if registered: + return True + + future = getattr(result, "_future", None) + add_done_callback = getattr(future, "add_done_callback", None) + if not callable(add_done_callback): + return False + + def done(completed: Any) -> None: + error: BaseException | None = None + try: + error = completed.exception() + if error is None: + error = _response_error(completed.result()) + except BaseException as exc: + error = exc + _record_logical_operation(owner, operation, collection, started_at, error, request_id) + + try: + add_done_callback(done) + except BaseException: + return False + return True + + +def telemetry_operation(operation: str) -> Callable[[Callable], Callable]: + """Record one completed public operation around retries and result processing.""" + + def decorate(func: Callable) -> Callable: + signature = inspect.signature(func) + + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + started_at = time.perf_counter() + collection, request_id = _logical_call_info(signature, args, kwargs) + depth = _LOGICAL_TELEMETRY_DEPTH.get() + token = _LOGICAL_TELEMETRY_DEPTH.set(depth + 1) + error: BaseException | None = None + try: + return await func(*args, **kwargs) + except BaseException as exc: + error = exc + raise + finally: + _LOGICAL_TELEMETRY_DEPTH.reset(token) + if depth == 0 and args: + _record_logical_operation( + args[0], operation, collection, started_at, error, request_id + ) + + return async_wrapper + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + started_at = time.perf_counter() + collection, request_id = _logical_call_info(signature, args, kwargs) + depth = _LOGICAL_TELEMETRY_DEPTH.get() + token = _LOGICAL_TELEMETRY_DEPTH.set(depth + 1) + error: BaseException | None = None + deferred = False + try: + result = func(*args, **kwargs) + except BaseException as exc: + error = exc + raise + else: + if depth == 0 and args: + deferred = _defer_sync_future_recording( + result, args[0], operation, collection, started_at, request_id + ) + return result + finally: + _LOGICAL_TELEMETRY_DEPTH.reset(token) + if depth == 0 and args and not deferred: + _record_logical_operation( + args[0], operation, collection, started_at, error, request_id + ) + + return wrapper + + return decorate + + +@dataclass +class TelemetryConfig: + """Configuration for client telemetry. + + ``client_id`` can be pinned to preserve identity across process restarts. When omitted, + a random UUID is generated for this process. + """ + + enabled: bool = True + heartbeat_interval: float = _DEFAULT_HEARTBEAT_INTERVAL + sampling_rate: float = 1.0 + error_max_count: int = 100 + client_id: str = "" + + @classmethod + def from_value(cls, value: Any) -> TelemetryConfig: + if value is None: + return cls() + if isinstance(value, cls): + # Managers apply server-pushed configuration in place. Give every + # manager its own snapshot so a push cannot mutate the caller's + # object or another connection that was initialized from it. + return cls( + enabled=value.enabled, + heartbeat_interval=value.heartbeat_interval, + sampling_rate=value.sampling_rate, + error_max_count=value.error_max_count, + client_id=value.client_id, + ) + if isinstance(value, Mapping): + aliases = { + "Enabled": "enabled", + "HeartbeatInterval": "heartbeat_interval", + "SamplingRate": "sampling_rate", + "ErrorMaxCount": "error_max_count", + "ClientID": "client_id", + } + normalized = {aliases.get(key, key): item for key, item in value.items()} + return cls(**normalized) + msg = "telemetry_config must be a TelemetryConfig or mapping" + raise TypeError(msg) + + def __post_init__(self) -> None: + try: + heartbeat_interval = float(self.heartbeat_interval) + except (TypeError, ValueError, OverflowError) as exc: + msg = "heartbeat_interval must be positive and fit in signed 64-bit milliseconds" + raise ValueError(msg) from exc + if ( + not math.isfinite(heartbeat_interval) + or heartbeat_interval <= 0 + or heartbeat_interval > _MAX_HEARTBEAT_INTERVAL + ): + msg = "heartbeat_interval must be positive and fit in signed 64-bit milliseconds" + raise ValueError(msg) + self.heartbeat_interval = heartbeat_interval + self.sampling_rate = min(1.0, max(0.0, float(self.sampling_rate))) + if self.error_max_count <= 0: + self.error_max_count = 100 + + +@dataclass +class Metrics: + request_count: int = 0 + success_count: int = 0 + error_count: int = 0 + avg_latency_ms: float = 0.0 + p99_latency_ms: float = 0.0 + max_latency_ms: float = 0.0 + + +@dataclass +class OperationMetrics: + operation: str + global_metrics: Metrics + collection_metrics: dict[str, Metrics] = field(default_factory=dict) + # History aggregation needs samples from the complete window, not an average of + # per-window percentiles. Keep a compact, private approximation of the global + # distribution; collection samples and this field never enter heartbeat/history JSON. + _global_latency_samples_us: array = field( + default_factory=lambda: array("q"), repr=False, compare=False + ) + + +@dataclass +class MetricsSnapshot: + timestamp: int + end_time: int + metrics: list[OperationMetrics] + + +@dataclass +class ClientCommand: + command_id: str + command_type: str + payload: bytes = b"" + create_time: int = 0 + persistent: bool = False + target_scope: str = "" + + +@dataclass +class CommandReply: + command_id: str + success: bool + error_message: str = "" + payload: bytes = b"" + + +@dataclass +class ErrorInfo: + timestamp: int + operation: str + error_msg: str + collection: str = "" + request_id: str = "" + + +class _MetricsBucket: + def __init__(self) -> None: + self.request_count = 0 + self.success_count = 0 + self.error_count = 0 + self.total_latency_us = 0 + self.max_latency_us = 0 + self.samples: deque[int] = deque(maxlen=_LATENCY_SAMPLE_SIZE) + + def record(self, latency_us: int, success: bool) -> None: + self.request_count += 1 + self.success_count += int(success) + self.error_count += int(not success) + self.total_latency_us += latency_us + self.max_latency_us = max(self.max_latency_us, latency_us) + self.samples.append(latency_us) + + def snapshot_and_reset( + self, *, retain_history_samples: bool = False + ) -> tuple[Metrics | None, array]: + if self.request_count == 0: + return None, array("q") + samples = sorted(self.samples) + p99 = samples[min(len(samples) - 1, int(len(samples) * 0.99))] if samples else 0 + history_samples = array("q") + if retain_history_samples: + if len(samples) <= _HISTORY_LATENCY_SAMPLE_SIZE: + history_samples = array("q", samples) + else: + # Evenly retain order statistics, including both endpoints. This keeps + # history memory bounded while preserving the whole observed range. + denominator = _HISTORY_LATENCY_SAMPLE_SIZE - 1 + last_index = len(samples) - 1 + history_samples = array( + "q", + ( + samples[(index * last_index + denominator // 2) // denominator] + for index in range(_HISTORY_LATENCY_SAMPLE_SIZE) + ), + ) + result = Metrics( + request_count=self.request_count, + success_count=self.success_count, + error_count=self.error_count, + avg_latency_ms=self.total_latency_us / self.request_count / 1000.0, + p99_latency_ms=p99 / 1000.0, + max_latency_ms=self.max_latency_us / 1000.0, + ) + self.__init__() + return result, history_samples + + +class _OperationCollector: + def __init__(self) -> None: + self.lock = threading.Lock() + self.global_bucket = _MetricsBucket() + self.collections: dict[str, _MetricsBucket] = {} + + def record(self, collection: str, latency_us: int, success: bool) -> None: + with self.lock: + self.global_bucket.record(latency_us, success) + if collection: + self.collections.setdefault(collection, _MetricsBucket()).record( + latency_us, success + ) + + def snapshot_and_reset(self, enabled_collections: set[str] | None) -> OperationMetrics | None: + with self.lock: + global_metrics, global_samples = self.global_bucket.snapshot_and_reset( + retain_history_samples=True + ) + if global_metrics is None: + return None + collection_metrics: dict[str, Metrics] = {} + for name, bucket in self.collections.items(): + metrics, _ = bucket.snapshot_and_reset() + if metrics is not None and ( + enabled_collections is None or name in enabled_collections + ): + collection_metrics[name] = metrics + self.collections = {} + return OperationMetrics("", global_metrics, collection_metrics, global_samples) + + +class ClientTelemetryManager: + """Collects metrics and exchanges commands over ``ClientHeartbeat``.""" + + def __init__( + self, + stub_provider: Callable[[], Any], + config: Any = None, + *, + user: str = "", + database_provider: Callable[[], str] | None = None, + config_provider: Callable[[], Mapping[str, Any]] | None = None, + owner_alive_provider: Callable[[], bool] | None = None, + runtime_client_id: str = "", + ) -> None: + self._stub_provider = stub_provider + # Rebinding a handler to another endpoint must fence the old endpoint's in-flight + # heartbeat. Keep the stub and its generation under one lock so an old response + # cannot clear replies/backoff or execute commands after the new endpoint commits. + self._endpoint_lock = threading.RLock() + self._bound_stub: Any = None + self._endpoint_bound = False + self._endpoint_generation = 0 + self._bound_database = "" + self._database_bound = False + self._transport_binding_token: Any = None + self._config = TelemetryConfig.from_value(config) + self._config_lock = threading.RLock() + self._user = user or "" + self._database_provider = database_provider or (lambda: "") + self._config_provider = config_provider or (dict) + self._owner_alive_provider = owner_alive_provider or (lambda: True) + self._client_id = self._config.client_id or runtime_client_id or str(uuid.uuid4()) + self._client_id_stable = bool(self._config.client_id) + + self._collectors: dict[str, _OperationCollector] = {} + self._collectors_lock = threading.Lock() + self._enabled_collections: set[str] = set() + self._all_collections_enabled = False + self._collections_lock = threading.RLock() + self._errors: deque[ErrorInfo] = deque(maxlen=self._config.error_max_count) + self._errors_lock = threading.Lock() + self._snapshots: deque[MetricsSnapshot] = deque() + self._snapshots_lock = threading.Lock() + + self._handlers: dict[str, Callable[[ClientCommand], CommandReply]] = {} + self._handlers_lock = threading.RLock() + self._command_batch_lock = threading.RLock() + self._pending_replies: list[common_pb2.CommandReply] = [] + self._pending_lock = threading.Lock() + self._executed_commands: dict[str, int] = {} + self._executed_lock = threading.Lock() + self._last_command_timestamp = 0 + self._config_hash = "" + # Carries the fractional sampling rate between calls, in _SAMPLING_SCALE units: + # each operation adds the rate and the one that pushes it past a whole unit is the + # one sampled. See _should_sample. + self._sampling_accum = 0 + self._state_lock = threading.RLock() + + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._ready = False + self._unsupported_streak = 0 + self._last_heartbeat_error: BaseException | None = None + self._last_snapshot_end = 0 + self._register_default_handlers() + + @property + def client_id(self) -> str: + return self._client_id + + @property + def config_hash(self) -> str: + with self._state_lock: + return self._config_hash + + @property + def last_command_timestamp(self) -> int: + with self._state_lock: + return self._last_command_timestamp + + @property + def ready(self) -> bool: + return self._ready + + def is_supported(self) -> bool: + with self._state_lock: + return self._unsupported_streak == 0 + + def last_heartbeat_error(self) -> BaseException | None: + with self._state_lock: + return self._last_heartbeat_error + + def start(self) -> None: + if self._thread is not None or self._ready: + return + self._ready = True + if not self._enabled(): + return + self._thread = threading.Thread( + target=self._heartbeat_loop, + name=f"pymilvus-telemetry-{self._client_id[:8]}", + daemon=True, + ) + self._thread.start() + + def stop(self) -> None: + self._stop_event.set() + thread = self._thread + if thread is None or thread is threading.current_thread(): + return + # Wait through the heartbeat RPC deadline. A custom handler can run longer; + # in that case retain the handle so a later stop can finish joining it. + thread.join(timeout=_HEARTBEAT_STOP_JOIN_TIMEOUT) + if not thread.is_alive() and self._thread is thread: + self._thread = None + + def register_command_handler( + self, command_type: str, handler: Callable[[ClientCommand], CommandReply] + ) -> None: + with self._handlers_lock: + self._handlers[command_type] = handler + + def rebind_stub(self, stub: Any, *, database: Any = _UNSET) -> None: + """Atomically bind heartbeat transport/identity and fence the old endpoint.""" + + with self._endpoint_lock: + self._transport_binding_token = None + self._bound_stub = stub + self._endpoint_bound = True + if database is not _UNSET: + self._bound_database = str(database or "") + self._database_bound = True + self._endpoint_generation += 1 + + def bind_transport(self, stub: Any, database: str, binding_token: Any) -> None: + """Bind this logical manager to a pooled handler using an opaque lease token.""" + + with self._endpoint_lock: + self._transport_binding_token = binding_token + self._bound_stub = stub + self._endpoint_bound = True + self._bound_database = database or "" + self._database_bound = True + self._endpoint_generation += 1 + + def rebind_transport(self, stub: Any, binding_token: Any) -> bool: + """Rebind a pooled transport only while its lease is still current.""" + + with self._endpoint_lock: + if binding_token is not self._transport_binding_token: + return False + self._bound_stub = stub + self._endpoint_bound = True + self._endpoint_generation += 1 + return True + + def unbind_transport(self, binding_token: Any) -> bool: + """Fence a detached transport without clearing a newer handler binding.""" + + with self._endpoint_lock: + if binding_token is not self._transport_binding_token: + return False + self._transport_binding_token = None + self._bound_stub = None + self._endpoint_bound = True + self._endpoint_generation += 1 + return True + + def _heartbeat_endpoint(self) -> tuple[Any, int, str]: + with self._endpoint_lock: + stub = self._bound_stub if self._endpoint_bound else self._stub_provider() + database = ( + self._bound_database if self._database_bound else self._database_provider() or "" + ) + return stub, self._endpoint_generation, database + + def record_operation( + self, + operation: str, + collection: str, + started_at: float, + error: BaseException | None = None, + request_id: str = "", + ) -> None: + with self._config_lock: + enabled = self._config.enabled + sampling_rate = self._config.sampling_rate + if not enabled or not self._should_sample(sampling_rate): + return + + latency_us = max(0, int((time.perf_counter() - started_at) * 1_000_000)) + with self._collections_lock: + collection_enabled = ( + self._all_collections_enabled or collection in self._enabled_collections + ) + collection_key = collection if collection_enabled else "" + + with self._collectors_lock: + collector = self._collectors.setdefault(operation, _OperationCollector()) + collector.record(collection_key, latency_us, error is None) + + if error is not None: + with self._errors_lock: + self._errors.append( + ErrorInfo( + timestamp=int(time.time() * 1000), + operation=operation, + error_msg=str(error), + collection=collection, + request_id=request_id, + ) + ) + + def get_recent_errors(self, max_count: int = 100) -> list[ErrorInfo]: + with self._errors_lock: + return list(reversed(self._errors))[:max_count] + + def get_metrics_snapshots(self) -> list[MetricsSnapshot]: + with self._snapshots_lock: + self._prune_snapshots_locked(int(time.time() * 1000)) + return list(self._snapshots) + + def process_commands( + self, commands: Iterable[Any], *, expected_generation: int | None = None + ) -> None: + commands = list(commands) + with self._command_batch_lock: + self._process_commands_locked(commands, expected_generation) + + def _generation_matches(self, expected_generation: int | None) -> bool: + if expected_generation is None: + return True + with self._endpoint_lock: + return expected_generation == self._endpoint_generation + + def _process_commands_locked( + self, commands: list[Any], expected_generation: int | None + ) -> None: + if not self._generation_matches(expected_generation): + return + with self._state_lock: + last_timestamp = self._last_command_timestamp + max_timestamp = last_timestamp + has_persistent = False + + for command in commands: + if not self._generation_matches(expected_generation): + return + local = ClientCommand( + command_id=command.command_id, + command_type=command.command_type, + payload=bytes(command.payload), + create_time=command.create_time, + persistent=command.persistent, + target_scope=command.target_scope, + ) + has_persistent = has_persistent or local.persistent + if not local.persistent: + max_timestamp = max(max_timestamp, local.create_time) + if local.create_time < last_timestamp: + self._queue_reply(CommandReply(local.command_id, True)) + if not self._generation_matches(expected_generation): + return + continue + with self._executed_lock: + already_executed = local.command_id in self._executed_commands + if already_executed: + self._queue_reply(CommandReply(local.command_id, True)) + if not self._generation_matches(expected_generation): + return + continue + + reply = self._handle_command(local) + if not local.persistent: + with self._executed_lock: + self._executed_commands[local.command_id] = local.create_time + if reply is not None: + self._queue_reply(reply) + + # Preserve the completed command's dedup state and ACK, but stop the old + # endpoint batch before applying any remaining commands or its cursor/hash. + if not self._generation_matches(expected_generation): + return + + def commit_batch() -> None: + with self._executed_lock: + # Timestamp filtering only rejects commands older than the cursor. + # Keep IDs at the new cursor timestamp so equal-timestamp + # redeliveries remain idempotent on every retry. + self._executed_commands = { + command_id: timestamp + for command_id, timestamp in self._executed_commands.items() + if timestamp >= max_timestamp + } + with self._state_lock: + # An empty persistent subset is not an authoritative empty snapshot: + # the current server also omits configs when our non-empty hash already + # matches. Keep the last accepted hash until the response protocol can + # explicitly distinguish those states, otherwise it oscillates every + # heartbeat and repeatedly reapplies the same persistent configs. + if has_persistent: + self._config_hash = self.calculate_config_hash(commands) + self._last_command_timestamp = max(self._last_command_timestamp, max_timestamp) + + if expected_generation is None: + commit_batch() + return + # Commit only while the response still belongs to the current endpoint. No user + # code runs in this short critical section. + with self._endpoint_lock: + if expected_generation != self._endpoint_generation: + return + commit_batch() + + @staticmethod + def calculate_config_hash(commands: Iterable[Any]) -> str: + persistent = sorted( + (command for command in commands if command.persistent), + key=lambda command: command.command_id, + ) + if not persistent: + return "" + digest = hashlib.sha256() + for command in persistent: + digest.update(command.command_id.encode()) + digest.update(command.command_type.encode()) + digest.update(bytes(command.payload)) + return digest.hexdigest()[:16] + + def _enabled(self) -> bool: + with self._config_lock: + return self._config.enabled + + def _heartbeat_interval(self) -> float: + with self._config_lock: + interval = self._config.heartbeat_interval + if not math.isfinite(interval) or interval <= 0 or interval > _MAX_HEARTBEAT_INTERVAL: + return _DEFAULT_HEARTBEAT_INTERVAL + return interval + + def _heartbeat_interval_ms(self) -> int: + # The max int64 millisecond value rounds up when represented as float seconds. + return min(_MAX_INT64, int(self._heartbeat_interval() * 1000)) + + def _wait_for_stop(self, delay: float) -> bool: + while delay > _MAX_WAIT_CHUNK_SECONDS: + if self._stop_event.wait(_MAX_WAIT_CHUNK_SECONDS): + return True + delay -= _MAX_WAIT_CHUNK_SECONDS + return self._stop_event.wait(max(0.0, delay)) + + def _next_heartbeat_delay(self) -> float: + interval = self._heartbeat_interval() + with self._state_lock: + unsupported_streak = self._unsupported_streak + if unsupported_streak <= 0: + return interval + if interval >= _MAX_UNSUPPORTED_BACKOFF: + return interval + delay = interval + for _ in range(unsupported_streak): + if delay >= _MAX_UNSUPPORTED_BACKOFF / 2: + return _MAX_UNSUPPORTED_BACKOFF + delay *= 2 + return delay + + def _owner_alive(self) -> bool: + try: + return bool(self._owner_alive_provider()) + except BaseException: + return False + + def owner_released(self) -> None: + """Wake the worker promptly when the logical client is garbage-collected.""" + + self._stop_event.set() + + def _heartbeat_loop(self) -> None: + while self._owner_alive() and not self._stop_event.is_set(): + try: + self._create_snapshot() + self._send_heartbeat() + except BaseException as exc: + with self._state_lock: + self._last_heartbeat_error = exc + if self._wait_for_stop(self._next_heartbeat_delay()): + return + + def _send_heartbeat(self) -> None: + stub, endpoint_generation, database = self._heartbeat_endpoint() + if stub is None: + return + + metrics_enabled = self._enabled() + with self._snapshots_lock: + latest = self._snapshots[-1] if metrics_enabled and self._snapshots else None + metrics = self._to_proto_metrics(latest.metrics if latest else []) + with self._pending_lock: + replies = list(self._pending_replies) + with self._state_lock: + config_hash = self._config_hash + last_timestamp = self._last_command_timestamp + + request = milvus_pb2.ClientHeartbeatRequest( + client_info=self._build_client_info(database), + report_timestamp=int(time.time() * 1000), + metrics=metrics, + command_replies=replies, + config_hash=config_hash, + last_command_timestamp=last_timestamp, + ) + try: + response = stub.ClientHeartbeat( + request, timeout=_HEARTBEAT_RPC_TIMEOUT, wait_for_ready=False + ) + except grpc.RpcError as exc: + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + with self._state_lock: + self._last_heartbeat_error = exc + if exc.code() == grpc.StatusCode.UNIMPLEMENTED: + self._unsupported_streak += 1 + return + except BaseException as exc: # best-effort background channel + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + with self._state_lock: + self._last_heartbeat_error = exc + return + + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + # Any server response proves the RPC exists, even when the application status + # is an error. Unsupported backoff is only for transport-level UNIMPLEMENTED. + with self._state_lock: + self._unsupported_streak = 0 + if response.status.code != 0 or response.status.error_code != 0: + self._last_heartbeat_error = RuntimeError( + response.status.reason or "client telemetry heartbeat failed" + ) + return + self._last_heartbeat_error = None + with self._pending_lock: + del self._pending_replies[: len(replies)] + # Custom handlers may reconnect. Run them without endpoint_lock and fence every + # command plus the final cursor/hash commit against this response's generation. + self.process_commands(response.commands, expected_generation=endpoint_generation) + + def _build_client_info(self, database: Any = _UNSET) -> common_pb2.ClientInfo: + reserved = { + "client_id": self._client_id, + "client_id_stable": str(self._client_id_stable).lower(), + } + if database is _UNSET: + with self._endpoint_lock: + database = ( + self._bound_database + if self._database_bound + else self._database_provider() or "" + ) + if database: + reserved["db_name"] = str(database) + return common_pb2.ClientInfo( + sdk_type="Python", + sdk_version=__version__, + local_time=time.ctime(), + user=self._user, + host=socket.gethostname(), + reserved=reserved, + ) + + def _should_sample(self, rate: float) -> bool: + """Decide whether this operation is recorded, spreading the sampled ones evenly. + + Each call adds the rate to an accumulator and samples on the call that carries it + across a whole unit: at 0.25 that is every fourth operation. What matters is that + the ratio holds over any stretch of calls, not only over a long one -- metrics are + reported per heartbeat window, and a window is tens or hundreds of operations. A + scheme that sampled a contiguous run and then dropped one would give the right + long-run ratio while making every individual window either complete or empty. + """ + if rate >= 1.0: + return True + if rate <= 0.0: + return False + # A rate too small to represent still means "sample rarely", never "sample never": + # silently disabling telemetry for a positive rate is the one outcome nobody could + # have intended. + step = max(1, int(rate * _SAMPLING_SCALE)) + with self._state_lock: + before = self._sampling_accum + self._sampling_accum = before + step + return self._sampling_accum // _SAMPLING_SCALE != before // _SAMPLING_SCALE + + def _create_snapshot(self) -> None: + if not self._enabled(): + return + with self._collections_lock: + enabled_collections = ( + None if self._all_collections_enabled else set(self._enabled_collections) + ) + metrics: list[OperationMetrics] = [] + with self._collectors_lock: + collectors = list(self._collectors.items()) + for operation, collector in collectors: + item = collector.snapshot_and_reset(enabled_collections) + if item is not None: + item.operation = operation + metrics.append(item) + + now = int(time.time() * 1000) + start = self._last_snapshot_end + if start == 0 or start > now: + start = now - self._heartbeat_interval_ms() + self._last_snapshot_end = now + with self._snapshots_lock: + self._snapshots.append(MetricsSnapshot(start, now, metrics)) + self._prune_snapshots_locked(now) + + def _prune_snapshots_locked(self, now_ms: int) -> None: + cutoff = now_ms - _SNAPSHOT_HISTORY_TTL_MS + while self._snapshots and self._snapshots[0].end_time < cutoff: + self._snapshots.popleft() + while len(self._snapshots) > _SNAPSHOT_HARD_LIMIT: + self._snapshots.popleft() + + def _to_proto_metrics( + self, items: Iterable[OperationMetrics] + ) -> list[common_pb2.OperationMetrics]: + with self._collections_lock: + all_collections_enabled = self._all_collections_enabled + enabled_collections = set(self._enabled_collections) + result = [] + for item in items: + result.append( + common_pb2.OperationMetrics( + **{ + "operation": item.operation, + "global": ClientTelemetryManager._metrics_proto(item.global_metrics), + "collection_metrics": { + name: ClientTelemetryManager._metrics_proto(metrics) + for name, metrics in item.collection_metrics.items() + if all_collections_enabled or name in enabled_collections + }, + } + ) + ) + return result + + @staticmethod + def _metrics_proto(metrics: Metrics) -> common_pb2.Metrics: + return common_pb2.Metrics( + request_count=metrics.request_count, + success_count=metrics.success_count, + error_count=metrics.error_count, + avg_latency_ms=metrics.avg_latency_ms, + p99_latency_ms=metrics.p99_latency_ms, + max_latency_ms=metrics.max_latency_ms, + ) + + def _queue_reply(self, reply: CommandReply) -> None: + with self._pending_lock: + self._pending_replies.append( + common_pb2.CommandReply( + command_id=reply.command_id, + success=reply.success, + error_message=reply.error_message, + payload=reply.payload, + ) + ) + + def _handle_command(self, command: ClientCommand) -> CommandReply: + with self._handlers_lock: + handler = self._handlers.get(command.command_type) + if handler is None: + return CommandReply( + command.command_id, + False, + error_message=f"unknown command type: {command.command_type}", + ) + try: + reply = handler(command) + if reply is None: + return CommandReply( + command.command_id, + False, + error_message="command handler returned no reply", + ) + return CommandReply( + command.command_id, + reply.success, + error_message=reply.error_message, + payload=reply.payload, + ) + except BaseException as exc: + return CommandReply( + command.command_id, + False, + error_message=_safe_exception_message(exc), + ) + + def _register_default_handlers(self) -> None: + self.register_command_handler("push_config", self._handle_push_config) + self.register_command_handler("collection_metrics", self._handle_collection_metrics) + self.register_command_handler("show_errors", self._handle_show_errors) + self.register_command_handler("show_latency_history", self._handle_latency_history) + self.register_command_handler("get_config", self._handle_get_config) + + @staticmethod + def _payload(command: ClientCommand) -> dict[str, Any]: + return _json_object(command.payload) + + def _handle_push_config(self, command: ClientCommand) -> CommandReply: + payload = self._payload(command) + # Validate every applied field before mutating the current config. This makes the + # command atomic when, for example, a valid enabled flag accompanies an invalid + # interval. Unknown fields (such as the server-side push API's ttl_seconds, which + # never reaches clients by design) are reported as ignored, never validated. + enabled = _optional_bool(payload, "enabled") + interval_ms = _optional_int(payload, "heartbeat_interval_ms") + sampling_rate = _optional_number(payload, "sampling_rate") + if interval_ms is not None and interval_ms <= 0: + msg = "heartbeat_interval_ms must be positive" + raise ValueError(msg) + + applied: list[str] = [] + if enabled is not None: + applied.append("enabled") + if interval_ms is not None: + applied.append("heartbeat_interval_ms") + if sampling_rate is not None: + applied.append("sampling_rate") + ignored = sorted(key for key in payload if key not in _PUSH_CONFIG_KEYS) + + with self._config_lock: + if enabled is not None: + self._config.enabled = enabled + if interval_ms is not None: + self._config.heartbeat_interval = interval_ms / 1000.0 + if sampling_rate is not None: + self._config.sampling_rate = min(1.0, max(0.0, sampling_rate)) + reply_payload: dict[str, Any] = {"applied": applied} + if ignored: + reply_payload["ignored"] = ignored + return CommandReply( + command.command_id, + True, + payload=json.dumps(reply_payload, separators=(",", ":")).encode(), + ) + + def _handle_collection_metrics(self, command: ClientCommand) -> CommandReply: + if not command.payload: + with self._collections_lock: + payload = { + "enabled_collections": sorted(self._enabled_collections), + "all_collections_enabled": self._all_collections_enabled, + } + return CommandReply(command.command_id, True, payload=json.dumps(payload).encode()) + + payload = self._payload(command) + collections = _optional_string_list(payload, "collections") or [] + enabled = _optional_bool(payload, "enabled") or False + _optional_string_list(payload, "metrics_types") + wildcard = "*" in collections + with self._collections_lock: + if enabled: + if not collections: + msg = "collections list cannot be empty when enabled=true" + raise ValueError(msg) + if wildcard: + self._all_collections_enabled = True + else: + self._enabled_collections.update(collections) + elif wildcard or not collections: + self._all_collections_enabled = False + self._enabled_collections.clear() + else: + self._enabled_collections.difference_update(collections) + return CommandReply(command.command_id, True) + + def _handle_show_errors(self, command: ClientCommand) -> CommandReply: + payload = self._payload(command) + max_count = _optional_int(payload, "max_count") + if max_count is None or max_count <= 0: + max_count = 100 + errors = [dict(vars(item)) for item in self.get_recent_errors(max_count)] + if not errors: + return CommandReply(command.command_id, True) + encoded = json.dumps(errors, separators=(",", ":")).encode() + while len(encoded) > _MAX_REPLY_PAYLOAD_SIZE and len(errors) > 1: + errors = errors[: max(1, len(errors) // 2)] + encoded = json.dumps(errors, separators=(",", ":")).encode() + if len(encoded) > _MAX_REPLY_PAYLOAD_SIZE and errors: + original_strings = { + key: value for key, value in errors[0].items() if isinstance(value, str) + } + limits = {key: len(value) for key, value in original_strings.items()} + while len(encoded) > _MAX_REPLY_PAYLOAD_SIZE and any(limits.values()): + key = max(limits, key=limits.get) + limit = limits[key] + limits[key] = limit // 2 if limit > 1 else 0 + errors[0][key] = original_strings[key][: limits[key]] + "...(truncated)" + encoded = json.dumps(errors, separators=(",", ":")).encode() + if len(encoded) > _MAX_REPLY_PAYLOAD_SIZE: + msg = "show_errors response exceeds the 1MB payload limit" + raise ValueError(msg) + return CommandReply(command.command_id, True, payload=encoded) + + def _handle_get_config(self, command: ClientCommand) -> CommandReply: + user_config = dict(self._config_provider()) + for key in ("password", "token", "api_key"): + user_config.pop(key, None) + with self._config_lock: + user_config.update( + telemetry_enabled=self._config.enabled, + telemetry_heartbeat_interval_ms=self._heartbeat_interval_ms(), + telemetry_sampling_rate=self._config.sampling_rate, + ) + with self._collections_lock: + user_config["enabled_collections"] = ( + ["*"] if self._all_collections_enabled else sorted(self._enabled_collections) + ) + user_config["all_collections_enabled"] = self._all_collections_enabled + return CommandReply( + command.command_id, + True, + payload=json.dumps({"user_config": user_config}, default=str).encode(), + ) + + def _handle_latency_history(self, command: ClientCommand) -> CommandReply: + payload = self._payload(command) + start_time = payload.get("start_time") + end_time = payload.get("end_time") + if start_time is not None and not isinstance(start_time, str): + msg = "start_time must be a string" + raise TypeError(msg) + if end_time is not None and not isinstance(end_time, str): + msg = "end_time must be a string" + raise TypeError(msg) + detail = _optional_bool(payload, "detail") or False + start_ms = _parse_rfc3339_ms(start_time) + end_ms = _parse_rfc3339_ms(end_time) + if end_ms < start_ms: + msg = "end_time must be after start_time" + raise ValueError(msg) + if end_ms - start_ms > 60 * 60 * 1000: + msg = "time range cannot exceed 1 hour" + raise ValueError(msg) + snapshots = [ + item + for item in self.get_metrics_snapshots() + if item.end_time >= start_ms and item.timestamp <= end_ms + ] + if detail: + body = { + "snapshots": [_snapshot_dict(item) for item in snapshots], + "total_snapshots": len(snapshots), + } + else: + body = _aggregate_snapshots(snapshots, start_ms, end_ms) + encoded = json.dumps(body, separators=(",", ":")).encode() + if len(encoded) > _MAX_REPLY_PAYLOAD_SIZE: + msg = "response too large, try a smaller time range" + raise ValueError(msg) + return CommandReply(command.command_id, True, payload=encoded) + + +class AsyncClientTelemetryManager(ClientTelemetryManager): + """Asyncio heartbeat variant used by ``AsyncGrpcHandler``.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._task: asyncio.Task | None = None + + def start(self) -> None: + if self._task is not None or self._ready: + return + self._ready = True + if not self._enabled(): + return + self._task = asyncio.get_running_loop().create_task(self._async_heartbeat_loop()) + + async def stop_async(self) -> None: + self._stop_event.set() + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + except BaseException as exc: + with self._state_lock: + self._last_heartbeat_error = exc + self._task = None + + def stop(self) -> None: + if self._task is not None: + self._task.cancel() + + def owner_released(self) -> None: + super().owner_released() + task = self._task + if task is None: + return + try: + loop = task.get_loop() + if not loop.is_closed(): + loop.call_soon_threadsafe(task.cancel) + except BaseException: + # Owner cleanup can run during event-loop or interpreter teardown. + return + + async def _async_heartbeat_loop(self) -> None: + while self._owner_alive(): + try: + self._create_snapshot() + await self._send_heartbeat_async() + except asyncio.CancelledError: + raise + except BaseException as exc: + with self._state_lock: + self._last_heartbeat_error = exc + if self._stop_event.is_set(): + return + await self._sleep_until_next_heartbeat(self._next_heartbeat_delay()) + + async def _sleep_until_next_heartbeat(self, delay: float) -> None: + while delay > 0: + chunk = min(delay, _MAX_WAIT_CHUNK_SECONDS) + await asyncio.sleep(chunk) + delay -= chunk + + async def _send_heartbeat_async(self) -> None: + stub, endpoint_generation, database = self._heartbeat_endpoint() + if stub is None: + return + metrics_enabled = self._enabled() + with self._snapshots_lock: + latest = self._snapshots[-1] if metrics_enabled and self._snapshots else None + with self._pending_lock: + replies = list(self._pending_replies) + with self._state_lock: + config_hash = self._config_hash + last_timestamp = self._last_command_timestamp + request = milvus_pb2.ClientHeartbeatRequest( + client_info=self._build_client_info(database), + report_timestamp=int(time.time() * 1000), + metrics=self._to_proto_metrics(latest.metrics if latest else []), + command_replies=replies, + config_hash=config_hash, + last_command_timestamp=last_timestamp, + ) + try: + response = await stub.ClientHeartbeat( + request, timeout=_HEARTBEAT_RPC_TIMEOUT, wait_for_ready=False + ) + except asyncio.CancelledError: + raise + except grpc.RpcError as exc: + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + with self._state_lock: + self._last_heartbeat_error = exc + if exc.code() == grpc.StatusCode.UNIMPLEMENTED: + self._unsupported_streak += 1 + return + except BaseException as exc: + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + with self._state_lock: + self._last_heartbeat_error = exc + return + with self._endpoint_lock: + if endpoint_generation != self._endpoint_generation: + return + with self._state_lock: + self._unsupported_streak = 0 + if response.status.code != 0 or response.status.error_code != 0: + self._last_heartbeat_error = RuntimeError( + response.status.reason or "client telemetry heartbeat failed" + ) + return + self._last_heartbeat_error = None + with self._pending_lock: + del self._pending_replies[: len(replies)] + self.process_commands(response.commands, expected_generation=endpoint_generation) + + +class TelemetryUnaryUnaryInterceptor(grpc.UnaryUnaryClientInterceptor): + def __init__(self, manager: ClientTelemetryManager) -> None: + self._manager = manager + + def intercept_unary_unary(self, continuation: Callable, call_details: Any, request: Any): + operation = _operation_from_method(call_details.method) + if operation is None or _LOGICAL_TELEMETRY_DEPTH.get() > 0: + return continuation(call_details, request) + started_at = time.perf_counter() + call = continuation(call_details, request) + + def done(completed: Any) -> None: + error = completed.exception() + if error is None: + try: + error = _response_error(completed.result()) + except BaseException as exc: + error = exc + self._manager.record_operation( + operation, + _collection_from_request(request), + started_at, + error, + _request_id_from_metadata(call_details.metadata), + ) + + call.add_done_callback(done) + return call + + +class AsyncTelemetryUnaryUnaryInterceptor(grpc.aio.UnaryUnaryClientInterceptor): + def __init__(self, manager: ClientTelemetryManager) -> None: + self._manager = manager + + async def intercept_unary_unary(self, continuation: Callable, call_details: Any, request: Any): + operation = _operation_from_method(call_details.method) + if operation is None or _LOGICAL_TELEMETRY_DEPTH.get() > 0: + return await continuation(call_details, request) + started_at = time.perf_counter() + call = await continuation(call_details, request) + + def done(completed: Any) -> None: + async def finalize() -> None: + error: BaseException | None = None + try: + response = await completed + error = _response_error(response) + except BaseException as exc: + error = exc + self._manager.record_operation( + operation, + _collection_from_request(request), + started_at, + error, + _request_id_from_metadata(call_details.metadata), + ) + + asyncio.get_running_loop().create_task(finalize()) + + call.add_done_callback(done) + return call + + +def _operation_from_method(method: Any) -> str | None: + name = method.decode() if isinstance(method, bytes) else str(method) + return _OPERATION_NAMES.get(name.rsplit("/", 1)[-1]) + + +def _collection_from_request(request: Any) -> str: + for name in ("collection_name", "collectionName"): + value = getattr(request, name, "") + if value: + return str(value) + return "" + + +def _request_id_from_metadata(metadata: Any) -> str: + for key, value in metadata or (): + if key in ("client_request_id", "client-request-id"): + request_id = value.decode() if isinstance(value, bytes) else str(value) + return request_id if is_valid_client_request_id(request_id) else "" + return "" + + +def _response_error(response: Any) -> BaseException | None: + """Return an error for Milvus failures carried in an otherwise-OK gRPC response.""" + + status = ( + response if isinstance(response, common_pb2.Status) else getattr(response, "status", None) + ) + if status is None: + return None + if int(getattr(status, "error_code", 0)) == 0 and int(getattr(status, "code", 0)) == 0: + return None + return RuntimeError(getattr(status, "reason", "") or "Milvus request failed") + + +def _parse_rfc3339_ms(value: Any) -> int: + if not value: + msg = "payload is required with start_time and end_time" + raise ValueError(msg) + from datetime import datetime # noqa: PLC0415 + + match = _RFC3339_PATTERN.fullmatch(value) if isinstance(value, str) else None + if match is None: + msg = "timestamp must use RFC3339 with seconds and timezone" + raise ValueError(msg) + # Python 3.9 rejects fractional seconds longer than microsecond precision, + # although RFC3339 permits any number of fractional digits. Telemetry stores + # milliseconds, so safely truncate only after the complete syntax is validated. + fraction = (match.group("fraction") or "")[:7] + timezone = "+00:00" if match.group("timezone") == "Z" else match.group("timezone") + normalized = f"{match.group('datetime')}{fraction}{timezone}" + return int(datetime.fromisoformat(normalized).timestamp() * 1000) + + +def _metrics_dict(metrics: Metrics) -> dict[str, Any]: + return { + "request_count": metrics.request_count, + "success_count": metrics.success_count, + "error_count": metrics.error_count, + "avg_latency_ms": metrics.avg_latency_ms, + "p99_latency_ms": metrics.p99_latency_ms, + "max_latency_ms": metrics.max_latency_ms, + } + + +def _snapshot_dict(snapshot: MetricsSnapshot) -> dict[str, Any]: + return { + "timestamp": snapshot.timestamp, + "end_time": snapshot.end_time, + "metrics": { + item.operation: _metrics_dict(item.global_metrics) for item in snapshot.metrics + }, + } + + +def _aggregate_snapshots( + snapshots: Iterable[MetricsSnapshot], start_ms: int, end_ms: int +) -> dict[str, Any]: + snapshots = list(snapshots) + totals: dict[str, dict[str, Any]] = {} + for snapshot in snapshots: + for item in snapshot.metrics: + metrics = item.global_metrics + total = totals.setdefault( + item.operation, + { + "request_count": 0, + "success_count": 0, + "error_count": 0, + "weighted_avg": 0.0, + "max_latency_ms": 0.0, + "latency_groups": [], + }, + ) + total["request_count"] += metrics.request_count + total["success_count"] += metrics.success_count + total["error_count"] += metrics.error_count + total["weighted_avg"] += metrics.avg_latency_ms * metrics.request_count + total["max_latency_ms"] = max(total["max_latency_ms"], metrics.max_latency_ms) + samples = item._global_latency_samples_us + if samples: + weight = metrics.request_count / len(samples) + total["latency_groups"].append((samples, weight)) + elif metrics.request_count: + # Runtime snapshots created before samples were retained remain usable. + total["latency_groups"].append( + ( + (metrics.p99_latency_ms * 1000.0,), + float(metrics.request_count), + ) + ) + result = {} + for operation, total in totals.items(): + count = int(total["request_count"]) + p99_latency_ms = 0.0 + groups = total["latency_groups"] + if count and groups: + threshold = count * 0.99 + cumulative = 0.0 + heap = [ + (samples[0], group_index, 0) + for group_index, (samples, _weight) in enumerate(groups) + if samples + ] + heapq.heapify(heap) + while heap: + latency_us, group_index, sample_index = heapq.heappop(heap) + samples, weight = groups[group_index] + cumulative += weight + p99_latency_ms = latency_us / 1000.0 + if cumulative > threshold: + break + next_index = sample_index + 1 + if next_index < len(samples): + heapq.heappush(heap, (samples[next_index], group_index, next_index)) + result[operation] = { + "request_count": count, + "success_count": int(total["success_count"]), + "error_count": int(total["error_count"]), + "avg_latency_ms": total["weighted_avg"] / count if count else 0.0, + "p99_latency_ms": p99_latency_ms, + "max_latency_ms": total["max_latency_ms"], + } + return { + "aggregated": {"start_time": start_ms, "end_time": end_ms, "metrics": result}, + "snapshot_count": len(snapshots), + } + + +def new_client_request_id() -> str: + """Return a lowercase 32-character OpenTelemetry TraceID.""" + + while True: + value = os.urandom(16) + if any(value): + return value.hex() diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py index b748d0c62..590f3271a 100644 --- a/pymilvus/milvus_client/async_milvus_client.py +++ b/pymilvus/milvus_client/async_milvus_client.py @@ -1,14 +1,18 @@ import asyncio import copy +import inspect +import logging import time import types -from typing import Dict, List, Optional, Type, Union +import weakref +from typing import Any, Dict, List, Optional, Type, Union from pymilvus.client import type_info from pymilvus.client.abstract import AnnSearchRequest, BaseRanker from pymilvus.client.connection_manager import AsyncConnectionManager, ConnectionConfig from pymilvus.client.constants import CLUSTER_ID, DEFAULT_CONSISTENCY_LEVEL from pymilvus.client.search_aggregation import SearchAggregation +from pymilvus.client.telemetry import AsyncClientTelemetryManager, telemetry_operation from pymilvus.client.types import ( ExceptionsMessage, FunctionType, @@ -43,6 +47,8 @@ from .index import IndexParam, IndexParams, extract_bound_index_param from .optimize_task import OptimizeResult, ProgressStage, parse_target_size +logger = logging.getLogger(__name__) + class AsyncMilvusClient(BaseMilvusClient): """AsyncMilvusClient is an EXPERIMENTAL class @@ -66,6 +72,8 @@ def __init__( # Store config for deferred connection self._dedicated = kwargs.pop("dedicated", False) kwargs.pop("cluster_id", None) + if user: + kwargs.setdefault("user", user) self._config = ConnectionConfig.from_uri( uri, token=final_token, @@ -78,6 +86,75 @@ def __init__( self._using = None self.is_self_hosted = None self._closed = False + self._lifecycle_lock: Optional[asyncio.Lock] = None + self._telemetry = self._new_telemetry_manager() + + def _new_telemetry_manager(self) -> AsyncClientTelemetryManager: + handler_kwargs = self._config.get_handler_kwargs() + client_ref = weakref.ref(self) + + def stub_provider(): + client = client_ref() + if client is None or client._handler is None: + return None + return getattr(client._handler, "telemetry_stub", None) + + def database_provider() -> str: + client = client_ref() + return client._config.db_name if client is not None else "" + + def config_provider() -> Dict[str, Any]: + client = client_ref() + config = client._config if client is not None else None + return { + # Keep the caller's logical endpoint stable across global-primary + # failover. ConnectionConfig.address contains no URI credentials. + "address": config.address if config is not None else "", + "username": handler_kwargs.get("user", "") or "", + "db_name": config.db_name if config is not None else "", + "secure": bool(handler_kwargs.get("secure", False)), + } + + def owner_alive_provider() -> bool: + return client_ref() is not None + + manager = AsyncClientTelemetryManager( + stub_provider, + handler_kwargs.get("telemetry_config"), + user=handler_kwargs.get("user", "") or "", + database_provider=database_provider, + config_provider=config_provider, + owner_alive_provider=owner_alive_provider, + runtime_client_id=handler_kwargs.get("_telemetry_client_id", ""), + ) + manager_ref = weakref.ref(manager) + + def owner_released(_client_ref: weakref.ReferenceType[Any]) -> None: + telemetry = manager_ref() + if telemetry is not None: + telemetry.owner_released() + + # The providers close over this cell, so replace the initial reference + # with a callback-bearing one after the manager exists. The callback + # retains only a weak manager reference and never keeps the client alive. + client_ref = weakref.ref(self, owner_released) + return manager + + def _bind_telemetry_handler(self, handler: Any, database: str) -> None: + register = getattr(handler, "register_client_telemetry", None) + if callable(register) and not inspect.iscoroutinefunction(register): + register(self._telemetry, database) + + def _unbind_telemetry_handler(self, handler: Any) -> None: + unregister = getattr(handler, "unregister_client_telemetry", None) + if callable(unregister) and not inspect.iscoroutinefunction(unregister): + unregister(self._telemetry) + + def _get_lifecycle_lock(self) -> asyncio.Lock: + """Create the lifecycle lock on the event loop that first uses the client.""" + if self._lifecycle_lock is None: + self._lifecycle_lock = asyncio.Lock() + return self._lifecycle_lock def session(self, cluster_id: str) -> "AsyncMilvusClientSession": """Create a lightweight client session pinned to a target cluster.""" @@ -100,18 +177,46 @@ async def __aexit__( async def _connect(self) -> None: """Establish the async connection. Call this before using the client.""" - if self._handler is not None: - return # Already connected + async with self._get_lifecycle_lock(): + if self._closed: + raise MilvusException(message="should create connection first") + if self._handler is not None: + return # Already connected + + manager = AsyncConnectionManager.get_instance() + candidate = await manager.get_or_create( + self._config, + dedicated=self._dedicated, + client=self, + timeout=self._timeout, + ) + try: + is_self_hosted = bool(candidate.get_server_type() == "milvus") + except BaseException: + try: + await manager.release(candidate, client=self) + except Exception: + logger.warning("Failed to release rejected async connection", exc_info=True) + raise - self._manager = AsyncConnectionManager.get_instance() - self._handler = await self._manager.get_or_create( - self._config, - dedicated=self._dedicated, - client=self, - timeout=self._timeout, - ) - self._using = f"cm-async-{id(self._handler)}" - self.is_self_hosted = bool(self._handler.get_server_type() == "milvus") + try: + self._bind_telemetry_handler(candidate, self._config.db_name) + self._manager = manager + self._handler = candidate + self._using = f"cm-async-{id(candidate)}" + self.is_self_hosted = is_self_hosted + self._telemetry.start() + except BaseException: + self._unbind_telemetry_handler(candidate) + self._manager = None + self._handler = None + self._using = None + self.is_self_hosted = None + try: + await manager.release(candidate, client=self) + except Exception: + logger.warning("Failed to release rejected async connection", exc_info=True) + raise async def _get_connection(self): """Return the handler for this client, auto-connecting if needed.""" @@ -429,6 +534,7 @@ async def list_partitions( **kwargs, ) + @telemetry_operation("Insert") async def insert( self, collection_name: str, @@ -468,6 +574,7 @@ async def insert( } ) + @telemetry_operation("Upsert") async def upsert( self, collection_name: str, @@ -537,6 +644,7 @@ async def upsert( } ) + @telemetry_operation("HybridSearch") async def hybrid_search( self, collection_name: str, @@ -563,6 +671,7 @@ async def hybrid_search( **kwargs, ) + @telemetry_operation("Search") async def search( self, collection_name: str, @@ -602,6 +711,7 @@ async def search( **kwargs, ) + @telemetry_operation("Query") async def query( self, collection_name: str, @@ -650,6 +760,7 @@ async def query( **kwargs, ) + @telemetry_operation("Query") async def get( self, collection_name: str, @@ -691,6 +802,7 @@ async def get( **kwargs, ) + @telemetry_operation("Delete") async def delete( self, collection_name: str, @@ -1277,10 +1389,28 @@ async def add_function_field( async def close(self): """Close the client and release the connection.""" - self._closed = True - if self._manager and self._handler: - await self._manager.release(self._handler, client=self) - self._handler = None + async with self._get_lifecycle_lock(): + self._closed = True + try: + await self._telemetry.stop_async() + except Exception: + logger.warning("Failed to stop client telemetry", exc_info=True) + finally: + if self._manager and self._handler: + try: + self._unbind_telemetry_handler(self._handler) + except Exception: + logger.warning("Failed to unbind client telemetry", exc_info=True) + try: + await self._manager.release(self._handler, client=self) + finally: + self._handler = None + + async def get_telemetry(self): + """Return this logical client's telemetry manager.""" + + await self._get_connection() + return self._telemetry async def list_indexes(self, collection_name: str, field_name: Optional[str] = "", **kwargs): conn = await self._get_connection() @@ -1409,8 +1539,28 @@ async def use_database(self, db_name: str, **kwargs): Raises: MilvusException: If the database does not exist (error code 800). """ - await self.describe_database(db_name, **kwargs) - self._config.db_name = db_name + # Establish deferred clients before entering the non-reentrant lifecycle + # section used for the database handoff. + await self._connect() + async with self._get_lifecycle_lock(): + await self.describe_database(db_name, **kwargs) + if db_name == self._config.db_name: + return + + # Database routing is carried by CallContext on every request. Keep the + # existing transport so operations that already captured it are not closed + # underneath another coroutine. Rebind only this logical client's telemetry + # database, retaining identity/history/server-pushed configuration. + old_config = self._config + new_config = copy.copy(self._config) + new_config.db_name = db_name + try: + self._bind_telemetry_handler(self._handler, db_name) + self._config = new_config + except BaseException: + self._config = old_config + self._bind_telemetry_handler(self._handler, old_config.db_name) + raise async def create_database( self, @@ -1999,6 +2149,7 @@ async def get_compaction_plans( job_id, timeout=timeout, context=self._generate_call_context(**kwargs), **kwargs ) + @telemetry_operation("RunAnalyzer") async def run_analyzer( self, texts: Union[str, List[str]], @@ -2394,6 +2545,7 @@ def remaining_timeout() -> Optional[float]: collection_name=collection_name, target_size=size_mb, timeout=remaining_timeout(), + context=self._generate_call_context(**kwargs), **kwargs, ) @@ -2822,18 +2974,22 @@ def _with_cluster_id(self, kwargs: Dict) -> Dict: async def close(self) -> None: self._closed = True + @telemetry_operation("Search") async def search(self, *args, **kwargs): self._ensure_open() return await self._parent.search(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("HybridSearch") async def hybrid_search(self, *args, **kwargs): self._ensure_open() return await self._parent.hybrid_search(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("Query") async def query(self, *args, **kwargs): self._ensure_open() return await self._parent.query(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("Query") async def get(self, *args, **kwargs): self._ensure_open() return await self._parent.get(*args, **self._with_cluster_id(kwargs)) diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py index 0a3d43baa..c19682d97 100644 --- a/pymilvus/milvus_client/milvus_client.py +++ b/pymilvus/milvus_client/milvus_client.py @@ -1,7 +1,9 @@ import copy import logging +import threading import time -from typing import Callable, Dict, List, Optional, Union +import weakref +from typing import Any, Callable, Dict, List, Optional, Union from pymilvus.client import type_info from pymilvus.client.abstract import AnnSearchRequest, BaseRanker @@ -11,6 +13,7 @@ from pymilvus.client.iterator import QueryIterator, SearchIterator, SearchIteratorV2 from pymilvus.client.search_aggregation import SearchAggregation from pymilvus.client.search_result import Hit, Hits +from pymilvus.client.telemetry import ClientTelemetryManager, telemetry_operation from pymilvus.client.types import ( CompactionPlans, ExceptionsMessage, @@ -86,8 +89,11 @@ def __init__( final_token = f"{user}:{password}" # Create config and get handler via ConnectionManager - dedicated = kwargs.pop("dedicated", False) + self._dedicated = kwargs.pop("dedicated", False) + self._timeout = timeout kwargs.pop("cluster_id", None) + if user: + kwargs.setdefault("user", user) self._config = ConnectionConfig.from_uri( uri, token=final_token, @@ -95,17 +101,86 @@ def __init__( **kwargs, ) self._manager = ConnectionManager.get_instance() + self._lifecycle_lock = threading.RLock() self._handler = self._manager.get_or_create( self._config, - dedicated=dedicated, + dedicated=self._dedicated, client=self, timeout=timeout, ) - - # Legacy compatibility - store alias for _using attribute + self._telemetry = self._new_telemetry_manager() self._using = f"cm-{id(self._handler)}" + try: + self._bind_telemetry_handler(self._handler, self._config.db_name) + self._telemetry.start() + self.is_self_hosted = bool(self.get_server_type() == "milvus") + except BaseException: + self._telemetry.stop() + self._unbind_telemetry_handler(self._handler) + self._manager.release(self._handler, client=self) + self._handler = None + raise + + def _new_telemetry_manager(self) -> ClientTelemetryManager: + handler_kwargs = self._config.get_handler_kwargs() + client_ref = weakref.ref(self) + + def stub_provider(): + client = client_ref() + if client is None or client._handler is None: + return None + return getattr(client._handler, "telemetry_stub", None) + + def database_provider() -> str: + client = client_ref() + return client._config.db_name if client is not None else "" + + def config_provider() -> Dict[str, Any]: + client = client_ref() + config = client._config if client is not None else None + return { + # Keep the caller's logical endpoint stable across global-primary + # failover. ConnectionConfig.address contains no URI credentials. + "address": config.address if config is not None else "", + "username": handler_kwargs.get("user", "") or "", + "db_name": config.db_name if config is not None else "", + "secure": bool(handler_kwargs.get("secure", False)), + } - self.is_self_hosted = bool(self.get_server_type() == "milvus") + def owner_alive_provider() -> bool: + return client_ref() is not None + + manager = ClientTelemetryManager( + stub_provider, + handler_kwargs.get("telemetry_config"), + user=handler_kwargs.get("user", "") or "", + database_provider=database_provider, + config_provider=config_provider, + owner_alive_provider=owner_alive_provider, + runtime_client_id=handler_kwargs.get("_telemetry_client_id", ""), + ) + manager_ref = weakref.ref(manager) + + def owner_released(_client_ref: weakref.ReferenceType[Any]) -> None: + telemetry = manager_ref() + if telemetry is not None: + telemetry.owner_released() + + # The providers close over this cell, so replace the initial reference + # with a callback-bearing one after the manager exists. The callback + # retains only a weak manager reference and never keeps the client alive. + client_ref = weakref.ref(self, owner_released) + return manager + + def _bind_telemetry_handler(self, handler: Any, database: str) -> None: + register = getattr(handler, "register_client_telemetry", None) + if callable(register): + register(self._telemetry, database) + + def _unbind_telemetry_handler(self, handler: Any) -> None: + unregister = getattr(handler, "unregister_client_telemetry", None) + if callable(unregister): + unregister(self._telemetry) def session(self, cluster_id: str) -> "MilvusClientSession": """Create a lightweight client session pinned to a target cluster.""" @@ -226,6 +301,7 @@ def _create_index( **kwargs, ) + @telemetry_operation("Insert") def insert( self, collection_name: str, @@ -283,6 +359,7 @@ def insert( } ) + @telemetry_operation("Upsert") def upsert( self, collection_name: str, @@ -354,6 +431,7 @@ def upsert( } ) + @telemetry_operation("HybridSearch") def hybrid_search( self, collection_name: str, @@ -417,6 +495,7 @@ def hybrid_search( **kwargs, ) + @telemetry_operation("Search") def search( self, collection_name: str, @@ -493,6 +572,7 @@ def search( **kwargs, ) + @telemetry_operation("Query") def query( self, collection_name: str, @@ -756,6 +836,7 @@ def search_iterator( rpc_options=kwargs, ) + @telemetry_operation("Query") def get( self, collection_name: str, @@ -810,6 +891,7 @@ def get( **kwargs, ) + @telemetry_operation("Delete") def delete( self, collection_name: str, @@ -1014,10 +1096,22 @@ def _get_connection(self): def close(self): """Close the client and release the connection.""" - if self._handler is None: - return - self._manager.release(self._handler, client=self) - self._handler = None + # A custom telemetry command runs on the heartbeat worker and may enter a + # lifecycle API such as use_database(). Do not hold this lock while waiting + # for that worker, or close and the command can wait on each other forever. + self._telemetry.stop() + with self._lifecycle_lock: + if self._handler is None: + return + self._unbind_telemetry_handler(self._handler) + self._manager.release(self._handler, client=self) + self._handler = None + + def get_telemetry(self): + """Return this logical client's telemetry manager.""" + + self._get_connection() + return self._telemetry def load_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): """Loads the collection.""" @@ -2037,8 +2131,25 @@ def use_database(self, db_name: str, **kwargs): MilvusException: If the database does not exist (error code 800). """ - self.describe_database(db_name, **kwargs) - self._config.db_name = db_name + with self._lifecycle_lock: + self.describe_database(db_name, **kwargs) + if db_name == self._config.db_name: + return + + # Database routing is carried by CallContext on every request. Keep the + # existing transport so operations that already captured it are not closed + # underneath another thread. Rebind only this logical client's telemetry + # database, retaining identity/history/server-pushed configuration. + old_config = self._config + new_config = copy.copy(self._config) + new_config.db_name = db_name + try: + self._bind_telemetry_handler(self._handler, db_name) + self._config = new_config + except BaseException: + self._config = old_config + self._bind_telemetry_handler(self._handler, old_config.db_name) + raise def create_database( self, @@ -2492,6 +2603,7 @@ def describe_replica( **kwargs, ) + @telemetry_operation("RunAnalyzer") def run_analyzer( self, texts: Union[str, List[str]], @@ -3591,14 +3703,17 @@ def _with_cluster_id(self, kwargs: Dict) -> Dict: def close(self) -> None: self._closed = True + @telemetry_operation("Search") def search(self, *args, **kwargs): self._ensure_open() return self._parent.search(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("HybridSearch") def hybrid_search(self, *args, **kwargs): self._ensure_open() return self._parent.hybrid_search(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("Query") def query(self, *args, **kwargs): self._ensure_open() return self._parent.query(*args, **self._with_cluster_id(kwargs)) @@ -3611,6 +3726,7 @@ def search_iterator(self, *args, **kwargs): self._ensure_open() return self._parent.search_iterator(*args, **self._with_cluster_id(kwargs)) + @telemetry_operation("Query") def get(self, *args, **kwargs): self._ensure_open() return self._parent.get(*args, **self._with_cluster_id(kwargs)) diff --git a/pymilvus/orm/connections.py b/pymilvus/orm/connections.py index c0ec815bf..6db21da3f 100644 --- a/pymilvus/orm/connections.py +++ b/pymilvus/orm/connections.py @@ -371,7 +371,7 @@ def connect( def connect_milvus(**kwargs): gh = GrpcHandler(**kwargs) if not _async else AsyncGrpcHandler(**kwargs) - exclude_keys = ["password", "token", "keep_alive"] + exclude_keys = ["password", "token", "keep_alive", "_telemetry_client_id"] if _unbind_with_db: exclude_keys.append("db_name") config_to_keep = {k: v for k, v in kwargs.items() if k not in exclude_keys} diff --git a/tests/e2e/__init__.py b/tests/e2e/__init__.py new file mode 100644 index 000000000..eb4e3fa5f --- /dev/null +++ b/tests/e2e/__init__.py @@ -0,0 +1 @@ +"""Opt-in end-to-end tests.""" diff --git a/tests/e2e/test_client_telemetry.py b/tests/e2e/test_client_telemetry.py new file mode 100644 index 000000000..ac05e542b --- /dev/null +++ b/tests/e2e/test_client_telemetry.py @@ -0,0 +1,301 @@ +import asyncio +import json +import os +import time +import urllib.parse +import urllib.request +import uuid + +import pytest +from pymilvus import ( + AsyncMilvusClient, + MilvusClient, + TelemetryConfig, + connections, + new_client_request_id, +) +from pymilvus.exceptions import MilvusException +from pymilvus.grpc_gen import common_pb2 + +pytestmark = pytest.mark.skipif( + os.getenv("MILVUS_TELEMETRY_E2E") != "true", + reason="set MILVUS_TELEMETRY_E2E=true to run against a local Milvus", +) + +MILVUS_URI = os.getenv("MILVUS_URI", "http://127.0.0.1:19530") +TELEMETRY_API = os.getenv("MILVUS_TELEMETRY_API", "http://127.0.0.1:9091/api/v1/_telemetry") + + +def test_default_telemetry_registers_automatically(): + client = MilvusClient(MILVUS_URI, dedicated=True) + try: + manager = client.get_telemetry() + assert manager.ready is True + _wait_for( + manager.client_id, + "default client registration", + lambda candidate: candidate["status"] == "active", + ) + finally: + client.close() + + +def test_legacy_connection_registers_automatically(): + alias = f"telemetry-e2e-{uuid.uuid4()}" + client_id = f"e2e-python-legacy-{uuid.uuid4()}" + try: + connections.connect( + alias=alias, + uri=MILVUS_URI, + telemetry_config=TelemetryConfig( + heartbeat_interval=0.5, + client_id=client_id, + ), + ) + _wait_for( + client_id, + "legacy client registration", + lambda candidate: candidate["status"] == "active", + ) + finally: + connections.disconnect(alias) + + +@pytest.mark.asyncio +async def test_async_client_metrics_and_request_id_round_trip(): + client_id = f"e2e-python-async-{uuid.uuid4()}" + client = AsyncMilvusClient( + MILVUS_URI, + telemetry_config=TelemetryConfig( + heartbeat_interval=0.5, + client_id=client_id, + ), + ) + try: + manager = await client.get_telemetry() + await asyncio.to_thread( + _wait_for, + client_id, + "async client registration", + lambda state: state["status"] == "active", + ) + + analyzer = await client.run_analyzer( + "hello async telemetry", analyzer_params={"type": "standard"} + ) + assert analyzer.tokens == ["hello", "async", "telemetry"] + await asyncio.to_thread( + _wait_for, + client_id, + "async RunAnalyzer metric", + lambda state: _has_metric(state, "RunAnalyzer", "success_count", 1), + ) + + collections_command = await asyncio.to_thread( + _push_command, + client_id, + "collection_metrics", + {"collections": ["*"], "enabled": True}, + ) + assert (await asyncio.to_thread(_wait_for_reply, client_id, collections_command))[ + "success" + ] is True + + request_id = new_client_request_id() + with pytest.raises(MilvusException): + await client.query( + "telemetry_e2e_async_missing", + filter="id > 0", + client_request_id=request_id, + ) + await asyncio.to_thread( + _wait_for, + client_id, + "async failed Query collection metric", + lambda state: _has_metric( + state, + "Query", + "error_count", + 1, + collection="telemetry_e2e_async_missing", + ), + ) + + errors_command = await asyncio.to_thread( + _push_command, client_id, "show_errors", {"max_count": 10} + ) + errors_reply = await asyncio.to_thread(_wait_for_reply, client_id, errors_command) + errors = json.loads(errors_reply["payload"]) + assert any(error.get("request_id") == request_id for error in errors) + assert manager.last_command_timestamp > 0 + finally: + await client.close() + + +def test_metrics_commands_config_and_request_id_round_trip(): + client_id = f"e2e-python-{uuid.uuid4()}" + config_command = None + client = MilvusClient( + MILVUS_URI, + telemetry_config=TelemetryConfig( + heartbeat_interval=0.5, + sampling_rate=1.0, + client_id=client_id, + ), + ) + try: + manager = client.get_telemetry() + assert manager.client_id == client_id + _wait_for(client_id, "client registration", lambda state: state["status"] == "active") + + analyzer = client.run_analyzer( + "hello milvus telemetry", analyzer_params={"type": "standard"} + ) + assert analyzer.tokens == ["hello", "milvus", "telemetry"] + _wait_for( + client_id, + "RunAnalyzer metric", + lambda state: _has_metric(state, "RunAnalyzer", "success_count", 1), + ) + + collections_command = _push_command( + client_id, + "collection_metrics", + {"collections": ["*"], "enabled": True}, + ) + assert _wait_for_reply(client_id, collections_command)["success"] is True + + request_id = new_client_request_id() + with pytest.raises(MilvusException): + client.query( + "telemetry_e2e_missing", + filter="id > 0", + client_request_id=request_id, + ) + _wait_for( + client_id, + "failed Query collection metric", + lambda state: _has_metric( + state, + "Query", + "error_count", + 1, + collection="telemetry_e2e_missing", + ), + ) + + errors_command = _push_command(client_id, "show_errors", {"max_count": 10}) + errors_reply = _wait_for_reply(client_id, errors_command) + assert errors_reply["success"] is True + errors = json.loads(errors_reply["payload"]) + assert any( + error.get("operation") == "Query" and error.get("request_id") == request_id + for error in errors + ) + + config_payload = {"sampling_rate": 0.75, "heartbeat_interval_ms": 600} + config_command = _push_command(client_id, "push_config", config_payload, persistent=True) + assert _wait_for_reply(client_id, config_command)["success"] is True + expected_hash = manager.calculate_config_hash( + [ + common_pb2.ClientCommand( + command_id=config_command, + command_type="push_config", + payload=json.dumps(config_payload, separators=(",", ":")).encode(), + persistent=True, + ) + ] + ) + assert manager.config_hash == expected_hash + assert manager.last_command_timestamp > 0 + + config_reply = _wait_for_reply(client_id, _push_command(client_id, "get_config", {})) + user_config = json.loads(config_reply["payload"])["user_config"] + assert user_config["telemetry_sampling_rate"] == 0.75 + assert user_config["telemetry_heartbeat_interval_ms"] == 600 + assert user_config["all_collections_enabled"] is True + finally: + try: + if config_command is not None: + _delete_command(config_command) + finally: + client.close() + + +def _request_json(url, method="GET", body=None): + parsed = urllib.parse.urlsplit(url) + if parsed.scheme not in {"http", "https"}: + raise ValueError(f"unsupported telemetry URL scheme: {parsed.scheme}") + data = None if body is None else json.dumps(body, separators=(",", ":")).encode() + request = urllib.request.Request(url, data=data, method=method) # noqa: S310 + if data is not None: + request.add_header("Content-Type", "application/json") + with urllib.request.urlopen(request, timeout=5) as response: # noqa: S310 + return json.load(response) + + +def _client_state(client_id): + query = urllib.parse.urlencode({"client_id": client_id, "include_metrics": "true"}) + clients = _request_json(f"{TELEMETRY_API}/clients?{query}").get("clients", []) + return clients[0] if clients else None + + +def _wait_for(client_id, label, predicate, timeout=15): + deadline = time.time() + timeout + last = None + while time.time() < deadline: + last = _client_state(client_id) + if last is not None and predicate(last): + return last + time.sleep(0.25) + raise AssertionError(f"timed out waiting for {label}; last={last}") + + +def _push_command(client_id, command_type, payload, persistent=False): + response = _request_json( + f"{TELEMETRY_API}/commands", + method="POST", + body={ + "command_type": command_type, + "target_client_id": client_id, + "payload": payload, + "ttl_seconds": 30, + "persistent": persistent, + }, + ) + return response["command_id"] + + +def _delete_command(command_id): + encoded_command_id = urllib.parse.quote(command_id, safe="") + _request_json(f"{TELEMETRY_API}/commands/{encoded_command_id}", method="DELETE") + + +def _wait_for_reply(client_id, command_id): + state = _wait_for( + client_id, + f"command reply {command_id}", + lambda candidate: _find_reply(candidate, command_id) is not None, + ) + return _find_reply(state, command_id) + + +def _find_reply(state, command_id): + return next( + ( + reply + for reply in state.get("command_replies") or [] + if reply.get("command_id") == command_id + ), + None, + ) + + +def _has_metric(state, operation, counter, minimum, collection=None): + for metric in state.get("metrics") or []: + if metric.get("operation") != operation: + continue + if metric.get("global", {}).get(counter, 0) < minimum: + continue + return collection is None or collection in metric.get("collection_metrics", {}) + return False diff --git a/tests/unit/async_grpc_handler/test_async_auth.py b/tests/unit/async_grpc_handler/test_async_auth.py index 20c404e47..351fddfdf 100644 --- a/tests/unit/async_grpc_handler/test_async_auth.py +++ b/tests/unit/async_grpc_handler/test_async_auth.py @@ -260,7 +260,8 @@ async def test_create_role_preserves_positional_context(self) -> None: mock_stub.CreateRole = AsyncMock(return_value=mock_status) handler._async_stub = mock_stub - context = CallContext(db_name="db1", client_request_id="req1") + request_id = "0123456789abcdef0123456789abcdef" + context = CallContext(db_name="db1", client_request_id=request_id) await handler.create_role("test_role", 30, context) req = mock_stub.CreateRole.call_args.args[0] @@ -268,7 +269,7 @@ async def test_create_role_preserves_positional_context(self) -> None: assert req.entity.description == "" assert kwargs["timeout"] == 30 assert ("dbname", "db1") in kwargs["metadata"] - assert ("client-request-id", "req1") in kwargs["metadata"] + assert ("client-request-id", request_id) in kwargs["metadata"] @pytest.mark.asyncio async def test_alter_role(self) -> None: diff --git a/tests/unit/async_grpc_handler/test_async_init.py b/tests/unit/async_grpc_handler/test_async_init.py index f90cb43bc..c84ac0138 100644 --- a/tests/unit/async_grpc_handler/test_async_init.py +++ b/tests/unit/async_grpc_handler/test_async_init.py @@ -220,7 +220,8 @@ def test_build_stub_reuses_authorization_and_one_time_log_level(self) -> None: handler._build_stub(next_channel) assert authorization_interceptor in next_channel._unary_unary_interceptors - assert len(next_channel._unary_unary_interceptors) == 2 + assert len(next_channel._unary_unary_interceptors) == 3 + assert next_channel._unary_unary_interceptors[-1] is handler._telemetry_interceptor assert handler._log_level is None def test_setup_secure_channel(self) -> None: diff --git a/tests/unit/client/__init__.py b/tests/unit/client/__init__.py new file mode 100644 index 000000000..a50b3410a --- /dev/null +++ b/tests/unit/client/__init__.py @@ -0,0 +1 @@ +"""Client unit tests.""" diff --git a/tests/unit/client/test_telemetry.py b/tests/unit/client/test_telemetry.py new file mode 100644 index 000000000..e229ed404 --- /dev/null +++ b/tests/unit/client/test_telemetry.py @@ -0,0 +1,2118 @@ +import asyncio +import json +import math +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor +from types import SimpleNamespace + +import grpc +import pymilvus.client.telemetry as telemetry_module +import pytest +from pymilvus.client.asynch import Future as PyMilvusFuture +from pymilvus.client.call_context import CallContext +from pymilvus.client.telemetry import ( + AsyncClientTelemetryManager, + AsyncTelemetryUnaryUnaryInterceptor, + ClientCommand, + ClientTelemetryManager, + CommandReply, + Metrics, + MetricsSnapshot, + OperationMetrics, + TelemetryConfig, + TelemetryUnaryUnaryInterceptor, + _request_id_from_metadata, + _response_error, + is_valid_client_request_id, + new_client_request_id, + telemetry_operation, +) +from pymilvus.decorators import retry_on_rpc_failure +from pymilvus.exceptions import MilvusException +from pymilvus.grpc_gen import common_pb2, milvus_pb2 + + +def test_config_hash_matches_server_algorithm(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=False)) + commands = [ + common_pb2.ClientCommand( + command_id="cfg-b", + command_type="push_config", + payload=b'{"sampling_rate":0.5}', + persistent=True, + ), + common_pb2.ClientCommand( + command_id="cfg-a", + command_type="push_config", + payload=b'{"heartbeat_interval_ms":5000}', + persistent=True, + ), + ] + + assert manager.calculate_config_hash(commands) == "a271ff0bb1941777" + + +def test_process_commands_is_idempotent_and_queues_replies(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + calls = [] + manager.register_command_handler( + "custom", + lambda command: ( + calls.append(command.command_id) + or CommandReply(command.command_id, True, payload=b"ok") + ), + ) + command = common_pb2.ClientCommand(command_id="cmd-1", command_type="custom", create_time=1000) + + manager.process_commands([command]) + manager.process_commands([command]) + + assert calls == ["cmd-1"] + assert [reply.command_id for reply in manager._pending_replies] == ["cmd-1", "cmd-1"] + + +def test_custom_command_replies_use_server_id_and_none_becomes_failure(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager.register_command_handler("missing", lambda _command: None) + + manager.process_commands( + [common_pb2.ClientCommand(command_id="missing-id", command_type="missing", create_time=1)] + ) + + missing = manager._pending_replies.pop() + assert missing.command_id == "missing-id" + assert missing.success is False + assert missing.error_message == "command handler returned no reply" + + manager.register_command_handler( + "wrong-id", lambda _command: CommandReply("other-id", False, "failed", b"payload") + ) + manager.process_commands( + [common_pb2.ClientCommand(command_id="server-id", command_type="wrong-id", create_time=2)] + ) + + canonical = manager._pending_replies.pop() + assert canonical.command_id == "server-id" + assert canonical.success is False + assert canonical.error_message == "failed" + assert canonical.payload == b"payload" + + +def test_builtin_config_and_error_commands(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager.record_operation("Search", "books", time.perf_counter(), RuntimeError("boom")) + commands = [ + common_pb2.ClientCommand( + command_id="cfg", + command_type="push_config", + payload=b'{"sampling_rate":0.5,"heartbeat_interval_ms":5000}', + create_time=1, + persistent=True, + ), + common_pb2.ClientCommand( + command_id="errors", + command_type="show_errors", + payload=b'{"max_count":1}', + create_time=2, + ), + ] + + manager.process_commands(commands) + + with manager._config_lock: + assert manager._config.sampling_rate == 0.5 + assert manager._config.heartbeat_interval == 5.0 + errors_reply = next(reply for reply in manager._pending_replies if reply.command_id == "errors") + assert json.loads(errors_reply.payload)[0]["error_msg"] == "boom" + + +def test_push_config_is_atomic_and_reports_applied_and_ignored_keys(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + with pytest.raises(ValueError, match="heartbeat_interval_ms must be positive"): + manager._handle_push_config( + ClientCommand( + command_id="bad", + command_type="push_config", + payload=b'{"enabled":false,"heartbeat_interval_ms":0}', + ) + ) + + with manager._config_lock: + assert manager._config.enabled is True + assert manager._config.heartbeat_interval == 10.0 + + reply = manager._handle_push_config( + ClientCommand( + command_id="good", + command_type="push_config", + payload=b'{"sampling_rate":2,"ttl_seconds":30,"future":"value"}', + ) + ) + assert json.loads(reply.payload) == { + "applied": ["sampling_rate"], + "ignored": ["future", "ttl_seconds"], + } + with manager._config_lock: + assert manager._config.sampling_rate == 1.0 + + +@pytest.mark.parametrize( + "payload,error", + [ + (b'{"enabled":"false"}', "enabled must be a boolean"), + (b'{"heartbeat_interval_ms":1.5}', "heartbeat_interval_ms must be an integer"), + (b'{"sampling_rate":true}', "sampling_rate must be a number"), + (b'{"sampling_rate":NaN}', "invalid JSON constant"), + ], +) +def test_push_config_rejects_wrong_json_types(payload, error): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + with pytest.raises((TypeError, ValueError), match=error): + manager._handle_push_config( + ClientCommand(command_id="bad", command_type="push_config", payload=payload) + ) + + +def test_process_command_batches_are_serialized(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + calls = 0 + calls_lock = threading.Lock() + second_lock_attempted = threading.Event() + + class ObservableLock: + def __init__(self): + self._lock = threading.RLock() + self._attempts = 0 + self._attempts_lock = threading.Lock() + + def __enter__(self): + with self._attempts_lock: + self._attempts += 1 + if self._attempts == 2: + second_lock_attempted.set() + self._lock.acquire() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self._lock.release() + + # Signal the second acquisition attempt before it blocks on the real lock. The first + # handler waits for that signal, so the test proves the two batches actually contend. + manager._command_batch_lock = ObservableLock() + + def handler(command): + nonlocal calls + with calls_lock: + calls += 1 + assert second_lock_attempted.wait(timeout=5) + return CommandReply(command.command_id, True) + + manager.register_command_handler("custom", handler) + command = common_pb2.ClientCommand( + command_id="same-command", command_type="custom", create_time=1000 + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(manager.process_commands, [command]) for _ in range(2)] + for future in futures: + future.result() + + assert second_lock_attempted.is_set() + assert calls == 1 + + +def test_equal_timestamp_command_stays_deduplicated_after_repeated_delivery(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + calls = 0 + + def handler(command): + nonlocal calls + calls += 1 + return CommandReply(command.command_id, True) + + manager.register_command_handler("custom", handler) + command = common_pb2.ClientCommand( + command_id="same-command", command_type="custom", create_time=1000 + ) + + manager.process_commands([command]) + manager.process_commands([command]) + manager.process_commands([command]) + + assert calls == 1 + + +def test_persistent_configs_bypass_one_time_command_state(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager.register_command_handler( + "custom", lambda command: CommandReply(command.command_id, True) + ) + manager.process_commands( + [common_pb2.ClientCommand(command_id="cursor", command_type="custom", create_time=20)] + ) + + older_config = common_pb2.ClientCommand( + command_id="database-config", + command_type="push_config", + payload=b'{"sampling_rate":0.25}', + create_time=10, + persistent=True, + ) + manager.process_commands([older_config]) + + assert manager.last_command_timestamp == 20 + assert manager.config_hash == manager.calculate_config_hash([older_config]) + with manager._config_lock: + assert manager._config.sampling_rate == 0.25 + + newer_config = common_pb2.ClientCommand( + command_id="newer-database-config", + command_type="push_config", + payload=b'{"sampling_rate":0.75}', + create_time=30, + persistent=True, + ) + manager.process_commands([newer_config]) + assert manager.last_command_timestamp == 20 + with manager._config_lock: + assert manager._config.sampling_rate == 0.75 + + # Returning to a previously delivered persistent set must reapply it even + # though its command ID has already been observed. + manager.process_commands([older_config]) + assert manager.last_command_timestamp == 20 + with manager._config_lock: + assert manager._config.sampling_rate == 0.25 + + +def test_empty_command_batch_preserves_hash_without_authoritative_config_snapshot(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + config = common_pb2.ClientCommand( + command_id="database-config", + command_type="push_config", + payload=b'{"sampling_rate":0.25}', + create_time=10, + persistent=True, + ) + + manager.process_commands([config]) + assert manager.config_hash == manager.calculate_config_hash([config]) + + manager.process_commands([]) + assert manager.config_hash == manager.calculate_config_hash([config]) + + +class _RetryableRpcError(grpc.RpcError): + def code(self): + return grpc.StatusCode.UNAVAILABLE + + +class _UnimplementedRpcError(grpc.RpcError): + def code(self): + return grpc.StatusCode.UNIMPLEMENTED + + +def test_logical_operation_counts_retries_once(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler: + def __init__(self): + self._telemetry = manager + self.attempts = 0 + + @telemetry_operation("Search") + @retry_on_rpc_failure(retry_times=2, initial_back_off=0, max_back_off=0) + def search(self, collection_name, context=None): + self.attempts += 1 + if self.attempts == 1: + raise _RetryableRpcError + return "parsed-result" + + handler = Handler() + + assert handler.search("books") == "parsed-result" + assert handler.attempts == 2 + collector = manager._collectors["Search"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + assert collector.global_bucket.success_count == 1 + assert collector.global_bucket.error_count == 0 + + +@pytest.mark.asyncio +async def test_async_logical_operation_counts_retries_once_and_records_final_error(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + request_id = new_client_request_id() + + class Handler: + def __init__(self): + self._telemetry = manager + self.attempts = 0 + + @telemetry_operation("Query") + @retry_on_rpc_failure(retry_times=1, initial_back_off=0, max_back_off=0) + async def query(self, collection_name, context=None): + self.attempts += 1 + raise _RetryableRpcError + + handler = Handler() + + with pytest.raises(MilvusException): + await handler.query("books", context=CallContext(client_request_id=request_id)) + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + assert collector.global_bucket.success_count == 0 + assert collector.global_bucket.error_count == 1 + assert manager.get_recent_errors(1)[0].request_id == request_id + + +def test_logical_operation_suppresses_per_rpc_interceptor_recording(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + interceptor = TelemetryUnaryUnaryInterceptor(manager) + + class ImmediateCall: + @staticmethod + def exception(): + return None + + @staticmethod + def result(): + return milvus_pb2.QueryResults(status=common_pb2.Status()) + + def add_done_callback(self, callback): + callback(self) + + class Handler: + _telemetry = manager + + @telemetry_operation("Query") + def query(self, collection_name, context=None): + details = SimpleNamespace( + method="/milvus.proto.milvus.MilvusService/Query", metadata=() + ) + request = SimpleNamespace(collection_name=collection_name) + return interceptor.intercept_unary_unary( + lambda *_args: ImmediateCall(), details, request + ) + + Handler().query("books") + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + + +def test_public_client_local_validation_is_included_in_logical_metric(): + from pymilvus import MilvusClient # noqa: PLC0415 + + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + client = MilvusClient.__new__(MilvusClient) + client._handler = SimpleNamespace(_telemetry=manager) + + with pytest.raises(TypeError, match="wrong type of argument 'data'"): + client.insert("books", "not-rows") + + collector = manager._collectors["Insert"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + assert collector.global_bucket.error_count == 1 + + +def test_sync_get_and_session_get_are_counted_once_as_query(): + from pymilvus import MilvusClient # noqa: PLC0415 + + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler: + def __init__(self): + self._telemetry = manager + self.query_calls = 0 + + def _get_schema(self, *_args, **_kwargs): + return ({"fields": []}, None) + + @telemetry_operation("Query") + def query(self, *_args, **_kwargs): + self.query_calls += 1 + return [] + + handler = Handler() + client = MilvusClient.__new__(MilvusClient) + client._handler = handler + client._config = SimpleNamespace(db_name="default") + client._pack_pks_expr = lambda _schema, _ids: "id in [1]" + + assert client.get("books", ids=[1]) == [] + assert client.session("cluster-1").get("books", ids=[2]) == [] + assert client.get("books", ids=[]) == [] + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 3 + assert collector.global_bucket.success_count == 3 + assert handler.query_calls == 2 + + +@pytest.mark.asyncio +async def test_async_get_and_session_get_are_counted_once_as_query(): + from pymilvus import AsyncMilvusClient # noqa: PLC0415 + + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler: + def __init__(self): + self._telemetry = manager + self.query_calls = 0 + + async def _get_schema(self, *_args, **_kwargs): + return ({"fields": []}, None) + + @telemetry_operation("Query") + async def query(self, *_args, **_kwargs): + self.query_calls += 1 + return [] + + handler = Handler() + client = AsyncMilvusClient.__new__(AsyncMilvusClient) + client._handler = handler + client._config = SimpleNamespace(db_name="default") + client._closed = False + client._pack_pks_expr = lambda _schema, _ids: "id in [1]" + + assert await client.get("books", ids=[1]) == [] + assert await client.session("cluster-1").get("books", ids=[2]) == [] + assert await client.get("books", ids=[]) == [] + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 3 + assert collector.global_bucket.success_count == 3 + assert handler.query_calls == 2 + + +def test_sync_session_keyword_collection_name_is_attributed(): + from pymilvus import MilvusClient # noqa: PLC0415 + + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._handle_collection_metrics( + ClientCommand( + command_id="enable-books", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":["books"]}', + ) + ) + + class Handler: + _telemetry = manager + + @staticmethod + def query(*_args, **_kwargs): + return [] + + client = MilvusClient.__new__(MilvusClient) + client._handler = Handler() + client._config = SimpleNamespace(db_name="") + + assert client.session("cluster-1").query(collection_name="books") == [] + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + assert collector.collections["books"].request_count == 1 + + +@pytest.mark.asyncio +async def test_async_session_keyword_collection_name_is_attributed(): + from pymilvus import AsyncMilvusClient # noqa: PLC0415 + + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._handle_collection_metrics( + ClientCommand( + command_id="enable-books", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":["books"]}', + ) + ) + + class Handler: + _telemetry = manager + + @staticmethod + async def query(*_args, **_kwargs): + return [] + + client = AsyncMilvusClient.__new__(AsyncMilvusClient) + client._closed = False + client._handler = Handler() + client._config = SimpleNamespace(db_name="") + + assert await client.session("cluster-1").query(collection_name="books") == [] + + collector = manager._collectors["Query"] + with collector.lock: + assert collector.global_bucket.request_count == 1 + assert collector.collections["books"].request_count == 1 + + +def test_collection_metrics_are_filtered_again_at_wire_time(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._handle_collection_metrics( + ClientCommand( + command_id="enable", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":["books"]}', + ) + ) + snapshot_metrics = [ + OperationMetrics( + "Search", + Metrics(request_count=1, success_count=1), + {"books": Metrics(request_count=1, success_count=1)}, + ) + ] + manager._handle_collection_metrics( + ClientCommand( + command_id="disable", + command_type="collection_metrics", + payload=b'{"enabled":false,"collections":["books"]}', + ) + ) + + assert dict(manager._to_proto_metrics(snapshot_metrics)[0].collection_metrics) == {} + + +def test_collection_metrics_rejects_coerced_json_types(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + with pytest.raises(TypeError, match="enabled must be a boolean"): + manager._handle_collection_metrics( + ClientCommand( + command_id="bad", + command_type="collection_metrics", + payload=b'{"enabled":"true","collections":["books"]}', + ) + ) + + +@pytest.mark.parametrize( + "start_time", + [ + "2026-08-23T12:34Z", + "2026-08-23T12:34:00", + "2026-08-23", + "2026-02-30T12:34:00Z", + ], +) +def test_latency_history_requires_strict_rfc3339(start_time): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + with pytest.raises(ValueError): + manager._handle_latency_history( + ClientCommand( + command_id="bad-history", + command_type="show_latency_history", + payload=json.dumps( + { + "start_time": start_time, + "end_time": "2026-08-23T12:35:00Z", + } + ).encode(), + ) + ) + + +def test_latency_history_accepts_rfc3339_offset_and_long_fraction(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + reply = manager._handle_latency_history( + ClientCommand( + command_id="valid-history", + command_type="show_latency_history", + payload=( + b'{"start_time":"2026-08-23T12:34:00.123456789123+08:00",' + b'"end_time":"2026-08-23T12:35:00.987654321987+08:00"}' + ), + ) + ) + + assert reply.success is True + + +def test_real_heartbeat_response_clears_unsupported_backoff_on_status_error(): + class Stub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(error_code=common_pb2.UnexpectedError, reason="failed") + ) + + manager = ClientTelemetryManager(Stub, TelemetryConfig(enabled=True)) + manager._unsupported_streak = 3 + + manager._send_heartbeat() + + assert manager.is_supported() is True + assert str(manager.last_heartbeat_error()) == "failed" + + +@pytest.mark.parametrize("outcome", ["response", "error"]) +def test_rebind_fences_stale_sync_heartbeat_state(outcome): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + original_error = RuntimeError("original") + manager._unsupported_streak = 3 + manager._last_heartbeat_error = original_error + manager._queue_reply(CommandReply("pending", True)) + + class OldStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + manager.rebind_stub(object()) + if outcome == "error": + raise _UnimplementedRpcError + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), + commands=[ + common_pb2.ClientCommand( + command_id="stale-config", + command_type="push_config", + payload=b'{"enabled":false}', + create_time=10, + persistent=True, + ) + ], + ) + + manager.rebind_stub(OldStub()) + manager._send_heartbeat() + + assert manager._unsupported_streak == 3 + assert manager.last_heartbeat_error() is original_error + assert [reply.command_id for reply in manager._pending_replies] == ["pending"] + assert manager.last_command_timestamp == 0 + with manager._config_lock: + assert manager._config.enabled is True + + +def test_reentrant_sync_rebind_stops_old_command_batch_until_redelivery(): + commands = [ + common_pb2.ClientCommand( + command_id="switch", command_type="switch_endpoint", create_time=10 + ), + common_pb2.ClientCommand( + command_id="config", + command_type="push_config", + payload=b'{"sampling_rate":0.5}', + create_time=20, + persistent=True, + ), + ] + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class NewStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status(), commands=commands) + + def switch_endpoint(_command): + manager.rebind_stub(NewStub()) + return CommandReply("wrong-id", True) + + manager.register_command_handler("switch_endpoint", switch_endpoint) + + class OldStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status(), commands=commands) + + manager.rebind_stub(OldStub()) + manager._send_heartbeat() + + assert [reply.command_id for reply in manager._pending_replies] == ["switch"] + assert manager.last_command_timestamp == 0 + with manager._config_lock: + assert manager._config.sampling_rate == 1.0 + + manager._send_heartbeat() + + assert [reply.command_id for reply in manager._pending_replies] == ["switch", "config"] + assert manager.last_command_timestamp == 10 + with manager._config_lock: + assert manager._config.sampling_rate == 0.5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["response", "error"]) +async def test_rebind_fences_stale_async_heartbeat_state(outcome): + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + original_error = RuntimeError("original") + manager._unsupported_streak = 3 + manager._last_heartbeat_error = original_error + manager._queue_reply(CommandReply("pending", True)) + + class OldStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + manager.rebind_stub(object()) + if outcome == "error": + raise _UnimplementedRpcError + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), + commands=[ + common_pb2.ClientCommand( + command_id="stale-config", + command_type="push_config", + payload=b'{"enabled":false}', + create_time=10, + persistent=True, + ) + ], + ) + + manager.rebind_stub(OldStub()) + await manager._send_heartbeat_async() + + assert manager._unsupported_streak == 3 + assert manager.last_heartbeat_error() is original_error + assert [reply.command_id for reply in manager._pending_replies] == ["pending"] + assert manager.last_command_timestamp == 0 + with manager._config_lock: + assert manager._config.enabled is True + + +@pytest.mark.asyncio +async def test_reentrant_async_rebind_stops_old_command_batch_until_redelivery(): + commands = [ + common_pb2.ClientCommand( + command_id="switch", command_type="switch_endpoint", create_time=10 + ), + common_pb2.ClientCommand( + command_id="config", + command_type="push_config", + payload=b'{"sampling_rate":0.25}', + create_time=20, + persistent=True, + ), + ] + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class NewStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status(), commands=commands) + + def switch_endpoint(_command): + manager.rebind_stub(NewStub()) + return CommandReply("", True) + + manager.register_command_handler("switch_endpoint", switch_endpoint) + + class OldStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status(), commands=commands) + + manager.rebind_stub(OldStub()) + await manager._send_heartbeat_async() + + assert [reply.command_id for reply in manager._pending_replies] == ["switch"] + assert manager.last_command_timestamp == 0 + with manager._config_lock: + assert manager._config.sampling_rate == 1.0 + + await manager._send_heartbeat_async() + + assert [reply.command_id for reply in manager._pending_replies] == ["switch", "config"] + assert manager.last_command_timestamp == 10 + with manager._config_lock: + assert manager._config.sampling_rate == 0.25 + + +def test_show_errors_uses_default_for_non_positive_count_and_empty_payload_for_no_errors(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + reply = manager._handle_show_errors( + ClientCommand(command_id="errors", command_type="show_errors", payload=b'{"max_count":0}') + ) + + assert reply.success is True + assert reply.payload == b"" + + +def test_new_client_request_id_is_valid_trace_id(): + request_id = new_client_request_id() + assert len(request_id) == 32 + assert request_id != "0" * 32 + int(request_id, 16) + assert is_valid_client_request_id(request_id) + + +@pytest.mark.parametrize( + "request_id", + ["", "0" * 32, "A" * 32, "g" * 32, "0123456789abcdef"], +) +def test_invalid_client_request_id_is_not_recorded(request_id): + assert not is_valid_client_request_id(request_id) + assert _request_id_from_metadata((("client-request-id", request_id),)) == "" + wire_metadata = dict(CallContext(client_request_id=request_id).to_grpc_metadata()) + if request_id: + assert wire_metadata["client-request-id"] == request_id + else: + assert "client-request-id" not in wire_metadata + + +def test_valid_client_request_id_is_forwarded_in_grpc_metadata(): + request_id = new_client_request_id() + + assert ( + dict(CallContext(client_request_id=request_id).to_grpc_metadata())["client-request-id"] + == request_id + ) + + +def test_response_error_detects_milvus_status_failure(): + response = milvus_pb2.QueryResults( + status=common_pb2.Status(error_code=common_pb2.CollectionNotExists, reason="missing") + ) + + error = _response_error(response) + + assert str(error) == "missing" + + +def test_response_error_accepts_success_status(): + response = milvus_pb2.QueryResults(status=common_pb2.Status()) + + assert _response_error(response) is None + + +def test_runtime_client_id_is_reused_without_becoming_stable(): + manager = ClientTelemetryManager( + lambda: None, + TelemetryConfig(enabled=False), + runtime_client_id="runtime-client-id", + ) + + assert manager.client_id == "runtime-client-id" + assert manager._client_id_stable is False + + +def test_show_errors_truncates_a_single_large_error(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager.record_operation( + "Query", + "books", + time.perf_counter(), + RuntimeError("x" * (2 * 1024 * 1024)), + ) + + reply = manager._handle_show_errors( + ClientCommand(command_id="errors", command_type="show_errors") + ) + + assert reply.success is True + assert len(reply.payload) <= 1024 * 1024 + + +def test_show_errors_truncates_oversized_non_message_fields_without_mutating_history(): + collection = "x" * (1024 * 1024 + 100) + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager.record_operation( + "Query", collection, time.perf_counter(), RuntimeError("short message") + ) + + reply = manager._handle_show_errors( + ClientCommand(command_id="errors", command_type="show_errors") + ) + + assert reply.success is True + assert len(reply.payload) <= 1024 * 1024 + assert json.loads(reply.payload)[0]["collection"].endswith("...(truncated)") + assert manager.get_recent_errors(1)[0].collection == collection + + +def test_metrics_snapshot_sampling_and_collection_scope(): + manager = ClientTelemetryManager( + lambda: None, + TelemetryConfig(enabled=True), + database_provider=lambda: "analytics", + ) + manager._handle_collection_metrics( + ClientCommand( + command_id="enable-books", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":["books"]}', + ) + ) + + assert manager._should_sample(0.0) is False + assert [manager._should_sample(0.25) for _ in range(4)] == [False, False, False, True] + + started_at = time.perf_counter() - 0.01 + manager.record_operation("Search", "books", started_at) + manager.record_operation("Search", "private", started_at, RuntimeError("failed")) + manager._create_snapshot() + + snapshot = manager.get_metrics_snapshots()[-1] + search = snapshot.metrics[0] + assert search.operation == "Search" + assert search.global_metrics.request_count == 2 + assert search.global_metrics.success_count == 1 + assert search.global_metrics.error_count == 1 + assert search.global_metrics.avg_latency_ms > 0 + assert search.collection_metrics["books"].request_count == 1 + assert "private" not in search.collection_metrics + + # A second window resets the collector and must not repeat the previous metrics. + manager._create_snapshot() + assert manager.get_metrics_snapshots()[-1].metrics == [] + assert manager._build_client_info().reserved["db_name"] == "analytics" + assert manager.config_hash == "" + assert manager.ready is False + + manager._unsupported_streak = 3 + assert manager._next_heartbeat_delay() == 80.0 + manager._unsupported_streak = 4096 + assert manager._next_heartbeat_delay() == 1800.0 + + disabled = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=False)) + disabled.record_operation("Query", "books", time.perf_counter()) + disabled._create_snapshot() + assert disabled.get_metrics_snapshots() == [] + + +def test_snapshot_history_uses_one_hour_ttl_with_independent_hard_cap(monkeypatch): + now_ms = 2_000_000_000_000 + monkeypatch.setattr(telemetry_module.time, "time", lambda: now_ms / 1000) + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + hour_ms = 60 * 60 * 1000 + manager._snapshots.extend( + [ + MetricsSnapshot(now_ms - hour_ms - 2, now_ms - hour_ms - 1, []), + MetricsSnapshot(now_ms - hour_ms, now_ms - hour_ms, []), + MetricsSnapshot(now_ms - hour_ms // 2, now_ms - hour_ms // 2, []), + ] + ) + + manager._handle_push_config( + ClientCommand( + command_id="interval", + command_type="push_config", + payload=b'{"heartbeat_interval_ms":600000}', + ) + ) + retained = manager.get_metrics_snapshots() + assert [snapshot.end_time for snapshot in retained] == [ + now_ms - hour_ms, + now_ms - hour_ms // 2, + ] + + manager._snapshots.clear() + manager._snapshots.extend( + MetricsSnapshot(now_ms + index, now_ms + index, []) for index in range(4097) + ) + retained = manager.get_metrics_snapshots() + assert len(retained) == 4096 + assert retained[0].timestamp == now_ms + 1 + + +def test_command_queries_redaction_collection_modes_and_failure_replies(): + manager = ClientTelemetryManager( + lambda: None, + TelemetryConfig(enabled=True), + config_provider=lambda: { + "password": "secret", + "token": "secret", + "api_key": "secret", + "address": "localhost:19530", + }, + ) + + initial = manager._handle_collection_metrics( + ClientCommand(command_id="state", command_type="collection_metrics") + ) + assert json.loads(initial.payload) == { + "enabled_collections": [], + "all_collections_enabled": False, + } + + with pytest.raises(ValueError, match="collections list cannot be empty"): + manager._handle_collection_metrics( + ClientCommand( + command_id="bad-enable", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":[]}', + ) + ) + + manager._handle_collection_metrics( + ClientCommand( + command_id="all", + command_type="collection_metrics", + payload=b'{"enabled":true,"collections":["*"]}', + ) + ) + config = manager._handle_get_config( + ClientCommand(command_id="config", command_type="get_config") + ) + user_config = json.loads(config.payload)["user_config"] + assert user_config["address"] == "localhost:19530" + assert user_config["enabled_collections"] == ["*"] + assert user_config["all_collections_enabled"] is True + assert not {"password", "token", "api_key"} & user_config.keys() + + manager._handle_collection_metrics( + ClientCommand( + command_id="none", + command_type="collection_metrics", + payload=b'{"enabled":false,"collections":[]}', + ) + ) + assert manager._all_collections_enabled is False + assert manager._enabled_collections == set() + + push = manager._handle_push_config( + ClientCommand( + command_id="disable", + command_type="push_config", + payload=b'{"enabled":false}', + ) + ) + assert json.loads(push.payload) == {"applied": ["enabled"]} + assert manager._config.enabled is False + + unknown = manager._handle_command(ClientCommand("unknown", "not-registered")) + assert unknown.success is False + assert "unknown command type" in unknown.error_message + + manager.register_command_handler("raises", lambda _command: 1 / 0) + failed = manager._handle_command(ClientCommand("failed", "raises")) + assert failed.success is False + assert "division by zero" in failed.error_message + + manager._last_command_timestamp = 10 + manager.process_commands( + [common_pb2.ClientCommand(command_id="stale", command_type="push_config", create_time=9)] + ) + assert manager._pending_replies[-1].command_id == "stale" + assert manager.calculate_config_hash([]) == "" + + +def test_unprintable_command_handler_error_does_not_stop_sync_heartbeat_loop(): + requests = [] + + class UnprintableError(Exception): + def __str__(self): + raise RuntimeError("formatting failed") + + class Stub: + @staticmethod + def ClientHeartbeat(request, **_kwargs): + requests.append(request) + if len(requests) == 1: + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), + commands=[ + common_pb2.ClientCommand( + command_id="unprintable", + command_type="custom", + create_time=1, + ) + ], + ) + manager._stop_event.set() + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + manager = ClientTelemetryManager(Stub, TelemetryConfig(enabled=True)) + + def raise_unprintable(_command): + raise UnprintableError + + manager.register_command_handler("custom", raise_unprintable) + manager._create_snapshot = lambda: None + manager._next_heartbeat_delay = lambda: 0 + + manager._heartbeat_loop() + + assert len(requests) == 2 + reply = requests[1].command_replies[0] + assert reply.command_id == "unprintable" + assert reply.success is False + assert reply.error_message == "UnprintableError (failed to format exception)" + + +@pytest.mark.parametrize( + "command,error", + [ + ( + ClientCommand("array", "push_config", payload=b"[]"), + "command payload must be a JSON object", + ), + ( + ClientCommand("finite", "push_config", payload=b'{"sampling_rate":1e309}'), + "sampling_rate must be finite", + ), + ( + ClientCommand( + "strings", + "collection_metrics", + payload=b'{"enabled":true,"collections":[1]}', + ), + "collections must be an array of strings", + ), + ], +) +def test_command_payload_rejects_strict_json_edge_cases(command, error): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + reply = manager._handle_command(command) + assert reply.success is False + assert error in reply.error_message + + +def test_config_validation_covers_invalid_type_interval_and_error_limit_default(): + with pytest.raises(TypeError, match="telemetry_config must be"): + TelemetryConfig.from_value("enabled") + with pytest.raises(ValueError, match="heartbeat_interval must be positive"): + TelemetryConfig(heartbeat_interval=0) + assert TelemetryConfig(error_max_count=0).error_max_count == 100 + + +def test_telemetry_config_is_copied_for_each_manager(): + supplied = TelemetryConfig(enabled=True, sampling_rate=0.75) + first = ClientTelemetryManager(lambda: None, supplied) + second = ClientTelemetryManager(lambda: None, supplied) + + first._handle_push_config( + ClientCommand( + command_id="disable", + command_type="push_config", + payload=b'{"enabled":false,"sampling_rate":0.25}', + ) + ) + + assert supplied.enabled is True + assert supplied.sampling_rate == 0.75 + assert second._config.enabled is True + assert second._config.sampling_rate == 0.75 + + +def test_latency_history_detail_aggregate_redaction_and_range_validation(monkeypatch): + monkeypatch.setattr(telemetry_module.time, "time", lambda: 3) + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._snapshots.append( + MetricsSnapshot( + timestamp=1000, + end_time=2000, + metrics=[ + OperationMetrics( + "Search", + Metrics( + request_count=2, + success_count=1, + error_count=1, + avg_latency_ms=4.0, + p99_latency_ms=7.0, + max_latency_ms=9.0, + ), + ) + ], + ) + ) + + def history(payload): + return manager._handle_latency_history( + ClientCommand( + command_id="history", + command_type="show_latency_history", + payload=json.dumps(payload).encode(), + ) + ) + + window = { + "start_time": "1970-01-01T00:00:00Z", + "end_time": "1970-01-01T00:00:03Z", + } + aggregate = json.loads(history(window).payload) + assert aggregate["snapshot_count"] == 1 + assert aggregate["aggregated"]["metrics"]["Search"] == { + "request_count": 2, + "success_count": 1, + "error_count": 1, + "avg_latency_ms": 4.0, + "p99_latency_ms": 7.0, + "max_latency_ms": 9.0, + } + + detail = json.loads(history({**window, "detail": True}).payload) + assert detail["total_snapshots"] == 1 + assert detail["snapshots"][0]["metrics"]["Search"]["request_count"] == 2 + + with pytest.raises(TypeError, match="start_time must be a string"): + history({"start_time": 1, "end_time": window["end_time"]}) + with pytest.raises(TypeError, match="end_time must be a string"): + history({"start_time": window["start_time"], "end_time": 3}) + with pytest.raises(ValueError, match="end_time must be after start_time"): + history( + { + "start_time": "1970-01-01T00:00:03Z", + "end_time": "1970-01-01T00:00:02Z", + } + ) + with pytest.raises(ValueError, match="time range cannot exceed 1 hour"): + history( + { + "start_time": "1970-01-01T00:00:00Z", + "end_time": "1970-01-01T01:00:01Z", + } + ) + with pytest.raises(ValueError, match="payload is required"): + history({}) + + monkeypatch.setattr(telemetry_module, "_MAX_REPLY_PAYLOAD_SIZE", 1) + with pytest.raises(ValueError, match="response too large"): + history(window) + + +def test_latency_history_aggregates_weighted_samples_without_flattening_sort(monkeypatch): + bucket = telemetry_module._MetricsBucket() + for latency_us in range(1000): + bucket.record(latency_us, True) + _, retained = bucket.snapshot_and_reset(retain_history_samples=True) + assert len(retained) == 128 + assert retained[0] == 0 + assert retained[-1] == 999 + + snapshots = [ + MetricsSnapshot( + timestamp=0, + end_time=1000, + metrics=[ + OperationMetrics( + "Search", + Metrics( + request_count=100, + success_count=100, + avg_latency_ms=1.0, + p99_latency_ms=1.0, + max_latency_ms=1.0, + ), + _global_latency_samples_us=telemetry_module.array("q", [1000]), + ) + ], + ), + MetricsSnapshot( + timestamp=1000, + end_time=2000, + metrics=[ + OperationMetrics( + "Search", + Metrics( + request_count=100, + success_count=100, + avg_latency_ms=100.0, + p99_latency_ms=100.0, + max_latency_ms=100.0, + ), + _global_latency_samples_us=telemetry_module.array("q", [100_000]), + ) + ], + ), + ] + + def reject_flattening_sort(*_args, **_kwargs): + pytest.fail("history aggregation must merge retained sorted samples without sorting a copy") + + monkeypatch.setattr(telemetry_module, "sorted", reject_flattening_sort, raising=False) + aggregate = telemetry_module._aggregate_snapshots(snapshots, 0, 2000) + search = aggregate["aggregated"]["metrics"]["Search"] + assert search["avg_latency_ms"] == 50.5 + assert search["p99_latency_ms"] == 100.0 + assert set(telemetry_module._snapshot_dict(snapshots[0])["metrics"]["Search"]) == { + "request_count", + "success_count", + "error_count", + "avg_latency_ms", + "p99_latency_ms", + "max_latency_ms", + } + + +@pytest.mark.parametrize("heartbeat_interval", [math.nan, math.inf, -math.inf]) +def test_config_rejects_non_finite_heartbeat_intervals(heartbeat_interval): + with pytest.raises(ValueError, match="fit in signed 64-bit milliseconds"): + TelemetryConfig(heartbeat_interval=heartbeat_interval) + + +def test_config_rejects_heartbeat_interval_beyond_signed_int64_milliseconds(): + with pytest.raises(ValueError, match="fit in signed 64-bit milliseconds"): + TelemetryConfig(heartbeat_interval=telemetry_module._MAX_HEARTBEAT_INTERVAL * 2) + + +def test_sync_long_heartbeat_interval_uses_cancellable_wait_chunks(monkeypatch): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + waits = [] + + class RecordingStopEvent: + def wait(self, delay): + waits.append(delay) + return len(waits) == 3 + + monkeypatch.setattr(telemetry_module, "_MAX_WAIT_CHUNK_SECONDS", 2.0) + manager._stop_event = RecordingStopEvent() + + assert manager._wait_for_stop(10.0) is True + assert waits == [2.0, 2.0, 2.0] + + +def test_max_int64_push_interval_remains_representable(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._handle_push_config( + ClientCommand( + command_id="max-interval", + command_type="push_config", + payload=f'{{"heartbeat_interval_ms":{2**63 - 1}}}'.encode(), + ) + ) + + assert manager._heartbeat_interval_ms() == 2**63 - 1 + reply = manager._handle_get_config(ClientCommand("get", "get_config")) + assert json.loads(reply.payload)["user_config"]["telemetry_heartbeat_interval_ms"] == ( + 2**63 - 1 + ) + + +def test_push_interval_beyond_int64_is_rejected_without_mutating_config(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + reply = manager._handle_command( + ClientCommand( + command_id="oversized-interval", + command_type="push_config", + payload=f'{{"heartbeat_interval_ms":{2**63}}}'.encode(), + ) + ) + + assert reply.success is False + assert "fit in a signed 64-bit integer" in reply.error_message + assert manager._heartbeat_interval() == telemetry_module._DEFAULT_HEARTBEAT_INTERVAL + + +def test_sync_heartbeat_loop_exits_before_work_when_owner_is_gone(): + manager = ClientTelemetryManager( + lambda: None, + TelemetryConfig(enabled=True), + owner_alive_provider=lambda: False, + ) + manager._stop_event.set() + + manager._heartbeat_loop() + + assert manager.get_metrics_snapshots() == [] + + +def test_sync_stop_waits_for_worker_exit_before_clearing_thread_handle(): + heartbeat_entered = threading.Event() + release_heartbeat = threading.Event() + stop_started = threading.Event() + stop_returned = threading.Event() + + class BlockingStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + heartbeat_entered.set() + assert release_heartbeat.wait(timeout=5) + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + manager = ClientTelemetryManager(BlockingStub, TelemetryConfig(enabled=True)) + manager.start() + assert heartbeat_entered.wait(timeout=1) + worker = manager._thread + assert worker is not None + + def stop_manager(): + stop_started.set() + manager.stop() + stop_returned.set() + + stopper = threading.Thread(target=stop_manager) + stopper.start() + + assert stop_started.wait(timeout=1) + assert not stop_returned.wait(timeout=0.05) + assert manager._thread is worker + assert worker.is_alive() + + release_heartbeat.set() + assert stop_returned.wait(timeout=1) + stopper.join(timeout=1) + assert not stopper.is_alive() + assert not worker.is_alive() + assert manager._thread is None + + +def test_sync_stop_retains_live_worker_handle_when_join_times_out(monkeypatch): + heartbeat_entered = threading.Event() + release_heartbeat = threading.Event() + + class BlockingStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + heartbeat_entered.set() + assert release_heartbeat.wait(timeout=5) + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + monkeypatch.setattr(telemetry_module, "_HEARTBEAT_STOP_JOIN_TIMEOUT", 0.01) + manager = ClientTelemetryManager(BlockingStub, TelemetryConfig(enabled=True)) + manager.start() + assert heartbeat_entered.wait(timeout=1) + worker = manager._thread + assert worker is not None + + manager.stop() + + assert worker.is_alive() + assert manager._thread is worker + + release_heartbeat.set() + worker.join(timeout=1) + assert not worker.is_alive() + manager.stop() + assert manager._thread is None + + +def test_sync_heartbeat_loop_isolates_unexpected_iteration_failure(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + calls = 0 + + def flaky_heartbeat(): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("transient telemetry failure") + manager._stop_event.set() + + manager._create_snapshot = lambda: None + manager._next_heartbeat_delay = lambda: 0 + manager._send_heartbeat = flaky_heartbeat + + manager._heartbeat_loop() + + assert calls == 2 + assert str(manager.last_heartbeat_error()) == "transient telemetry failure" + + +def test_sync_owner_release_wakes_long_wait(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + manager.owner_released() + + assert manager._stop_event.is_set() + + +@pytest.mark.asyncio +async def test_async_long_heartbeat_interval_uses_cancellable_sleep_chunks(monkeypatch): + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + sleeps = [] + + async def record_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(telemetry_module, "_MAX_WAIT_CHUNK_SECONDS", 2.0) + monkeypatch.setattr(telemetry_module.asyncio, "sleep", record_sleep) + + await manager._sleep_until_next_heartbeat(5.5) + + assert sleeps == [2.0, 2.0, 1.5] + + +@pytest.mark.asyncio +async def test_async_heartbeat_loop_exits_before_work_when_owner_is_gone(): + manager = AsyncClientTelemetryManager( + lambda: None, + TelemetryConfig(enabled=True), + owner_alive_provider=lambda: False, + ) + manager._stop_event.set() + + await manager._async_heartbeat_loop() + + assert manager.get_metrics_snapshots() == [] + + +@pytest.mark.asyncio +async def test_async_owner_release_cancels_long_sleep(): + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + manager._task = asyncio.create_task(asyncio.Event().wait()) + + manager.owner_released() + with pytest.raises(asyncio.CancelledError): + await manager._task + + assert manager._stop_event.is_set() + assert manager._task.cancelled() + + +def test_sync_heartbeat_success_unimplemented_and_lifecycle_short_circuits(): + captured = [] + + class SuccessStub: + @staticmethod + def ClientHeartbeat(request, **_kwargs): + captured.append(request) + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), + commands=[ + common_pb2.ClientCommand( + command_id="server-config", + command_type="push_config", + payload=b'{"sampling_rate":0.5}', + create_time=7, + persistent=True, + ) + ], + ) + + manager = ClientTelemetryManager(SuccessStub, TelemetryConfig(enabled=True)) + manager._unsupported_streak = 2 + manager._queue_reply(CommandReply("pending", True)) + manager._send_heartbeat() + + assert captured[0].command_replies[0].command_id == "pending" + assert manager._pending_replies[0].command_id == "server-config" + assert manager._unsupported_streak == 0 + assert manager.last_heartbeat_error() is None + assert manager.last_command_timestamp == 0 + + class UnsupportedStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + raise _UnimplementedRpcError + + unsupported = ClientTelemetryManager( + UnsupportedStub, TelemetryConfig(enabled=True, heartbeat_interval=0.5) + ) + unsupported._send_heartbeat() + assert unsupported._unsupported_streak == 1 + assert isinstance(unsupported.last_heartbeat_error(), _UnimplementedRpcError) + assert unsupported._next_heartbeat_delay() == 1.0 + + ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=False))._send_heartbeat() + ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True))._send_heartbeat() + + disabled = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=False)) + disabled.start() + disabled.start() + assert disabled.ready is True + assert disabled._thread is None + + joined = [] + + def record_join(timeout): + joined.append(timeout) + + disabled._thread = SimpleNamespace(join=record_join, is_alive=lambda: False) + disabled.stop() + assert joined == [telemetry_module._HEARTBEAT_STOP_JOIN_TIMEOUT] + + +def test_sync_disabled_metrics_keeps_control_plane_and_can_be_reenabled(): + captured = [] + + class Stub: + @staticmethod + def ClientHeartbeat(request, **_kwargs): + captured.append(request) + if len(captured) == 1: + command = common_pb2.ClientCommand( + command_id="disable", + command_type="push_config", + payload=b'{"enabled":false}', + create_time=1, + persistent=True, + ) + elif len(captured) == 2: + command = common_pb2.ClientCommand( + command_id="enable", + command_type="push_config", + payload=b'{"enabled":true}', + create_time=2, + persistent=True, + ) + else: + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), commands=[command] + ) + + manager = ClientTelemetryManager(Stub, TelemetryConfig(enabled=True)) + manager.record_operation("Query", "books", time.perf_counter() - 0.001) + manager._create_snapshot() + + manager._send_heartbeat() + disable_hash = manager.config_hash + with manager._config_lock: + assert manager._config.enabled is False + assert captured[0].metrics + + manager.record_operation("Query", "books", time.perf_counter() - 0.001) + manager._create_snapshot() + assert len(manager.get_metrics_snapshots()) == 1 + manager._send_heartbeat() + + assert not captured[1].metrics + assert [reply.command_id for reply in captured[1].command_replies] == ["disable"] + assert captured[1].config_hash == disable_hash + assert [reply.command_id for reply in manager._pending_replies] == ["enable"] + with manager._config_lock: + assert manager._config.enabled is True + + manager._create_snapshot() + enable_hash = manager.config_hash + manager._send_heartbeat() + assert [reply.command_id for reply in captured[2].command_replies] == ["enable"] + assert captured[2].config_hash == enable_hash + assert manager._pending_replies == [] + + +@pytest.mark.asyncio +async def test_async_heartbeat_success_errors_and_lifecycle_short_circuits(): + class SuccessStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + manager = AsyncClientTelemetryManager(SuccessStub, TelemetryConfig(enabled=True)) + manager._unsupported_streak = 2 + manager._queue_reply(CommandReply("pending", True)) + await manager._send_heartbeat_async() + assert manager._unsupported_streak == 0 + assert manager.last_heartbeat_error() is None + assert manager._pending_replies == [] + + class StatusErrorStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(error_code=common_pb2.UnexpectedError, reason="bad") + ) + + status_error = AsyncClientTelemetryManager(StatusErrorStub, TelemetryConfig(enabled=True)) + status_error._unsupported_streak = 2 + await status_error._send_heartbeat_async() + assert status_error._unsupported_streak == 0 + assert str(status_error.last_heartbeat_error()) == "bad" + + class UnsupportedStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + raise _UnimplementedRpcError + + unsupported = AsyncClientTelemetryManager(UnsupportedStub, TelemetryConfig(enabled=True)) + await unsupported._send_heartbeat_async() + assert unsupported._unsupported_streak == 1 + assert isinstance(unsupported.last_heartbeat_error(), _UnimplementedRpcError) + + class StaleFailureStub: + def __init__(self, owner): + self.owner = owner + + async def ClientHeartbeat(self, *_args, **_kwargs): + self.owner.rebind_stub(object()) + raise RuntimeError("stale") + + stale = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + stale.rebind_stub(StaleFailureStub(stale)) + await stale._send_heartbeat_async() + assert stale.last_heartbeat_error() is None + + disabled = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=False)) + await disabled._send_heartbeat_async() + enabled_without_stub = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + await enabled_without_stub._send_heartbeat_async() + + disabled.start() + disabled.start() + assert disabled.ready is True + assert disabled._task is None + + blocker = asyncio.Event() + task_manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + task_manager._task = asyncio.create_task(blocker.wait()) + await task_manager.stop_async() + assert task_manager._task is None + + task_manager._task = asyncio.create_task(blocker.wait()) + task_manager.stop() + await asyncio.sleep(0) + assert task_manager._task.cancelled() + + +@pytest.mark.asyncio +async def test_async_heartbeat_loop_isolates_failures_but_propagates_cancellation(): + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + calls = 0 + + async def flaky_heartbeat(): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("transient telemetry failure") + manager._stop_event.set() + + manager._create_snapshot = lambda: None + manager._next_heartbeat_delay = lambda: 0 + manager._send_heartbeat_async = flaky_heartbeat + await manager._async_heartbeat_loop() + assert calls == 2 + assert str(manager.last_heartbeat_error()) == "transient telemetry failure" + + cancelled = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + entered = asyncio.Event() + blocker = asyncio.Event() + + async def blocking_heartbeat(): + entered.set() + await blocker.wait() + + cancelled._create_snapshot = lambda: None + cancelled._send_heartbeat_async = blocking_heartbeat + task = asyncio.create_task(cancelled._async_heartbeat_loop()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + send_entered = asyncio.Event() + + class BlockingStub: + @staticmethod + async def ClientHeartbeat(*_args, **_kwargs): + send_entered.set() + await asyncio.Event().wait() + + send_manager = AsyncClientTelemetryManager(BlockingStub, TelemetryConfig(enabled=True)) + send_task = asyncio.create_task(send_manager._send_heartbeat_async()) + await send_entered.wait() + send_task.cancel() + with pytest.raises(asyncio.CancelledError): + await send_task + + +@pytest.mark.asyncio +async def test_async_stop_swallows_completed_telemetry_task_failure(): + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + async def fail(): + raise RuntimeError("background telemetry failed") + + manager._task = asyncio.create_task(fail()) + await asyncio.sleep(0) + await manager.stop_async() + + assert manager._task is None + assert str(manager.last_heartbeat_error()) == "background telemetry failed" + + +@pytest.mark.asyncio +async def test_async_disabled_metrics_keeps_control_plane_and_can_be_reenabled(): + captured = [] + + class Stub: + @staticmethod + async def ClientHeartbeat(request, **_kwargs): + captured.append(request) + if len(captured) == 1: + command = common_pb2.ClientCommand( + command_id="disable", + command_type="push_config", + payload=b'{"enabled":false}', + create_time=1, + persistent=True, + ) + elif len(captured) == 2: + command = common_pb2.ClientCommand( + command_id="enable", + command_type="push_config", + payload=b'{"enabled":true}', + create_time=2, + persistent=True, + ) + else: + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), commands=[command] + ) + + manager = AsyncClientTelemetryManager(Stub, TelemetryConfig(enabled=True)) + manager.record_operation("Query", "books", time.perf_counter() - 0.001) + manager._create_snapshot() + + await manager._send_heartbeat_async() + disable_hash = manager.config_hash + with manager._config_lock: + assert manager._config.enabled is False + assert captured[0].metrics + + manager.record_operation("Query", "books", time.perf_counter() - 0.001) + manager._create_snapshot() + assert len(manager.get_metrics_snapshots()) == 1 + await manager._send_heartbeat_async() + + assert not captured[1].metrics + assert [reply.command_id for reply in captured[1].command_replies] == ["disable"] + assert captured[1].config_hash == disable_hash + assert [reply.command_id for reply in manager._pending_replies] == ["enable"] + with manager._config_lock: + assert manager._config.enabled is True + + manager._create_snapshot() + enable_hash = manager.config_hash + await manager._send_heartbeat_async() + assert [reply.command_id for reply in captured[2].command_replies] == ["enable"] + assert captured[2].config_hash == enable_hash + assert manager._pending_replies == [] + + +def test_sync_interceptor_and_deferred_future_record_final_status(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + interceptor = TelemetryUnaryUnaryInterceptor(manager) + request_id = new_client_request_id() + + class ImmediateCall: + @staticmethod + def exception(): + return None + + @staticmethod + def result(): + return milvus_pb2.QueryResults( + status=common_pb2.Status( + error_code=common_pb2.CollectionNotExists, reason="missing" + ) + ) + + def add_done_callback(self, callback): + callback(self) + + details = SimpleNamespace( + method="/milvus.proto.milvus.MilvusService/Query", + metadata=(("client-request-id", request_id),), + ) + result = interceptor.intercept_unary_unary( + lambda *_args: ImmediateCall(), details, SimpleNamespace(collectionName="books") + ) + assert isinstance(result, ImmediateCall) + assert manager.get_recent_errors(1)[0].request_id == request_id + + future = Future() + + class Handler: + _telemetry = manager + + @telemetry_operation("Search") + def search(self, collection_name): + return SimpleNamespace(_exception=None, _future=future) + + Handler().search("books") + assert "Search" not in manager._collectors + future.set_result(milvus_pb2.SearchResults(status=common_pb2.Status())) + assert manager._collectors["Search"].global_bucket.success_count == 1 + + +def test_sync_pymilvus_future_records_after_result_processing(monkeypatch): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + clock = {"now": 10.0} + monkeypatch.setattr(telemetry_module.time, "perf_counter", lambda: clock["now"]) + + class ParsingFuture(PyMilvusFuture): + def on_response(self, _response): + clock["now"] = 10.04 + raise ValueError("response parser failed") + + raw_future = Future() + + class Handler: + _telemetry = manager + + @telemetry_operation("Search") + def search(self, collection_name): + return ParsingFuture(raw_future) + + result = Handler().search("books") + raw_future.set_result("wire response") + + assert "Search" not in manager._collectors + with pytest.raises(ValueError, match="response parser failed"): + result.result() + + bucket = manager._collectors["Search"].global_bucket + assert bucket.request_count == 1 + assert bucket.success_count == 0 + assert bucket.error_count == 1 + assert bucket.total_latency_us == pytest.approx(40_000, abs=1) + assert manager.get_recent_errors(1)[0].error_msg == "response parser failed" + + +def test_sync_pymilvus_future_done_records_response_processing_error(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class ParsingFuture(PyMilvusFuture): + def on_response(self, _response): + raise ValueError("done parser failed") + + raw_future = Future() + + class Handler: + _telemetry = manager + + @telemetry_operation("Search") + def search(self, collection_name): + return ParsingFuture(raw_future) + + result = Handler().search("books") + raw_future.set_result("wire response") + + assert "Search" not in manager._collectors + result.done() + + bucket = manager._collectors["Search"].global_bucket + assert bucket.request_count == 1 + assert bucket.error_count == 1 + assert str(result._exception) == "done parser failed" + + +@pytest.mark.parametrize("failure_stage", ["callback", "on_response"]) +def test_sync_pymilvus_future_concurrent_consumers_record_first_failure_once(failure_stage): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + processing_entered = threading.Event() + release_processing = threading.Event() + done_returned = threading.Event() + parse_calls = 0 + callback_calls = 0 + + class ParsingFuture(PyMilvusFuture): + def on_response(self, _response): + nonlocal parse_calls + parse_calls += 1 + if failure_stage == "on_response": + processing_entered.set() + assert release_processing.wait(timeout=2) + raise ValueError("on_response failed") + return "parsed" + + def user_callback(_result): + nonlocal callback_calls + callback_calls += 1 + if failure_stage == "callback": + processing_entered.set() + assert release_processing.wait(timeout=2) + raise ValueError("callback failed") + + raw_future = Future() + raw_future.set_result("wire response") + + class Handler: + _telemetry = manager + + @telemetry_operation("Search") + def search(self, collection_name): + return ParsingFuture(raw_future, user_callback) + + result = Handler().search("books") + result_errors = [] + done_errors = [] + + def consume_result(): + try: + result.result() + except BaseException as exc: + result_errors.append(exc) + + def consume_done(): + try: + result.done() + except BaseException as exc: + done_errors.append(exc) + finally: + done_returned.set() + + result_thread = threading.Thread(target=consume_result) + result_thread.start() + assert processing_entered.wait(timeout=2) + + done_thread = threading.Thread(target=consume_done) + done_thread.start() + assert not done_returned.wait(timeout=0.05) + + release_processing.set() + result_thread.join(timeout=2) + done_thread.join(timeout=2) + + assert not result_thread.is_alive() + assert not done_thread.is_alive() + assert done_errors == [] + assert len(result_errors) == 1 + assert str(result_errors[0]) == f"{failure_stage} failed" + assert parse_calls == 1 + assert callback_calls == (1 if failure_stage == "callback" else 0) + + bucket = manager._collectors["Search"].global_bucket + assert bucket.request_count == 1 + assert bucket.success_count == 0 + assert bucket.error_count == 1 + assert len(manager.get_recent_errors()) == 1 + assert manager.get_recent_errors()[0].error_msg == f"{failure_stage} failed" + + +@pytest.mark.asyncio +async def test_async_interceptor_callback_records_response_and_transport_error(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + interceptor = AsyncTelemetryUnaryUnaryInterceptor(manager) + + class AsyncCall: + def __init__(self, response=None, error=None): + self.response = response + self.error = error + + def __await__(self): + async def complete(): + if self.error is not None: + raise self.error + return self.response + + return complete().__await__() + + def add_done_callback(self, callback): + callback(self) + + details = SimpleNamespace(method=b"/milvus.proto.milvus.MilvusService/Search", metadata=()) + + async def continuation(_details, _request): + return AsyncCall(milvus_pb2.SearchResults(status=common_pb2.Status())) + + call = await interceptor.intercept_unary_unary( + continuation, details, SimpleNamespace(collection_name="books") + ) + assert isinstance(call, AsyncCall) + await asyncio.sleep(0) + await asyncio.sleep(0) + + async def failing_continuation(_details, _request): + return AsyncCall(error=RuntimeError("transport")) + + await interceptor.intercept_unary_unary( + failing_continuation, details, SimpleNamespace(collection_name="books") + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + + collector = manager._collectors["Search"] + assert collector.global_bucket.request_count == 2 + assert collector.global_bucket.success_count == 1 + assert collector.global_bucket.error_count == 1 + assert _request_id_from_metadata(()) == "" + assert _response_error(SimpleNamespace()) is None diff --git a/tests/unit/grpc_handler/test_auth.py b/tests/unit/grpc_handler/test_auth.py index 864c5a410..78852217b 100644 --- a/tests/unit/grpc_handler/test_auth.py +++ b/tests/unit/grpc_handler/test_auth.py @@ -71,14 +71,15 @@ def test_create_role_with_description(self, handler): def test_create_role_preserves_positional_context(self, handler): handler._stub.CreateRole.return_value = make_status() - context = CallContext(db_name="db1", client_request_id="req1") + request_id = "0123456789abcdef0123456789abcdef" + context = CallContext(db_name="db1", client_request_id=request_id) handler.create_role("role", 30, context) req = handler._stub.CreateRole.call_args.args[0] kwargs = handler._stub.CreateRole.call_args.kwargs assert req.entity.description == "" assert kwargs["timeout"] == 30 assert ("dbname", "db1") in kwargs["metadata"] - assert ("client-request-id", "req1") in kwargs["metadata"] + assert ("client-request-id", request_id) in kwargs["metadata"] def test_alter_role(self, handler): handler._stub.AlterRole.return_value = make_status() diff --git a/tests/unit/grpc_handler/test_init.py b/tests/unit/grpc_handler/test_init.py index 51011f3ff..642658841 100644 --- a/tests/unit/grpc_handler/test_init.py +++ b/tests/unit/grpc_handler/test_init.py @@ -1,6 +1,6 @@ """Tests for GrpcHandler initialization and connection management.""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import grpc import pytest @@ -46,6 +46,23 @@ def test_reconnect_on_idle_non_idle_clears_state(self): handler.reconnect_on_idle(MagicMock(value=(None, "ready"))) assert handler.is_idle_state is False + def test_idle_reconnect_preserves_existing_handler_and_telemetry(self): + mock_conns = MagicMock() + old_handler = MagicMock() + telemetry = object() + old_handler.telemetry = telemetry + mock_conns._fetch_handler.return_value = old_handler + handler = ReconnectHandler(mock_conns, "test", {"timeout": 1}) + handler.is_idle_state = True + + with patch("pymilvus.client.grpc_handler.time.sleep"): + handler.check_state_and_reconnect_later() + + old_handler.reconnect.assert_called_once_with(timeout=1) + assert old_handler.telemetry is telemetry + mock_conns.disconnect.assert_not_called() + mock_conns.connect.assert_not_called() + class TestGrpcHandlerInit: """Tests for GrpcHandler initialization.""" @@ -331,7 +348,10 @@ def test_set_onetime_loglevel_rebuilds_stub_with_loglevel_interceptor(self): handler.set_onetime_loglevel("debug") mock_header.assert_called_once_with(["log_level"], ["debug"]) - mock_intercept.assert_called_once_with(handler._channel, log_interceptor) + assert mock_intercept.call_args_list == [ + call(handler._channel, log_interceptor), + call(final_channel, handler._telemetry_interceptor), + ] assert handler._log_level is None assert handler._stub is stub @@ -417,6 +437,7 @@ def test_internal_register_forwards_timeout_to_connect(self): def test_wait_for_channel_ready_register_rpc_error_reports_connect_failed(self): handler = GrpcHandler(channel=MagicMock()) + handler._stub = MagicMock() handler._stub.Connect.side_effect = UnavailableRpcError() with pytest.raises(MilvusException) as exc_info: diff --git a/tests/unit/test_async_milvus_client_ops.py b/tests/unit/test_async_milvus_client_ops.py index 0873acda8..e5ef80a4b 100644 --- a/tests/unit/test_async_milvus_client_ops.py +++ b/tests/unit/test_async_milvus_client_ops.py @@ -201,8 +201,12 @@ async def test_list_aliases(self): async def test_using_database(self): client, handler = _make_client() handler.describe_database.return_value = {"db_name": "mydb"} + client._manager = AsyncMock() await client.using_database("mydb") assert client._config.db_name == "mydb" + assert client._handler is handler + client._manager.get_or_create.assert_not_awaited() + client._manager.release.assert_not_awaited() @pytest.mark.asyncio async def test_list_resource_groups(self): @@ -496,6 +500,8 @@ async def test_wait_for_compaction_failed_raises(self): @pytest.mark.asyncio async def test_execute_optimize_success(self): client, handler = _make_client() + handler.describe_database.return_value = {"db_name": "analytics"} + await client.use_database("analytics") handler.compact.return_value = 99 state = MagicMock() state.state = 2 @@ -511,6 +517,7 @@ async def test_execute_optimize_success(self): task._target_size = None result = await client._execute_optimize(task, "col", None, None) assert result.collection_name == "col" + assert handler.compact.await_args.kwargs["context"].get_db_name() == "analytics" @pytest.mark.asyncio async def test_optimize_wait_true(self): diff --git a/tests/unit/test_connection_manager.py b/tests/unit/test_connection_manager.py index a9101ca1a..0ddd290a3 100644 --- a/tests/unit/test_connection_manager.py +++ b/tests/unit/test_connection_manager.py @@ -2,8 +2,11 @@ """Tests for ConnectionManager and related classes.""" import asyncio +import contextlib +import gc import threading import time +import weakref from unittest.mock import AsyncMock, Mock, patch import grpc @@ -25,7 +28,15 @@ ) from pymilvus.client.global_topology import ClusterInfo from pymilvus.client.grpc_handler import GrpcHandler +from pymilvus.client.telemetry import ( + AsyncClientTelemetryManager, + ClientCommand, + ClientTelemetryManager, + CommandReply, + TelemetryConfig, +) from pymilvus.exceptions import ConnectionConfigException, MilvusException +from pymilvus.grpc_gen import common_pb2, milvus_pb2 # ============================================================================= # Module-level helpers @@ -68,6 +79,59 @@ def code(self): return grpc.StatusCode.NOT_FOUND +def _pooled_sync_handler(db_name: str = "") -> GrpcHandler: + return GrpcHandler( + uri="http://localhost:19530", + address="localhost:19530", + db_name=db_name, + channel=Mock(), + _client_owned_telemetry=True, + ) + + +def _pooled_async_handler(db_name: str = "") -> AsyncGrpcHandler: + channel = Mock() + channel._unary_unary_interceptors = [] + channel.close = AsyncMock() + return AsyncGrpcHandler( + uri="http://localhost:19530", + address="localhost:19530", + db_name=db_name, + channel=channel, + _client_owned_telemetry=True, + ) + + +class _LogicalSyncPool: + def __init__(self, handlers): + self.handlers = handlers + self.fail_database = None + self.releases = [] + + def get_or_create(self, config, **_kwargs): + if config.db_name == self.fail_database: + raise RuntimeError("candidate failed") + return self.handlers[config.db_name] + + def release(self, handler, client=None): + self.releases.append((handler, client)) + + +class _LogicalAsyncPool: + def __init__(self, handlers): + self.handlers = handlers + self.fail_database = None + self.releases = [] + + async def get_or_create(self, config, **_kwargs): + if config.db_name == self.fail_database: + raise RuntimeError("candidate failed") + return self.handlers[config.db_name] + + async def release(self, handler, client=None): + self.releases.append((handler, client)) + + def _make_topology(version, cluster_id, endpoint, capability=3): """Create a single-cluster GlobalTopology for tests.""" return GlobalTopology( @@ -202,11 +266,22 @@ def test_from_uri(self, uri, token, db_name, expected): assert config.token == expected["token"] assert config.db_name == expected["db_name"] + def test_uri_username_is_forwarded_for_telemetry_identity(self): + config = ConnectionConfig.from_uri("https://alice:secret@host:19530") + + assert config.get_handler_kwargs()["user"] == "alice" + + def test_uri_token_is_not_forwarded_as_telemetry_identity(self): + config = ConnectionConfig.from_uri("https://secret-token@host:19530") + + assert config.token == "secret-token" + assert "user" not in config.get_handler_kwargs() + @pytest.mark.parametrize( "uri,expected_key", [ - ("https://user:pass@host:19530", "host:19530|user:pass"), - ("http://localhost:19530", "localhost:19530|"), + ("https://user:pass@host:19530", ("host:19530", "user:pass", "", "")), + ("http://localhost:19530", ("localhost:19530", "", "", "")), ], ) def test_key_property(self, uri, expected_key): @@ -214,6 +289,60 @@ def test_key_property(self, uri, expected_key): config = ConnectionConfig.from_uri(uri) assert config.key == expected_key + def test_default_telemetry_keeps_historical_pool_key(self): + implicit = ConnectionConfig.from_uri("http://localhost:19530") + explicit = ConnectionConfig.from_uri( + "http://localhost:19530", + telemetry_config={ + "enabled": True, + "heartbeat_interval": 10.0, + "sampling_rate": 1.0, + "error_max_count": 100, + "client_id": "", + }, + ) + + assert explicit.key == implicit.key + + def test_custom_telemetry_config_does_not_split_transport_pool(self): + first = ConnectionConfig.from_uri( + "http://localhost:19530", + telemetry_config={"client_id": "client-a"}, + ) + second = ConnectionConfig.from_uri( + "http://localhost:19530", + telemetry_config={"client_id": "client-b"}, + ) + + assert first.key == second.key + + def test_database_gets_dedicated_pool_key(self): + first = ConnectionConfig.from_uri( + "http://localhost:19530", + token="token", + db_name="database-a", + ) + second = ConnectionConfig.from_uri( + "http://localhost:19530", + token="token", + db_name="database-b", + ) + + assert first.key != second.key + + def test_pool_key_has_unambiguous_token_and_database_boundaries(self): + token_with_delimiter = ConnectionConfig.from_uri( + "http://localhost:19530", + token="token|db=database-a", + ) + separate_database = ConnectionConfig.from_uri( + "http://localhost:19530", + token="token", + db_name="database-a", + ) + + assert token_with_delimiter.key != separate_database.key + @pytest.mark.parametrize( "uri,expected_is_global", [ @@ -349,6 +478,7 @@ def test_handler_kwargs_forwarded_to_grpc_handler(self): token=config.token, db_name=config.db_name, secure=True, + _client_owned_telemetry=True, ) def test_handler_kwargs_forwarded_async(self): @@ -365,6 +495,7 @@ def test_handler_kwargs_forwarded_async(self): token=config.token, db_name=config.db_name, secure=True, + _client_owned_telemetry=True, ) def test_empty_handler_kwargs(self): @@ -396,6 +527,7 @@ def test_create_handler(self): token="mytoken", db_name="mydb", secure=True, + _client_owned_telemetry=True, ) assert handler is mock_handler_cls.return_value @@ -473,6 +605,7 @@ def test_create_handler_fetches_topology( token="mytoken", db_name="", secure=True, + _client_owned_telemetry=True, ) assert strategy.get_topology() is sample_topology assert handler is mock_handler_cls.return_value @@ -597,6 +730,19 @@ def test_get_or_create_modes(self, dedicated): assert h1 is h2 assert mock_handler_cls.call_count == 1 + def test_failed_connection_candidate_is_closed_before_publish(self): + config = ConnectionConfig.from_uri("http://localhost:19530", token="test") + mgr = ConnectionManager.get_instance() + handler = _make_sync_handler() + handler._wait_for_channel_ready.side_effect = RuntimeError("not ready") + + with patch("pymilvus.client.grpc_handler.GrpcHandler", return_value=handler): + with pytest.raises(RuntimeError, match="not ready"): + mgr.get_or_create(config, client=Mock()) + + handler.close.assert_called_once() + assert mgr._registry == {} + def test_release_removes_client_reference(self, mock_grpc_handler): """Test release removes client from managed connection.""" config = ConnectionConfig.from_uri("http://localhost:19530", token="test") @@ -611,6 +757,28 @@ def test_release_removes_client_reference(self, mock_grpc_handler): mgr.release(handler, client=client) assert client not in managed.clients + assert mgr._get_managed(handler) is None + mock_grpc_handler.close.assert_called_once() + + def test_failed_recovery_does_not_add_candidate_owner(self): + config = ConnectionConfig.from_uri("http://localhost:19530", token="test") + mgr = ConnectionManager.get_instance() + original_owner = Mock() + candidate_owner = Mock() + + with patch("pymilvus.client.grpc_handler.GrpcHandler") as handler_cls: + handler_cls.return_value = _make_sync_handler() + handler = mgr.get_or_create(config, client=original_owner) + managed = mgr._get_managed(handler) + managed.last_used_at = time.time() - IDLE_THRESHOLD_SECONDS - 1 + with patch.object(mgr, "_check_health", return_value=False), patch.object( + mgr, "_recover", side_effect=RuntimeError("recovery failed") + ): + with pytest.raises(RuntimeError, match="recovery failed"): + mgr.get_or_create(config, client=candidate_owner) + + assert original_owner in managed.clients + assert candidate_owner not in managed.clients def test_close_all(self): """Test close_all closes all connections.""" @@ -841,6 +1009,46 @@ def bump_gen_then_return_true(m): # reconnect should NOT have been called - concurrent recovery detected handler.reconnect.assert_not_called() + def test_handle_error_does_not_resurrect_last_client_released_handler(self): + """A recovery decision made outside the lock cannot revive a removed handler.""" + mgr = ConnectionManager.get_instance() + config = ConnectionConfig.from_uri("http://localhost:19530", token="test") + handler = _make_sync_handler() + client = Mock() + entered = threading.Event() + proceed = threading.Event() + result = {} + + with patch("pymilvus.client.grpc_handler.GrpcHandler", return_value=handler): + mgr.get_or_create(config, client=client) + managed = mgr._get_managed(handler) + + def wait_until_released(_managed): + entered.set() + assert proceed.wait(timeout=2) + return True + + managed.strategy.on_unavailable = wait_until_released + + recovery = threading.Thread( + target=lambda: result.setdefault( + "retry", mgr.handle_error(handler, _MockRpcError()) + ) + ) + recovery.start() + assert entered.wait(timeout=2) + + mgr.release(handler, client=client) + assert mgr._get_managed(handler) is None + handler.close.assert_called_once() + + proceed.set() + recovery.join(timeout=2) + + assert not recovery.is_alive() + assert result == {"retry": False} + handler.reconnect.assert_not_called() + def test_dedicated_connection_findable_after_recovery(self): """Test that after in-place recovery, dedicated connection is still findable.""" mgr = ConnectionManager.get_instance() @@ -955,6 +1163,7 @@ def test_create_handler_fetches_topology_and_starts_refresher( token="mytoken", db_name="", secure=True, + _client_owned_telemetry=True, ) mock_refresher.start.assert_called_once() assert handler is mock_handler_cls.return_value @@ -1081,6 +1290,20 @@ async def test_get_or_create_modes(self, mock_async_handler, dedicated): else: assert h1 is h2 + @pytest.mark.asyncio + async def test_failed_async_connection_candidate_is_closed_before_publish(self): + config = ConnectionConfig.from_uri("http://localhost:19530", token="test") + mgr = AsyncConnectionManager.get_instance() + handler = _make_async_handler() + handler.ensure_channel_ready.side_effect = RuntimeError("not ready") + + with patch("pymilvus.client.async_grpc_handler.AsyncGrpcHandler", return_value=handler): + with pytest.raises(RuntimeError, match="not ready"): + await mgr.get_or_create(config, client=Mock()) + + handler.close.assert_awaited_once() + assert mgr._registry == {} + @pytest.mark.asyncio async def test_release(self, mock_async_handler): """Test release removes client reference.""" @@ -1098,6 +1321,31 @@ async def test_release(self, mock_async_handler): await mgr.release(handler, client=client) assert client not in managed.clients + assert mgr._get_managed(handler) is None + mock_async_handler.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_failed_async_recovery_does_not_add_candidate_owner(self): + config = ConnectionConfig.from_uri("http://localhost:19530", token="test") + mgr = AsyncConnectionManager.get_instance() + original_owner = Mock() + candidate_owner = Mock() + + with patch("pymilvus.client.async_grpc_handler.AsyncGrpcHandler") as handler_cls: + handler_cls.return_value = _make_async_handler() + handler = await mgr.get_or_create(config, client=original_owner) + managed = mgr._get_managed(handler) + managed.last_used_at = time.time() - IDLE_THRESHOLD_SECONDS - 1 + with patch.object(mgr, "_check_health", AsyncMock(return_value=False)), patch.object( + mgr, + "_recover", + AsyncMock(side_effect=RuntimeError("recovery failed")), + ): + with pytest.raises(RuntimeError, match="recovery failed"): + await mgr.get_or_create(config, client=candidate_owner) + + assert original_owner in managed.clients + assert candidate_owner not in managed.clients @pytest.mark.asyncio async def test_get_or_create_adds_client_to_existing(self): @@ -2461,6 +2709,7 @@ def test_user_password_token_construction(self): ) client = MilvusClient(uri="http://localhost:19530", user="admin", password="secret") assert client._config.token == "admin:secret" + assert client._config.get_handler_kwargs()["user"] == "admin" client.close() def test_explicit_token_overrides_user_password(self): @@ -2499,16 +2748,22 @@ def test_get_connection_raises_when_closed(self): client._get_connection() def test_use_database_updates_config(self): - """Test use_database updates _config.db_name.""" + """Test use_database updates routing without replacing the live transport.""" with patch("pymilvus.client.grpc_handler.GrpcHandler") as mock_handler_cls: - mock_handler_cls.return_value = _make_sync_handler( + original_handler = _make_sync_handler( get_server_type=Mock(return_value="milvus"), describe_database=Mock(return_value={}), ) + mock_handler_cls.return_value = original_handler client = MilvusClient(uri="http://localhost:19530") + original_using = client._using assert client._config.db_name == "" client.use_database("mydb") assert client._config.db_name == "mydb" + assert client._handler is original_handler + assert client._using == original_using + assert mock_handler_cls.call_count == 1 + original_handler.reset_db_name.assert_not_called() client.close() def test_dedicated_kwarg(self): @@ -2537,12 +2792,16 @@ class TestAsyncMilvusClientChangedCode: def test_init_deferred_state(self): """Test __init__ sets deferred state without connecting.""" - client = AsyncMilvusClient(uri="http://localhost:19530", token="test") + with patch("pymilvus.milvus_client.async_milvus_client.asyncio.Lock") as mock_lock: + client = AsyncMilvusClient(uri="http://localhost:19530", token="test") + + mock_lock.assert_not_called() assert client._handler is None assert client._manager is None assert client._using is None assert client.is_self_hosted is None assert client._closed is False + assert client._lifecycle_lock is None assert client._config.address == "localhost:19530" assert client._config.token == "test" @@ -2550,6 +2809,7 @@ def test_user_password_token_construction(self): """Test that user/password are combined into token when no token given.""" client = AsyncMilvusClient(uri="http://localhost:19530", user="admin", password="secret") assert client._config.token == "admin:secret" + assert client._config.get_handler_kwargs()["user"] == "admin" def test_explicit_token_overrides_user_password(self): """Test that explicit token takes precedence over user/password.""" @@ -2603,6 +2863,50 @@ async def test_connect_idempotent(self): await client.close() + @pytest.mark.asyncio + async def test_concurrent_dedicated_connect_publishes_one_candidate(self): + client = AsyncMilvusClient(uri="http://localhost:19530", dedicated=True) + manager = Mock() + candidate = Mock(get_server_type=Mock(return_value="milvus")) + manager.get_or_create = AsyncMock(return_value=candidate) + manager.release = AsyncMock() + + with patch.object(AsyncConnectionManager, "get_instance", return_value=manager): + await asyncio.gather(client._connect(), client._connect()) + + manager.get_or_create.assert_awaited_once() + assert client._handler is candidate + await client.close() + + @pytest.mark.asyncio + async def test_close_waits_for_connect_and_releases_published_candidate(self): + client = AsyncMilvusClient(uri="http://localhost:19530", dedicated=True) + manager = Mock() + candidate = Mock(get_server_type=Mock(return_value="milvus")) + started = asyncio.Event() + allow_ready = asyncio.Event() + + async def get_or_create(*args, **kwargs): + started.set() + await allow_ready.wait() + return candidate + + manager.get_or_create = AsyncMock(side_effect=get_or_create) + manager.release = AsyncMock() + + with patch.object(AsyncConnectionManager, "get_instance", return_value=manager): + connect_task = asyncio.create_task(client._connect()) + await started.wait() + close_task = asyncio.create_task(client.close()) + await asyncio.sleep(0) + assert not close_task.done() + allow_ready.set() + await asyncio.gather(connect_task, close_task) + + assert client._closed is True + assert client._handler is None + manager.release.assert_awaited_once_with(candidate, client=client) + @pytest.mark.asyncio async def test_context_manager(self): """Test async with AsyncMilvusClient works.""" @@ -2695,19 +2999,26 @@ async def test_close_without_connect_is_safe(self): @pytest.mark.asyncio async def test_use_database_updates_config(self): - """Test use_database updates _config.db_name.""" + """Test use_database updates routing without replacing the live transport.""" client = AsyncMilvusClient(uri="http://localhost:19530") with patch("pymilvus.client.async_grpc_handler.AsyncGrpcHandler") as mock_handler_cls: - mock_handler_cls.return_value = _make_async_handler( + original_handler = _make_async_handler( get_server_type=Mock(return_value="milvus"), describe_database=AsyncMock(return_value={}), + reset_db_name=Mock(), ) + mock_handler_cls.return_value = original_handler await client._connect() + original_using = client._using assert client._config.db_name == "" await client.use_database("mydb") assert client._config.db_name == "mydb" + assert client._handler is original_handler + assert client._using == original_using + assert mock_handler_cls.call_count == 1 + original_handler.reset_db_name.assert_not_called() await client.close() @@ -2733,3 +3044,650 @@ async def mock_get_or_create(*args, **kwargs): await client.close() mock_manager.release.assert_called_once() + + +def test_sync_logical_clients_share_transport_but_isolate_telemetry_and_close(monkeypatch): + handler = _pooled_sync_handler() + pool = _LogicalSyncPool({"": handler}) + captured = [] + + class HeartbeatStub: + def ClientHeartbeat(self, request, **_kwargs): + captured.append(request) + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + monkeypatch.setattr( + ClientTelemetryManager, "start", lambda manager: setattr(manager, "_ready", True) + ) + with patch.object(ConnectionManager, "get_instance", return_value=pool): + client_a = MilvusClient(telemetry_config={"client_id": "logical-a", "sampling_rate": 1.0}) + client_b = MilvusClient(telemetry_config={"client_id": "logical-b", "sampling_rate": 1.0}) + + assert client_a._handler is handler is client_b._handler + assert handler.telemetry is None + assert handler._telemetry_interceptor is None + assert len(handler._client_telemetry_bindings) == 2 + assert client_a.get_telemetry() is client_a._telemetry + assert client_b.get_telemetry() is client_b._telemetry + assert client_a._telemetry is not client_b._telemetry + + with pytest.raises(TypeError): + client_a.insert("books", "invalid") + assert client_a._telemetry._collectors["Insert"].global_bucket.request_count == 1 + assert "Insert" not in client_b._telemetry._collectors + + client_a._telemetry._queue_reply(CommandReply("only-a", True)) + handler._rebind_telemetry_stub(HeartbeatStub()) + client_a._telemetry._send_heartbeat() + client_b._telemetry._send_heartbeat() + assert [request.client_info.reserved["client_id"] for request in captured] == [ + "logical-a", + "logical-b", + ] + assert [reply.command_id for reply in captured[0].command_replies] == ["only-a"] + assert not captured[1].command_replies + + manager_a = client_a._telemetry + client_a.close() + assert len(handler._client_telemetry_bindings) == 1 + assert manager_a._heartbeat_endpoint()[0] is None + + captured.clear() + handler._rebind_telemetry_stub(HeartbeatStub()) + client_b._telemetry._send_heartbeat() + assert captured[0].client_info.reserved["client_id"] == "logical-b" + assert client_b._telemetry.ready + client_b.close() + handler.close() + + +def test_sync_close_does_not_hold_lifecycle_lock_while_joining_telemetry_worker(): + client = MilvusClient.__new__(MilvusClient) + client._lifecycle_lock = threading.RLock() + client._config = ConnectionConfig.from_uri("http://localhost:19530") + client._handler = object() + client._manager = Mock() + client._bind_telemetry_handler = Mock() + client._unbind_telemetry_handler = Mock() + client.describe_database = Mock(return_value={}) + + command_entered = threading.Event() + allow_lifecycle_call = threading.Event() + command = common_pb2.ClientCommand( + command_id="lifecycle-command", + command_type="custom-lifecycle", + create_time=1, + ) + + class CommandStub: + @staticmethod + def ClientHeartbeat(*_args, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), commands=[command] + ) + + manager = ClientTelemetryManager(CommandStub, TelemetryConfig(enabled=True)) + client._telemetry = manager + + def lifecycle_handler(local_command): + command_entered.set() + assert allow_lifecycle_call.wait(timeout=2) + client.use_database("analytics") + return CommandReply(local_command.command_id, True) + + manager.register_command_handler("custom-lifecycle", lifecycle_handler) + manager.start() + assert command_entered.wait(timeout=1) + + close_thread = threading.Thread(target=client.close, daemon=True) + close_thread.start() + assert manager._stop_event.wait(timeout=1) + allow_lifecycle_call.set() + close_thread.join(timeout=2) + + assert not close_thread.is_alive() + assert client._handler is None + assert client._config.db_name == "analytics" + client._manager.release.assert_called_once() + + +def test_sync_client_gc_wakes_telemetry_worker_without_retaining_owner(): + client = MilvusClient.__new__(MilvusClient) + client._config = ConnectionConfig.from_uri("http://localhost:19530") + client._handler = None + manager = client._new_telemetry_manager() + client._telemetry = manager + owner = weakref.ref(client) + worker_waiting = threading.Event() + original_wait = manager._wait_for_stop + + def wait_for_stop(delay): + worker_waiting.set() + return original_wait(delay) + + manager._wait_for_stop = wait_for_stop + + manager.start() + worker = manager._thread + assert worker_waiting.wait(timeout=2) + del client + gc.collect() + worker.join(timeout=2) + + assert owner() is None + assert not worker.is_alive() + + +def test_sync_use_database_retains_logical_state_and_failed_candidate_is_atomic(monkeypatch): + old_handler = _pooled_sync_handler() + pool = _LogicalSyncPool({"": old_handler}) + monkeypatch.setattr( + ClientTelemetryManager, "start", lambda manager: setattr(manager, "_ready", True) + ) + with patch.object(ConnectionManager, "get_instance", return_value=pool): + client_a = MilvusClient(telemetry_config={"client_id": "stable-a"}) + client_b = MilvusClient(telemetry_config={"client_id": "stable-b"}) + + telemetry = client_a._telemetry + telemetry.record_operation("Search", "books", time.perf_counter()) + telemetry._create_snapshot() + snapshots = telemetry.get_metrics_snapshots() + telemetry._handle_push_config( + ClientCommand("rate", "push_config", payload=b'{"sampling_rate":0.25}') + ) + client_a.describe_database = Mock(return_value={}) + + client_a.use_database("analytics") + assert client_a._telemetry is telemetry + assert telemetry.client_id == "stable-a" + assert telemetry.get_metrics_snapshots() == snapshots + assert telemetry._config.sampling_rate == 0.25 + assert telemetry._build_client_info().reserved["db_name"] == "analytics" + assert client_a._handler is old_handler is client_b._handler + assert client_b._telemetry.client_id == "stable-b" + assert len(old_handler._client_telemetry_bindings) == 2 + + bound_stub, _, bound_database = telemetry._heartbeat_endpoint() + original_bind = telemetry.bind_transport + + def partially_bind_then_reject(stub, database, binding_token): + original_bind(stub, database, binding_token) + if database == "broken": + raise RuntimeError("binding failed") + + monkeypatch.setattr(telemetry, "bind_transport", partially_bind_then_reject) + with pytest.raises(RuntimeError, match="binding failed"): + client_a.use_database("broken") + assert client_a._handler is old_handler + assert client_a._config.db_name == "analytics" + assert client_a._telemetry is telemetry + restored_stub, _, restored_database = telemetry._heartbeat_endpoint() + assert restored_stub is bound_stub + assert bound_database == restored_database == "analytics" + binding_token, binding_database = old_handler._client_telemetry_bindings[telemetry] + assert binding_token is telemetry._transport_binding_token + assert binding_database == "analytics" + assert len(old_handler._client_telemetry_bindings) == 2 + + client_a.close() + client_b.close() + old_handler.close() + + +def test_sync_use_database_keeps_transport_for_in_flight_operation(monkeypatch): + handler = _pooled_sync_handler() + pool = _LogicalSyncPool({"": handler}) + started = threading.Event() + proceed = threading.Event() + observed = [] + + def insert_rows(_collection, _data, **kwargs): + observed.append(kwargs["context"].get_db_name()) + if len(observed) == 1: + started.set() + assert proceed.wait(timeout=2) + return Mock(insert_count=1, primary_keys=[1], cost=0) + + monkeypatch.setattr(handler, "insert_rows", insert_rows) + monkeypatch.setattr( + ClientTelemetryManager, "start", lambda manager: setattr(manager, "_ready", True) + ) + with patch.object(ConnectionManager, "get_instance", return_value=pool): + client = MilvusClient(telemetry_config={"enabled": False}) + client.describe_database = Mock(return_value={}) + result = {} + + operation = threading.Thread( + target=lambda: result.setdefault("value", client.insert("books", [{"id": 1}])) + ) + operation.start() + assert started.wait(timeout=2) + + client.use_database("analytics") + assert client._handler is handler + assert pool.releases == [] + new_result = client.insert("books", [{"id": 2}]) + assert new_result["insert_count"] == 1 + assert operation.is_alive() + assert observed == ["", "analytics"] + handler._channel.close.assert_not_called() + + proceed.set() + operation.join(timeout=2) + assert not operation.is_alive() + assert result["value"]["insert_count"] == 1 + assert observed == ["", "analytics"] + + client.close() + assert pool.releases == [(handler, client)] + handler.close() + + +def test_transport_binding_token_rejects_late_old_reconnect(): + old_handler = _pooled_sync_handler() + new_handler = _pooled_sync_handler("analytics") + old_stub = object() + new_stub = object() + late_old_stub = object() + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + old_handler._rebind_telemetry_stub(old_stub) + new_handler._rebind_telemetry_stub(new_stub) + old_handler.register_client_telemetry(manager, "") + new_handler.register_client_telemetry(manager, "analytics") + old_handler._rebind_telemetry_stub(late_old_stub) + old_handler.unregister_client_telemetry(manager) + + stub, _, database = manager._heartbeat_endpoint() + assert stub is new_stub + assert database == "analytics" + new_handler.close() + old_handler.close() + + +def _assert_registration_serializes_with_rebind(handler, manager): + old_stub = object() + replacement_stub = object() + bind_entered = threading.Event() + release_bind = threading.Event() + rebind_attempted = threading.Event() + rebind_done = threading.Event() + errors = [] + original_bind = manager.bind_transport + + def blocking_bind(stub, database, binding_token): + bind_entered.set() + if not release_bind.wait(timeout=2): + raise TimeoutError("test did not release telemetry registration") + original_bind(stub, database, binding_token) + + def register(): + try: + handler.register_client_telemetry(manager, "books") + except BaseException as exc: + errors.append(exc) + + def rebind(): + rebind_attempted.set() + try: + handler._rebind_telemetry_stub(replacement_stub) + except BaseException as exc: + errors.append(exc) + finally: + rebind_done.set() + + handler._rebind_telemetry_stub(old_stub) + manager.bind_transport = blocking_bind + register_thread = threading.Thread(target=register) + rebind_thread = threading.Thread(target=rebind) + register_thread.start() + assert bind_entered.wait(timeout=2) + rebind_thread.start() + assert rebind_attempted.wait(timeout=2) + try: + # Manager callbacks run outside the handler lock, so reconnect cannot + # deadlock behind a blocked register callback. Register's generation + # stabilization must still observe this completed publication afterward. + assert rebind_done.wait(timeout=2) + finally: + release_bind.set() + register_thread.join(timeout=2) + rebind_thread.join(timeout=2) + + assert not register_thread.is_alive() + assert not rebind_thread.is_alive() + assert not errors + stub, _, database = manager._heartbeat_endpoint() + assert stub is replacement_stub + assert database == "books" + handler.unregister_client_telemetry(manager) + + +def test_sync_registration_and_reconnect_publish_one_atomic_binding(): + handler = _pooled_sync_handler() + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + _assert_registration_serializes_with_rebind(handler, manager) + + handler.close() + + +@pytest.mark.asyncio +async def test_async_registration_and_reconnect_publish_one_atomic_binding(): + handler = _pooled_async_handler() + manager = AsyncClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + _assert_registration_serializes_with_rebind(handler, manager) + + await handler.close() + + +def test_custom_command_lifecycle_does_not_invert_handler_and_endpoint_locks(): + handler = _pooled_sync_handler() + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + command_entered = threading.Event() + allow_lifecycle = threading.Event() + reconnect_called_manager = threading.Event() + errors = [] + command = common_pb2.ClientCommand( + command_id="lifecycle", + command_type="custom-lifecycle", + create_time=1, + ) + + class CommandStub: + def ClientHeartbeat(self, _request, **_kwargs): + return milvus_pb2.ClientHeartbeatResponse( + status=common_pb2.Status(), commands=[command] + ) + + replacement_stub = object() + handler._rebind_telemetry_stub(CommandStub()) + handler.register_client_telemetry(manager, "books") + + def lifecycle_handler(local_command): + command_entered.set() + if not allow_lifecycle.wait(timeout=2): + raise TimeoutError("test did not allow lifecycle callback") + handler.unregister_client_telemetry(manager) + handler.register_client_telemetry(manager, "books") + return CommandReply(local_command.command_id, True) + + manager.register_command_handler("custom-lifecycle", lifecycle_handler) + original_rebind = manager.rebind_transport + + def tracking_rebind(stub, binding_token): + reconnect_called_manager.set() + return original_rebind(stub, binding_token) + + manager.rebind_transport = tracking_rebind + + def heartbeat(): + try: + manager._send_heartbeat() + except BaseException as exc: + errors.append(exc) + + def reconnect(): + try: + handler._rebind_telemetry_stub(replacement_stub) + except BaseException as exc: + errors.append(exc) + + heartbeat_thread = threading.Thread(target=heartbeat) + reconnect_thread = threading.Thread(target=reconnect) + heartbeat_thread.start() + reconnect_started = False + try: + assert command_entered.wait(timeout=2) + reconnect_thread.start() + reconnect_started = True + assert reconnect_called_manager.wait(timeout=2) + finally: + allow_lifecycle.set() + heartbeat_thread.join(timeout=2) + if reconnect_started: + reconnect_thread.join(timeout=2) + + assert not heartbeat_thread.is_alive() + assert not reconnect_started or not reconnect_thread.is_alive() + assert not errors + assert manager._heartbeat_endpoint()[0] is replacement_stub + handler.unregister_client_telemetry(manager) + handler.close() + + +def test_global_client_telemetry_keeps_logical_identity_across_failover(monkeypatch): + handler = _pooled_sync_handler("sales") + handler._address = "physical-primary.example.com:19530" + first_stub = object() + replacement_stub = object() + handler._rebind_telemetry_stub(first_stub) + pool = _LogicalSyncPool({"sales": handler}) + monkeypatch.setattr( + ClientTelemetryManager, "start", lambda manager: setattr(manager, "_ready", True) + ) + + with patch.object(ConnectionManager, "get_instance", return_value=pool): + client = MilvusClient( + uri="https://alice:secret@global-cluster.example.com:19530/sales", + telemetry_config={"client_id": "global-logical"}, + ) + + config = client._telemetry._config_provider() + assert config["address"] == "global-cluster.example.com:19530" + assert config["username"] == "alice" + assert config["db_name"] == "sales" + assert config["secure"] is True + assert all("secret" not in str(value) for value in config.values()) + info = client._telemetry._build_client_info() + assert info.user == "alice" + assert info.reserved["db_name"] == "sales" + + _, generation, _ = client._telemetry._heartbeat_endpoint() + client._bind_telemetry_handler(handler, "sales") + assert client._telemetry._heartbeat_endpoint()[1] == generation + + handler._rebind_telemetry_stub(replacement_stub) + rebound_stub, rebound_generation, database = client._telemetry._heartbeat_endpoint() + assert rebound_stub is replacement_stub + assert rebound_generation == generation + 1 + assert database == "sales" + assert client._telemetry._config_provider()["address"] == ("global-cluster.example.com:19530") + + client.close() + handler.close() + + +@pytest.mark.asyncio +async def test_async_logical_clients_share_transport_use_database_and_close(monkeypatch): + old_handler = _pooled_async_handler() + pool = _LogicalAsyncPool({"": old_handler}) + captured = [] + + class HeartbeatStub: + async def ClientHeartbeat(self, request, **_kwargs): + captured.append(request) + return milvus_pb2.ClientHeartbeatResponse(status=common_pb2.Status()) + + monkeypatch.setattr( + AsyncClientTelemetryManager, + "start", + lambda manager: setattr(manager, "_ready", True), + ) + with patch.object(AsyncConnectionManager, "get_instance", return_value=pool): + client_a = AsyncMilvusClient(telemetry_config={"client_id": "async-a"}) + client_b = AsyncMilvusClient(telemetry_config={"client_id": "async-b"}) + await asyncio.gather(client_a._connect(), client_b._connect()) + + assert client_a._handler is old_handler is client_b._handler + assert old_handler.telemetry is None + assert len(old_handler._client_telemetry_bindings) == 2 + with pytest.raises(TypeError): + await client_a.insert("books", "invalid") + assert client_a._telemetry._collectors["Insert"].global_bucket.request_count == 1 + assert "Insert" not in client_b._telemetry._collectors + + old_handler._rebind_telemetry_stub(HeartbeatStub()) + await client_a._telemetry._send_heartbeat_async() + await client_b._telemetry._send_heartbeat_async() + assert [request.client_info.reserved["client_id"] for request in captured] == [ + "async-a", + "async-b", + ] + + telemetry = client_a._telemetry + telemetry.record_operation("Search", "books", time.perf_counter()) + telemetry._create_snapshot() + snapshots = telemetry.get_metrics_snapshots() + client_a.describe_database = AsyncMock(return_value={}) + await client_a.use_database("analytics") + assert client_a._telemetry is telemetry + assert telemetry.client_id == "async-a" + assert telemetry.get_metrics_snapshots() == snapshots + assert telemetry._build_client_info().reserved["db_name"] == "analytics" + assert client_a._handler is old_handler is client_b._handler + assert len(old_handler._client_telemetry_bindings) == 2 + + bound_stub, _, bound_database = telemetry._heartbeat_endpoint() + original_bind = telemetry.bind_transport + + def partially_bind_then_reject(stub, database, binding_token): + original_bind(stub, database, binding_token) + if database == "broken": + raise RuntimeError("binding failed") + + monkeypatch.setattr(telemetry, "bind_transport", partially_bind_then_reject) + with pytest.raises(RuntimeError, match="binding failed"): + await client_a.use_database("broken") + assert client_a._handler is old_handler + assert client_a._config.db_name == "analytics" + restored_stub, _, restored_database = telemetry._heartbeat_endpoint() + assert restored_stub is bound_stub + assert bound_database == restored_database == "analytics" + binding_token, binding_database = old_handler._client_telemetry_bindings[telemetry] + assert binding_token is telemetry._transport_binding_token + assert binding_database == "analytics" + + await client_a.close() + assert len(old_handler._client_telemetry_bindings) == 1 + old_handler._rebind_telemetry_stub(HeartbeatStub()) + captured.clear() + await client_b._telemetry._send_heartbeat_async() + assert captured[0].client_info.reserved["client_id"] == "async-b" + await client_b.close() + await old_handler.close() + + +@pytest.mark.asyncio +async def test_async_client_gc_cancels_telemetry_worker_without_retaining_owner(): + client = AsyncMilvusClient.__new__(AsyncMilvusClient) + client._config = ConnectionConfig.from_uri("http://localhost:19530") + client._handler = None + manager = client._new_telemetry_manager() + client._telemetry = manager + owner = weakref.ref(client) + worker_waiting = asyncio.Event() + original_sleep = manager._sleep_until_next_heartbeat + + async def sleep_until_next_heartbeat(delay): + worker_waiting.set() + await original_sleep(delay) + + manager._sleep_until_next_heartbeat = sleep_until_next_heartbeat + + manager.start() + worker = manager._task + await asyncio.wait_for(worker_waiting.wait(), timeout=2) + del client + gc.collect() + with contextlib.suppress(asyncio.CancelledError): + await asyncio.wait_for(worker, timeout=2) + + assert owner() is None + assert worker.done() + + +@pytest.mark.asyncio +async def test_async_use_database_keeps_transport_for_in_flight_operation(monkeypatch): + handler = _pooled_async_handler() + pool = _LogicalAsyncPool({"": handler}) + started = asyncio.Event() + proceed = asyncio.Event() + observed = [] + + async def insert_rows(_collection, _data, **kwargs): + observed.append(kwargs["context"].get_db_name()) + if len(observed) == 1: + started.set() + await proceed.wait() + return Mock(insert_count=1, primary_keys=[1], cost=0) + + monkeypatch.setattr(handler, "insert_rows", insert_rows) + monkeypatch.setattr( + AsyncClientTelemetryManager, + "start", + lambda manager: setattr(manager, "_ready", True), + ) + with patch.object(AsyncConnectionManager, "get_instance", return_value=pool): + client = AsyncMilvusClient(telemetry_config={"enabled": False}) + await client._connect() + client.describe_database = AsyncMock(return_value={}) + + operation = asyncio.create_task(client.insert("books", [{"id": 1}])) + await asyncio.wait_for(started.wait(), timeout=2) + + await client.use_database("analytics") + assert client._handler is handler + assert pool.releases == [] + new_result = await client.insert("books", [{"id": 2}]) + assert new_result["insert_count"] == 1 + assert not operation.done() + assert observed == ["", "analytics"] + handler._async_channel.close.assert_not_awaited() + + proceed.set() + result = await asyncio.wait_for(operation, timeout=2) + assert result["insert_count"] == 1 + assert observed == ["", "analytics"] + + await client.close() + assert pool.releases == [(handler, client)] + await handler.close() + + +@pytest.mark.asyncio +async def test_legacy_direct_handlers_keep_one_fallback_manager_and_interceptor(): + sync_handler = GrpcHandler(channel=Mock()) + async_channel = Mock() + async_channel._unary_unary_interceptors = [] + async_channel.close = AsyncMock() + async_handler = AsyncGrpcHandler(channel=async_channel) + + assert sync_handler.telemetry is not None + assert sync_handler._telemetry_interceptor is not None + assert async_handler.telemetry is not None + assert async_handler._telemetry_interceptor is not None + assert len(sync_handler._client_telemetry_bindings) == 0 + assert len(async_handler._client_telemetry_bindings) == 0 + + sync_handler.close() + await async_handler.close() + + +@pytest.mark.asyncio +async def test_async_client_close_releases_transport_after_telemetry_stop_failure(): + client = AsyncMilvusClient(uri="http://localhost:19530") + manager = Mock() + manager.release = AsyncMock() + handler = Mock() + handler.unregister_client_telemetry = Mock(side_effect=RuntimeError("unbind failed")) + client._manager = manager + client._handler = handler + client._telemetry.stop_async = AsyncMock(side_effect=RuntimeError("telemetry failed")) + + await client.close() + + handler.unregister_client_telemetry.assert_called_once_with(client._telemetry) + manager.release.assert_awaited_once_with(handler, client=client) + assert client._handler is None diff --git a/tests/unit/test_grpc_channel_options.py b/tests/unit/test_grpc_channel_options.py index eabac5a9b..cfc6d31a2 100644 --- a/tests/unit/test_grpc_channel_options.py +++ b/tests/unit/test_grpc_channel_options.py @@ -7,6 +7,10 @@ from grpc._cython import cygrpc from pymilvus.client.async_grpc_handler import AsyncGrpcHandler from pymilvus.client.grpc_handler import GrpcHandler +from pymilvus.client.telemetry import ( + AsyncTelemetryUnaryUnaryInterceptor, + TelemetryUnaryUnaryInterceptor, +) # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -20,6 +24,7 @@ def _create_sync_handler(grpc_options=None, secure=False): handler._address = "localhost:19530" handler._log_level = None handler._authorization_interceptor = None + handler._telemetry_interceptor = TelemetryUnaryUnaryInterceptor(MagicMock()) if secure: handler._server_name = "" handler._server_pem_path = "" @@ -38,6 +43,7 @@ def _create_async_handler(grpc_options=None): handler._address = "localhost:19530" handler._log_level = None handler._async_authorization_interceptor = None + handler._telemetry_interceptor = AsyncTelemetryUnaryUnaryInterceptor(MagicMock()) return handler diff --git a/tests/unit/test_iterator_ownership.py b/tests/unit/test_iterator_ownership.py index 6f3bcecc2..26694df18 100644 --- a/tests/unit/test_iterator_ownership.py +++ b/tests/unit/test_iterator_ownership.py @@ -25,6 +25,11 @@ ) from pymilvus.client.iterator import QueryIterator as ClientQueryIterator from pymilvus.client.iterator import QueryIteratorCursor as ClientQueryIteratorCursor +from pymilvus.client.telemetry import ( + ClientTelemetryManager, + TelemetryConfig, + telemetry_operation, +) from pymilvus.client.types import DataType from pymilvus.exceptions import MilvusException, ServerVersionIncompatibleException from pymilvus.milvus_client import milvus_client as milvus_client_module @@ -185,6 +190,35 @@ def test_query_iterator_forwards_context_during_offset_seek(): assert handler.query_calls[1][1]["iterator"] == "False" +def test_query_iterator_internal_setup_seek_and_pages_do_not_count_as_query_operations(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler(_SwitchingQueryHandler): + def __init__(self): + super().__init__() + self._telemetry = manager + + @telemetry_operation("Query") + def query(self, collection_name, **kwargs): + return super().query(collection_name, **kwargs) + + handler = Handler() + iterator = ClientQueryIterator( + handler=handler, + context=CallContext(db_name="db"), + collection_name="collection", + batch_size=10, + expr="pk > 0", + output_fields=["pk"], + schema=_QUERY_SCHEMA, + rpc_options={OFFSET: 1}, + ) + + assert iterator.next() == [{"pk": 2, "transport": "before"}] + assert len(handler.query_calls) == 3 + assert "Query" not in manager._collectors + + def test_milvus_client_public_seam_returns_shared_iterator_page(): handler = _SwitchingQueryHandler() context = CallContext(db_name="db", client_request_id="request") @@ -356,6 +390,28 @@ def test_search_v1_keeps_handler_identity_and_context_across_pages(): assert handler.close_calls == 0 +def test_search_v1_internal_init_and_pages_do_not_count_as_search_operations(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler(_SwitchingSearchV1Handler): + def __init__(self): + super().__init__() + self._telemetry = manager + + @telemetry_operation("Search") + def search(self, collection_name, **kwargs): + return super().search(collection_name, **kwargs) + + handler = Handler() + iterator = _new_search_v1_iterator(handler, CallContext(db_name="db")) + + iterator.next() + iterator.next() + + assert len(handler.search_calls) == 2 + assert "Search" not in manager._collectors + + def test_collection_search_iterator_uses_shared_v1_owner(): handler = _SwitchingSearchV1Handler() context = CallContext(db_name="db", client_request_id="request") @@ -525,6 +581,27 @@ def test_search_v2_keeps_handler_identity_and_context_across_probe_and_pages(): assert handler.close_calls == 0 +def test_search_v2_internal_probe_and_pages_do_not_count_as_search_operations(): + manager = ClientTelemetryManager(lambda: None, TelemetryConfig(enabled=True)) + + class Handler(_SwitchingSearchV2Handler): + def __init__(self): + super().__init__() + self._telemetry = manager + + @telemetry_operation("Search") + def search(self, **kwargs): + return super().search(**kwargs) + + handler = Handler() + iterator = _new_search_v2_iterator(handler, CallContext(db_name="db")) + + iterator.next() + + assert len(handler.search_calls) == 2 + assert "Search" not in manager._collectors + + def test_milvus_client_uses_shared_search_v2_owner_and_one_context(): handler = _SwitchingSearchV2Handler() context = CallContext(db_name="db", client_request_id="request") diff --git a/tests/unit/test_repro_issues.py b/tests/unit/test_repro_issues.py index 5ea0c47b0..90354d962 100644 --- a/tests/unit/test_repro_issues.py +++ b/tests/unit/test_repro_issues.py @@ -412,6 +412,9 @@ async def close(self, grace=None): async def channel_ready(self): return None + def unary_unary(self, *args, **kwargs): + return MagicMock() + class _Issue3541AsyncStub: def __init__(self, channel, rpc_entered=None, release_rpc=None):