更新

middleware/traceId.go

@@ -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)
 }