From 9941055eaa611cdc5e163336f31750d9874586db Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Fri, 2 Dec 2022 11:00:44 +0800
Subject: [PATCH] feat: add trace.SpanIDFromContext and
 trace.TraceIDFromContext (#2654)

---
 core/logx/richlogger.go      | 24 +++---------------------
 core/trace/utils.go          |  8 ++++++++
 internal/trace/trace.go      | 25 +++++++++++++++++++++++++
 internal/trace/trace_test.go | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+), 21 deletions(-)
 create mode 100644 internal/trace/trace.go
 create mode 100644 internal/trace/trace_test.go

diff --git a/core/logx/richlogger.go b/core/logx/richlogger.go
index e88c2f75..bacceb1b 100644
--- a/core/logx/richlogger.go
+++ b/core/logx/richlogger.go
@@ -6,7 +6,7 @@ import (
 	"time"
 
 	"github.com/zeromicro/go-zero/core/timex"
-	"go.opentelemetry.io/otel/trace"
+	"github.com/zeromicro/go-zero/internal/trace"
 )
 
 // WithCallerSkip returns a Logger with given caller skip.
@@ -136,12 +136,12 @@ func (l *richLogger) buildFields(fields ...LogField) []LogField {
 		return fields
 	}
 
-	traceID := traceIdFromContext(l.ctx)
+	traceID := trace.TraceIDFromContext(l.ctx)
 	if len(traceID) > 0 {
 		fields = append(fields, Field(traceKey, traceID))
 	}
 
-	spanID := spanIdFromContext(l.ctx)
+	spanID := trace.SpanIDFromContext(l.ctx)
 	if len(spanID) > 0 {
 		fields = append(fields, Field(spanKey, spanID))
 	}
@@ -179,21 +179,3 @@ func (l *richLogger) slow(v interface{}, fields ...LogField) {
 		getWriter().Slow(v, l.buildFields(fields...)...)
 	}
 }
-
-func spanIdFromContext(ctx context.Context) string {
-	spanCtx := trace.SpanContextFromContext(ctx)
-	if spanCtx.HasSpanID() {
-		return spanCtx.SpanID().String()
-	}
-
-	return ""
-}
-
-func traceIdFromContext(ctx context.Context) string {
-	spanCtx := trace.SpanContextFromContext(ctx)
-	if spanCtx.HasTraceID() {
-		return spanCtx.TraceID().String()
-	}
-
-	return ""
-}
diff --git a/core/trace/utils.go b/core/trace/utils.go
index b039f875..ce83e863 100644
--- a/core/trace/utils.go
+++ b/core/trace/utils.go
@@ -5,6 +5,7 @@ import (
 	"net"
 	"strings"
 
+	ztrace "github.com/zeromicro/go-zero/internal/trace"
 	"go.opentelemetry.io/otel/attribute"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	"google.golang.org/grpc/peer"
@@ -12,6 +13,13 @@ import (
 
 const localhost = "127.0.0.1"
 
+var (
+	// SpanIDFromContext returns the span id from ctx.
+	SpanIDFromContext = ztrace.SpanIDFromContext
+	// TraceIDFromContext returns the trace id from ctx.
+	TraceIDFromContext = ztrace.TraceIDFromContext
+)
+
 // PeerFromCtx returns the peer from ctx.
 func PeerFromCtx(ctx context.Context) string {
 	p, ok := peer.FromContext(ctx)
diff --git a/internal/trace/trace.go b/internal/trace/trace.go
new file mode 100644
index 00000000..3304b043
--- /dev/null
+++ b/internal/trace/trace.go
@@ -0,0 +1,25 @@
+package trace
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/trace"
+)
+
+func SpanIDFromContext(ctx context.Context) string {
+	spanCtx := trace.SpanContextFromContext(ctx)
+	if spanCtx.HasSpanID() {
+		return spanCtx.SpanID().String()
+	}
+
+	return ""
+}
+
+func TraceIDFromContext(ctx context.Context) string {
+	spanCtx := trace.SpanContextFromContext(ctx)
+	if spanCtx.HasTraceID() {
+		return spanCtx.TraceID().String()
+	}
+
+	return ""
+}
diff --git a/internal/trace/trace_test.go b/internal/trace/trace_test.go
new file mode 100644
index 00000000..e40339cb
--- /dev/null
+++ b/internal/trace/trace_test.go
@@ -0,0 +1,32 @@
+package trace
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+func TestSpanIDFromContext(t *testing.T) {
+	tracer := sdktrace.NewTracerProvider().Tracer("test")
+	ctx, span := tracer.Start(
+		context.Background(),
+		"foo",
+		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
+		oteltrace.WithAttributes(semconv.HTTPClientAttributesFromHTTPRequest(httptest.NewRequest(http.MethodGet, "/", nil))...),
+	)
+	defer span.End()
+
+	assert.NotEmpty(t, TraceIDFromContext(ctx))
+	assert.NotEmpty(t, SpanIDFromContext(ctx))
+}
+
+func TestSpanIDFromContextEmpty(t *testing.T) {
+	assert.Empty(t, TraceIDFromContext(context.Background()))
+	assert.Empty(t, SpanIDFromContext(context.Background()))
+}