From 60bfc77a505dd0072c5d22eb59e5d2faa2ed9f65 Mon Sep 17 00:00:00 2001 From: xiaofanluan Date: Tue, 25 Aug 2026 12:19:21 -0700 Subject: [PATCH] feat: add client telemetry support Signed-off-by: xiaofanluan --- milvus/grpc/BaseClient.ts | 14 +- milvus/grpc/Data.ts | 93 +- milvus/grpc/GrpcClient.ts | 472 +++++-- milvus/index.ts | 1 + milvus/telemetry/ClientTelemetry.ts | 1285 ++++++++++++++++++++ milvus/telemetry/index.ts | 1 + milvus/types/Client.ts | 4 + milvus/utils/Function.ts | 292 ++++- milvus/utils/Grpc.ts | 3 +- milvus/utils/Search.ts | 20 +- test/grpc/Basic.spec.ts | 12 +- test/grpc/MilvusClient.spec.ts | 11 + test/telemetry/ClientTelemetry.e2e.spec.ts | 214 ++++ test/telemetry/ClientTelemetry.spec.ts | 946 ++++++++++++++ test/telemetry/GrpcClientTelemetry.spec.ts | 68 ++ test/utils/Collection.spec.ts | 16 +- test/utils/Data.spec.ts | 205 ++++ test/utils/Function.spec.ts | 386 +++++- test/utils/GlobalConnection.spec.ts | 273 ++++- 19 files changed, 4058 insertions(+), 258 deletions(-) create mode 100644 milvus/telemetry/ClientTelemetry.ts create mode 100644 milvus/telemetry/index.ts create mode 100644 test/telemetry/ClientTelemetry.e2e.spec.ts create mode 100644 test/telemetry/ClientTelemetry.spec.ts create mode 100644 test/telemetry/GrpcClientTelemetry.spec.ts diff --git a/milvus/grpc/BaseClient.ts b/milvus/grpc/BaseClient.ts index 32f8a9ba..c7d81a13 100644 --- a/milvus/grpc/BaseClient.ts +++ b/milvus/grpc/BaseClient.ts @@ -60,7 +60,7 @@ export class BaseClient { // Mutex flag to serialize failover reconnections protected isReconnecting: boolean = false; // Promise that concurrent callers can await during reconnection - protected reconnectingPromise: Promise | null = null; + protected reconnectingPromise: Promise | null = null; // ChannelCredentials object used for authenticating the client on the gRPC channel. protected creds!: ChannelCredentials; @@ -207,18 +207,18 @@ export class BaseClient { const rootCertBuff: Buffer | null = rootCert ? rootCert : rootCertPath - ? readFileSync(rootCertPath) - : null; + ? readFileSync(rootCertPath) + : null; const privateKeyBuff: Buffer | null = privateKey ? privateKey : privateKeyPath - ? readFileSync(privateKeyPath) - : null; + ? readFileSync(privateKeyPath) + : null; const certChainBuff: Buffer | null = certChain ? certChain : certChainPath - ? readFileSync(certChainPath) - : null; + ? readFileSync(certChainPath) + : null; this.creds = credentials.createSsl( rootCertBuff, privateKeyBuff, diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index a45e246e..62b63b2b 100644 --- a/milvus/grpc/Data.ts +++ b/milvus/grpc/Data.ts @@ -83,6 +83,9 @@ import { FieldPartialUpdateOpType, FieldPartialUpdateOp, CLUSTER_ID, + isHybridSearchRequest, + withTelemetryLogicalOperation, + withTelemetrySuppressed, } from '../'; import { Collection } from './Collection'; @@ -128,14 +131,18 @@ export class Data extends Collection { * Upsert data into Milvus, view _insert for detail */ async upsert(data: UpsertReq): Promise { - return this._insert(data, true); + return withTelemetryLogicalOperation(this.channelPool, 'Upsert', data, () => + this._insert(data, true) + ); } /** * Insert data into Milvus, view _insert for detail */ async insert(data: InsertReq): Promise { - return this._insert(data); + return withTelemetryLogicalOperation(this.channelPool, 'Insert', data, () => + this._insert(data) + ); } /** @@ -632,6 +639,14 @@ export class Data extends Collection { * ``` */ async deleteEntities(data: DeleteEntitiesReq): Promise { + return withTelemetryLogicalOperation(this.channelPool, 'Delete', data, () => + this._deleteEntities(data) + ); + } + + private async _deleteEntities( + data: DeleteEntitiesReq + ): Promise { if (!data || !data.collection_name) { throw new Error(ERROR_REASONS.COLLECTION_NAME_IS_REQUIRED); } @@ -695,6 +710,12 @@ export class Data extends Collection { * ``` */ async delete(data: DeleteReq): Promise { + return withTelemetryLogicalOperation(this.channelPool, 'Delete', data, () => + this._delete(data) + ); + } + + private async _delete(data: DeleteReq): Promise { if (!data || !data.collection_name) { throw new Error(ERROR_REASONS.COLLECTION_NAME_IS_REQUIRED); } @@ -798,6 +819,18 @@ export class Data extends Collection { async search( params: T ): Promise> { + const operation = isHybridSearchRequest(params) ? 'HybridSearch' : 'Search'; + return withTelemetryLogicalOperation( + this.channelPool, + operation, + params, + () => this._search(params) + ); + } + + private async _search< + T extends SearchReq | SearchSimpleReq | HybridSearchReq, + >(params: T): Promise> { // default collection request const describeCollectionRequest = { collection_name: params.collection_name, @@ -893,13 +926,16 @@ export class Data extends Collection { async searchIterator(param: SearchIteratorReq): Promise { const client = this; - // Get available count - const count = await client.count({ - collection_name: param.collection_name, - expr: param.expr || param.filter || '', - db_name: param.db_name, - cluster_id: param.cluster_id, - }); + // Iterators are not logical operations: suppress telemetry for the setup + // count and for every per-page internal search below. + const count = await withTelemetrySuppressed(() => + client.count({ + collection_name: param.collection_name, + expr: param.expr || param.filter || '', + db_name: param.db_name, + cluster_id: param.cluster_id, + }) + ); // get collection Info const collectionInfo = await this.describeCollection({ @@ -949,11 +985,13 @@ export class Data extends Collection { } try { - const batchRes = await client.search({ - ...param, - params, - limit: batchSize, - }); + const batchRes = await withTelemetrySuppressed(() => + client.search({ + ...param, + params, + limit: batchSize, + }) + ); // update current total and batch size currentTotal += batchRes.results.length; @@ -1019,13 +1057,16 @@ export class Data extends Collection { const pkField = isElementFilter ? collectionInfo!.schema.fields.find(field => field.is_primary_key)! : await this.getPkField(data); - // get count - const count = await client.count({ - collection_name: data.collection_name, - expr: userExpr, - db_name: data.db_name, - cluster_id: data.cluster_id, - }); + // Iterators are not logical operations: suppress telemetry for the setup + // count and for every per-page internal query below. + const count = await withTelemetrySuppressed(() => + client.count({ + collection_name: data.collection_name, + expr: userExpr, + db_name: data.db_name, + cluster_id: data.cluster_id, + }) + ); // remove filter field to avoid conflict with expr in query method const queryData = { ...data }; delete queryData.filter; @@ -1088,7 +1129,9 @@ export class Data extends Collection { } // search data - const res = await client.query(queryData as QueryReq); + const res = await withTelemetrySuppressed(() => + client.query(queryData as QueryReq) + ); if (!res.data.length) { return { done: true, value: null }; @@ -1236,6 +1279,12 @@ export class Data extends Collection { * ``` */ async query(data: QueryReq): Promise { + return withTelemetryLogicalOperation(this.channelPool, 'Query', data, () => + this._query(data) + ); + } + + private async _query(data: QueryReq): Promise { checkCollectionName(data); // Set up limits and offset for the query diff --git a/milvus/grpc/GrpcClient.ts b/milvus/grpc/GrpcClient.ts index 46970c08..2cfd5134 100644 --- a/milvus/grpc/GrpcClient.ts +++ b/milvus/grpc/GrpcClient.ts @@ -38,6 +38,9 @@ import { getPrimaryCluster, TopologyRefresher, setPoolFailoverHandler, + setPoolTelemetryManager, + ClientTelemetryManager, + withTelemetryLogicalOperation, } from '../'; import { User } from './User'; @@ -56,6 +59,23 @@ export const LOADER_OPTIONS = { export class GRPCClient extends User { // Store the gRPC service constructor for pool rebuild on failover private _MilvusService!: ServiceClientConstructor; + private _ClientTelemetryService!: ServiceClientConstructor; + private telemetryClient?: Client; + private telemetryChannelOptions!: ChannelOptions; + // Incremented whenever the live telemetry endpoint changes. The telemetry + // manager uses this epoch to discard a response that arrives from an old + // endpoint after a successful global-cluster failover. + private telemetryEndpointEpoch = 0; + private readonly telemetry: ClientTelemetryManager; + // Operation RPCs use the effective default database in their metadata, but telemetry + // must distinguish an omitted database from an explicitly selected database named + // "default". The server reserves an absent db_name for the former. + private telemetryDatabaseExplicit: boolean; + // Async topology discovery and candidate validation can complete after close(). Fence + // their publication so a released candidate can never resurrect a closed client. + private lifecycleGeneration = 0; + private closed = false; + private closePromise?: Promise; // Store sdkVersion for reconnection private _sdkVersion: string = ''; @@ -84,6 +104,31 @@ export class GRPCClient extends User { }, { ...LOADER_OPTIONS, ...this.config.loaderOptions } ); + this._ClientTelemetryService = getGRPCService( + { + serviceName: 'milvus.proto.milvus.ClientTelemetryService', + }, + { ...LOADER_OPTIONS, ...this.config.loaderOptions } + ); + this.telemetryDatabaseExplicit = Boolean(this.config.database); + + this.telemetry = new ClientTelemetryManager({ + sender: request => this.sendTelemetryHeartbeat(request), + config: this.config.telemetry, + userProvider: () => this.config.username || '', + databaseProvider: () => + this.telemetryDatabaseExplicit + ? this.metadata.get(METADATA.DATABASE) || DEFAULT_DB + : '', + configProvider: () => ({ + address: this.config.address, + username: this.config.username, + database: this.metadata.get(METADATA.DATABASE) || DEFAULT_DB, + ssl: this.config.ssl, + timeout: this.config.timeout, + }), + senderEpochProvider: () => this.telemetryEndpointEpoch, + }); // setup auth if necessary const auth = getAuthString(this.config); @@ -124,11 +169,16 @@ export class GRPCClient extends User { interceptors.push(getTraceInterceptor()); } - // add retry interceptor - interceptors.push(retryInterceptor); + // Heartbeats are best-effort: keep auth/request metadata/trace propagation, but do not + // let the ordinary RPC retry interceptor turn one heartbeat into a retry burst. The + // telemetry manager owns its own next-heartbeat/backoff policy. + this.telemetryChannelOptions = { + ...this.channelOptions, + interceptors: [...interceptors], + }; - // add interceptors - this.channelOptions.interceptors = interceptors; + // Ordinary Milvus RPCs retain the configured retry behavior. + this.channelOptions.interceptors = [...interceptors, retryInterceptor]; // For global cluster, skip pool creation here — pool will be created // in connect() after topology is fetched and primary endpoint is resolved. @@ -139,13 +189,28 @@ export class GRPCClient extends User { // create a grpc service client(connect) connect(sdkVersion: string) { + if (this.closed) { + return; + } + this.lifecycleGeneration += 1; + const lifecycleGeneration = this.lifecycleGeneration; this._sdkVersion = sdkVersion; + this.telemetry.setSdkVersion(sdkVersion); if (this.isGlobal) { // For global cluster: fetch topology → create pool → connect - this.connectPromise = this._initGlobalConnection(sdkVersion); + this.connectPromise = this._initGlobalConnection( + sdkVersion, + lifecycleGeneration + ); } else { // Normal connection - this.connectPromise = this._getServerInfo(sdkVersion); + this.replaceTelemetryClient(); + this.connectPromise = this._getServerInfo( + sdkVersion, + this.channelPool, + true, + lifecycleGeneration + ); } } @@ -153,7 +218,10 @@ export class GRPCClient extends User { * Initializes a global cluster connection. * Fetches topology, resolves primary endpoint, creates pool, starts refresher. */ - private async _initGlobalConnection(sdkVersion: string) { + private async _initGlobalConnection( + sdkVersion: string, + lifecycleGeneration = this.lifecycleGeneration + ) { const token = this.config.token || ''; logger.debug( @@ -162,6 +230,9 @@ export class GRPCClient extends User { // Fetch topology to discover primary cluster const topology = await fetchTopology(this.globalEndpoint, token); + if (!this.isLifecycleCurrent(lifecycleGeneration)) { + return; + } this.globalTopology = topology; // Resolve primary endpoint and create pool @@ -173,6 +244,7 @@ export class GRPCClient extends User { ); this.channelPool = this.createChannelPool(); + this.replaceTelemetryClient(); this._attachFailoverHandler(); // Start background topology refresher @@ -186,8 +258,18 @@ export class GRPCClient extends User { }); this.topologyRefresher.start(); - // Now connect to the primary - return this._getServerInfo(sdkVersion); + // Now connect to the primary. Validate without side effects first because close may + // run while the RPC is awaiting a response. + const serverInfo = await this._getServerInfo( + sdkVersion, + this.channelPool, + false + ); + if (!this.isLifecycleCurrent(lifecycleGeneration)) { + return serverInfo; + } + this.applyServerInfo(serverInfo); + return serverInfo; } /** @@ -196,21 +278,30 @@ export class GRPCClient extends User { * @returns true if primary changed and reconnection happened, false if primary unchanged */ async reconnectToPrimary(): Promise { + if (this.closed) { + return false; + } // Serialize concurrent failover attempts if (this.isReconnecting) { logger.debug( `\x1b[36m[Global]\x1b[0m Reconnect already in progress, waiting for completion` ); if (this.reconnectingPromise) { - await this.reconnectingPromise; + return this.reconnectingPromise; } - return true; + return false; } let primaryChanged = false; + const lifecycleGeneration = this.lifecycleGeneration; this.isReconnecting = true; this.reconnectingPromise = (async () => { + let candidatePool: + | ReturnType + | undefined; + let candidateTelemetryClient: Client | undefined; + let candidateRefresher: TopologyRefresher | undefined; try { const token = this.config.token || ''; @@ -220,6 +311,9 @@ export class GRPCClient extends User { // Fetch fresh topology const newTopology = await fetchTopology(this.globalEndpoint, token); + if (!this.isLifecycleCurrent(lifecycleGeneration)) { + return false; + } const newPrimary = getPrimaryCluster(newTopology); // Check if primary actually changed @@ -228,46 +322,41 @@ export class GRPCClient extends User { `\x1b[36m[Global]\x1b[0m Primary unchanged (${this.config.address}), no reconnect needed` ); this.globalTopology = newTopology; - return; // Primary hasn't changed, no reconnect needed + return false; // Primary hasn't changed, no reconnect needed } - primaryChanged = true; - logger.info( `Global cluster failover: ${this.config.address} -> ${newPrimary.endpoint}` ); - // Create new pool BEFORE draining old pool, so if creation fails - // the old pool is still usable + // Build and validate a complete candidate lifecycle without mutating the live + // address, pool, telemetry transport, topology, or connection status. In + // particular, pool factories must capture the candidate address rather than + // reading this.config.address later when generic-pool creates a client. logger.debug( `\x1b[36m[Global]\x1b[0m Creating new channel pool for ${newPrimary.endpoint}` ); - const oldPool = this.channelPool; - this.config.address = newPrimary.endpoint; - this.channelPool = this.createChannelPool(); - this._attachFailoverHandler(); - - // Now drain old pool (non-critical, best-effort) - if (oldPool) { - logger.debug( - `\x1b[36m[Global]\x1b[0m Draining old channel pool` - ); + candidatePool = this.createChannelPool(newPrimary.endpoint); + candidateTelemetryClient = this.createTelemetryClient( + newPrimary.endpoint + ); + const candidateServerInfo = await this._getServerInfo( + this._sdkVersion, + candidatePool, + false + ); + if (!this.isLifecycleCurrent(lifecycleGeneration)) { try { - await oldPool.drain(); - await oldPool.clear(); + candidateTelemetryClient.close(); } catch { - // ignore cleanup errors on old pool + // ignore cleanup errors on an unpublished stale candidate } + candidateTelemetryClient = undefined; + await this.disposeChannelPool(candidatePool); + candidatePool = undefined; + return false; } - - // Update state - this.globalTopology = newTopology; - - // Update topology refresher - if (this.topologyRefresher) { - this.topologyRefresher.stop(); - } - this.topologyRefresher = new TopologyRefresher({ + candidateRefresher = new TopologyRefresher({ globalEndpoint: this.globalEndpoint, token, topology: newTopology, @@ -275,28 +364,106 @@ export class GRPCClient extends User { this.globalTopology = t; }, }); - this.topologyRefresher.start(); + // Attaching the handler only mutates the isolated candidate pool and makes it + // ready for publication without exposing it to ordinary operations yet. + this._attachFailoverHandler(candidatePool); + + const oldAddress = this.config.address; + const oldPool = this.channelPool; + const oldTelemetryClient = this.telemetryClient; + const oldRefresher = this.topologyRefresher; - // Re-establish server info - this.connectStatus = CONNECT_STATUS.CONNECTING; - await this._getServerInfo(this._sdkVersion); + // JavaScript runs this assignment block without interleaving another promise + // continuation. Advance the epoch before publishing the candidate telemetry + // client so any late response from the old endpoint is ignored as a whole. + this.telemetryEndpointEpoch += 1; + this.config.address = newPrimary.endpoint; + this.channelPool = candidatePool; + this.telemetryClient = candidateTelemetryClient; + this.applyServerInfo(candidateServerInfo, false); + this.globalTopology = newTopology; + this.topologyRefresher = candidateRefresher; + primaryChanged = true; + + // Candidate resources are now owned by the live lifecycle and must not be + // cleaned up by the failure path. + candidatePool = undefined; + candidateTelemetryClient = undefined; + candidateRefresher = undefined; + + // Everything below is post-commit cleanup/startup. None of it may enter the + // candidate-validation catch path and pretend the already-published lifecycle + // was rolled back. + try { + this.topologyRefresher?.start(); + } catch (error: any) { + logger.warn(`Failed to start topology refresher: ${error.message}`); + } + if (candidateServerInfo?.identifier) { + try { + this.telemetry.start(); + } catch (error: any) { + logger.warn(`Failed to start client telemetry: ${error.message}`); + } + } + try { + oldRefresher?.stop(); + } catch (error: any) { + logger.warn( + `Failed to stop old topology refresher: ${error.message}` + ); + } + try { + oldTelemetryClient?.close(); + } catch (error: any) { + logger.warn(`Failed to close old telemetry client: ${error.message}`); + } + + // Existing operations may still hold a client from the old pool. drain() waits + // for those borrowers before clear() closes the channels, so late operations can + // finish while all new operations use the newly published pool. + if (oldPool) { + logger.debug( + `\x1b[36m[Global]\x1b[0m Draining old channel pool for ${oldAddress}` + ); + await this.disposeChannelPool(oldPool); + } } catch (e: any) { + if (primaryChanged) { + // Publication is the commit point. An unexpected post-commit cleanup error + // must not close the new live resources or report a rollback that did not + // happen. + logger.warn( + `Global cluster failover committed with a cleanup error: ${e.message}` + ); + return true; + } logger.warn(`Global cluster failover failed: ${e.message}`); - // Clean up resources created during failed failover - if (this.topologyRefresher) { - this.topologyRefresher.stop(); - this.topologyRefresher = null; + // The live lifecycle was not touched unless the synchronous publication block + // completed. Candidate validation failures therefore leave the old address, + // pool, telemetry manager/client/state, topology refresher, and status usable. + try { + candidateRefresher?.stop(); + } catch { + // ignore cleanup errors on an unpublished candidate + } + try { + candidateTelemetryClient?.close(); + } catch { + // ignore cleanup errors on an unpublished candidate + } + if (candidatePool) { + await this.disposeChannelPool(candidatePool); } - this.connectStatus = CONNECT_STATUS.SHUTDOWN; throw e; } + return primaryChanged; })(); try { - await this.reconnectingPromise; - return primaryChanged; + return await this.reconnectingPromise; } finally { this.isReconnecting = false; this.reconnectingPromise = null; @@ -308,8 +475,8 @@ export class GRPCClient extends User { * When promisify encounters a gRPC UNAVAILABLE error after all retries, * this handler triggers topology refresh and pool rebuild. */ - private _attachFailoverHandler() { - setPoolFailoverHandler(this.channelPool, async () => { + private _attachFailoverHandler(pool = this.channelPool) { + setPoolFailoverHandler(pool, async () => { // Trigger topology refresh if (this.topologyRefresher) { this.topologyRefresher.triggerRefresh(); @@ -324,14 +491,15 @@ export class GRPCClient extends User { * Creates a pool of gRPC service clients. * @returns {Pool} - A pool of gRPC service clients. */ - private createChannelPool() { + private createChannelPool(address = this.config.address) { const ServiceClientConstructor = this._MilvusService; - return createPool( + const formattedAddress = formatAddress(address); + const pool = createPool( { create: async () => { // Create a new gRPC service client return new ServiceClientConstructor( - formatAddress(this.config.address), // format the address + formattedAddress, this.creds, this.channelOptions ); @@ -349,6 +517,57 @@ export class GRPCClient extends User { max: DEFAULT_POOL_MAX, } ); + setPoolTelemetryManager(pool, this.telemetry); + return pool; + } + + private createTelemetryClient(address: string): Client { + return new this._ClientTelemetryService( + formatAddress(address), + this.creds, + this.telemetryChannelOptions + ); + } + + private replaceTelemetryClient(address = this.config.address) { + const replacement = this.createTelemetryClient(address); + const previous = this.telemetryClient; + this.telemetryEndpointEpoch += 1; + this.telemetryClient = replacement; + previous?.close(); + } + + private async disposeChannelPool( + pool: ReturnType + ) { + try { + await pool.drain(); + await pool.clear(); + } catch { + // Pool cleanup is best-effort. A cleanup failure must neither roll back a + // successful publication nor hide the original candidate validation error. + } + } + + private sendTelemetryHeartbeat(request: Record) { + return new Promise((resolve, reject) => { + if (!this.telemetryClient) { + reject(new Error('telemetry client is not connected')); + return; + } + (this.telemetryClient as any).ClientHeartbeat( + request, + new Metadata(), + { deadline: new Date(Date.now() + 10_000) }, + (error: any, response: any) => + error ? reject(error) : resolve(response) + ); + }); + } + + /** Returns the telemetry manager for inspection and custom command handlers. */ + getTelemetry() { + return this.telemetry; } /** @@ -377,6 +596,7 @@ export class GRPCClient extends User { `No database name provided, using default database: ${DEFAULT_DB}` ); } + this.telemetryDatabaseExplicit = Boolean(data?.db_name); // update database this.metadata.set( METADATA.DATABASE, @@ -394,7 +614,12 @@ export class GRPCClient extends User { * @param {string} sdkVersion - The version of the SDK being used. * @returns {Promise} - A Promise that resolves when the server information has been retrieved. */ - private async _getServerInfo(sdkVersion: string) { + private async _getServerInfo( + sdkVersion: string, + pool = this.channelPool, + apply = true, + lifecycleGeneration = this.lifecycleGeneration + ) { // build user info const userInfo = { client_info: { @@ -406,26 +631,28 @@ export class GRPCClient extends User { }, }; - // update connect status - this.connectStatus = CONNECT_STATUS.CONNECTING; + if (apply) { + this.connectStatus = CONNECT_STATUS.CONNECTING; + } - return promisify(this.channelPool, 'Connect', userInfo, this.timeout).then( - f => { - // add new identifier interceptor - if (f && f.identifier) { - // update identifier - this.metadata.set(METADATA.CLIENT_ID, f.identifier); + const response = await promisify(pool, 'Connect', userInfo, this.timeout); + if (apply && this.isLifecycleCurrent(lifecycleGeneration)) { + this.applyServerInfo(response); + } + return response; + } - // setup identifier - this.serverInfo = f.server_info; - } - // update connect status - this.connectStatus = - f && f.identifier - ? CONNECT_STATUS.CONNECTED - : CONNECT_STATUS.UNIMPLEMENTED; - } - ); + private applyServerInfo(response: any, startTelemetry = true) { + if (response?.identifier) { + this.metadata.set(METADATA.CLIENT_ID, response.identifier); + this.serverInfo = response.server_info; + } + this.connectStatus = response?.identifier + ? CONNECT_STATUS.CONNECTED + : CONNECT_STATUS.UNIMPLEMENTED; + if (response?.identifier && startTelemetry) { + this.telemetry.start(); + } } /** @@ -434,24 +661,54 @@ export class GRPCClient extends User { * @returns {Promise} The updated connection status. */ async closeConnection() { - // Stop topology refresher if running (global cluster) - if (this.topologyRefresher) { - logger.debug( - `\x1b[36m[Global]\x1b[0m Stopping topology refresher on connection close` - ); - this.topologyRefresher.stop(); - this.topologyRefresher = null; + if (this.closePromise) { + return this.closePromise; } - // Close all connections in the pool - if (this.channelPool) { - await this.channelPool.drain(); - await this.channelPool.clear(); + // Publish the close fence before touching resources. Any topology/candidate promise + // released from this point on observes a different generation and may only clean up. + this.closed = true; + this.lifecycleGeneration += 1; + this.telemetryEndpointEpoch += 1; + this.connectStatus = CONNECT_STATUS.SHUTDOWN; - // update status - this.connectStatus = CONNECT_STATUS.SHUTDOWN; - } - return this.connectStatus; + const telemetryClient = this.telemetryClient; + this.telemetryClient = undefined; + const topologyRefresher = this.topologyRefresher; + this.topologyRefresher = null; + const channelPool = this.channelPool; + + this.closePromise = (async () => { + try { + this.telemetry.stop(); + } catch { + // best-effort telemetry shutdown must not prevent channel cleanup + } + try { + telemetryClient?.close(); + } catch { + // best-effort cleanup + } + if (topologyRefresher) { + logger.debug( + `\x1b[36m[Global]\x1b[0m Stopping topology refresher on connection close` + ); + try { + topologyRefresher.stop(); + } catch { + // best-effort cleanup + } + } + if (channelPool) { + await this.disposeChannelPool(channelPool); + } + return this.connectStatus; + })(); + return this.closePromise; + } + + private isLifecycleCurrent(generation: number) { + return !this.closed && generation === this.lifecycleGeneration; } /** @@ -480,24 +737,31 @@ export class GRPCClient extends User { * @returns {Promise} - A Promise that resolves with the analyzer response. */ async runAnalyzer(data: RunAnalyzerRequest): Promise { - return await promisify( + return withTelemetryLogicalOperation( this.channelPool, 'RunAnalyzer', - { - analyzer_params: data.analyzer_params - ? JSON.stringify(data.analyzer_params) - : '', - placeholder: (Array.isArray(data.text) ? data.text : [data.text]).map( - d => new TextEncoder().encode(String(d)) - ), - with_detail: data.with_detail, - with_hash: data.with_hash, - db_name: data.db_name, - collection_name: data.collection_name, - field_name: data.field_name, - analyzer_names: data.analyzer_names, - }, - this.timeout + data, + async () => + promisify( + this.channelPool, + 'RunAnalyzer', + { + analyzer_params: data.analyzer_params + ? JSON.stringify(data.analyzer_params) + : '', + placeholder: (Array.isArray(data.text) + ? data.text + : [data.text] + ).map(d => new TextEncoder().encode(String(d))), + with_detail: data.with_detail, + with_hash: data.with_hash, + db_name: data.db_name, + collection_name: data.collection_name, + field_name: data.field_name, + analyzer_names: data.analyzer_names, + }, + this.timeout + ) ); } diff --git a/milvus/index.ts b/milvus/index.ts index e026aa24..206b817a 100644 --- a/milvus/index.ts +++ b/milvus/index.ts @@ -4,6 +4,7 @@ export * from './const'; export * from './utils'; // types export * from './types'; +export * from './telemetry'; // clients export * from './grpc/GrpcClient'; export * from './MilvusClient'; diff --git a/milvus/telemetry/ClientTelemetry.ts b/milvus/telemetry/ClientTelemetry.ts new file mode 100644 index 00000000..bb3b8a08 --- /dev/null +++ b/milvus/telemetry/ClientTelemetry.ts @@ -0,0 +1,1285 @@ +import crypto from 'crypto'; +import os from 'os'; +import { AsyncLocalStorage } from 'async_hooks'; +import { status as grpcStatus } from '@grpc/grpc-js'; + +const MAX_UNIMPLEMENTED_BACKOFF_MS = 30 * 60 * 1000; +// Node turns delays larger than a signed 32-bit integer into a 1ms timer. Split long +// heartbeat intervals into safe, cancellable chunks so a valid server-pushed interval +// cannot turn into a tight heartbeat loop. +const MAX_TIMER_DELAY_MS = 2_147_483_647; +const SAMPLE_BUFFER_SIZE = 1000; +// History needs enough distribution shape to aggregate percentiles across windows, +// but retaining the full live p99 ring for every snapshot would exceed 200 MiB at +// the one-hour/one-second hard cap. Keep 128 sorted, equidistant quantiles (including +// min/max): seven operations x 4096 windows is about 28 MiB of raw doubles. +const HISTORY_SAMPLE_BUFFER_SIZE = 128; +const HISTORY_RETENTION_MS = 60 * 60 * 1000; +// Preserve the 3601 boundary-inclusive windows produced by a full hour at a one-second +// heartbeat while still bounding memory if a server pushes a sub-second interval. +const MAX_HISTORY_SNAPSHOTS = 4096; +const MAX_REPLY_BYTES = 1024 * 1024; +// 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". +const SAMPLING_SCALE = 1_000_000_000; +const PUSH_CONFIG_KEYS = new Set([ + 'enabled', + 'heartbeat_interval_ms', + 'sampling_rate', +]); + +export interface TelemetryConfig { + enabled?: boolean; + heartbeatIntervalMs?: number; + samplingRate?: number; + errorMaxCount?: number; + /** Pins telemetry identity across process restarts. */ + clientId?: string; +} + +export interface TelemetryMetric { + request_count: number; + success_count: number; + error_count: number; + avg_latency_ms: number; + p99_latency_ms: number; + max_latency_ms: number; +} + +export interface TelemetryOperationMetrics { + operation: string; + global: TelemetryMetric; + collection_metrics: Record; +} + +export interface TelemetrySnapshot { + timestamp: number; + end_time: number; + metrics: TelemetryOperationMetrics[]; +} + +export interface TelemetryError { + timestamp: number; + operation: string; + error_msg: string; + collection?: string; + request_id?: string; +} + +export interface ClientCommand { + command_id: string; + command_type: string; + payload?: Uint8Array | Buffer | string; + create_time?: number | string; + persistent?: boolean; + target_scope?: string; +} + +export interface CommandReply { + command_id: string; + success: boolean; + error_message?: string; + payload?: Buffer; +} + +export type CommandHandler = ( + command: ClientCommand +) => CommandReply | Promise; + +export interface OperationRecord { + operation: string; + collection: string; + startTime: number; + error?: unknown; + requestId?: string; +} + +type HeartbeatSender = (request: Record) => Promise; + +class MetricBucket { + requests = 0; + successes = 0; + failures = 0; + totalLatencyMs = 0; + maxLatencyMs = 0; + samples: number[] = []; + + record(latencyMs: number, success: boolean) { + this.requests += 1; + this.successes += success ? 1 : 0; + this.failures += success ? 0 : 1; + this.totalLatencyMs += latencyMs; + this.maxLatencyMs = Math.max(this.maxLatencyMs, latencyMs); + this.samples.push(latencyMs); + if (this.samples.length > SAMPLE_BUFFER_SIZE) { + this.samples.shift(); + } + } + + snapshot(): TelemetryMetric | undefined { + if (this.requests === 0) { + return undefined; + } + const sorted = [...this.samples].sort((left, right) => left - right); + const p99 = sorted.length + ? sorted[Math.min(sorted.length - 1, Math.floor(sorted.length * 0.99))] + : 0; + return { + request_count: this.requests, + success_count: this.successes, + error_count: this.failures, + avg_latency_ms: this.totalLatencyMs / this.requests, + p99_latency_ms: p99, + max_latency_ms: this.maxLatencyMs, + }; + } +} + +class OperationCollector { + global = new MetricBucket(); + collections = new Map(); + + record(collection: string, latencyMs: number, success: boolean) { + this.global.record(latencyMs, success); + if (collection) { + let bucket = this.collections.get(collection); + if (!bucket) { + bucket = new MetricBucket(); + this.collections.set(collection, bucket); + } + bucket.record(latencyMs, success); + } + } + + snapshot(operation: string): + | { + metric: TelemetryOperationMetrics; + globalLatencySamples: Float64Array; + } + | undefined { + const global = this.global.snapshot(); + if (!global) { + return undefined; + } + const globalLatencySamples = retainQuantileSamples(this.global.samples); + const collectionMetrics: Record = {}; + for (const [name, bucket] of this.collections) { + const metric = bucket.snapshot(); + if (metric) { + collectionMetrics[name] = metric; + } + } + this.global = new MetricBucket(); + this.collections.clear(); + return { + metric: { + operation, + global, + collection_metrics: collectionMetrics, + }, + globalLatencySamples, + }; + } +} + +export class ClientTelemetryManager { + public readonly clientId: string; + public readonly stableClientId: boolean; + public ready = false; + public configHash = ''; + public lastCommandTimestamp = 0; + public lastHeartbeatError: unknown; + + private readonly sender: HeartbeatSender; + private sdkVersion: string; + private readonly userProvider: () => string; + private readonly databaseProvider: () => string; + private readonly configProvider: () => Record; + private readonly senderEpochProvider: () => number; + private readonly collectors = new Map(); + private readonly errors: TelemetryError[] = []; + private readonly snapshots: TelemetrySnapshot[] = []; + // Kept outside TelemetrySnapshot so samples never enter heartbeat/detail/public JSON. + // Pruning explicitly deletes matching entries to enforce the hard memory bound. + private readonly snapshotLatencySamples = new Map< + TelemetrySnapshot, + Record + >(); + private readonly pendingReplies: CommandReply[] = []; + private readonly executedCommands = new Map(); + private readonly handlers = new Map(); + private readonly enabledCollections = new Set(); + private allCollectionsEnabled = false; + private enabled: boolean; + private heartbeatIntervalMs: number; + private samplingRate: number; + private readonly errorMaxCount: number; + // 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 shouldSample. + private samplingAccum = 0; + private unsupportedStreak = 0; + private timer?: NodeJS.Timeout; + private stopped = false; + private lastSnapshotEnd = 0; + // Command handlers may be asynchronous and ProcessCommands is public. Keep whole batches + // on one queue so two concurrent heartbeats/callers cannot both execute the same ID after + // observing it as absent. + private commandQueue: Promise = Promise.resolve(); + // Awaiting processCommands() from a handler would enqueue work behind the batch that is + // currently awaiting that handler. Track the active handler context so recursive use can + // fail as a correlated command reply instead of creating a promise cycle. The mutable + // flag lets work scheduled for after a completed handler use the public queue normally. + private readonly commandHandlerScope = new AsyncLocalStorage<{ + active: boolean; + }>(); + + constructor(options: { + sender: HeartbeatSender; + config?: TelemetryConfig; + sdkVersion?: string; + userProvider?: () => string; + databaseProvider?: () => string; + configProvider?: () => Record; + senderEpochProvider?: () => number; + }) { + const config = options.config || {}; + this.sender = options.sender; + this.sdkVersion = options.sdkVersion || ''; + this.userProvider = options.userProvider || (() => ''); + this.databaseProvider = options.databaseProvider || (() => ''); + this.configProvider = options.configProvider || (() => ({})); + this.senderEpochProvider = options.senderEpochProvider || (() => 0); + if (config.enabled !== undefined && typeof config.enabled !== 'boolean') { + throw new Error('enabled must be a boolean'); + } + this.enabled = config.enabled ?? true; + // Milliseconds 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. + this.heartbeatIntervalMs = Number(config.heartbeatIntervalMs ?? 10_000); + if ( + !Number.isFinite(this.heartbeatIntervalMs) || + this.heartbeatIntervalMs <= 0 + ) { + throw new Error('heartbeatIntervalMs must be a finite positive number'); + } + const samplingRate = config.samplingRate ?? 1; + if (typeof samplingRate !== 'number' || !Number.isFinite(samplingRate)) { + throw new Error('samplingRate must be a finite number'); + } + this.samplingRate = clamp(samplingRate, 0, 1); + const errorMaxCount = config.errorMaxCount ?? 100; + if ( + typeof errorMaxCount !== 'number' || + !Number.isSafeInteger(errorMaxCount) || + errorMaxCount <= 0 + ) { + throw new Error('errorMaxCount must be a positive integer'); + } + this.errorMaxCount = errorMaxCount; + this.stableClientId = Boolean(config.clientId); + this.clientId = config.clientId || crypto.randomUUID(); + this.registerDefaultHandlers(); + } + + start() { + if (this.ready) { + return; + } + this.ready = true; + if (!this.enabled) { + return; + } + void this.heartbeatLoop(); + } + + setSdkVersion(version: string) { + this.sdkVersion = version; + } + + stop() { + this.stopped = true; + if (this.timer) { + clearTimeout(this.timer); + this.timer = undefined; + } + } + + isSupported() { + return this.unsupportedStreak === 0; + } + + getConfig(): Required { + return { + enabled: this.enabled, + heartbeatIntervalMs: this.heartbeatIntervalMs, + samplingRate: this.samplingRate, + errorMaxCount: this.errorMaxCount, + clientId: this.clientId, + }; + } + + registerCommandHandler(type: string, handler: CommandHandler) { + this.handlers.set(type, handler); + } + + recordOperation(record: OperationRecord) { + if (!this.enabled || !this.shouldSample()) { + return; + } + const latencyMs = Math.max(0, performance.now() - record.startTime); + const collection = + record.collection && + (this.allCollectionsEnabled || + this.enabledCollections.has(record.collection)) + ? record.collection + : ''; + let collector = this.collectors.get(record.operation); + if (!collector) { + collector = new OperationCollector(); + this.collectors.set(record.operation, collector); + } + collector.record(collection, latencyMs, !record.error); + if (record.error) { + this.errors.push({ + timestamp: Date.now(), + operation: record.operation, + error_msg: errorMessage(record.error), + collection: record.collection || undefined, + request_id: record.requestId || undefined, + }); + while (this.errors.length > this.errorMaxCount) { + this.errors.shift(); + } + } + } + + getRecentErrors(maxCount = 100): TelemetryError[] { + return [...this.errors].reverse().slice(0, maxCount); + } + + getMetricsSnapshots(): TelemetrySnapshot[] { + this.pruneSnapshots(Date.now()); + return [...this.snapshots]; + } + + processCommands(commands: ClientCommand[]): Promise { + if (this.commandHandlerScope.getStore()?.active) { + const rejected = Promise.reject( + new Error( + 'processCommands cannot be called recursively from a command handler' + ) + ); + // Preserve the rejection for callers that await it, but attach a handler immediately + // so fire-and-forget recursion cannot surface as an unhandled rejection. + void rejected.catch(() => undefined); + return rejected; + } + return this.enqueueCommandBatch(commands); + } + + private enqueueCommandBatch( + commands: ClientCommand[], + expectedSenderEpoch?: number + ): Promise { + const queued = this.commandQueue.then(() => + this.processCommandBatch(commands, expectedSenderEpoch) + ); + // A custom handler is isolated by handleCommand, but keep the queue usable even if a + // future batch-level change throws unexpectedly. + this.commandQueue = queued.catch(() => undefined); + return queued; + } + + private async processCommandBatch( + commands: ClientCommand[], + expectedSenderEpoch?: number + ) { + if ( + expectedSenderEpoch !== undefined && + expectedSenderEpoch !== this.senderEpochProvider() + ) { + return; + } + const previousTimestamp = this.lastCommandTimestamp; + let maxTimestamp = previousTimestamp; + let hasPersistent = false; + for (const command of commands) { + const createTime = Number(command.create_time || 0); + maxTimestamp = Math.max(maxTimestamp, createTime); + hasPersistent ||= Boolean(command.persistent); + if (createTime < previousTimestamp) { + this.pendingReplies.push(successReply(command.command_id)); + continue; + } + if (this.executedCommands.has(command.command_id)) { + this.pendingReplies.push(successReply(command.command_id)); + continue; + } + const reply = await this.handleCommand(command); + this.executedCommands.set(command.command_id, createTime); + this.pendingReplies.push(reply); + if ( + expectedSenderEpoch !== undefined && + expectedSenderEpoch !== this.senderEpochProvider() + ) { + // The completed handler may itself have changed the live endpoint. Retain its ID + // and correlated reply so redelivery cannot repeat a non-idempotent prefix, but do + // not apply the retired endpoint's remaining commands or commit its cursor/hash. + return; + } + } + // Keep IDs at the new cursor timestamp: timestamp filtering only rejects commands + // strictly older than the cursor, so equal-timestamp redeliveries still need ID-based + // deduplication. Entries below the new cursor can be discarded safely. + for (const [id, timestamp] of this.executedCommands) { + if (timestamp < maxTimestamp) { + this.executedCommands.delete(id); + } + } + if (hasPersistent) { + this.configHash = ClientTelemetryManager.calculateConfigHash(commands); + } + this.lastCommandTimestamp = Math.max( + this.lastCommandTimestamp, + maxTimestamp + ); + } + + static calculateConfigHash(commands: ClientCommand[]): string { + const persistent = commands + .filter(command => command.persistent) + .sort((left, right) => compareUtf8(left.command_id, right.command_id)); + if (!persistent.length) { + return ''; + } + const hash = crypto.createHash('sha256'); + for (const command of persistent) { + hash.update(command.command_id); + hash.update(command.command_type); + hash.update(payloadBuffer(command.payload)); + } + return hash.digest('hex').slice(0, 16); + } + + private async heartbeatLoop() { + try { + this.createSnapshot(); + await this.sendHeartbeat(); + } catch (error) { + // Telemetry is an optional background control plane. An unexpected collector, + // serializer, command, or transport failure must not terminate the loop forever. + this.lastHeartbeatError = error; + } finally { + if (this.stopped) { + return; + } + this.scheduleNextHeartbeat(this.nextHeartbeatDelay()); + } + } + + private scheduleNextHeartbeat(remainingDelayMs: number) { + const chunk = Math.min(remainingDelayMs, MAX_TIMER_DELAY_MS); + this.timer = setTimeout(() => { + this.timer = undefined; + if (this.stopped) { + return; + } + const remaining = Math.max(0, remainingDelayMs - chunk); + if (remaining > 0) { + this.scheduleNextHeartbeat(remaining); + return; + } + void this.heartbeatLoop(); + }, chunk); + // Telemetry is best-effort background work and must not keep an otherwise idle Node + // process alive. + this.timer.unref?.(); + } + + private nextHeartbeatDelay() { + if (this.unsupportedStreak <= 0) { + return this.heartbeatIntervalMs; + } + return Math.max( + this.heartbeatIntervalMs, + Math.min( + MAX_UNIMPLEMENTED_BACKOFF_MS, + this.heartbeatIntervalMs * 2 ** this.unsupportedStreak + ) + ); + } + + private async sendHeartbeat() { + const latest = this.snapshots[this.snapshots.length - 1]; + const replies = [...this.pendingReplies]; + const reserved: Record = { + client_id: this.clientId, + client_id_stable: String(this.stableClientId), + }; + const database = this.databaseProvider(); + if (database) { + reserved.db_name = database; + } + let response: any; + const senderEpoch = this.senderEpochProvider(); + try { + response = await this.sender({ + client_info: { + sdk_type: 'nodejs', + sdk_version: this.sdkVersion, + local_time: new Date().toISOString(), + user: this.userProvider(), + host: os.hostname(), + reserved, + }, + report_timestamp: Date.now(), + // Do not resend the last enabled snapshot while collection is disabled. Replies, + // config hash, cursor and commands remain active as the control plane. + metrics: this.enabled + ? this.filterMetricsForWire(latest?.metrics || []) + : [], + command_replies: replies, + config_hash: this.configHash, + last_command_timestamp: this.lastCommandTimestamp, + }); + } catch (error: any) { + if (senderEpoch !== this.senderEpochProvider()) { + return; + } + this.lastHeartbeatError = error; + if (error?.code === grpcStatus.UNIMPLEMENTED) { + this.unsupportedStreak += 1; + } + return; + } + + // A global-cluster failover can publish a new telemetry transport while this + // heartbeat is awaiting the old endpoint. Treat that old response as if it never + // arrived: do not acknowledge replies, reset backoff/errors, execute commands, or + // mutate any command/config history. + if (senderEpoch !== this.senderEpochProvider()) { + return; + } + + // Any real response proves the RPC exists. Reset an old UNIMPLEMENTED streak before + // checking the business status, which may fail while the service is still starting. + this.unsupportedStreak = 0; + if (!responseSucceeded(response)) { + this.lastHeartbeatError = new Error( + response?.status?.reason || 'client heartbeat failed' + ); + return; + } + this.pendingReplies.splice(0, replies.length); + this.lastHeartbeatError = undefined; + await this.enqueueCommandBatch(response?.commands || [], senderEpoch); + } + + /** + * Decide whether this operation is recorded, spreading the sampled ones evenly rather + * than in runs. + * + * 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. The ratio has to hold + * 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, so sampling a + * contiguous run would make each window either complete or empty. + */ + private shouldSample() { + if (this.samplingRate >= 1) { + return true; + } + if (this.samplingRate <= 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. + const step = Math.max(1, Math.floor(this.samplingRate * SAMPLING_SCALE)); + this.samplingAccum += step; + if (this.samplingAccum < SAMPLING_SCALE) { + return false; + } + // Keep only the remainder rather than letting the accumulator grow without bound: a + // JavaScript number is exact only to 2^53, which at a few hundred million units per + // operation is reached within a day of steady traffic, after which the comparison + // would start losing operations silently. + this.samplingAccum -= SAMPLING_SCALE; + return true; + } + + private createSnapshot() { + if (!this.enabled) { + return; + } + const metrics: TelemetryOperationMetrics[] = []; + const latencySamples: Record = {}; + for (const [operation, collector] of this.collectors) { + const collected = collector.snapshot(operation); + if (collected) { + metrics.push(collected.metric); + latencySamples[operation] = collected.globalLatencySamples; + } + } + const now = Date.now(); + const start = + !this.lastSnapshotEnd || this.lastSnapshotEnd > now + ? now - this.heartbeatIntervalMs + : this.lastSnapshotEnd; + this.lastSnapshotEnd = now; + const snapshot = { timestamp: start, end_time: now, metrics }; + this.snapshots.push(snapshot); + this.snapshotLatencySamples.set(snapshot, latencySamples); + this.pruneSnapshots(now); + } + + private pruneSnapshots(now: number) { + const cutoff = now - HISTORY_RETENTION_MS; + while (this.snapshots.length && this.snapshots[0].end_time < cutoff) { + const expired = this.snapshots.shift(); + if (expired) { + this.snapshotLatencySamples.delete(expired); + } + } + while (this.snapshots.length > MAX_HISTORY_SNAPSHOTS) { + const expired = this.snapshots.shift(); + if (expired) { + this.snapshotLatencySamples.delete(expired); + } + } + } + + private filterMetricsForWire( + metrics: TelemetryOperationMetrics[] + ): TelemetryOperationMetrics[] { + return metrics.map(operation => { + const collectionMetrics: Record = {}; + for (const [collection, metric] of Object.entries( + operation.collection_metrics + )) { + if ( + this.allCollectionsEnabled || + this.enabledCollections.has(collection) + ) { + collectionMetrics[collection] = metric; + } + } + return { ...operation, collection_metrics: collectionMetrics }; + }); + } + + private async handleCommand(command: ClientCommand): Promise { + const handler = this.handlers.get(command.command_type); + if (!handler) { + return failedReply( + command.command_id, + `unknown command type: ${command.command_type}` + ); + } + const handlerScope = { active: true }; + try { + const reply = await this.commandHandlerScope.run(handlerScope, () => + handler(command) + ); + if (!reply || typeof reply !== 'object') { + return failedReply( + command.command_id, + 'command handler returned no reply' + ); + } + // command_id is the server's correlation key. A custom handler may customize the + // result, but it may never acknowledge a different command. + return { ...reply, command_id: command.command_id }; + } catch (error) { + return failedReply(command.command_id, errorMessage(error)); + } finally { + handlerScope.active = false; + } + } + + private registerDefaultHandlers() { + this.registerCommandHandler('push_config', command => { + const payload = parsePayload(command); + const enabled = optionalBoolean(payload, 'enabled'); + const heartbeatIntervalMs = optionalInteger( + payload, + 'heartbeat_interval_ms' + ); + const samplingRate = optionalFiniteNumber(payload, 'sampling_rate'); + // Validate the complete payload before changing any field. A failed command must not + // leave an earlier key applied while reporting that the whole command failed. + if (heartbeatIntervalMs !== undefined && heartbeatIntervalMs <= 0) { + throw new Error('heartbeat_interval_ms must be a positive integer'); + } + + const applied: string[] = []; + if (enabled !== undefined) { + applied.push('enabled'); + } + if (heartbeatIntervalMs !== undefined) { + applied.push('heartbeat_interval_ms'); + } + if (samplingRate !== undefined) { + applied.push('sampling_rate'); + } + const ignored = Object.keys(payload) + .filter(key => !PUSH_CONFIG_KEYS.has(key)) + .sort(compareUtf8); + + if (enabled !== undefined) { + this.enabled = enabled; + } + if (heartbeatIntervalMs !== undefined) { + this.heartbeatIntervalMs = heartbeatIntervalMs; + } + if (samplingRate !== undefined) { + this.samplingRate = clamp(samplingRate, 0, 1); + } + + return successReply( + command.command_id, + Buffer.from( + JSON.stringify({ + applied, + ...(ignored.length ? { ignored } : {}), + }) + ) + ); + }); + + this.registerCommandHandler('collection_metrics', command => { + if (!payloadBuffer(command.payload).length) { + return successReply( + command.command_id, + Buffer.from( + JSON.stringify({ + enabled_collections: [...this.enabledCollections].sort(), + all_collections_enabled: this.allCollectionsEnabled, + }) + ) + ); + } + const payload = parsePayload(command); + const enabled = optionalBoolean(payload, 'enabled') ?? false; + const collections = optionalStringArray(payload, 'collections') ?? []; + // metrics_types is not acted on yet, but it is part of the protocol payload and must + // still have the same typed-JSON behavior as the Go SDK. + optionalStringArray(payload, 'metrics_types'); + const wildcard = collections.includes('*'); + if (enabled) { + if (!collections.length) { + throw new Error('collections list cannot be empty when enabled=true'); + } + if (wildcard) { + this.allCollectionsEnabled = true; + } else { + collections.forEach(name => this.enabledCollections.add(name)); + } + } else if (wildcard || !collections.length) { + this.allCollectionsEnabled = false; + this.enabledCollections.clear(); + } else { + collections.forEach(name => this.enabledCollections.delete(name)); + } + return successReply(command.command_id); + }); + + this.registerCommandHandler('show_errors', command => { + const payload = parsePayload(command); + const configuredMaxCount = optionalInteger(payload, 'max_count'); + const maxCount = + configuredMaxCount !== undefined && configuredMaxCount > 0 + ? configuredMaxCount + : 100; + let errors = this.getRecentErrors(maxCount).map(error => ({ ...error })); + if (!errors.length) { + return successReply(command.command_id); + } + let encoded = Buffer.from(JSON.stringify(errors)); + while (encoded.length > MAX_REPLY_BYTES && errors.length > 1) { + errors = errors.slice(0, Math.max(1, Math.floor(errors.length / 2))); + encoded = Buffer.from(JSON.stringify(errors)); + } + if (encoded.length > MAX_REPLY_BYTES && errors.length === 1) { + encoded = truncateSingleErrorReply(errors[0], MAX_REPLY_BYTES); + } + if (encoded.length > MAX_REPLY_BYTES) { + throw new Error('show_errors response exceeds the 1MB payload limit'); + } + return successReply(command.command_id, encoded); + }); + + this.registerCommandHandler('get_config', command => { + const userConfig = { ...this.configProvider() }; + ['password', 'token', 'api_key', 'authorization'].forEach( + key => delete userConfig[key] + ); + Object.assign(userConfig, { + telemetry_enabled: this.enabled, + telemetry_heartbeat_interval_ms: this.heartbeatIntervalMs, + telemetry_sampling_rate: this.samplingRate, + enabled_collections: this.allCollectionsEnabled + ? ['*'] + : [...this.enabledCollections].sort(), + all_collections_enabled: this.allCollectionsEnabled, + }); + return successReply( + command.command_id, + Buffer.from(JSON.stringify({ user_config: userConfig })) + ); + }); + + this.registerCommandHandler('show_latency_history', command => { + const payload = parsePayload(command); + if ( + typeof payload.start_time !== 'string' || + typeof payload.end_time !== 'string' + ) { + throw new Error('payload is required with start_time and end_time'); + } + if ('detail' in payload && typeof payload.detail !== 'boolean') { + throw new Error('detail must be a boolean'); + } + const start = parseRfc3339(payload.start_time, 'start_time'); + const end = parseRfc3339(payload.end_time, 'end_time'); + if (end < start) { + throw new Error('end_time must be after start_time'); + } + if (end - start > 60 * 60 * 1000) { + throw new Error('time range cannot exceed 1 hour'); + } + this.pruneSnapshots(Date.now()); + const snapshots = this.snapshots.filter( + snapshot => snapshot.end_time >= start && snapshot.timestamp <= end + ); + const body = + payload.detail === true + ? { + snapshots: detailSnapshots(snapshots), + total_snapshots: snapshots.length, + } + : aggregateSnapshots( + snapshots, + start, + end, + snapshot => this.snapshotLatencySamples.get(snapshot) || {} + ); + const encoded = Buffer.from(JSON.stringify(body)); + if (encoded.length > MAX_REPLY_BYTES) { + throw new Error('response too large, try a smaller time range'); + } + return successReply(command.command_id, encoded); + }); + } +} + +export function newClientRequestId(): string { + let value: Buffer; + do { + value = crypto.randomBytes(16); + } while (value.every(byte => byte === 0)); + return value.toString('hex'); +} + +function payloadBuffer(payload?: Uint8Array | Buffer | string): Buffer { + if (!payload) { + return Buffer.alloc(0); + } + return typeof payload === 'string' + ? Buffer.from(payload) + : Buffer.from(payload); +} + +function parsePayload(command: ClientCommand): Record { + const payload = payloadBuffer(command.payload); + if (!payload.length) { + return {}; + } + const parsed = JSON.parse(payload.toString()); + if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) { + throw new Error('command payload must be a JSON object'); + } + return parsed; +} + +function truncateSingleErrorReply( + error: TelemetryError, + maxBytes: number +): Buffer { + const mutableError = error as unknown as Record; + const fields = ['error_msg', 'collection', 'request_id', 'operation'].sort( + (left, right) => + Buffer.byteLength(String(mutableError[right] || '')) - + Buffer.byteLength(String(mutableError[left] || '')) + ); + const suffix = '...(truncated)'; + let encoded = Buffer.from(JSON.stringify([error])); + + for (const field of fields) { + if (encoded.length <= maxBytes) { + return encoded; + } + const original = mutableError[field]; + if (typeof original !== 'string' || original.length === 0) { + continue; + } + + // Halve an original prefix, never the already suffixed value. The prefix length strictly + // decreases to zero, so this terminates even when another field caused the overflow. + let prefixLength = original.length; + while (encoded.length > maxBytes && prefixLength > 0) { + prefixLength = Math.floor(prefixLength / 2); + const prefix = unicodeSafePrefix(original, prefixLength); + mutableError[field] = prefix ? `${prefix}${suffix}` : ''; + encoded = Buffer.from(JSON.stringify([error])); + } + if (encoded.length <= maxBytes) { + return encoded; + } + + if (field === 'collection' || field === 'request_id') { + delete mutableError[field]; + } + encoded = Buffer.from(JSON.stringify([error])); + } + return encoded; +} + +function unicodeSafePrefix(value: string, length: number): string { + let end = Math.min(value.length, Math.max(0, length)); + if ( + end > 0 && + end < value.length && + value.charCodeAt(end - 1) >= 0xd800 && + value.charCodeAt(end - 1) <= 0xdbff && + value.charCodeAt(end) >= 0xdc00 && + value.charCodeAt(end) <= 0xdfff + ) { + end -= 1; + } + return value.slice(0, end); +} + +function optionalBoolean( + payload: Record, + key: string +): boolean | undefined { + if (!(key in payload)) { + return undefined; + } + if (typeof payload[key] !== 'boolean') { + throw new Error(`${key} must be a boolean`); + } + return payload[key]; +} + +function optionalFiniteNumber( + payload: Record, + key: string +): number | undefined { + if (!(key in payload)) { + return undefined; + } + const value = payload[key]; + if (typeof value !== 'number' || !Number.isFinite(value)) { + throw new Error(`${key} must be a finite number`); + } + return value; +} + +function optionalInteger( + payload: Record, + key: string +): number | undefined { + const value = optionalFiniteNumber(payload, key); + if (value !== undefined && !Number.isSafeInteger(value)) { + throw new Error(`${key} must be an integer`); + } + return value; +} + +function optionalStringArray( + payload: Record, + key: string +): string[] | undefined { + if (!(key in payload)) { + return undefined; + } + const value = payload[key]; + if (!Array.isArray(value) || value.some(item => typeof item !== 'string')) { + throw new Error(`${key} must be an array of strings`); + } + return value; +} + +function compareUtf8(left: string, right: string): number { + return Buffer.compare(Buffer.from(left), Buffer.from(right)); +} + +const RFC3339_PATTERN = + /^(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})(?:\.(\d+))?(?:Z|([+-])(\d{2}):(\d{2}))$/; + +function parseRfc3339(value: string, field: string): number { + const match = RFC3339_PATTERN.exec(value); + if (!match) { + throw new Error(`invalid ${field} format, expected RFC3339`); + } + const [, year, month, day, hour, minute, second, , , zoneHour, zoneMinute] = + match; + const numericYear = Number(year); + const numericMonth = Number(month); + const numericDay = Number(day); + const numericHour = Number(hour); + const numericMinute = Number(minute); + const numericSecond = Number(second); + const numericZoneHour = zoneHour === undefined ? 0 : Number(zoneHour); + const numericZoneMinute = zoneMinute === undefined ? 0 : Number(zoneMinute); + const leapYear = + numericYear % 4 === 0 && + (numericYear % 100 !== 0 || numericYear % 400 === 0); + const daysInMonth = [ + 31, + leapYear ? 29 : 28, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ][numericMonth - 1]; + if ( + numericMonth < 1 || + numericMonth > 12 || + numericDay < 1 || + numericDay > (daysInMonth || 0) || + numericHour > 23 || + numericMinute > 59 || + numericSecond > 59 || + numericZoneHour > 23 || + numericZoneMinute > 59 + ) { + throw new Error(`invalid ${field} format, expected RFC3339`); + } + const parsed = Date.parse(value); + if (!Number.isFinite(parsed)) { + throw new Error(`invalid ${field} format, expected RFC3339`); + } + return parsed; +} + +function successReply(commandId: string, payload = Buffer.alloc(0)) { + return { command_id: commandId, success: true, payload }; +} + +function failedReply(commandId: string, error: string) { + return { command_id: commandId, success: false, error_message: error }; +} + +function clamp(value: number, minimum: number, maximum: number) { + return Math.max(minimum, Math.min(maximum, value)); +} + +function errorMessage(error: unknown) { + return error instanceof Error ? error.message : String(error); +} + +function responseSucceeded(response: any) { + const status = response?.status; + if (!status) { + return true; + } + const code = Number(status.code || 0); + const errorCode = status.error_code; + return ( + code === 0 && + (errorCode === undefined || + errorCode === 0 || + errorCode === '0' || + errorCode === 'Success' || + errorCode === 'SUCCESS') + ); +} + +function retainQuantileSamples(samples: number[]) { + const sorted = [...samples].sort((left, right) => left - right); + if (sorted.length <= HISTORY_SAMPLE_BUFFER_SIZE) { + return Float64Array.from(sorted); + } + return Float64Array.from( + { length: HISTORY_SAMPLE_BUFFER_SIZE }, + (_, index) => { + const sourceIndex = Math.round( + (index * (sorted.length - 1)) / (HISTORY_SAMPLE_BUFFER_SIZE - 1) + ); + return sorted[sourceIndex]; + } + ); +} + +function weightedPercentile( + groups: Array<{ samples: Float64Array; weight: number }>, + target: number +) { + type Cursor = { + samples: Float64Array; + weight: number; + index: number; + latency: number; + }; + const heap: Cursor[] = []; + const push = (cursor: Cursor) => { + heap.push(cursor); + let index = heap.length - 1; + while (index > 0) { + const parent = Math.floor((index - 1) / 2); + if (heap[parent].latency <= cursor.latency) { + break; + } + heap[index] = heap[parent]; + index = parent; + } + heap[index] = cursor; + }; + const pop = () => { + const first = heap[0]; + const last = heap.pop(); + if (heap.length && last) { + let index = 0; + while (true) { + const left = index * 2 + 1; + if (left >= heap.length) { + break; + } + const right = left + 1; + const child = + right < heap.length && heap[right].latency < heap[left].latency + ? right + : left; + if (heap[child].latency >= last.latency) { + break; + } + heap[index] = heap[child]; + index = child; + } + heap[index] = last; + } + return first; + }; + + for (const group of groups) { + if (group.samples.length) { + push({ ...group, index: 0, latency: group.samples[0] }); + } + } + let cumulative = 0; + let percentile = 0; + while (heap.length) { + const cursor = pop(); + cumulative += cursor.weight; + percentile = cursor.latency; + if (cumulative > target) { + return percentile; + } + cursor.index += 1; + if (cursor.index < cursor.samples.length) { + cursor.latency = cursor.samples[cursor.index]; + push(cursor); + } + } + return percentile; +} + +function aggregateSnapshots( + snapshots: TelemetrySnapshot[], + start: number, + end: number, + latencySamplesFor: ( + snapshot: TelemetrySnapshot + ) => Record +) { + const totals: Record< + string, + { + request_count: number; + success_count: number; + error_count: number; + weighted_avg: number; + max_latency_ms: number; + latency_sample_groups: Array<{ + samples: Float64Array; + weight: number; + }>; + } + > = {}; + for (const snapshot of snapshots) { + const latencySamples = latencySamplesFor(snapshot); + for (const operation of snapshot.metrics) { + const metric = operation.global; + const total = (totals[operation.operation] ||= { + request_count: 0, + success_count: 0, + error_count: 0, + weighted_avg: 0, + max_latency_ms: 0, + latency_sample_groups: [], + }); + total.request_count += metric.request_count; + total.success_count += metric.success_count; + total.error_count += metric.error_count; + total.weighted_avg += metric.avg_latency_ms * metric.request_count; + total.max_latency_ms = Math.max( + total.max_latency_ms, + metric.max_latency_ms + ); + const samples = latencySamples[operation.operation]; + if (samples?.length) { + const weight = metric.request_count / samples.length; + total.latency_sample_groups.push({ samples, weight }); + } + } + } + const metrics: Record = {}; + for (const [operation, total] of Object.entries(totals)) { + const p99 = weightedPercentile( + total.latency_sample_groups, + 0.99 * total.request_count + ); + metrics[operation] = { + request_count: total.request_count, + success_count: total.success_count, + error_count: total.error_count, + avg_latency_ms: total.request_count + ? total.weighted_avg / total.request_count + : 0, + p99_latency_ms: p99, + max_latency_ms: total.max_latency_ms, + }; + } + return { + aggregated: { start_time: start, end_time: end, metrics }, + snapshot_count: snapshots.length, + }; +} + +function detailSnapshots(snapshots: TelemetrySnapshot[]) { + return snapshots.map(snapshot => { + const metrics: Record = {}; + for (const operation of snapshot.metrics) { + metrics[operation.operation] = operation.global; + } + return { + timestamp: snapshot.timestamp, + end_time: snapshot.end_time, + metrics, + }; + }); +} diff --git a/milvus/telemetry/index.ts b/milvus/telemetry/index.ts new file mode 100644 index 00000000..f85806ae --- /dev/null +++ b/milvus/telemetry/index.ts @@ -0,0 +1 @@ +export * from './ClientTelemetry'; diff --git a/milvus/types/Client.ts b/milvus/types/Client.ts index 48cf48e6..7a19a753 100644 --- a/milvus/types/Client.ts +++ b/milvus/types/Client.ts @@ -2,6 +2,7 @@ import { ChannelOptions } from '@grpc/grpc-js'; import { Options as LoaderOption } from '@grpc/proto-loader'; import { Options } from 'generic-pool'; import { GrpcTimeOut, ResStatus } from './Common'; +import type { TelemetryConfig } from '../telemetry'; /** * Configuration options for the Milvus client. @@ -72,6 +73,9 @@ export interface ClientConfig { // enable trace trace?: boolean; + // Client metrics, heartbeat, and server-pushed command configuration. + telemetry?: TelemetryConfig; + // Explicitly enable/disable global cluster mode. // When true, the SDK fetches topology from the endpoint and routes to the primary cluster. // When omitted, auto-detected from the address URI (looks for 'global-cluster'). diff --git a/milvus/utils/Function.ts b/milvus/utils/Function.ts index 93da96b1..61937bbe 100644 --- a/milvus/utils/Function.ts +++ b/milvus/utils/Function.ts @@ -10,6 +10,7 @@ import { import { logger } from './logger'; import { Pool } from 'generic-pool'; import { Metadata, status as grpcStatus } from '@grpc/grpc-js'; +import { AsyncLocalStorage } from 'async_hooks'; /** * Failover handler type for global cluster support. @@ -21,6 +22,25 @@ export type FailoverHandler = (error: any) => Promise | null>; /** Well-known property key for attaching a failover handler to a pool. */ export const FAILOVER_HANDLER_KEY = '__failoverHandler'; +export const TELEMETRY_MANAGER_KEY = '__telemetryManager'; + +type TelemetryRecorder = { + recordOperation(record: { + operation: string; + collection: string; + startTime: number; + error?: unknown; + requestId?: string; + }): void; +}; + +/** Attach the telemetry recorder used by the common RPC path. */ +export function setPoolTelemetryManager( + pool: Pool, + manager: TelemetryRecorder +): void { + (pool as any)[TELEMETRY_MANAGER_KEY] = manager; +} /** * Attach a failover handler to a pool for global cluster support. @@ -59,22 +79,17 @@ function executeCall( finalRequestMetadata = extractRequestMetadata(params); } - const metadata = finalRequestMetadata ? new Metadata() : undefined; - if (metadata && finalRequestMetadata) { - const clientRequestId = getClientRequestId(finalRequestMetadata); - if (clientRequestId) { - metadata.add(METADATA.CLIENT_REQUEST_ID, String(clientRequestId)); - } + const clientRequestId = getClientRequestId(finalRequestMetadata); + let metadata: Metadata | undefined; + if (clientRequestId) { + metadata = new Metadata(); + metadata.add(METADATA.CLIENT_REQUEST_ID, clientRequestId); } return new Promise((resolve, reject) => { try { const callOptions: any = { deadline: new Date(Date.now() + t) }; - if (metadata) { - callOptions.metadata = metadata; - } - - client[target](params, callOptions, (err: any, result: any) => { + const callback = (err: any, result: any) => { if (err) { reject(err); } else { @@ -83,7 +98,16 @@ function executeCall( if (client) { pool.release(client); } - }); + }; + // grpc-js unary calls accept either (request, options, callback) or + // (request, metadata, options, callback). Metadata is not an option field: + // putting it inside callOptions silently drops it on a real channel even though + // loose unit-test doubles may appear to accept it. + if (metadata) { + client[target](params, metadata, callOptions, callback); + } else { + client[target](params, callOptions, callback); + } } catch (e: any) { reject(e); if (client) { @@ -94,6 +118,128 @@ function executeCall( })(); } +const TELEMETRY_OPERATIONS = new Set([ + 'Insert', + 'Delete', + 'Upsert', + 'Search', + 'HybridSearch', + 'Query', + 'RunAnalyzer', +]); + +type TelemetryLogicalScope = { + operation: string; + // Set by withTelemetrySuppressed: internal SDK machinery (iterator setup + // queries and per-page fetches) must not emit telemetry at any level. + suppress?: boolean; +}; + +// A process-wide store is safe here: AsyncLocalStorage isolates concurrent promise +// chains, while nested public helpers (for example delete -> deleteEntities) inherit +// the current logical operation and therefore do not emit a second measurement. +const telemetryLogicalScope = new AsyncLocalStorage(); + +function telemetryOperation(target: string): string | undefined { + return TELEMETRY_OPERATIONS.has(target) ? target : undefined; +} + +function getBusinessError(result: any): Error | undefined { + const responseStatus = result?.status || result; + if (!responseStatus) { + return undefined; + } + const code = Number(responseStatus.code || 0); + const errorCode = responseStatus.error_code; + const success = + code === 0 && + (errorCode === undefined || + errorCode === 0 || + errorCode === '0' || + errorCode === 'Success' || + errorCode === 'SUCCESS'); + return success + ? undefined + : new Error(responseStatus.reason || 'Milvus operation failed'); +} + +/** + * Measure one complete public SDK operation, including validation, request + * preprocessing, retries, response formatting, and postprocessing. + * + * promisify remains instrumented as a fallback for direct/internal callers, but + * suppresses its RPC-level measurement while the same logical operation is active. + */ +export async function withTelemetryLogicalOperation( + pool: Pool | undefined, + operation: string, + params: any, + call: () => Promise +): Promise { + const activeScope = telemetryLogicalScope.getStore(); + if (activeScope?.suppress || activeScope?.operation === operation) { + return call(); + } + + const telemetry: TelemetryRecorder | undefined = pool + ? (pool as any)[TELEMETRY_MANAGER_KEY] + : undefined; + const startTime = performance.now(); + const requestId = getTelemetryRequestId(extractRequestMetadata(params)); + let result: T | undefined; + let finalError: unknown; + let failed = false; + + try { + result = await telemetryLogicalScope.run({ operation }, call); + } catch (error) { + failed = true; + finalError = error; + } + + if (telemetry) { + try { + telemetry.recordOperation({ + operation, + collection: + operation === 'RunAnalyzer' + ? '' + : String(params?.collection_name || ''), + startTime, + error: failed ? finalError : getBusinessError(result), + requestId, + }); + } catch { + // Optional telemetry must never replace the business result, including when a + // thrown value has a hostile coercion hook or logging itself is unavailable. + } + } + + if (failed) { + throw finalError; + } + return result as T; +} + +/** + * Run internal SDK machinery without emitting telemetry. Iterators are not + * logical operations, so neither their setup queries nor their per-page + * internal RPCs may be measured: wrapping them here suppresses both the + * logical-operation wrapper and the promisify RPC-level fallback. + */ +export async function withTelemetrySuppressed( + call: () => Promise +): Promise { + const activeScope = telemetryLogicalScope.getStore(); + if (activeScope?.suppress) { + return call(); + } + return telemetryLogicalScope.run( + { operation: activeScope?.operation ?? '', suppress: true }, + call + ); +} + /** * Promisify a function call with optional timeout, metadata, and global cluster failover. * @param pool - The pool of gRPC clients @@ -110,36 +256,91 @@ export async function promisify( timeout: number, requestMetadata?: { 'client-request-id'?: string; client_request_id?: string } ): Promise { + const operation = telemetryOperation(target); + const activeScope = telemetryLogicalScope.getStore(); + const recordRpcOperation = + operation !== undefined && + !activeScope?.suppress && + activeScope?.operation !== operation; + const telemetry: TelemetryRecorder | undefined = (pool as any)[ + TELEMETRY_MANAGER_KEY + ]; + const startTime = performance.now(); + const finalRequestMetadata = + requestMetadata || (params ? extractRequestMetadata(params) : undefined); + const requestId = getTelemetryRequestId(finalRequestMetadata); + let result: any; + let finalError: any; + let failed = false; + try { - return await executeCall(pool, target, params, timeout, requestMetadata); - } catch (error: any) { - // Check for global cluster failover handler - const handler: FailoverHandler | undefined = (pool as any)[ - FAILOVER_HANDLER_KEY - ]; + try { + result = await executeCall( + pool, + target, + params, + timeout, + finalRequestMetadata + ); + } catch (error: any) { + // Check for global cluster failover handler. Instrumentation stays outside both + // attempts so one logical SDK call contributes exactly one telemetry outcome. + const handler: FailoverHandler | undefined = (pool as any)[ + FAILOVER_HANDLER_KEY + ]; + + if (!handler || !isUnavailableError(error)) { + throw error; + } - if (handler && isUnavailableError(error)) { logger.debug( `\x1b[36m[Global]\x1b[0m UNAVAILABLE error on \x1b[1m${target}\x1b[0m, triggering failover handler` ); const newPool = await handler(error); - if (newPool) { - logger.debug( - `\x1b[36m[Global]\x1b[0m Failover complete, retrying \x1b[1m${target}\x1b[0m with new pool` - ); - // Retry once with the new pool (after failover) - return await executeCall( - newPool, - target, - params, - timeout, - requestMetadata - ); + if (!newPool) { + throw error; } + logger.debug( + `\x1b[36m[Global]\x1b[0m Failover complete, retrying \x1b[1m${target}\x1b[0m with new pool` + ); + result = await executeCall( + newPool, + target, + params, + timeout, + finalRequestMetadata + ); + } + } catch (error: any) { + failed = true; + finalError = error; + } + + if (recordRpcOperation && operation && telemetry) { + const businessError = failed ? finalError : getBusinessError(result); + try { + telemetry.recordOperation({ + operation, + // Go intentionally records RunAnalyzer globally even though its request happens to + // carry a collection_name field. + collection: + operation === 'RunAnalyzer' + ? '' + : String(params?.collection_name || ''), + startTime, + error: businessError, + requestId, + }); + } catch { + // Optional telemetry must never replace the business result, including when a + // thrown value has a hostile coercion hook or logging itself is unavailable. } + } - throw error; + if (failed) { + throw finalError; } + return result; } export const findKeyValue = (obj: KeyValuePair[], key: string) => @@ -238,8 +439,27 @@ const getClientRequestId = (metadata?: { if (!metadata) { return undefined; } - // Priority: client_request_id > client-request-id (JavaScript/TypeScript convention) - return metadata.client_request_id || metadata['client-request-id']; + // Preserve the established wire contract: arbitrary non-empty string IDs were + // documented and sent before client telemetry existed. New telemetry correlation has + // stricter OpenTelemetry requirements, but must not silently remove metadata from + // existing applications. + const value = metadata.client_request_id ?? metadata['client-request-id']; + return typeof value === 'string' && value.length > 0 ? value : undefined; +}; + +/** Returns whether value is a non-zero lowercase 128-bit OpenTelemetry trace ID. */ +export const isValidClientRequestId = (value: unknown): value is string => + typeof value === 'string' && + /^[0-9a-f]{32}$/.test(value) && + value !== '00000000000000000000000000000000'; + +/** Returns a request ID only when it is safe to expose as an OTel trace ID. */ +const getTelemetryRequestId = (metadata?: { + 'client-request-id'?: string; + client_request_id?: string; +}): string | undefined => { + const value = getClientRequestId(metadata); + return isValidClientRequestId(value) ? value : undefined; }; /** @@ -257,7 +477,5 @@ export const extractRequestMetadata = ( } | undefined => { const clientRequestId = getClientRequestId(data); - return clientRequestId - ? { 'client-request-id': String(clientRequestId) } - : undefined; + return clientRequestId ? { 'client-request-id': clientRequestId } : undefined; }; diff --git a/milvus/utils/Grpc.ts b/milvus/utils/Grpc.ts index 3da4ad3f..a2936c60 100644 --- a/milvus/utils/Grpc.ts +++ b/milvus/utils/Grpc.ts @@ -160,7 +160,8 @@ export const getRetryInterceptor = ({ // fall back to the legacy AddCollectionField RPC. if ( status.code === grpcStatus.UNIMPLEMENTED && - methodName !== 'AlterCollectionSchema' + methodName !== 'AlterCollectionSchema' && + methodName !== 'ClientHeartbeat' ) { savedReceiveMessage = {}; status.code = grpcStatus.OK; diff --git a/milvus/utils/Search.ts b/milvus/utils/Search.ts index 076db37d..0fb51bc7 100644 --- a/milvus/utils/Search.ts +++ b/milvus/utils/Search.ts @@ -259,6 +259,19 @@ export const createFunctionScore = ( * @returns {Number[][]} return.searchVectors - The search vectors used in the operation. * @returns {number} return.round_decimal - The score precision. */ +export const isHybridSearchRequest = ( + params: SearchReq | SearchSimpleReq | HybridSearchReq | undefined +): boolean => { + const data = (params as HybridSearchReq | undefined)?.data; + const firstRequest = Array.isArray(data) ? data[0] : undefined; + return !!( + firstRequest && + typeof firstRequest === 'object' && + 'anns_field' in firstRequest && + firstRequest.anns_field + ); +}; + export const buildSearchRequest = ( params: SearchReq | SearchSimpleReq | HybridSearchReq, collectionInfo: DescribeCollectionResponse, @@ -274,12 +287,7 @@ export const buildSearchRequest = ( const requests: FormatedSearchRequest[] = []; // detect if the request is hybrid search request - const isHybridSearch = !!( - searchHybridReq.data && - searchHybridReq.data.length && - typeof searchHybridReq.data[0] === 'object' && - searchHybridReq.data[0].anns_field - ); + const isHybridSearch = isHybridSearchRequest(searchHybridReq); const searchAggregation = searchSimpleReq.search_aggregation || searchReq.search_aggregation; diff --git a/test/grpc/Basic.spec.ts b/test/grpc/Basic.spec.ts index 59486941..63664e78 100644 --- a/test/grpc/Basic.spec.ts +++ b/test/grpc/Basic.spec.ts @@ -1,4 +1,10 @@ -import { MilvusClient, ErrorCode, DataType, FieldType } from '../../milvus'; +import { + MilvusClient, + ErrorCode, + DataType, + FieldType, + newClientRequestId, +} from '../../milvus'; import { IP, GENERATE_NAME, generateInsertData } from '../tools'; const milvusClient = new MilvusClient({ @@ -162,7 +168,7 @@ describe(`Basic API without database`, () => { }); it(`Create collection with traceid should be successful`, async () => { - const traceId = 'test-trace-id-' + Date.now(); + const traceId = newClientRequestId(); const res = await milvusClient.createCollection({ collection_name: COLLECTION_NAME + '_trace', fields: schema, @@ -172,7 +178,7 @@ describe(`Basic API without database`, () => { }); it(`Create collection with traceid (alternative format) should be successful`, async () => { - const traceId = 'test-trace-id-alt-' + Date.now(); + const traceId = newClientRequestId(); const res = await milvusClient.createCollection({ collection_name: COLLECTION_NAME + '_trace_alt', fields: schema, diff --git a/test/grpc/MilvusClient.spec.ts b/test/grpc/MilvusClient.spec.ts index d98ac144..4bf977c2 100644 --- a/test/grpc/MilvusClient.spec.ts +++ b/test/grpc/MilvusClient.spec.ts @@ -252,6 +252,17 @@ describe(`Milvus client`, () => { expect(client.channelOptions.interceptors.length).toBeGreaterThanOrEqual(3); }); + it(`should omit the ordinary retry interceptor from telemetry heartbeats`, async () => { + const client = new MilvusClient({ + address: IP, + __SKIP_CONNECT__: true, + }); + + expect((client as any).telemetryChannelOptions.interceptors).toHaveLength( + client.channelOptions.interceptors.length - 1 + ); + }); + it(`Expect get node sdk info`, async () => { expect(MilvusClient.sdkInfo.version).toEqual(sdkInfo.version); expect(MilvusClient.sdkInfo.recommendMilvus).toEqual(sdkInfo.milvusVersion); diff --git a/test/telemetry/ClientTelemetry.e2e.spec.ts b/test/telemetry/ClientTelemetry.e2e.spec.ts new file mode 100644 index 00000000..577ab616 --- /dev/null +++ b/test/telemetry/ClientTelemetry.e2e.spec.ts @@ -0,0 +1,214 @@ +import crypto from 'crypto'; +import fetch from 'node-fetch'; +import { MilvusClient, newClientRequestId } from '../../milvus'; + +const runE2E = + process.env.MILVUS_TELEMETRY_E2E === 'true' ? describe : describe.skip; +const address = process.env.MILVUS_ADDRESS || '127.0.0.1:19530'; +const telemetryApi = + process.env.MILVUS_TELEMETRY_API || 'http://127.0.0.1:9091/api/v1/_telemetry'; + +runE2E('client telemetry local E2E', () => { + jest.setTimeout(30_000); + + it('registers default telemetry automatically', async () => { + const client = new MilvusClient({ address }); + try { + const manager = client.getTelemetry(); + expect(manager.stableClientId).toBe(false); + await waitFor( + 'default client registration', + manager.clientId, + state => state.status === 'active' + ); + } finally { + await client.closeConnection(); + } + }); + + it('round-trips metrics, commands, config, and request IDs', async () => { + const clientId = `e2e-node-${crypto.randomUUID()}`; + const client = new MilvusClient({ + address, + telemetry: { + heartbeatIntervalMs: 500, + samplingRate: 1, + clientId, + }, + }); + + try { + const manager = client.getTelemetry(); + expect(manager.clientId).toBe(clientId); + await waitFor( + 'client registration', + clientId, + state => state.status === 'active' + ); + + const analyzer = await client.runAnalyzer({ + text: 'hello milvus telemetry', + analyzer_params: { type: 'standard' }, + with_detail: true, + }); + expect( + analyzer.results[0].tokens.map(token => + Buffer.from(token.token as any).toString() + ) + ).toEqual(['hello', 'milvus', 'telemetry']); + await waitFor('RunAnalyzer metric', clientId, state => + hasMetric(state, 'RunAnalyzer', 'success_count', 1) + ); + + const collectionsCommand = await pushCommand( + clientId, + 'collection_metrics', + { collections: ['*'], enabled: true } + ); + expect((await waitForReply(clientId, collectionsCommand)).success).toBe( + true + ); + + const requestId = newClientRequestId(); + const query = await client.query({ + collection_name: 'telemetry_e2e_missing', + filter: 'id > 0', + client_request_id: requestId, + }); + expect(query.status.error_code).not.toBe('Success'); + await waitFor('failed Query collection metric', clientId, state => + hasMetric(state, 'Query', 'error_count', 1, 'telemetry_e2e_missing') + ); + + const errorsCommand = await pushCommand(clientId, 'show_errors', { + max_count: 10, + }); + const errorsReply = await waitForReply(clientId, errorsCommand); + expect(errorsReply.success).toBe(true); + expect(JSON.parse(errorsReply.payload)).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + operation: 'Query', + request_id: requestId, + }), + ]) + ); + + const configPayload = { + sampling_rate: 0.75, + heartbeat_interval_ms: 600, + }; + const configCommand = await pushCommand( + clientId, + 'push_config', + configPayload, + true + ); + expect((await waitForReply(clientId, configCommand)).success).toBe(true); + expect(manager.configHash).toBe( + crypto + .createHash('sha256') + .update(configCommand) + .update('push_config') + .update(JSON.stringify(configPayload)) + .digest('hex') + .slice(0, 16) + ); + expect(manager.lastCommandTimestamp).toBeGreaterThan(0); + + const getConfigReply = await waitForReply( + clientId, + await pushCommand(clientId, 'get_config', {}) + ); + const userConfig = JSON.parse(getConfigReply.payload).user_config; + expect(userConfig.telemetry_sampling_rate).toBe(0.75); + expect(userConfig.telemetry_heartbeat_interval_ms).toBe(600); + expect(userConfig.all_collections_enabled).toBe(true); + } finally { + await client.closeConnection(); + } + }); +}); + +async function clientState(clientId: string): Promise { + const query = new URLSearchParams({ + client_id: clientId, + include_metrics: 'true', + }); + const response = await fetch(`${telemetryApi}/clients?${query}`); + expect(response.ok).toBe(true); + const body = (await response.json()) as any; + return body.clients?.[0]; +} + +async function waitFor( + label: string, + clientId: string, + predicate: (state: any) => boolean +): Promise { + const deadline = Date.now() + 15_000; + let last: any; + while (Date.now() < deadline) { + last = await clientState(clientId); + if (last && predicate(last)) { + return last; + } + await new Promise(resolve => setTimeout(resolve, 250)); + } + throw new Error( + `Timed out waiting for ${label}; last=${JSON.stringify(last)}` + ); +} + +async function pushCommand( + clientId: string, + commandType: string, + payload: Record, + persistent = false +): Promise { + const response = await fetch(`${telemetryApi}/commands`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + command_type: commandType, + target_client_id: clientId, + payload, + ttl_seconds: 30, + persistent, + }), + }); + expect(response.ok).toBe(true); + return ((await response.json()) as any).command_id; +} + +async function waitForReply(clientId: string, commandId: string): Promise { + const state = await waitFor( + `command reply ${commandId}`, + clientId, + candidate => + candidate.command_replies?.some( + (reply: any) => reply.command_id === commandId + ) + ); + return state.command_replies.find( + (reply: any) => reply.command_id === commandId + ); +} + +function hasMetric( + state: any, + operation: string, + counter: string, + minimum: number, + collection?: string +): boolean { + return (state.metrics || []).some((metric: any) => { + if ( + metric.operation !== operation || + Number(metric.global?.[counter] || 0) < minimum + ) { + return false; + } + return !collection || Boolean(metric.collection_metrics?.[collection]); + }); +} diff --git a/test/telemetry/ClientTelemetry.spec.ts b/test/telemetry/ClientTelemetry.spec.ts new file mode 100644 index 00000000..2c397e1f --- /dev/null +++ b/test/telemetry/ClientTelemetry.spec.ts @@ -0,0 +1,946 @@ +import { + ClientTelemetryManager, + newClientRequestId, +} from '../../milvus/telemetry'; +import { status as grpcStatus } from '@grpc/grpc-js'; + +describe('ClientTelemetryManager', () => { + it('matches the cross-SDK persistent command hash vector', () => { + expect( + ClientTelemetryManager.calculateConfigHash([ + { + command_id: 'cfg-b', + command_type: 'push_config', + payload: Buffer.from('{"sampling_rate":0.5}'), + persistent: true, + }, + { + command_id: 'cfg-a', + command_type: 'push_config', + payload: Buffer.from('{"heartbeat_interval_ms":5000}'), + persistent: true, + }, + ]) + ).toBe('a271ff0bb1941777'); + }); + + it('sorts persistent command IDs by UTF-8 bytes, not host locale', () => { + expect( + ClientTelemetryManager.calculateConfigHash([ + { + command_id: 'a', + command_type: 'push_config', + payload: Buffer.from('A'), + persistent: true, + }, + { + command_id: 'Z', + command_type: 'push_config', + payload: Buffer.from('B'), + persistent: true, + }, + ]) + ).toBe('0e793afed772d6a5'); + }); + + it('applies built-in commands and deduplicates command IDs', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + let calls = 0; + manager.registerCommandHandler('custom', command => { + calls += 1; + return { command_id: command.command_id, success: true }; + }); + + await manager.processCommands([ + { + command_id: 'config', + command_type: 'push_config', + payload: Buffer.from( + '{"heartbeat_interval_ms":5000,"sampling_rate":0.25}' + ), + create_time: 1, + persistent: true, + }, + { + command_id: 'custom', + command_type: 'custom', + create_time: 2, + }, + ]); + await manager.processCommands([ + { + command_id: 'custom', + command_type: 'custom', + create_time: 2, + }, + ]); + await manager.processCommands([ + { + command_id: 'custom', + command_type: 'custom', + create_time: 2, + }, + ]); + + expect(manager.getConfig().heartbeatIntervalMs).toBe(5000); + expect(manager.getConfig().samplingRate).toBe(0.25); + expect(manager.lastCommandTimestamp).toBe(2); + expect(manager.configHash).not.toBe(''); + expect(calls).toBe(1); + manager.stop(); + }); + + it('retains command IDs at the cursor timestamp for exact deduplication', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + const executed: string[] = []; + manager.registerCommandHandler('custom', command => { + executed.push(command.command_id); + return { command_id: command.command_id, success: true }; + }); + + await manager.processCommands([ + { + command_id: 'equal-a', + command_type: 'custom', + create_time: 10, + }, + ]); + await manager.processCommands([ + { + command_id: 'equal-b', + command_type: 'custom', + create_time: 10, + }, + { + command_id: 'equal-a', + command_type: 'custom', + create_time: 10, + }, + ]); + + expect(executed).toEqual(['equal-a', 'equal-b']); + expect(manager.lastCommandTimestamp).toBe(10); + expect([...(manager as any).executedCommands.keys()]).toEqual([ + 'equal-a', + 'equal-b', + ]); + manager.stop(); + }); + + it.each([ + ['wrong', { command_id: 'wrong-id', success: true }], + ['empty', { command_id: '', success: true }], + ['missing', undefined], + ])( + 'keeps the server command ID when a custom handler returns a %s reply', + async (_case, customReply) => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + manager.registerCommandHandler('custom', (() => customReply) as any); + + await manager.processCommands([ + { + command_id: 'server-command-id', + command_type: 'custom', + create_time: 1, + }, + ]); + + expect((manager as any).pendingReplies[0]).toEqual( + expect.objectContaining({ + command_id: 'server-command-id', + success: customReply ? true : false, + }) + ); + expect(manager.lastCommandTimestamp).toBe(1); + manager.stop(); + } + ); + + it('rejects invalid heartbeat intervals without changing the current interval', async () => { + expect( + () => + new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { heartbeatIntervalMs: Number.NaN }, + }) + ).toThrow('heartbeatIntervalMs must be a finite positive number'); + + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false, heartbeatIntervalMs: 5000 }, + }); + const reply = await (manager as any).handleCommand({ + command_id: 'invalid-config', + command_type: 'push_config', + payload: Buffer.from('{"heartbeat_interval_ms":"not-a-number"}'), + }); + + expect(reply.success).toBe(false); + expect(reply.error_message).toContain('heartbeat_interval_ms'); + expect(manager.getConfig().heartbeatIntervalMs).toBe(5000); + manager.stop(); + }); + + it('applies push_config atomically and reports applied and ignored keys', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { + enabled: true, + heartbeatIntervalMs: 5000, + samplingRate: 0.75, + }, + }); + + const rejected = await (manager as any).handleCommand({ + command_id: 'invalid-combined-config', + command_type: 'push_config', + payload: Buffer.from( + '{"enabled":false,"heartbeat_interval_ms":-1,"sampling_rate":0.1}' + ), + }); + expect(rejected.success).toBe(false); + expect(manager.getConfig()).toEqual( + expect.objectContaining({ + enabled: true, + heartbeatIntervalMs: 5000, + samplingRate: 0.75, + }) + ); + + // ttl_seconds belongs to the server-side push API and never reaches clients by + // design: a stray value is reported as ignored, never validated. + const strayTtl = await (manager as any).handleCommand({ + command_id: 'stray-ttl', + command_type: 'push_config', + payload: Buffer.from('{"enabled":false,"ttl_seconds":"60"}'), + }); + expect(strayTtl.success).toBe(true); + expect(JSON.parse(strayTtl.payload.toString())).toEqual({ + applied: ['enabled'], + ignored: ['ttl_seconds'], + }); + expect(manager.getConfig().enabled).toBe(false); + + const applied = await (manager as any).handleCommand({ + command_id: 'valid-config', + command_type: 'push_config', + payload: Buffer.from( + '{"sampling_rate":1.5,"enabled":false,"ttl_seconds":60,"future_key":1}' + ), + }); + expect(applied.success).toBe(true); + expect(JSON.parse(applied.payload.toString())).toEqual({ + applied: ['enabled', 'sampling_rate'], + ignored: ['future_key', 'ttl_seconds'], + }); + expect(manager.getConfig()).toEqual( + expect.objectContaining({ enabled: false, samplingRate: 1 }) + ); + manager.stop(); + }); + + it('serializes concurrent command batches before deduplicating IDs', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + let calls = 0; + let release!: () => void; + const gate = new Promise(resolve => { + release = resolve; + }); + manager.registerCommandHandler('slow', async command => { + calls += 1; + await gate; + return { command_id: command.command_id, success: true }; + }); + const command = { + command_id: 'same-id', + command_type: 'slow', + create_time: 1, + }; + + const first = manager.processCommands([command]); + const second = manager.processCommands([command]); + release(); + await Promise.all([first, second]); + + expect(calls).toBe(1); + manager.stop(); + }); + + it('fails recursive command processing without deadlocking the queue', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + let innerCalls = 0; + let laterCalls = 0; + manager.registerCommandHandler('outer', async command => { + await manager.processCommands([ + { + command_id: 'inner-id', + command_type: 'inner', + create_time: 1, + }, + ]); + return { command_id: command.command_id, success: true }; + }); + manager.registerCommandHandler('inner', command => { + innerCalls += 1; + return { command_id: command.command_id, success: true }; + }); + manager.registerCommandHandler('later', command => { + laterCalls += 1; + return { command_id: command.command_id, success: true }; + }); + + await manager.processCommands([ + { command_id: 'outer-id', command_type: 'outer', create_time: 1 }, + ]); + await manager.processCommands([ + { command_id: 'later-id', command_type: 'later', create_time: 2 }, + ]); + + expect(innerCalls).toBe(0); + expect(laterCalls).toBe(1); + expect((manager as any).pendingReplies).toEqual([ + expect.objectContaining({ + command_id: 'outer-id', + success: false, + error_message: + 'processCommands cannot be called recursively from a command handler', + }), + expect.objectContaining({ command_id: 'later-id', success: true }), + ]); + manager.stop(); + }); + + it('strictly validates collection_metrics payloads', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + config: { enabled: false }, + }); + const malformed = await (manager as any).handleCommand({ + command_id: 'bad-collections', + command_type: 'collection_metrics', + payload: Buffer.from('{"enabled":false,"collections":"books"}'), + }); + + expect(malformed.success).toBe(false); + expect(malformed.error_message).toContain('array of strings'); + manager.stop(); + }); + + it('returns operation-keyed latency detail metrics', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now() - 5, + }); + (manager as any).createSnapshot(); + const snapshot = manager.getMetricsSnapshots()[0]; + + const reply = await (manager as any).handleCommand({ + command_id: 'latency-detail', + command_type: 'show_latency_history', + payload: Buffer.from( + JSON.stringify({ + start_time: new Date(snapshot.timestamp - 1).toISOString(), + end_time: new Date(snapshot.end_time + 1).toISOString(), + detail: true, + }) + ), + }); + const body = JSON.parse(reply.payload.toString()); + + expect(reply.success).toBe(true); + expect(body.total_snapshots).toBe(1); + expect(Array.isArray(body.snapshots[0].metrics)).toBe(false); + expect(body.snapshots[0].metrics.Search).toEqual( + expect.objectContaining({ + request_count: 1, + success_count: 1, + error_count: 0, + }) + ); + manager.stop(); + }); + + it('aggregates p99 from bounded cross-window latency samples', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + for (let index = 0; index < 256; index += 1) { + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now() - 1, + }); + } + (manager as any).createSnapshot(); + for (let index = 0; index < 256; index += 1) { + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now() - 100, + }); + } + (manager as any).createSnapshot(); + + const snapshots = manager.getMetricsSnapshots(); + const retainedSamples = snapshots.map( + snapshot => (manager as any).snapshotLatencySamples.get(snapshot).Search + ); + expect( + retainedSamples.every(samples => samples instanceof Float64Array) + ).toBe(true); + expect(retainedSamples.map(samples => samples.length)).toEqual([128, 128]); + expect(JSON.stringify(snapshots)).not.toContain('latency_samples'); + + const reply = await (manager as any).handleCommand({ + command_id: 'latency-aggregate', + command_type: 'show_latency_history', + payload: Buffer.from( + JSON.stringify({ + start_time: new Date(snapshots[0].timestamp - 1).toISOString(), + end_time: new Date( + snapshots[snapshots.length - 1].end_time + 1 + ).toISOString(), + detail: false, + }) + ), + }); + const body = JSON.parse(reply.payload.toString()); + const search = body.aggregated.metrics.Search; + + expect(reply.success).toBe(true); + expect(body.snapshot_count).toBe(2); + expect(search.request_count).toBe(512); + expect(search.avg_latency_ms).toBeGreaterThan(40); + expect(search.avg_latency_ms).toBeLessThan(70); + // Averaging the two window p99s would land near 50 ms. The combined p99 + // belongs to the equally sized slow window. + expect(search.p99_latency_ms).toBeGreaterThan(90); + manager.stop(); + }); + + it('generates a lowercase 128-bit client request ID', () => { + expect(newClientRequestId()).toMatch(/^[0-9a-f]{32}$/); + }); + + it('backs off when the telemetry service is unimplemented', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => { + throw Object.assign(new Error('unimplemented'), { + code: grpcStatus.UNIMPLEMENTED, + }); + }, + }); + + await (manager as any).sendHeartbeat(); + + expect(manager.isSupported()).toBe(false); + manager.stop(); + }); + + it('clears unsupported backoff on a real response with a business error', async () => { + let calls = 0; + const manager = new ClientTelemetryManager({ + sender: async () => { + calls += 1; + if (calls === 1) { + throw Object.assign(new Error('unimplemented'), { + code: grpcStatus.UNIMPLEMENTED, + }); + } + return { status: { code: 1, reason: 'not ready' } }; + }, + }); + + await (manager as any).sendHeartbeat(); + expect(manager.isSupported()).toBe(false); + await (manager as any).sendHeartbeat(); + + expect(manager.isSupported()).toBe(true); + expect(manager.lastHeartbeatError).toEqual( + expect.objectContaining({ message: 'not ready' }) + ); + manager.stop(); + }); + + it('keeps control-plane heartbeats alive while operation telemetry is disabled', async () => { + const disableCommand = { + command_id: 'disable', + command_type: 'push_config', + payload: Buffer.from('{"enabled":false}'), + create_time: 1, + persistent: true, + }; + const enableCommand = { + command_id: 'enable', + command_type: 'push_config', + payload: Buffer.from('{"enabled":true}'), + create_time: 2, + persistent: true, + }; + const requests: any[] = []; + let releaseEnable!: () => void; + const enableGate = new Promise(resolve => { + releaseEnable = resolve; + }); + let secondHeartbeat!: () => void; + const secondHeartbeatReceived = new Promise(resolve => { + secondHeartbeat = resolve; + }); + let thirdHeartbeat!: () => void; + const thirdHeartbeatReceived = new Promise(resolve => { + thirdHeartbeat = resolve; + }); + const manager = new ClientTelemetryManager({ + config: { heartbeatIntervalMs: 100 }, + sender: async request => { + requests.push(request); + if (requests.length === 1) { + return { + status: { error_code: 'Success' }, + commands: [disableCommand], + }; + } + if (requests.length === 2) { + secondHeartbeat(); + await enableGate; + return { + status: { error_code: 'Success' }, + commands: [enableCommand], + }; + } + thirdHeartbeat(); + return { status: { error_code: 'Success' } }; + }, + }); + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now() - 1, + }); + + manager.start(); + while (manager.getConfig().enabled) { + await new Promise(resolve => setTimeout(resolve, 0)); + } + await secondHeartbeatReceived; + expect(manager.getConfig().enabled).toBe(false); + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now() - 1, + }); + releaseEnable(); + await thirdHeartbeatReceived; + manager.stop(); + + expect(requests[0].metrics).toHaveLength(1); + expect(requests[1].metrics).toEqual([]); + expect(requests[1].command_replies).toEqual([ + expect.objectContaining({ command_id: 'disable', success: true }), + ]); + expect(requests[1].config_hash).toBe( + ClientTelemetryManager.calculateConfigHash([disableCommand]) + ); + expect(requests[2].metrics).toEqual([]); + expect(requests[2].command_replies).toEqual([ + expect.objectContaining({ command_id: 'enable', success: true }), + ]); + expect(requests[2].config_hash).toBe( + ClientTelemetryManager.calculateConfigHash([enableCommand]) + ); + expect(manager.getConfig().enabled).toBe(true); + }); + + it('does not heartbeat when the initial config explicitly opts out', async () => { + const sender = jest.fn(async () => ({ + status: { error_code: 'Success' }, + })); + const manager = new ClientTelemetryManager({ + sender, + config: { enabled: false, heartbeatIntervalMs: 1 }, + }); + + manager.start(); + await new Promise(resolve => setTimeout(resolve, 10)); + manager.stop(); + + expect(sender).not.toHaveBeenCalled(); + }); + + it('chunks heartbeat delays beyond the Node timer limit', async () => { + jest.useFakeTimers(); + const maxTimerDelayMs = 2_147_483_647; + const sender = jest.fn(async () => ({ + status: { error_code: 'Success' }, + })); + const manager = new ClientTelemetryManager({ + sender, + config: { heartbeatIntervalMs: maxTimerDelayMs + 1_000 }, + }); + try { + (manager as any).scheduleNextHeartbeat(maxTimerDelayMs + 1_000); + + jest.advanceTimersByTime(maxTimerDelayMs); + expect(sender).not.toHaveBeenCalled(); + + jest.advanceTimersByTime(1_000); + await Promise.resolve(); + expect(sender).toHaveBeenCalledTimes(1); + } finally { + manager.stop(); + jest.useRealTimers(); + } + }); + + it('continues the heartbeat loop after an unexpected collector failure', async () => { + let heartbeatObserved!: () => void; + const heartbeat = new Promise(resolve => { + heartbeatObserved = resolve; + }); + const sender = jest.fn(async () => { + heartbeatObserved(); + return { status: { error_code: 'Success' } }; + }); + const manager = new ClientTelemetryManager({ + sender, + config: { heartbeatIntervalMs: 1 }, + }); + const createSnapshot = (manager as any).createSnapshot.bind(manager); + let attempts = 0; + (manager as any).createSnapshot = jest.fn(() => { + attempts += 1; + if (attempts === 1) { + throw new Error('collector failed'); + } + createSnapshot(); + }); + + manager.start(); + await heartbeat; + manager.stop(); + + expect(attempts).toBeGreaterThanOrEqual(2); + expect(sender).toHaveBeenCalled(); + }); + + it('retains at most one hour of snapshots with a hard memory cap', () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + const now = Date.now(); + const snapshots = (manager as any).snapshots as any[]; + const latencySamples = (manager as any).snapshotLatencySamples as Map< + object, + object + >; + const appendSnapshot = (snapshot: any) => { + snapshots.push(snapshot); + latencySamples.set(snapshot, { Search: new Float64Array([1]) }); + }; + appendSnapshot({ + timestamp: now - 60 * 60 * 1000 - 2, + end_time: now - 60 * 60 * 1000 - 1, + metrics: [], + }); + for (let index = 0; index < 4098; index += 1) { + appendSnapshot({ + timestamp: now - 1000 + index, + end_time: now - 999 + index, + metrics: [], + index, + }); + } + + const retained = manager.getMetricsSnapshots() as any[]; + + expect(retained).toHaveLength(4096); + expect(retained[0].index).toBe(2); + expect(latencySamples.size).toBe(4096); + expect( + retained.every(snapshot => snapshot.end_time >= now - 60 * 60 * 1000) + ).toBe(true); + manager.stop(); + }); + + it('drops a stale endpoint response without mutating telemetry state', async () => { + let endpointEpoch = 1; + let resolveHeartbeat!: (response: any) => void; + const sender = jest.fn( + () => + new Promise(resolve => { + resolveHeartbeat = resolve; + }) + ); + const manager = new ClientTelemetryManager({ + sender, + senderEpochProvider: () => endpointEpoch, + }); + manager.registerCommandHandler('baseline', command => ({ + command_id: command.command_id, + success: true, + })); + await manager.processCommands([ + { + command_id: 'baseline-command', + command_type: 'baseline', + payload: Buffer.from('baseline'), + create_time: 1, + persistent: true, + }, + ]); + (manager as any).unsupportedStreak = 1; + const oldError = new Error('old endpoint was unsupported'); + manager.lastHeartbeatError = oldError; + + const before = { + config: manager.getConfig(), + configHash: manager.configHash, + lastCommandTimestamp: manager.lastCommandTimestamp, + pendingReplyIds: (manager as any).pendingReplies.map( + (reply: any) => reply.command_id + ), + executedCommandIds: [...(manager as any).executedCommands.keys()], + snapshots: manager.getMetricsSnapshots(), + }; + + const heartbeat = (manager as any).sendHeartbeat(); + await Promise.resolve(); + expect(sender).toHaveBeenCalledTimes(1); + + // Publish a new endpoint before the old transport completes. + endpointEpoch += 1; + resolveHeartbeat({ + status: { error_code: 'Success' }, + commands: [ + { + command_id: 'stale-config', + command_type: 'push_config', + payload: Buffer.from( + '{"enabled":false,"heartbeat_interval_ms":1234,"sampling_rate":0.1}' + ), + create_time: 2, + persistent: true, + }, + ], + }); + await heartbeat; + + expect(manager.getConfig()).toEqual(before.config); + expect(manager.configHash).toBe(before.configHash); + expect(manager.lastCommandTimestamp).toBe(before.lastCommandTimestamp); + expect( + (manager as any).pendingReplies.map((reply: any) => reply.command_id) + ).toEqual(before.pendingReplyIds); + expect([...(manager as any).executedCommands.keys()]).toEqual( + before.executedCommandIds + ); + expect(manager.getMetricsSnapshots()).toEqual(before.snapshots); + expect(manager.isSupported()).toBe(false); + expect(manager.lastHeartbeatError).toBe(oldError); + manager.stop(); + }); + + it('stops an old endpoint batch when a handler changes sender epoch', async () => { + let endpointEpoch = 1; + const commands = [ + { + command_id: 'switch-generation', + command_type: 'switch-generation', + create_time: 1, + persistent: true, + payload: Buffer.from('switch'), + }, + { + command_id: 'followup', + command_type: 'followup', + create_time: 2, + persistent: true, + payload: Buffer.from('followup'), + }, + ]; + const manager = new ClientTelemetryManager({ + sender: async () => ({ + status: { error_code: 'Success' }, + commands, + }), + senderEpochProvider: () => endpointEpoch, + }); + let switchCalls = 0; + let followupCalls = 0; + manager.registerCommandHandler('switch-generation', async () => { + switchCalls += 1; + await Promise.resolve(); + endpointEpoch += 1; + return { command_id: 'wrong-id', success: true }; + }); + manager.registerCommandHandler('followup', command => { + followupCalls += 1; + return { command_id: command.command_id, success: true }; + }); + + await (manager as any).sendHeartbeat(); + + expect(switchCalls).toBe(1); + expect(followupCalls).toBe(0); + expect(manager.lastCommandTimestamp).toBe(0); + expect(manager.configHash).toBe(''); + expect((manager as any).pendingReplies).toEqual([ + expect.objectContaining({ + command_id: 'switch-generation', + success: true, + }), + ]); + + await (manager as any).sendHeartbeat(); + + expect(switchCalls).toBe(1); + expect(followupCalls).toBe(1); + expect(manager.lastCommandTimestamp).toBe(2); + expect(manager.configHash).toBe( + ClientTelemetryManager.calculateConfigHash(commands) + ); + manager.stop(); + }); + + it('filters snapshot collection metrics against the current wire scope', async () => { + let heartbeat: any; + const manager = new ClientTelemetryManager({ + sender: async request => { + heartbeat = request; + return { status: { error_code: 'Success' } }; + }, + }); + await manager.processCommands([ + { + command_id: 'scope-on', + command_type: 'collection_metrics', + payload: Buffer.from('{"enabled":true,"collections":["books"]}'), + create_time: 1, + }, + ]); + manager.recordOperation({ + operation: 'Search', + collection: 'books', + startTime: performance.now(), + }); + (manager as any).createSnapshot(); + await manager.processCommands([ + { + command_id: 'scope-off', + command_type: 'collection_metrics', + payload: Buffer.from('{"enabled":false,"collections":["books"]}'), + create_time: 2, + }, + ]); + + await (manager as any).sendHeartbeat(); + + expect(heartbeat.metrics[0].collection_metrics).toEqual({}); + manager.stop(); + }); + + it('requires RFC3339 timestamps and a boolean history detail flag', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + const missingZone = await (manager as any).handleCommand({ + command_id: 'missing-zone', + command_type: 'show_latency_history', + payload: Buffer.from( + '{"start_time":"2026-01-01T00:00:00","end_time":"2026-01-01T00:01:00Z"}' + ), + }); + const stringDetail = await (manager as any).handleCommand({ + command_id: 'string-detail', + command_type: 'show_latency_history', + payload: Buffer.from( + '{"start_time":"2026-01-01T00:00:00Z","end_time":"2026-01-01T00:01:00Z","detail":"false"}' + ), + }); + + expect(missingZone.success).toBe(false); + expect(missingZone.error_message).toContain('RFC3339'); + expect(stringDetail.success).toBe(false); + expect(stringDetail.error_message).toContain('boolean'); + manager.stop(); + }); + + it('truncates a single oversized error reply', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + manager.recordOperation({ + operation: 'Query', + collection: 'books', + startTime: performance.now(), + error: new Error('x'.repeat(2 * 1024 * 1024)), + }); + + const reply = await (manager as any).handleCommand({ + command_id: 'errors', + command_type: 'show_errors', + }); + + expect(reply.success).toBe(true); + expect(reply.payload.length).toBeLessThanOrEqual(1024 * 1024); + manager.stop(); + }); + + it('bounds an oversized non-message error field without looping', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + manager.recordOperation({ + operation: 'Query', + collection: 'c'.repeat(2 * 1024 * 1024), + startTime: performance.now(), + error: new Error('xy'), + }); + + const reply = await (manager as any).handleCommand({ + command_id: 'errors-with-large-collection', + command_type: 'show_errors', + }); + const errors = JSON.parse(reply.payload.toString()); + + expect(reply.success).toBe(true); + expect(reply.payload.length).toBeLessThanOrEqual(1024 * 1024); + expect(errors[0].error_msg).toBe('xy'); + expect(errors[0].collection).toContain('...(truncated)'); + manager.stop(); + }); + + it('returns an empty payload when there are no recent errors', async () => { + const manager = new ClientTelemetryManager({ + sender: async () => ({ status: { error_code: 'Success' } }), + }); + + const reply = await (manager as any).handleCommand({ + command_id: 'no-errors', + command_type: 'show_errors', + }); + + expect(reply.success).toBe(true); + expect(reply.payload).toHaveLength(0); + manager.stop(); + }); +}); diff --git a/test/telemetry/GrpcClientTelemetry.spec.ts b/test/telemetry/GrpcClientTelemetry.spec.ts new file mode 100644 index 00000000..176f3bf9 --- /dev/null +++ b/test/telemetry/GrpcClientTelemetry.spec.ts @@ -0,0 +1,68 @@ +import { Metadata } from '@grpc/grpc-js'; +import { MilvusClient } from '../../milvus'; + +describe('GRPCClient telemetry transport', () => { + function installHeartbeatCapture(client: MilvusClient, requests: any[]) { + (client as any).telemetryClient = { + ClientHeartbeat: ( + request: any, + metadata: Metadata, + options: any, + callback: Function + ) => { + requests.push({ request, metadata, options }); + callback(null, { status: { error_code: 'Success' } }); + }, + close: jest.fn(), + }; + } + + it('uses the grpc-js metadata overload for heartbeat unary calls', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const requests: any[] = []; + installHeartbeatCapture(client, requests); + + await (client.getTelemetry() as any).sendHeartbeat(); + + expect(requests[0].metadata).toBeInstanceOf(Metadata); + expect(requests[0].options.deadline).toBeInstanceOf(Date); + await client.closeConnection(); + }); + + it('omits an unset database but reports an explicitly selected default', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const requests: any[] = []; + installHeartbeatCapture(client, requests); + + await (client.getTelemetry() as any).sendHeartbeat(); + await client.useDatabase({ db_name: 'default' }); + await (client.getTelemetry() as any).sendHeartbeat(); + + expect(requests[0].request.client_info.reserved).not.toHaveProperty( + 'db_name' + ); + expect(requests[1].request.client_info.reserved.db_name).toBe('default'); + await client.closeConnection(); + }); + + it('reports an explicitly configured default database', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + database: 'default', + __SKIP_CONNECT__: true, + }); + const requests: any[] = []; + installHeartbeatCapture(client, requests); + + await (client.getTelemetry() as any).sendHeartbeat(); + + expect(requests[0].request.client_info.reserved.db_name).toBe('default'); + await client.closeConnection(); + }); +}); diff --git a/test/utils/Collection.spec.ts b/test/utils/Collection.spec.ts index 8711d58a..4e4a4ca6 100644 --- a/test/utils/Collection.spec.ts +++ b/test/utils/Collection.spec.ts @@ -46,12 +46,16 @@ const externalCollectionAlterUnsupportedStatus: ResStatus = { }; const respondWith = - (response: any) => (_request: any, _options: any, callback: RpcCallback) => + (response: any) => (...args: any[]) => { + const callback = args[args.length - 1] as RpcCallback; callback(null, response); + }; const failWith = - (error: any) => (_request: any, _options: any, callback: RpcCallback) => + (error: any) => (...args: any[]) => { + const callback = args[args.length - 1] as RpcCallback; callback(error); + }; const createTestClient = () => { const rpcClient = { @@ -104,7 +108,7 @@ describe('collection schema alteration', () => { const result = await client.addCollectionField({ collection_name: COLLECTION_NAME, db_name: 'db1', - client_request_id: 'request-1', + client_request_id: '11111111111111111111111111111111', field: { name: 'age', data_type: DataType.Int64, @@ -150,10 +154,8 @@ describe('collection schema alteration', () => { .default_value ).toEqual(expect.objectContaining({ long_data: '42' })); expect( - rpcClient.AlterCollectionSchema.mock.calls[0][1].metadata.get( - 'client-request-id' - ) - ).toEqual(['request-1']); + rpcClient.AlterCollectionSchema.mock.calls[0][1].get('client-request-id') + ).toEqual(['11111111111111111111111111111111']); expect(cache.has(cacheKey('db1'))).toBe(false); }); diff --git a/test/utils/Data.spec.ts b/test/utils/Data.spec.ts index 8d3fdc49..01034710 100644 --- a/test/utils/Data.spec.ts +++ b/test/utils/Data.spec.ts @@ -18,9 +18,214 @@ import { findKeyValue, FieldPartialUpdateOpType, CLUSTER_ID, + setPoolTelemetryManager, } from '../../milvus'; describe('utils/Data', () => { + it('records public validation failures before any RPC is attempted', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const pool: any = { + acquire: jest.fn(), + release: jest.fn(), + }; + const telemetry = { recordOperation: jest.fn() }; + (client as any).channelPool = pool; + setPoolTelemetryManager(pool, telemetry); + + await expect( + client.insert({ collection_name: 'test_collection' } as any) + ).rejects.toThrow(ERROR_REASONS.INSERT_CHECK_FIELD_DATA_IS_REQUIRED); + + expect(pool.acquire).not.toHaveBeenCalled(); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Insert', + collection: 'test_collection', + error: expect.any(Error), + }) + ); + }); + + it.each([ + ['Insert', 'insert'], + ['Upsert', 'upsert'], + ] as const)( + 'records one %s outcome across a SchemaMismatch rebuild', + async (operation, method) => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + (client as any).describeCollection = jest.fn().mockResolvedValue({ + status: { error_code: ErrorCode.SUCCESS, reason: '' }, + schema: { + enable_dynamic_field: false, + fields: [ + { + name: 'id', + data_type: DataType.Int64, + is_primary_key: true, + autoID: false, + nullable: false, + element_type: DataType.None, + }, + ], + }, + properties: [], + }); + let rpcCalls = 0; + const rpc = jest.fn((params: any, options: any, callback: any) => { + rpcCalls += 1; + callback(null, { + status: { + error_code: + rpcCalls === 1 ? ErrorCode.SchemaMismatch : ErrorCode.SUCCESS, + reason: rpcCalls === 1 ? 'stale schema' : '', + }, + }); + }); + const pool: any = { + acquire: jest.fn().mockResolvedValue({ [operation]: rpc }), + release: jest.fn(), + }; + const telemetry = { recordOperation: jest.fn() }; + (client as any).channelPool = pool; + setPoolTelemetryManager(pool, telemetry); + + await (client as any)[method]({ + collection_name: 'test_collection', + fields_data: [{ id: 1 }], + }); + + expect(rpc).toHaveBeenCalledTimes(2); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation, + collection: 'test_collection', + error: undefined, + }) + ); + } + ); + + it('records response postprocessing failures at the Query boundary', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const pool: any = { + acquire: jest.fn().mockResolvedValue({ + Query: (params: any, options: any, callback: any) => + callback(null, { + status: { error_code: ErrorCode.SUCCESS, reason: '' }, + }), + }), + release: jest.fn(), + }; + const telemetry = { recordOperation: jest.fn() }; + (client as any).channelPool = pool; + setPoolTelemetryManager(pool, telemetry); + + await expect( + client.query({ collection_name: 'test_collection' }) + ).rejects.toThrow(); + + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Query', + collection: 'test_collection', + error: expect.any(Error), + }) + ); + }); + + it('classifies public Search and HybridSearch logical operations', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const pool: any = {}; + const telemetry = { recordOperation: jest.fn() }; + (client as any).channelPool = pool; + (client as any)._search = jest.fn().mockResolvedValue({ + status: { error_code: ErrorCode.SUCCESS, reason: '' }, + results: [], + }); + setPoolTelemetryManager(pool, telemetry); + + await client.search({ + collection_name: 'dense_collection', + data: [[0.1, 0.2]], + } as any); + await client.hybridSearch({ + collection_name: 'hybrid_collection', + data: [{ anns_field: 'vector', data: [[0.1, 0.2]] }], + } as any); + + expect(telemetry.recordOperation).toHaveBeenCalledTimes(2); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Search', + collection: 'dense_collection', + }) + ); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'HybridSearch', + collection: 'hybrid_collection', + }) + ); + }); + + it('records Delete validation and RunAnalyzer preprocessing at public boundaries', async () => { + const client = new MilvusClient({ + address: 'localhost:19530', + __SKIP_CONNECT__: true, + }); + const runAnalyzer = jest.fn((params: any, options: any, callback: any) => + callback(null, { status: { error_code: ErrorCode.SUCCESS, reason: '' } }) + ); + const pool: any = { + acquire: jest.fn().mockResolvedValue({ RunAnalyzer: runAnalyzer }), + release: jest.fn(), + }; + const telemetry = { recordOperation: jest.fn() }; + (client as any).channelPool = pool; + setPoolTelemetryManager(pool, telemetry); + + await expect( + client.deleteEntities({ collection_name: 'test_collection' } as any) + ).rejects.toThrow(ERROR_REASONS.FILTER_EXPR_REQUIRED); + await client.runAnalyzer({ + collection_name: 'test_collection', + text: 'hello', + analyzer_params: { tokenizer: 'standard' }, + }); + + expect(runAnalyzer).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(2); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Delete', + collection: 'test_collection', + error: expect.any(Error), + }) + ); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'RunAnalyzer', + collection: '', + error: undefined, + }) + ); + }); + it('should pass query grouping and order by fields through query params', async () => { const client = new MilvusClient({ address: 'localhost:19530', diff --git a/test/utils/Function.spec.ts b/test/utils/Function.spec.ts index 9c8522dd..d0e7f3ef 100644 --- a/test/utils/Function.spec.ts +++ b/test/utils/Function.spec.ts @@ -11,8 +11,19 @@ import { DataType, getValidDataArray, extractRequestMetadata, + setPoolFailoverHandler, + setPoolTelemetryManager, + withTelemetryLogicalOperation, + withTelemetrySuppressed, + getGRPCService, + LOADER_OPTIONS, } from '../../milvus'; -import { Metadata } from '@grpc/grpc-js'; +import { + credentials, + Metadata, + Server, + ServerCredentials, +} from '@grpc/grpc-js'; describe('Function API testing', () => { let pool: any; @@ -278,15 +289,15 @@ describe('Function API testing', () => { describe('promisify with traceid', () => { it('should extract client_request_id from params automatically', async () => { - const traceId = 'test-trace-id-123'; + const traceId = '11111111111111111111111111111111'; const params = { collection_name: 'test_collection', client_request_id: traceId, }; let capturedMetadata: Metadata | undefined; - client.testFunction = jest.fn((params, options, callback) => { - capturedMetadata = options.metadata; + client.testFunction = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; callback(null, 'success'); }); @@ -297,15 +308,15 @@ describe('Function API testing', () => { }); it('should extract client-request-id from params automatically', async () => { - const traceId = 'test-trace-id-456'; + const traceId = '22222222222222222222222222222222'; const params = { collection_name: 'test_collection', 'client-request-id': traceId, }; let capturedMetadata: Metadata | undefined; - client.testFunction = jest.fn((params, options, callback) => { - capturedMetadata = options.metadata; + client.testFunction = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; callback(null, 'success'); }); @@ -318,13 +329,13 @@ describe('Function API testing', () => { it('should prefer client_request_id over client-request-id when both exist', async () => { const params = { collection_name: 'test_collection', - client_request_id: 'preferred-id', - 'client-request-id': 'alternative-id', + client_request_id: '33333333333333333333333333333333', + 'client-request-id': '44444444444444444444444444444444', }; let capturedMetadata: Metadata | undefined; - client.testFunction = jest.fn((params, options, callback) => { - capturedMetadata = options.metadata; + client.testFunction = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; callback(null, 'success'); }); @@ -332,22 +343,22 @@ describe('Function API testing', () => { expect(capturedMetadata).toBeDefined(); expect(capturedMetadata?.get('client-request-id')).toEqual([ - 'preferred-id', + '33333333333333333333333333333333', ]); }); it('should use explicit requestMetadata when provided', async () => { const params = { collection_name: 'test_collection', - client_request_id: 'params-id', + client_request_id: '55555555555555555555555555555555', }; const explicitMetadata = { - 'client-request-id': 'explicit-id', + 'client-request-id': '66666666666666666666666666666666', }; let capturedMetadata: Metadata | undefined; - client.testFunction = jest.fn((params, options, callback) => { - capturedMetadata = options.metadata; + client.testFunction = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; callback(null, 'success'); }); @@ -355,7 +366,7 @@ describe('Function API testing', () => { expect(capturedMetadata).toBeDefined(); expect(capturedMetadata?.get('client-request-id')).toEqual([ - 'explicit-id', + '66666666666666666666666666666666', ]); }); @@ -374,35 +385,120 @@ describe('Function API testing', () => { expect(capturedMetadata).toBeUndefined(); }); + + it('preserves documented legacy request IDs on the wire', async () => { + for (const clientRequestId of [ + 'not-a-trace-id', + 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', + '00000000000000000000000000000000', + ]) { + let capturedMetadata: Metadata | undefined; + client.testFunction = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; + callback(null, 'success'); + }); + + await promisify( + pool, + 'testFunction', + { client_request_id: clientRequestId }, + 1000 + ); + + expect(capturedMetadata?.get('client-request-id')).toEqual([ + clientRequestId, + ]); + } + }); + + it('sends a legacy request ID through a real grpc-js channel', async () => { + const Service = getGRPCService( + { serviceName: 'milvus.proto.milvus.MilvusService' }, + LOADER_OPTIONS + ); + const server = new Server(); + let wireRequestId: string[] = []; + server.addService((Service as any).service, { + GetVersion: (call: any, callback: Function) => { + wireRequestId = call.metadata.get('client-request-id'); + callback(null, { + status: { error_code: 'Success', reason: '' }, + version: 'test', + }); + }, + }); + const port = await new Promise((resolve, reject) => { + server.bindAsync( + '127.0.0.1:0', + ServerCredentials.createInsecure(), + (error, boundPort) => (error ? reject(error) : resolve(boundPort)) + ); + }); + const grpcClient = new Service( + `127.0.0.1:${port}`, + credentials.createInsecure() + ); + const grpcPool: any = { + acquire: jest.fn().mockResolvedValue(grpcClient), + release: jest.fn(), + }; + + try { + await promisify( + grpcPool, + 'GetVersion', + { client_request_id: 'legacy-wire-id' }, + 1000 + ); + expect(wireRequestId).toEqual(['legacy-wire-id']); + } finally { + grpcClient.close(); + await new Promise(resolve => server.tryShutdown(() => resolve())); + } + }); }); describe('extractRequestMetadata', () => { it('should extract client_request_id from data', () => { const data = { collection_name: 'test', - client_request_id: 'trace-123', + client_request_id: '77777777777777777777777777777777', }; const result = extractRequestMetadata(data); - expect(result).toEqual({ 'client-request-id': 'trace-123' }); + expect(result).toEqual({ + 'client-request-id': '77777777777777777777777777777777', + }); + }); + + it('should preserve a documented arbitrary string request ID', () => { + expect( + extractRequestMetadata({ + client_request_id: 'insert-trace-123', + }) + ).toEqual({ 'client-request-id': 'insert-trace-123' }); }); it('should extract client-request-id from data', () => { const data = { collection_name: 'test', - 'client-request-id': 'trace-456', + 'client-request-id': '88888888888888888888888888888888', }; const result = extractRequestMetadata(data); - expect(result).toEqual({ 'client-request-id': 'trace-456' }); + expect(result).toEqual({ + 'client-request-id': '88888888888888888888888888888888', + }); }); it('should prefer client_request_id over client-request-id', () => { const data = { collection_name: 'test', - client_request_id: 'preferred', - 'client-request-id': 'alternative', + client_request_id: '99999999999999999999999999999999', + 'client-request-id': 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', }; const result = extractRequestMetadata(data); - expect(result).toEqual({ 'client-request-id': 'preferred' }); + expect(result).toEqual({ + 'client-request-id': '99999999999999999999999999999999', + }); }); it('should return undefined when no traceid is provided', () => { @@ -419,13 +515,251 @@ describe('Function API testing', () => { expect(extractRequestMetadata({})).toBeUndefined(); }); - it('should convert traceid to string', () => { + it('should reject non-string trace IDs', () => { const data = { collection_name: 'test', client_request_id: 12345, }; const result = extractRequestMetadata(data); - expect(result).toEqual({ 'client-request-id': '12345' }); + expect(result).toBeUndefined(); + }); + }); + + it('records one telemetry outcome across global failover', async () => { + const unavailable = Object.assign(new Error('unavailable'), { + code: 14, + }); + client.Search = jest.fn((params, options, callback) => + callback(unavailable) + ); + const failoverClient = { + Search: jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ), + }; + const failoverPool: any = { + acquire: jest.fn().mockResolvedValue(failoverClient), + release: jest.fn(), + }; + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + setPoolTelemetryManager(failoverPool, telemetry); + setPoolFailoverHandler(pool, async () => failoverPool); + + await promisify(pool, 'Search', { collection_name: 'books' }, 1000); + + expect(client.Search).toHaveBeenCalledTimes(1); + expect(failoverClient.Search).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Search', + collection: 'books', + error: undefined, + }) + ); + }); + + it('keeps a legacy wire ID out of OTel telemetry correlation', async () => { + let capturedMetadata: Metadata | undefined; + client.Search = jest.fn((params, metadata, options, callback) => { + capturedMetadata = metadata; + callback(null, { status: { error_code: 'Success' } }); }); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await promisify( + pool, + 'Search', + { + collection_name: 'books', + client_request_id: 'insert-trace-123', + }, + 1000 + ); + + expect(capturedMetadata?.get('client-request-id')).toEqual([ + 'insert-trace-123', + ]); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ requestId: undefined }) + ); + }); + + it('records RunAnalyzer globally even when the request carries a collection', async () => { + client.RunAnalyzer = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await promisify(pool, 'RunAnalyzer', { collection_name: 'books' }, 1000); + + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ operation: 'RunAnalyzer', collection: '' }) + ); + }); + + it('does not let a telemetry recorder failure change the logical result', async () => { + const hostileFailure = { + toString: () => { + throw new Error('coercion must never run'); + }, + }; + const telemetry = { + recordOperation: jest.fn(() => { + throw hostileFailure; + }), + }; + setPoolTelemetryManager(pool, telemetry); + + await expect( + withTelemetryLogicalOperation( + pool, + 'Search', + { collection_name: 'books' }, + async () => 'business-result' + ) + ).resolves.toBe('business-result'); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + }); + + it('records a monitored RPC once at the enclosing public logical boundary', async () => { + client.Search = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await withTelemetryLogicalOperation( + pool, + 'Search', + { collection_name: 'books' }, + () => promisify(pool, 'Search', { collection_name: 'books' }, 1000) + ); + + expect(client.Search).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'Search', + collection: 'books', + error: undefined, + }) + ); + }); + + it('coalesces nested helpers for the same logical operation', async () => { + client.Delete = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await withTelemetryLogicalOperation( + pool, + 'Delete', + { collection_name: 'books' }, + () => + withTelemetryLogicalOperation( + pool, + 'Delete', + { collection_name: 'books' }, + () => promisify(pool, 'Delete', { collection_name: 'books' }, 1000) + ) + ); + + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + }); + + it('isolates concurrent logical operations', async () => { + client.Search = jest.fn((params, options, callback) => + setTimeout(() => callback(null, { status: { error_code: 'Success' } }), 5) + ); + client.Query = jest.fn((params, options, callback) => + setTimeout(() => callback(null, { status: { error_code: 'Success' } }), 1) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await Promise.all([ + withTelemetryLogicalOperation( + pool, + 'Search', + { collection_name: 'books' }, + () => promisify(pool, 'Search', { collection_name: 'books' }, 1000) + ), + withTelemetryLogicalOperation( + pool, + 'Query', + { collection_name: 'authors' }, + () => promisify(pool, 'Query', { collection_name: 'authors' }, 1000) + ), + ]); + + expect(telemetry.recordOperation).toHaveBeenCalledTimes(2); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ operation: 'Search', collection: 'books' }) + ); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ operation: 'Query', collection: 'authors' }) + ); + }); + + it('suppresses logical-operation and RPC-level telemetry inside iterators', async () => { + client.Search = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await withTelemetrySuppressed(() => + withTelemetryLogicalOperation( + pool, + 'Search', + { collection_name: 'books' }, + () => promisify(pool, 'Search', { collection_name: 'books' }, 1000) + ) + ); + + expect(client.Search).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).not.toHaveBeenCalled(); + }); + + it('suppresses the promisify RPC-level fallback for internal page fetches', async () => { + client.Query = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await withTelemetrySuppressed(() => + promisify(pool, 'Query', { collection_name: 'books' }, 1000) + ); + + expect(client.Query).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).not.toHaveBeenCalled(); + }); + + it('resumes recording after the suppressed section exits', async () => { + client.Search = jest.fn((params, options, callback) => + callback(null, { status: { error_code: 'Success' } }) + ); + const telemetry = { recordOperation: jest.fn() }; + setPoolTelemetryManager(pool, telemetry); + + await withTelemetrySuppressed(() => + withTelemetrySuppressed(() => + promisify(pool, 'Search', { collection_name: 'books' }, 1000) + ) + ); + expect(telemetry.recordOperation).not.toHaveBeenCalled(); + + await promisify(pool, 'Search', { collection_name: 'books' }, 1000); + expect(telemetry.recordOperation).toHaveBeenCalledTimes(1); + expect(telemetry.recordOperation).toHaveBeenCalledWith( + expect.objectContaining({ operation: 'Search', collection: 'books' }) + ); }); }); diff --git a/test/utils/GlobalConnection.spec.ts b/test/utils/GlobalConnection.spec.ts index fe4d59fb..ee0b5b2d 100644 --- a/test/utils/GlobalConnection.spec.ts +++ b/test/utils/GlobalConnection.spec.ts @@ -33,12 +33,8 @@ describe('Global connection lifecycle', () => { }); it('should detect global cluster from address', () => { - expect( - isGlobalEndpoint('https://glo-xxx.global-cluster.xyz') - ).toBe(true); - expect( - isGlobalEndpoint('https://in01-xxx.zilliz.com') - ).toBe(false); + expect(isGlobalEndpoint('https://glo-xxx.global-cluster.xyz')).toBe(true); + expect(isGlobalEndpoint('https://in01-xxx.zilliz.com')).toBe(false); }); it('should resolve primary endpoint during initialization', async () => { @@ -56,9 +52,7 @@ describe('Global connection lifecycle', () => { // The client should be detected as global expect(client.isGlobal).toBe(true); - expect(client.globalEndpoint).toBe( - 'https://glo-xxx.global-cluster.xyz' - ); + expect(client.globalEndpoint).toBe('https://glo-xxx.global-cluster.xyz'); }); it('should not treat regular addresses as global', () => { @@ -318,7 +312,11 @@ describe('Global connection lifecycle', () => { data: { version: '2', clusters: [ - { clusterId: 'c3', endpoint: 'new-primary:19530', capability: 3 }, + { + clusterId: 'c3', + endpoint: 'new-primary:19530', + capability: 3, + }, ], }, }), @@ -332,27 +330,76 @@ describe('Global connection lifecycle', () => { __SKIP_CONNECT__: true, }); - // Mock _getServerInfo BEFORE init so it doesn't hang on real gRPC - (client as any)._getServerInfo = jest.fn().mockResolvedValue(undefined); + const oldPool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + const candidatePool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + (client as any).createChannelPool = jest + .fn() + .mockReturnValueOnce(oldPool) + .mockReturnValueOnce(candidatePool); + const oldTelemetryClient = { close: jest.fn() } as any; + const candidateTelemetryClient = { close: jest.fn() } as any; + (client as any).createTelemetryClient = jest + .fn() + .mockReturnValueOnce(oldTelemetryClient) + .mockReturnValueOnce(candidateTelemetryClient); + (client as any)._getServerInfo = jest + .fn() + .mockResolvedValueOnce(undefined) + .mockResolvedValueOnce({ + identifier: 'new-client-id', + server_info: { build_tags: 'new-primary' }, + }); await (client as any)._initGlobalConnection('test'); expect(client.config.address).toBe('primary-host:19530'); - - // Mock old pool drain/clear - const oldPool = client.channelPool; - oldPool.drain = jest.fn().mockResolvedValue(undefined); - oldPool.clear = jest.fn().mockResolvedValue(undefined); + expect(client.channelPool).toBe(oldPool); + const oldRefresher = client.topologyRefresher; + const telemetry = client.getTelemetry(); + await telemetry.processCommands([ + { + command_id: 'preserved-state', + command_type: 'push_config', + payload: Buffer.from('{"sampling_rate":0.5}'), + create_time: 1, + persistent: true, + }, + ]); + const oldTelemetryState = { + config: telemetry.getConfig(), + configHash: telemetry.configHash, + lastCommandTimestamp: telemetry.lastCommandTimestamp, + }; const changed = await client.reconnectToPrimary(); expect(changed).toBe(true); expect(client.config.address).toBe('new-primary:19530'); + expect(client.channelPool).toBe(candidatePool); + expect((client as any).telemetryClient).toBe(candidateTelemetryClient); + expect(client.connectStatus).toBe(CONNECT_STATUS.CONNECTED); + expect(client.serverInfo).toEqual({ build_tags: 'new-primary' }); + expect(client.topologyRefresher).not.toBe(oldRefresher); + expect(client.topologyRefresher?.isRunning()).toBe(true); + expect(oldRefresher?.isRunning()).toBe(false); + expect(oldTelemetryClient.close).toHaveBeenCalledTimes(1); + expect(candidateTelemetryClient.close).not.toHaveBeenCalled(); + expect(oldPool.drain).toHaveBeenCalledTimes(1); + expect(oldPool.clear).toHaveBeenCalledTimes(1); + expect(candidatePool.drain).not.toHaveBeenCalled(); + expect(telemetry.getConfig()).toEqual(oldTelemetryState.config); + expect(telemetry.configHash).toBe(oldTelemetryState.configHash); + expect(telemetry.lastCommandTimestamp).toBe( + oldTelemetryState.lastCommandTimestamp + ); // Clean up client.topologyRefresher?.stop(); - if (client.channelPool) { - try { await client.channelPool.drain(); await client.channelPool.clear(); } catch {} - } }); it('should return false from reconnectToPrimary when primary unchanged', async () => { @@ -381,7 +428,7 @@ describe('Global connection lifecycle', () => { client.topologyRefresher?.stop(); }); - it('should clean up on failover failure', async () => { + it('should preserve the live lifecycle when candidate validation fails', async () => { let fetchCount = 0; global.fetch = jest.fn().mockImplementation(() => { fetchCount++; @@ -399,7 +446,11 @@ describe('Global connection lifecycle', () => { data: { version: '2', clusters: [ - { clusterId: 'c3', endpoint: 'failing-host:19530', capability: 3 }, + { + clusterId: 'c3', + endpoint: 'failing-host:19530', + capability: 3, + }, ], }, }), @@ -413,33 +464,165 @@ describe('Global connection lifecycle', () => { __SKIP_CONNECT__: true, }); - // Mock _getServerInfo: succeed on init, fail on reconnect - let getServerInfoCallCount = 0; - (client as any)._getServerInfo = jest.fn().mockImplementation(() => { - getServerInfoCallCount++; - if (getServerInfoCallCount <= 1) { - return Promise.resolve(undefined); - } - return Promise.reject(new Error('connection failed')); - }); + const oldPool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + const candidatePool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + (client as any).createChannelPool = jest + .fn() + .mockReturnValueOnce(oldPool) + .mockReturnValueOnce(candidatePool); + const oldTelemetryClient = { close: jest.fn() } as any; + const candidateTelemetryClient = { close: jest.fn() } as any; + (client as any).createTelemetryClient = jest + .fn() + .mockReturnValueOnce(oldTelemetryClient) + .mockReturnValueOnce(candidateTelemetryClient); + (client as any)._getServerInfo = jest + .fn() + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('connection failed')); await (client as any)._initGlobalConnection('test'); + client.connectStatus = CONNECT_STATUS.CONNECTED; + client.serverInfo = { build_tags: 'old-primary' }; + const oldAddress = client.config.address; + const oldTopology = client.globalTopology; + const oldRefresher = client.topologyRefresher; + const oldEpoch = (client as any).telemetryEndpointEpoch; + const telemetry = client.getTelemetry(); + await telemetry.processCommands([ + { + command_id: 'preserved-state', + command_type: 'push_config', + payload: Buffer.from('{"sampling_rate":0.5}'), + create_time: 1, + persistent: true, + }, + ]); + const oldTelemetryState = { + config: telemetry.getConfig(), + configHash: telemetry.configHash, + lastCommandTimestamp: telemetry.lastCommandTimestamp, + }; + + await expect(client.reconnectToPrimary()).rejects.toThrow( + 'connection failed' + ); - // Mock old pool drain/clear - const oldPool = client.channelPool; - oldPool.drain = jest.fn().mockResolvedValue(undefined); - oldPool.clear = jest.fn().mockResolvedValue(undefined); + expect(client.config.address).toBe(oldAddress); + expect(client.channelPool).toBe(oldPool); + expect((client as any).telemetryClient).toBe(oldTelemetryClient); + expect((client as any).telemetryEndpointEpoch).toBe(oldEpoch); + expect(client.globalTopology).toBe(oldTopology); + expect(client.topologyRefresher).toBe(oldRefresher); + expect(oldRefresher?.isRunning()).toBe(true); + expect(client.connectStatus).toBe(CONNECT_STATUS.CONNECTED); + expect(client.serverInfo).toEqual({ build_tags: 'old-primary' }); + expect(oldTelemetryClient.close).not.toHaveBeenCalled(); + expect(oldPool.drain).not.toHaveBeenCalled(); + expect(oldPool.clear).not.toHaveBeenCalled(); + expect(candidateTelemetryClient.close).toHaveBeenCalledTimes(1); + expect(candidatePool.drain).toHaveBeenCalledTimes(1); + expect(candidatePool.clear).toHaveBeenCalledTimes(1); + expect(telemetry.getConfig()).toEqual(oldTelemetryState.config); + expect(telemetry.configHash).toBe(oldTelemetryState.configHash); + expect(telemetry.lastCommandTimestamp).toBe( + oldTelemetryState.lastCommandTimestamp + ); - // Reconnect will fail at _getServerInfo (second call) - try { - await client.reconnectToPrimary(); - } catch { - // Expected + oldRefresher?.stop(); + }); + + it('discards and cleans a blocked failover candidate released after close', async () => { + global.fetch = jest.fn().mockResolvedValue({ + ok: true, + json: () => + Promise.resolve({ + code: 0, + data: { + version: '2', + clusters: [ + { + clusterId: 'new', + endpoint: 'new-primary:19530', + capability: 3, + }, + ], + }, + }), + }) as any; + const client = new MilvusClient({ + address: 'https://glo-xxx.global-cluster.xyz', + token: 'test-token', + pool: { min: 0, max: 1 }, + __SKIP_CONNECT__: true, + }); + client.config.address = 'old-primary:19530'; + const oldTopology = { + version: '1', + clusters: [ + { + clusterId: 'old', + endpoint: 'old-primary:19530', + capability: 3, + }, + ], + } as any; + client.globalTopology = oldTopology; + + const oldPool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + const candidatePool = { + drain: jest.fn().mockResolvedValue(undefined), + clear: jest.fn().mockResolvedValue(undefined), + } as any; + client.channelPool = oldPool; + (client as any).createChannelPool = jest.fn(() => candidatePool); + const oldTelemetryClient = { close: jest.fn() } as any; + const candidateTelemetryClient = { close: jest.fn() } as any; + (client as any).telemetryClient = oldTelemetryClient; + (client as any).createTelemetryClient = jest.fn( + () => candidateTelemetryClient + ); + let releaseCandidate!: (value: any) => void; + const candidateValidation = new Promise(resolve => { + releaseCandidate = resolve; + }); + (client as any)._getServerInfo = jest.fn(() => candidateValidation); + const oldEpoch = (client as any).telemetryEndpointEpoch; + + const reconnect = client.reconnectToPrimary(); + while (!(client as any)._getServerInfo.mock.calls.length) { + await Promise.resolve(); } + const concurrentReconnect = client.reconnectToPrimary(); + await client.closeConnection(); + releaseCandidate({ + identifier: 'candidate-client', + server_info: { build_tags: 'candidate' }, + }); - // After failure, refresher should be cleaned up - expect(client.topologyRefresher).toBeNull(); + await expect(reconnect).resolves.toBe(false); + await expect(concurrentReconnect).resolves.toBe(false); expect(client.connectStatus).toBe(CONNECT_STATUS.SHUTDOWN); + expect(client.config.address).toBe('old-primary:19530'); + expect(client.channelPool).toBe(oldPool); + expect(client.globalTopology).toBe(oldTopology); + expect((client as any).telemetryClient).toBeUndefined(); + expect((client as any).telemetryEndpointEpoch).toBe(oldEpoch + 1); + expect(oldTelemetryClient.close).toHaveBeenCalledTimes(1); + expect(oldPool.drain).toHaveBeenCalledTimes(1); + expect(oldPool.clear).toHaveBeenCalledTimes(1); + expect(candidateTelemetryClient.close).toHaveBeenCalledTimes(1); + expect(candidatePool.drain).toHaveBeenCalledTimes(1); + expect(candidatePool.clear).toHaveBeenCalledTimes(1); }); it('should serialize concurrent reconnect via isReconnecting flag', async () => { @@ -469,14 +652,14 @@ describe('Global connection lifecycle', () => { // Simulate an ongoing reconnect by setting the flag c.isReconnecting = true; - let resolveReconnect!: () => void; - c.reconnectingPromise = new Promise((resolve: any) => { + let resolveReconnect!: (changed: boolean) => void; + c.reconnectingPromise = new Promise(resolve => { resolveReconnect = resolve; }); // Second call should wait for the existing promise, not start a new one const waitPromise = client.reconnectToPrimary(); - resolveReconnect(); + resolveReconnect(true); const result = await waitPromise; // Should return true (reconnect was handled by the first caller)