modelRT/middleware/trace.go

// Package middleware defines gin framework middlewares and OTel tracing infrastructure.
package middleware

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"strings"
	"time"

	"modelRT/config"
	"modelRT/constants"
	"modelRT/logger"

	"github.com/gin-gonic/gin"
	"go.opentelemetry.io/contrib/propagators/b3"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/propagation"
	sdkresource "go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	oteltrace "go.opentelemetry.io/otel/trace"
)

// InitTracerProvider creates an OTLP TracerProvider and registers it as the global provider.
// It also registers the B3 propagator to stay compatible with existing B3 infrastructure.
// The caller is responsible for calling Shutdown on the returned provider during graceful shutdown.
func InitTracerProvider(ctx context.Context, cfg config.ModelRTConfig) (*sdktrace.TracerProvider, error) {
	opts := []otlptracehttp.Option{
		otlptracehttp.WithEndpoint(cfg.OtelConfig.Endpoint),
	}
	if cfg.OtelConfig.Insecure {
		opts = append(opts, otlptracehttp.WithInsecure())
	}

	exporter, err := otlptracehttp.New(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("create OTLP exporter: %w", err)
	}

	res := sdkresource.NewSchemaless(
		attribute.String("service.name", cfg.ServiceName),
		attribute.String("deployment.environment", cfg.DeployEnv),
	)

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exporter),
		sdktrace.WithResource(res),
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
	)

	otel.SetTracerProvider(tp)
	otel.SetTextMapPropagator(b3.New())

	return tp, nil
}

// StartTrace extracts upstream B3 trace context from request headers and starts a server span.
// Typed context keys are also injected for backward compatibility with the existing logger
// until the logger is migrated to read from the OTel span context (Step 6).
func StartTrace() gin.HandlerFunc {
	tracer := otel.Tracer("modelRT/http")
	return func(c *gin.Context) {
		// Extract upstream trace context from B3 headers (X-B3-TraceId etc.)
		ctx := otel.GetTextMapPropagator().Extract(
			c.Request.Context(),
			propagation.HeaderCarrier(c.Request.Header),
		)

		spanName := c.FullPath()
		if spanName == "" {
			spanName = c.Request.URL.Path
		}
		ctx, span := tracer.Start(ctx, spanName,
			oteltrace.WithSpanKind(oteltrace.SpanKindServer),
		)
		defer span.End()

		// backward compat: inject typed keys so existing logger reads work until Step 6
		spanCtx := span.SpanContext()
		ctx = context.WithValue(ctx, constants.CtxKeyTraceID, spanCtx.TraceID().String())
		ctx = context.WithValue(ctx, constants.CtxKeySpanID, spanCtx.SpanID().String())

		c.Request = c.Request.WithContext(ctx)

		// set in gin context for accessLog (logger.New(c) reads via gin.Context.Value)
		c.Set(constants.HeaderTraceID, spanCtx.TraceID().String())
		c.Set(constants.HeaderSpanID, spanCtx.SpanID().String())

		c.Next()
	}
}

type bodyLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// TODO 包装一下 gin.ResponseWriter，通过这种方式拦截写响应
// 让gin写响应的时候先写到 bodyLogWriter 再写gin.ResponseWriter ，
// 这样利用中间件里输出访问日志时就能拿到响应了
// https://stackoverflow.com/questions/38501325/how-to-log-response-body-in-gin
func (w bodyLogWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// TODO 用于访问request与reponse的日志记录
// LogAccess define func of log access info
func LogAccess() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 保存body
		var reqBody []byte
		contentType := c.GetHeader("Content-Type")
		// multipart/form-data 文件上传请求, 不在日志里记录body
		if !strings.Contains(contentType, "multipart/form-data") {
			reqBody, _ = io.ReadAll(c.Request.Body)
			c.Request.Body = io.NopCloser(bytes.NewReader(reqBody))

			// var request map[string]interface{}
			// if err := c.ShouldBindBodyWith(&request, binding.JSON); err != nil {
			// 	c.JSON(400, gin.H{"error": err.Error()})
			// 	return
			// }
		}
		start := time.Now()
		blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
		c.Writer = blw

		accessLog(c, "access_start", time.Since(start), reqBody, nil)
		defer func() {
			var responseLogging string
			if c.Writer.Size() > 10*1024 { // 响应大于10KB 不记录
				responseLogging = "Response data size is too Large to log"
			} else {
				responseLogging = blw.body.String()
			}
			accessLog(c, "access_end", time.Since(start), reqBody, responseLogging)
		}()
		c.Next()
	}
}

func accessLog(c *gin.Context, accessType string, dur time.Duration, body []byte, dataOut interface{}) {
	req := c.Request
	bodyStr := string(body)
	query := req.URL.RawQuery
	path := req.URL.Path
	// TODO: 实现Token认证后再把访问日志里也加上token记录
	// token := c.Request.Header.Get("token")
	logger.New(c).Info("AccessLog",
		"type", accessType,
		"ip", c.ClientIP(),
		//"token", token,
		"method", req.Method,
		"path", path,
		"query", query,
		"body", bodyStr,
		"output", dataOut,
		"time(ms)", int64(dur/time.Millisecond))
}