loggerx/middleware/traceId.go

package middleware

import (
	"context"

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

// traceCtxKey 是放在 context 里的 trace id 的键类型
//
// 刻意用自定义类型而不是裸 string：go vet 会直接报 SA1029，
// 而且裸 string 作 key 有和第三方库撞键的风险
type traceCtxKey struct{ field string }

// TraceHeader 默认的 trace 请求头/响应头名字
const TraceHeader = "X-Trace-Id"

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

// 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(c *gin.Context) {
		var traceId string
		if header != "" {
			traceId = c.Request.Header.Get(header)
		}
		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()
	}
}

// SetGinTraceIdByLogger 兼容旧名字：用 logger 的 trace 字段名生成中间件
func SetGinTraceIdByLogger(logger *loggerx.Logger) gin.HandlerFunc {
	return SetGinTraceId(logger)
}