Support trace context propagation and traceID visibility (#412)

* Support trace context propagation and traceID visibility Signed-off-by: ryjiang <[email protected]> * fix build Signed-off-by: ryjiang <[email protected]> --------- Signed-off-by: ryjiang <[email protected]>
milvus-io · Jan 28, 2025 · 98a23bc · 98a23bc
1 parent 3f97a5e
commit 98a23bc
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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"