From 98a23bc9ebcb5334c57ee0a2cdcd6018aacf5246 Mon Sep 17 00:00:00 2001 From: ryjiang Date: Tue, 28 Jan 2025 08:11:52 +0800 Subject: [PATCH] Support trace context propagation and traceID visibility (#412) * Support trace context propagation and traceID visibility Signed-off-by: ryjiang * fix build Signed-off-by: ryjiang --------- Signed-off-by: ryjiang --- milvus/grpc/GrpcClient.ts | 12 +++++++++- milvus/types/Client.ts | 3 +++ milvus/utils/Grpc.ts | 44 ++++++++++++++++++++++++++++++++++ package.json | 1 + test/grpc/MilvusClient.spec.ts | 17 ++++++++++++- yarn.lock | 5 ++++ 6 files changed, 80 insertions(+), 2 deletions(-) diff --git a/milvus/grpc/GrpcClient.ts b/milvus/grpc/GrpcClient.ts index 688efcfc..64951565 100644 --- a/milvus/grpc/GrpcClient.ts +++ b/milvus/grpc/GrpcClient.ts @@ -17,6 +17,7 @@ import { getAuthString, getRetryInterceptor, getMetaInterceptor, + getTraceInterceptor, ErrorCode, DEFAULT_DB, METADATA, @@ -96,7 +97,16 @@ export class GRPCClient extends User { }); // interceptors - const interceptors = [metaInterceptor, retryInterceptor]; + const interceptors = [metaInterceptor]; + + // add trace if necessary + if (this.config.trace) { + // add trace interceptor + interceptors.push(getTraceInterceptor()); + } + + // add retry interceptor + interceptors.push(retryInterceptor); // add interceptors this.channelOptions.interceptors = interceptors; diff --git a/milvus/types/Client.ts b/milvus/types/Client.ts index abf23647..a9002433 100644 --- a/milvus/types/Client.ts +++ b/milvus/types/Client.ts @@ -71,6 +71,9 @@ export interface ClientConfig { // Number converts int64 to number, it will lose precision // String converts int64 to string, it's the default behavior loaderOptions?: LoaderOption; + + // enable trace + trace?: boolean; } export interface ServerInfo { diff --git a/milvus/utils/Grpc.ts b/milvus/utils/Grpc.ts index 4945ca91..23bc914b 100644 --- a/milvus/utils/Grpc.ts +++ b/milvus/utils/Grpc.ts @@ -15,7 +15,13 @@ import { isInvalidMessage, logger, } from '.'; +import { context, propagation, Span, trace } from '@opentelemetry/api'; +interface Carrier { + traceparent?: string; + tracestate?: string; +} import { DEFAULT_DB } from '../const'; +import sdkInfo from '../../sdk.json'; interface IServiceDetails { protoPath: string; // file to your proto file @@ -225,3 +231,41 @@ export const getRetryInterceptor = ({ }; return new InterceptingCall(nextCall(options), requester); }; + +/** + * Returns a gRPC interceptor function that adds trace context to outgoing requests. + */ +/* istanbul ignore next */ +export const getTraceInterceptor = () => { + // Get the name and version of the client. + const name = 'milvus-node-client'; + const version = sdkInfo.version; + // Get the tracer. + const tracer = trace.getTracer(name, version); + + return function (options: any, nextCall: any) { + // Create a new InterceptingCall object with nextCall(options) as its first parameter. + return new InterceptingCall(nextCall(options), { + // Define the start method of the InterceptingCall object. + start: function (metadata, listener, next) { + tracer.startActiveSpan('grpc-intercept', (span: Span) => { + // Set the span context. + const output: Carrier = {}; + // Inject the span context into the metadata. + propagation.inject(context.active(), output); + // Add the traceparent and tracestate to the metadata. + const { traceparent, tracestate } = output; + if (traceparent) { + metadata.add('traceparent', traceparent); + } + if (tracestate) { + metadata.add('tracestate', tracestate); + } + span.end(); + }); + // Call next(metadata, listener) to continue the call with the modified metadata. + next(metadata, listener); + }, + }); + }; +}; diff --git a/package.json b/package.json index 76496d02..e4e8ddc5 100644 --- a/package.json +++ b/package.json @@ -22,6 +22,7 @@ "dependencies": { "@grpc/grpc-js": "^1.12.1", "@grpc/proto-loader": "^0.7.10", + "@opentelemetry/api": "^1.9.0", "@petamoriken/float16": "^3.8.6", "dayjs": "^1.11.7", "generic-pool": "^3.9.0", diff --git a/test/grpc/MilvusClient.spec.ts b/test/grpc/MilvusClient.spec.ts index 5759ab9e..4ac10199 100644 --- a/test/grpc/MilvusClient.spec.ts +++ b/test/grpc/MilvusClient.spec.ts @@ -72,7 +72,7 @@ describe(`Milvus client`, () => { const m1u = new MilvusClient({ address: IP, tls: { - skipCertCheck : true + skipCertCheck: true, }, id: '1', __SKIP_CONNECT__: true, @@ -210,6 +210,21 @@ describe(`Milvus client`, () => { } }); + it(`should add trace interceptor if enableTrace is true`, async () => { + const nonTraceClient = new MilvusClient({ + address: IP, + __SKIP_CONNECT__: true, + }); + expect(nonTraceClient.channelOptions.interceptors.length).toEqual(2); + const traceClient = new MilvusClient({ + address: IP, + trace: true, + __SKIP_CONNECT__: true, + }); + + expect(traceClient.channelOptions.interceptors.length).toEqual(3); + }); + it(`Expect get node sdk info`, async () => { expect(MilvusClient.sdkInfo.version).toEqual(sdkInfo.version); expect(MilvusClient.sdkInfo.recommendMilvus).toEqual(sdkInfo.milvusVersion); diff --git a/yarn.lock b/yarn.lock index a5302d72..c7d0504f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -560,6 +560,11 @@ resolved "https://registry.npmjs.org/@js-sdsl/ordered-map/-/ordered-map-4.4.2.tgz" integrity sha512-iUKgm52T8HOE/makSxjqoWhe95ZJA1/G1sYsGev2JDKUSS14KAgg1LHb+Ba+IPow0xflbnSkOsZcO08C7w1gYw== +"@opentelemetry/api@^1.9.0": + version "1.9.0" + resolved "https://registry.yarnpkg.com/@opentelemetry/api/-/api-1.9.0.tgz#d03eba68273dc0f7509e2a3d5cba21eae10379fe" + integrity sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg== + "@petamoriken/float16@^3.8.6": version "3.8.6" resolved "https://registry.npmjs.org/@petamoriken/float16/-/float16-3.8.6.tgz"