Add trace ut (#7445)

Signed-off-by: godchen <qingxiang.chen@zilliz.com>

internal/util/trace/util.go

@@ -25,26 +25,15 @@ import (
 )
 
 func InitTracing(serviceName string) io.Closer {
-	if opentracing.IsGlobalTracerRegistered() {
+	cfg := &config.Configuration{
-		return nil
+		ServiceName: serviceName,
+		Sampler: &config.SamplerConfig{
+			Type:  "const",
+			Param: 0,
+		},
 	}
-	var cfg *config.Configuration
-	var err error
 	if true {
-		cfg, err = config.FromEnv()
+		cfg = InitFromEnv(serviceName)
-		if err != nil {
-			log.Error(err)
-			return nil
-		}
-		cfg.ServiceName = serviceName
-	} else {
-		cfg = &config.Configuration{
-			ServiceName: serviceName,
-			Sampler: &config.SamplerConfig{
-				Type:  "const",
-				Param: 1,
-			},
-		}
 	}
 	tracer, closer, err := cfg.NewTracer()
 	if err != nil {
@@ -55,13 +44,27 @@ func InitTracing(serviceName string) io.Closer {
 	return closer
 }
 
+func InitFromEnv(serviceName string) *config.Configuration {
+	cfg, err := config.FromEnv()
+	if err != nil {
+		log.Error(err)
+		return nil
+	}
+	cfg.ServiceName = serviceName
+	return cfg
+}
+
 func StartSpanFromContext(ctx context.Context, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
+	return StartSpanFromContextWithSkip(ctx, 2, opts...)
+}
+
+func StartSpanFromContextWithSkip(ctx context.Context, skip int, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
 	if ctx == nil {
 		return NoopSpan(), ctx
 	}
 
 	var pcs [1]uintptr
-	n := runtime.Callers(2, pcs[:])
+	n := runtime.Callers(skip, pcs[:])
 	if n < 1 {
 		span, ctx := opentracing.StartSpanFromContext(ctx, "unknown", opts...)
 		span.LogFields(log.Error(errors.New("runtime.Callers failed")))
@@ -85,12 +88,16 @@ func StartSpanFromContext(ctx context.Context, opts ...opentracing.StartSpanOpti
 }
 
 func StartSpanFromContextWithOperationName(ctx context.Context, operationName string, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
+	return StartSpanFromContextWithOperationNameWithSkip(ctx, operationName, 2, opts...)
+}
+
+func StartSpanFromContextWithOperationNameWithSkip(ctx context.Context, operationName string, skip int, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
 	if ctx == nil {
 		return NoopSpan(), ctx
 	}
 
 	var pcs [1]uintptr
-	n := runtime.Callers(2, pcs[:])
+	n := runtime.Callers(skip, pcs[:])
 	if n < 1 {
 		span, ctx := opentracing.StartSpanFromContext(ctx, operationName, opts...)
 		span.LogFields(log.Error(errors.New("runtime.Callers failed")))
@@ -130,17 +137,21 @@ func LogError(span opentracing.Span, err error) error {
 }
 
 func InfoFromSpan(span opentracing.Span) (traceID string, sampled bool, found bool) {
-	if spanContext, ok := span.Context().(jaeger.SpanContext); ok {
+	if span != nil {
-		traceID = spanContext.TraceID().String()
+		if spanContext, ok := span.Context().(jaeger.SpanContext); ok {
-		sampled = spanContext.IsSampled()
+			traceID = spanContext.TraceID().String()
-		return traceID, sampled, true
+			sampled = spanContext.IsSampled()
+			return traceID, sampled, true
+		}
 	}
 	return "", false, false
 }
 
 func InfoFromContext(ctx context.Context) (traceID string, sampled bool, found bool) {
-	if span := opentracing.SpanFromContext(ctx); span != nil {
+	if ctx != nil {
-		return InfoFromSpan(span)
+		if span := opentracing.SpanFromContext(ctx); span != nil {
+			return InfoFromSpan(span)
+		}
 	}
 	return "", false, false
 }

internal/util/trace/util_test.go

@@ -18,7 +18,9 @@ import (
 
 	"errors"
 
+	"github.com/opentracing/opentracing-go"
 	oplog "github.com/opentracing/opentracing-go/log"
+	"github.com/stretchr/testify/assert"
 )
 
 type simpleStruct struct {
@@ -26,6 +28,11 @@ type simpleStruct struct {
 	value string
 }
 
+func TestInit(t *testing.T) {
+	cfg := InitFromEnv("test")
+	assert.NotNil(t, cfg)
+}
+
 func TestTracing(t *testing.T) {
 	//Already Init in each framework, this can be ignored in debug
 	closer := InitTracing("test")
@@ -37,6 +44,8 @@ func TestTracing(t *testing.T) {
 	//start span
 	//default use function name for operation name
 	sp, ctx := StartSpanFromContext(ctx)
+	id, sampled, found := InfoFromContext(ctx)
+	fmt.Printf("traceID = %s, sampled = %t, found = %t", id, sampled, found)
 	sp.SetTag("tag1", "tag1")
 	// use self-defined operation name for span
 	// sp, ctx := StartSpanFromContextWithOperationName(ctx, "self-defined name")
@@ -77,3 +86,56 @@ func caller(ctx context.Context) error {
 	}
 	return nil
 }
+
+func TestInject(t *testing.T) {
+	//Already Init in each framework, this can be ignored in debug
+	closer := InitTracing("test")
+	defer closer.Close()
+
+	// context normally can be propagated through func params
+	ctx := context.Background()
+
+	//start span
+	//default use function name for operation name
+	sp, ctx := StartSpanFromContext(ctx)
+	id, sampled, found := InfoFromContext(ctx)
+	fmt.Printf("traceID = %s, sampled = %t, found = %t", id, sampled, found)
+	pp := PropertiesReaderWriter{PpMap: map[string]string{}}
+	InjectContextToPulsarMsgProperties(sp.Context(), pp.PpMap)
+	tracer := opentracing.GlobalTracer()
+	sc, _ := tracer.Extract(opentracing.TextMap, pp)
+	assert.NotNil(t, sc)
+
+}
+
+func TestTraceError(t *testing.T) {
+	//Already Init in each framework, this can be ignored in debug
+	closer := InitTracing("test")
+	defer closer.Close()
+
+	// context normally can be propagated through func params
+	sp, ctx := StartSpanFromContext(nil)
+	assert.Nil(t, ctx)
+	assert.NotNil(t, sp)
+
+	sp, ctx = StartSpanFromContextWithOperationName(nil, "test")
+	assert.Nil(t, ctx)
+	assert.NotNil(t, sp)
+
+	//Will Cause span log error
+	StartSpanFromContextWithOperationNameWithSkip(context.Background(), "test", 10000)
+
+	//Will Cause span log error
+	StartSpanFromContextWithSkip(context.Background(), 10000)
+
+	id, sampled, found := InfoFromSpan(nil)
+	assert.Equal(t, id, "")
+	assert.Equal(t, sampled, false)
+	assert.Equal(t, found, false)
+
+	id, sampled, found = InfoFromContext(nil)
+	assert.Equal(t, id, "")
+	assert.Equal(t, sampled, false)
+	assert.Equal(t, found, false)
+
+}