更新

2026-09-14 00:14:30 +08:00
parent 316c6420a7
commit 7e1b474823
20 changed files with 1885 additions and 166 deletions
@@ -4,45 +4,91 @@ import (
 	"context"

 	"github.com/gin-gonic/gin"
-	uuid "github.com/satori/go.uuid"
+	"github.com/google/uuid"
 	"github.com/yuninks/loggerx"
 )

-// 设置普通的traceId
-func SetTraceIdByKey(ctx context.Context, traceKey string) context.Context {
-	if traceKey == "" {
-		traceKey = "trace_id"
-	}
+// traceCtxKey 是放在 context 里的 trace id 的键类型
+//
+// 刻意用自定义类型而不是裸 string：go vet 会直接报 SA1029，
+// 而且裸 string 作 key 有和第三方库撞键的风险
+type traceCtxKey struct{ field string }

-	val := ctx.Value(traceKey)
-	if val == nil {
-		ctx = context.WithValue(ctx, traceKey, uuid.NewV4().String())
-	}
-	return ctx
-}
+// TraceHeader 默认的 trace 请求头/响应头名字
+const TraceHeader = "X-Trace-Id"

-// 设置logger的traceId
+// SetTraceId 给 ctx 打上 logger 需要的 trace id，返回带值的 ctx
+// 已经存在时不覆盖，保证同一个请求内链路一致
 func SetTraceId(ctx context.Context, logger *loggerx.Logger) context.Context {
 	return SetTraceIdByKey(ctx, logger.GetTraceField())
 }

-// 设置Gin的traceId
-func SetGinTraceIdByKey(traceKey string) gin.HandlerFunc {
+// SetTraceIdByKey 按指定字段名给 ctx 打 trace id
+// traceKey 为空时用 "trace_id"
+func SetTraceIdByKey(ctx context.Context, traceKey string) context.Context {
+	if traceKey == "" {
+		traceKey = "trace_id"
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if v, _ := ctx.Value(traceCtxKey{traceKey}).(string); v != "" {
+		return ctx
+	}
+	return context.WithValue(ctx, traceCtxKey{traceKey}, uuid.NewString())
+}

+// GetTraceId 从 ctx 里取 trace id（取不到返回空串）
+func GetTraceId(ctx context.Context, traceKey string) string {
+	if traceKey == "" {
+		traceKey = "trace_id"
+	}
+	if ctx == nil {
+		return ""
+	}
+	v, _ := ctx.Value(traceCtxKey{traceKey}).(string)
+	return v
+}
+
+// SetGinTraceId 生成/透传 trace id 的 Gin 中间件
+//
+// 行为：
+//   - 优先取请求头 X-Trace-Id（方便上游透传，跨服务串起同一条链路）
+//   - 没有就生成一个
+//   - 写回响应头，下游/客户端能拿到同一个 id，便于排障
+//   - 同时写进 gin.Context（日志里能取到）与 request context
+func SetGinTraceId(logger *loggerx.Logger) gin.HandlerFunc {
+	return SetGinTraceIdByKey(logger.GetTraceField(), TraceHeader)
+}
+
+// SetGinTraceIdByKey 与 SetGinTraceId 相同，但可自定义字段名
+// header 为空表示不读也不写请求/响应头
+func SetGinTraceIdByKey(traceKey, header string) gin.HandlerFunc {
 	if traceKey == "" {
 		traceKey = "trace_id"
 	}

-	return func(ctx *gin.Context) {
-		traceId := ctx.Request.Header.Get(traceKey)
-		if traceId == "" {
-			traceId = uuid.NewV4().String()
+	return func(c *gin.Context) {
+		var traceId string
+		if header != "" {
+			traceId = c.Request.Header.Get(header)
 		}
-		ctx.Set(traceKey, traceId)
+		if traceId == "" {
+			traceId = uuid.NewString()
+		}
+
+		c.Set(traceKey, traceId)
+		ctx := context.WithValue(c.Request.Context(), traceCtxKey{traceKey}, traceId)
+		c.Request = c.Request.WithContext(ctx)
+
+		if header != "" {
+			c.Writer.Header().Set(header, traceId)
+		}
+		c.Next()
 	}
 }

-// 设置Gin的traceId
+// SetGinTraceIdByLogger 兼容旧名字：用 logger 的 trace 字段名生成中间件
 func SetGinTraceIdByLogger(logger *loggerx.Logger) gin.HandlerFunc {
-	return SetGinTraceIdByKey(logger.GetTraceField())
+	return SetGinTraceId(logger)
 }
@@ -0,0 +1,154 @@
+package middleware_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"github.com/yuninks/loggerx"
+	"github.com/yuninks/loggerx/middleware"
+)
+
+func newTestEngine(log *loggerx.Logger) *gin.Engine {
+	gin.SetMode(gin.TestMode)
+	g := gin.New()
+	g.Use(middleware.SetGinTraceId(log))
+	return g
+}
+
+// trace id 应当写回响应头，并在 handler 里可读
+func TestGinTraceIdSetsResponseHeader(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	var seen string
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		seen = middleware.GetTraceId(c.Request.Context(), log.GetTraceField())
+		c.Status(200)
+	})
+
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ping", nil))
+
+	header := w.Header().Get(middleware.TraceHeader)
+	if header == "" {
+		t.Fatalf("响应头 %s 没有写回，跨服务排障拿不到同一个 id", middleware.TraceHeader)
+	}
+	if _, err := uuid.Parse(header); err != nil {
+		t.Errorf("生成的 trace id 不是合法 uuid: %q", header)
+	}
+	if seen != header {
+		t.Errorf("handler 里读到的 trace id (%q) 与响应头 (%q) 不一致", seen, header)
+	}
+}
+
+// 上游带 trace id 时必须透传，不能重新生成
+func TestGinTraceIdPropagatesFromHeader(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	const upstream = "8f3d1c22-1111-4222-8333-444455556666"
+	var seen string
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		seen = middleware.GetTraceId(c.Request.Context(), log.GetTraceField())
+		c.Status(200)
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
+	req.Header.Set(middleware.TraceHeader, upstream)
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, req)
+
+	if seen != upstream {
+		t.Errorf("上游 trace id 未透传: 期望 %q 实际 %q", upstream, seen)
+	}
+	if got := w.Header().Get(middleware.TraceHeader); got != upstream {
+		t.Errorf("响应头应回传上游 trace id，实际 %q", got)
+	}
+}
+
+// 日志里应当带上 middleware 注入的 trace id（gin.Context 与 request ctx 两条路径）
+func TestTraceIdAppearsInLog(t *testing.T) {
+	dir := t.TempDir()
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(dir))
+	defer log.Close()
+
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		log.Info(c, "via-gin-ctx")
+		log.Info(c.Request.Context(), "via-request-ctx")
+		c.Status(200)
+	})
+
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ping", nil))
+	traceId := w.Header().Get(middleware.TraceHeader)
+	if traceId == "" {
+		t.Fatal("没有 trace id")
+	}
+	if err := log.MustSync(); err != nil {
+		t.Fatal(err)
+	}
+
+	content := readAll(t, dir)
+	if !contains(content, traceId) {
+		t.Errorf("日志里没有出现 trace id %q:\n%s", traceId, content)
+	}
+	if !contains(content, "via-gin-ctx") || !contains(content, "via-request-ctx") {
+		t.Errorf("两条日志都应当写入:\n%s", content)
+	}
+}
+
+// SetTraceId 打标：已存在时不覆盖，取不到时返回空串而不是 panic
+func TestSetTraceIdByKeySemantics(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	ctx := middleware.SetTraceId(context.Background(), log)
+	first := middleware.GetTraceId(ctx, log.GetTraceField())
+	if first == "" {
+		t.Fatal("SetTraceId 应当生成一个 id")
+	}
+
+	again := middleware.SetTraceId(ctx, log)
+	if got := middleware.GetTraceId(again, log.GetTraceField()); got != first {
+		t.Errorf("已有 trace id 时不该覆盖: %q -> %q", first, got)
+	}
+
+	if got := middleware.GetTraceId(nil, "trace_id"); got != "" {
+		t.Errorf("nil ctx 应返回空串，实际 %q", got)
+	}
+	if got := middleware.GetTraceId(context.Background(), "trace_id"); got != "" {
+		t.Errorf("没有值时该返回空串，实际 %q", got)
+	}
+}
+
+func readAll(t *testing.T, dir string) string {
+	t.Helper()
+	var sb []byte
+	files, _ := filepath.Glob(filepath.Join(dir, "*.log"))
+	for _, f := range files {
+		b, err := os.ReadFile(f)
+		if err != nil {
+			t.Fatalf("读取 %s: %v", f, err)
+		}
+		sb = append(sb, b...)
+	}
+	return string(sb)
+}
+
+func contains(s, sub string) bool {
+	for i := 0; i+len(sub) <= len(s); i++ {
+		if s[i:i+len(sub)] == sub {
+			return true
+		}
+	}
+	return false
+}