54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"log/slog"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
const traceIDKey = "trace_id"
|
|
|
|
// Trace 透传或生成 trace_id，并注入可携带 trace_id 的 slog 日志句柄，
|
|
// 供后续日志统一携带链路标识。
|
|
func Trace() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
tid := c.GetHeader("X-Trace-ID")
|
|
if tid == "" {
|
|
tid = newTraceID()
|
|
}
|
|
c.Set(traceIDKey, tid)
|
|
c.Set("slog_ctx", slog.With(slog.String("trace_id", tid)))
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func newTraceID() string {
|
|
b := make([]byte, 16)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return "unknown"
|
|
}
|
|
return hex.EncodeToString(b)
|
|
}
|
|
|
|
// TraceID 从 gin.Context 取出 trace_id。
|
|
func TraceID(c *gin.Context) string {
|
|
if v, ok := c.Get(traceIDKey); ok {
|
|
if s, ok := v.(string); ok {
|
|
return s
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// SlogCtx 返回携带 trace_id 的 slog 日志句柄（无则回退默认）。
|
|
func SlogCtx(c *gin.Context) *slog.Logger {
|
|
if v, ok := c.Get("slog_ctx"); ok {
|
|
if l, ok := v.(*slog.Logger); ok {
|
|
return l
|
|
}
|
|
}
|
|
return slog.Default()
|
|
}
|