95 lines
2.7 KiB
Go
95 lines
2.7 KiB
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)
|
||
}
|