66 lines
1.1 KiB
Go
66 lines
1.1 KiB
Go
package trace
|
|
|
|
import (
|
|
"context"
|
|
"loveuer/utodo/pkg/logger"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/gofiber/fiber/v3"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type Option func(*option)
|
|
|
|
type option struct {
|
|
traceKey string
|
|
traceFn func(c fiber.Ctx) string
|
|
}
|
|
|
|
func WithTraceKey(key string) Option {
|
|
return func(o *option) {
|
|
if key != "" {
|
|
o.traceKey = key
|
|
}
|
|
}
|
|
}
|
|
|
|
func WithTraceFn(fn func(c fiber.Ctx) string) Option {
|
|
return func(o *option) {
|
|
if fn != nil {
|
|
o.traceFn = fn
|
|
}
|
|
}
|
|
}
|
|
|
|
func New(opts ...Option) fiber.Handler {
|
|
opt := &option{
|
|
traceKey: "X-Trace-Id",
|
|
traceFn: func(c fiber.Ctx) string {
|
|
uid, err := uuid.NewV7()
|
|
if err != nil {
|
|
return strconv.FormatInt(time.Now().UnixNano(), 10)
|
|
}
|
|
|
|
return uid.String()
|
|
},
|
|
}
|
|
|
|
for _, o := range opts {
|
|
o(opt)
|
|
}
|
|
|
|
return func(c fiber.Ctx) error {
|
|
traceId := c.Get(opt.traceKey)
|
|
if traceId == "" {
|
|
traceId = opt.traceFn(c)
|
|
c.Request().Header.Set(opt.traceKey, traceId)
|
|
}
|
|
|
|
c.SetContext(context.WithValue(c.Context(), logger.CtxKey, traceId))
|
|
c.Set(opt.traceKey, traceId)
|
|
|
|
return c.Next()
|
|
}
|
|
}
|