Support passing trace ctx and displaying traceID in case of request e…

…rrors Signed-off-by: tinswzy <[email protected]>
milvus-io · Dec 20, 2024 · d5e9964 · d5e9964
1 parent 5a960e9
commit d5e9964
Show file tree

Hide file tree

Showing 19 changed files with 744 additions and 89 deletions.
diff --git a/client/admin.go b/client/admin.go
@@ -14,12 +14,17 @@ package client
 import (
 	"context"
 
+	"go.opentelemetry.io/otel"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
 // GetVersion returns milvus server version information.
 func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
+	method := "GetVersion"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
 	if c.Service == nil {
 		return "", ErrClientNotReady
 	}
@@ -32,6 +37,9 @@ func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
 
 // CheckHealth returns milvus state
 func (c *GrpcClient) CheckHealth(ctx context.Context) (*entity.MilvusState, error) {
+	method := "CheckHealth"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
 	if c.Service == nil {
 		return nil, ErrClientNotReady
 	}

diff --git a/client/alias.go b/client/alias.go
@@ -18,12 +18,19 @@ package client
 
 import (
 	"context"
+	"log"
+
+	"go.opentelemetry.io/otel"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 )
 
 // CreateAlias creates an alias for collection
 func (c *GrpcClient) CreateAlias(ctx context.Context, collName string, alias string) error {
+	method := "CreateAlias"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -36,17 +43,23 @@ func (c *GrpcClient) CreateAlias(ctx context.Context, collName string, alias str
 
 	resp, err := c.Service.CreateAlias(ctx, req)
 	if err != nil {
+		log.Fatalf("create alias failed, collName:%s, alias:%s, traceID:%s err: %v", collName, alias, traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("create alias failed, collName:%s, alias:%s, traceID:%s err: %v", collName, alias, traceID, err)
 		return err
 	}
 	return nil
 }
 
 // DropAlias drops the specified Alias
 func (c *GrpcClient) DropAlias(ctx context.Context, alias string) error {
+	method := "DropAlias"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -58,17 +71,23 @@ func (c *GrpcClient) DropAlias(ctx context.Context, alias string) error {
 
 	resp, err := c.Service.DropAlias(ctx, req)
 	if err != nil {
+		log.Fatalf("drop alias failed, alias:%s, traceID:%s err: %v", alias, traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("drop alias failed, alias:%s, traceID:%s err: %v", alias, traceID, err)
 		return err
 	}
 	return nil
 }
 
 // AlterAlias changes collection alias to provided alias
 func (c *GrpcClient) AlterAlias(ctx context.Context, collName string, alias string) error {
+	method := "AlterAlias"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -81,10 +100,12 @@ func (c *GrpcClient) AlterAlias(ctx context.Context, collName string, alias stri
 
 	resp, err := c.Service.AlterAlias(ctx, req)
 	if err != nil {
+		log.Fatalf("alter alias failed, collName:%s, alias:%s, traceID:%s err: %v", collName, alias, traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("alter alias failed, collName:%s, alias:%s, traceID:%s err: %v", collName, alias, traceID, err)
 		return err
 	}
 	return nil

diff --git a/client/authentication.go b/client/authentication.go
@@ -2,13 +2,20 @@ package client
 
 import (
 	"context"
+	"log"
+
+	"go.opentelemetry.io/otel"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/internal/utils/crypto"
 )
 
 // CreateCredential create new user and password
 func (c *GrpcClient) CreateCredential(ctx context.Context, username string, password string) error {
+	method := "CreateCredential"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -18,17 +25,23 @@ func (c *GrpcClient) CreateCredential(ctx context.Context, username string, pass
 	}
 	resp, err := c.Service.CreateCredential(ctx, req)
 	if err != nil {
+		log.Fatalf("create credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("create credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	return nil
 }
 
 // UpdateCredential update password for a user
 func (c *GrpcClient) UpdateCredential(ctx context.Context, username string, oldPassword string, newPassword string) error {
+	method := "UpdateCredential"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -39,17 +52,23 @@ func (c *GrpcClient) UpdateCredential(ctx context.Context, username string, oldP
 	}
 	resp, err := c.Service.UpdateCredential(ctx, req)
 	if err != nil {
+		log.Fatalf("update credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("update credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	return nil
 }
 
 // DeleteCredential delete a user
 func (c *GrpcClient) DeleteCredential(ctx context.Context, username string) error {
+	method := "DeleteCredential"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -58,27 +77,35 @@ func (c *GrpcClient) DeleteCredential(ctx context.Context, username string) erro
 	}
 	resp, err := c.Service.DeleteCredential(ctx, req)
 	if err != nil {
+		log.Fatalf("delete credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	err = handleRespStatus(resp)
 	if err != nil {
+		log.Fatalf("delete credential failed, traceID:%s err: %v", traceID, err)
 		return err
 	}
 	return nil
 }
 
 // ListCredUsers list all usernames
 func (c *GrpcClient) ListCredUsers(ctx context.Context) ([]string, error) {
+	method := "ListCredUsers"
+	ctx, span := otel.Tracer("client").Start(ctx, method)
+	defer span.End()
+	traceID := span.SpanContext().TraceID().String()
 	if c.Service == nil {
 		return nil, ErrClientNotReady
 	}
 	req := &milvuspb.ListCredUsersRequest{}
 	resp, err := c.Service.ListCredUsers(ctx, req)
 	if err != nil {
+		log.Fatalf("list credential users failed, traceID:%s err: %v", traceID, err)
 		return nil, err
 	}
 	err = handleRespStatus(resp.Status)
 	if err != nil {
+		log.Fatalf("list credential users failed, traceID:%s err: %v", traceID, err)
 		return nil, err
 	}
 	return resp.Usernames, nil

diff --git a/client/client.go b/client/client.go
@@ -19,7 +19,6 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
-
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
@@ -277,6 +276,8 @@ func NewClient(ctx context.Context, config Config) (Client, error) {
 
 	// Parse grpc options
 	options := c.config.getDialOption()
+	traceOpt := grpc.WithStatsHandler(getDynamicClientStatsHandler())
+	options = append(options, traceOpt)
 
 	// Connect the grpc server.
 	if err := c.connect(ctx, addr, options...); err != nil {