observe-server/common/middleware/request_id.go

package middleware

import (
	"net/http"
	"regexp"
	"strings"

	"github.com/go-admin-team/go-admin-core/logger"
	"github.com/go-admin-team/go-admin-core/sdk/pkg"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"

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

// RequestId 自动增加requestId
func RequestId(trafficKey string) gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Request.Method == http.MethodOptions {
			c.Next()
			return
		}
		prop := otel.GetTextMapPropagator()
		ctx := prop.Extract(c.Request.Context(), propagation.HeaderCarrier(c.Request.Header))

		span := trace.SpanFromContext(ctx)
		requestId := ""
		if span.SpanContext().IsValid() {
			requestId = span.SpanContext().TraceID().String()
		}
		if requestId == "" {
			requestId = c.GetHeader(trafficKey)
			if requestId == "" {
				requestId = c.GetHeader(strings.ToLower(trafficKey))
			}
			if requestId == "" {
				requestId = uuid.New().String()
			}
		}
		if valid, _ := regexp.Match("[0-9A-Za-z]+", []byte(requestId)); valid {
			c.Request.Header.Set(trafficKey, requestId)
			c.Set(trafficKey, requestId)
			c.Set(pkg.LoggerKey,
				logger.NewHelper(logger.DefaultLogger).
					WithFields(map[string]interface{}{
						trafficKey: requestId,
					}))

		}
		c.Next()
	}
}