diff --git a/net/gtrace/gtrace.go b/net/gtrace/gtrace.go index 31c48452e..0eb96ebac 100644 --- a/net/gtrace/gtrace.go +++ b/net/gtrace/gtrace.go @@ -12,8 +12,6 @@ import ( "os" "strings" - "github.com/gogf/gf/v2/errors/gcode" - "github.com/gogf/gf/v2/errors/gerror" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/propagation" @@ -22,16 +20,18 @@ import ( "github.com/gogf/gf/v2/container/gmap" "github.com/gogf/gf/v2/container/gvar" + "github.com/gogf/gf/v2/errors/gcode" + "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/internal/command" "github.com/gogf/gf/v2/net/gipv4" "github.com/gogf/gf/v2/net/gtrace/internal/provider" + "github.com/gogf/gf/v2/text/gstr" "github.com/gogf/gf/v2/util/gconv" ) const ( tracingCommonKeyIpIntranet = `ip.intranet` tracingCommonKeyIpHostname = `hostname` - commandEnvKeyForTraceEnabled = "gf.trace.enabled" // Main switch for tracing feature. commandEnvKeyForMaxContentLogSize = "gf.gtrace.max.content.log.size" // To avoid too big tracing content. commandEnvKeyForTracingInternal = "gf.gtrace.tracing.internal" // For detailed controlling for tracing content. ) @@ -146,6 +146,11 @@ func GetBaggageVar(ctx context.Context, key string) *gvar.Var { return NewBaggage(ctx).GetVar(key) } +// WithUUID injects custom trace id with UUID into context to propagate. +func WithUUID(ctx context.Context, uuid string) (context.Context, error) { + return WithTraceID(ctx, gstr.Replace(uuid, "-", "")) +} + // WithTraceID injects custom trace id into context to propagate. func WithTraceID(ctx context.Context, traceID string) (context.Context, error) { generatedTraceID, err := trace.TraceIDFromHex(traceID) diff --git a/net/gtrace/gtrace_z_unit_test.go b/net/gtrace/gtrace_z_unit_test.go index 76f86d706..97a4d2ee6 100644 --- a/net/gtrace/gtrace_z_unit_test.go +++ b/net/gtrace/gtrace_z_unit_test.go @@ -17,16 +17,16 @@ import ( func TestWithTraceID(t *testing.T) { var ( - ctx = context.Background() - requestId = `a323f910-f690-11ec-963d-79c0b7fcf119` + ctx = context.Background() + uuid = `a323f910-f690-11ec-963d-79c0b7fcf119` ) gtest.C(t, func(t *gtest.T) { - newCtx, err := gtrace.WithTraceID(ctx, requestId) + newCtx, err := gtrace.WithTraceID(ctx, uuid) t.AssertNE(err, nil) t.Assert(newCtx, ctx) }) gtest.C(t, func(t *gtest.T) { - var traceId = gstr.Replace(requestId, "-", "") + var traceId = gstr.Replace(uuid, "-", "") newCtx, err := gtrace.WithTraceID(ctx, traceId) t.AssertNil(err) t.AssertNE(newCtx, ctx) @@ -34,3 +34,22 @@ func TestWithTraceID(t *testing.T) { t.Assert(gtrace.GetTraceID(newCtx), traceId) }) } + +func TestWithUUID(t *testing.T) { + var ( + ctx = context.Background() + uuid = `a323f910-f690-11ec-963d-79c0b7fcf119` + ) + gtest.C(t, func(t *gtest.T) { + newCtx, err := gtrace.WithTraceID(ctx, uuid) + t.AssertNE(err, nil) + t.Assert(newCtx, ctx) + }) + gtest.C(t, func(t *gtest.T) { + newCtx, err := gtrace.WithUUID(ctx, uuid) + t.AssertNil(err) + t.AssertNE(newCtx, ctx) + t.Assert(gtrace.GetTraceID(ctx), "") + t.Assert(gtrace.GetTraceID(newCtx), gstr.Replace(uuid, "-", "")) + }) +} diff --git a/net/gtrace/internal/provider/provider_idgenerator.go b/net/gtrace/internal/provider/provider_idgenerator.go index 83ffa845e..6c5baec3b 100644 --- a/net/gtrace/internal/provider/provider_idgenerator.go +++ b/net/gtrace/internal/provider/provider_idgenerator.go @@ -9,12 +9,15 @@ package provider import ( "context" - "github.com/gogf/gf/v2/internal/tracing" "go.opentelemetry.io/otel/trace" + + "github.com/gogf/gf/v2/internal/tracing" ) +// IDGenerator is a trace ID generator. type IDGenerator struct{} +// NewIDGenerator returns a new IDGenerator. func NewIDGenerator() *IDGenerator { return &IDGenerator{} }