common/jaeger/jaeger.go

package jaeger

import (
	"context"
	"encoding/json"
	"strconv"
	"strings"
	"sync"

	"github.com/gogf/gf/contrib/trace/otlphttp/v2"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/net/gtrace"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
)

var (
	ShutDown func(ctx context.Context)
	initOnce sync.Once
)

// Init 初始化 Jaeger 链路追踪（延迟初始化，首次调用时执行）
func Init() {
	initOnce.Do(func() {
		ctx := context.Background()
		jaegerAgent := g.Cfg().MustGet(ctx, "jaeger.addr").String()
		serverName := g.Cfg().MustGet(ctx, "server.name").String()

		if jaegerAgent == "" {
			g.Log().Warning(ctx, "⚠️  Jaeger 配置未找到，跳过初始化")
			ShutDown = func(ctx context.Context) {} // 空函数，避免 nil panic
			return
		}

		shutdown, err := otlphttp.Init(serverName, jaegerAgent, "/v1/traces")
		if err != nil {
			g.Log().Errorf(ctx, "Jaeger 初始化失败: %v", err)
			ShutDown = func(ctx context.Context) {}
			return
		}
		ShutDown = shutdown
		g.Log().Infof(ctx, "✅ Jaeger 初始化成功: %s", jaegerAgent)
	})
}

func init() {
	// 默认自动初始化（保持向后兼容）
	Init()
}

// NewSpan 创建新的链路追踪 Span
// spanName: Span 名称，用于在 Jaeger UI 中标识
// 返回带有 Span 的 context 和 Span 对象，调用方需 defer span.End()
func NewSpan(ctx context.Context, spanName string, opts ...trace.SpanStartOption) (context.Context, *gtrace.Span) {
	return gtrace.NewSpan(ctx, spanName, opts...)
}

// RecordError 统一错误记录方法
// 功能：
//  1. 控制台输出错误（带完整堆栈 %+v）
//  2. Jaeger 链路追踪记录错误
//  3. 设置 Span 错误状态
//
// 使用示例：
//
//	jaeger.RecordError(ctx, err, "保存数据失败")
//
// 参数：
//   - ctx: 包含 trace span 的上下文
//   - err: 错误对象（支持 gerror 堆栈）
//   - msg: 可选的错误描述（用于日志和 Jaeger 显示）
func RecordError(ctx context.Context, err error, msg ...string) {
	if err == nil {
		return
	}
	// 1. 控制台输出（%+v 打印完整堆栈）
	if len(msg) > 0 && msg[0] != "" {
		g.Log().Errorf(ctx, "%s: %+v", msg[0], err)
	} else {
		g.Log().Errorf(ctx, "%+v", err)
	}
	// 2. Jaeger 记录（从 context 获取当前 span）
	span := trace.SpanFromContext(ctx)
	if span == nil || !span.IsRecording() {
		return
	}
	// 3. 记录错误到 span
	span.RecordError(err)
	span.SetAttributes(
		attribute.Bool("error", true),
		attribute.String("error.message", err.Error()),
	)
	// 4. 设置 span 状态为错误
	if len(msg) > 0 && msg[0] != "" {
		span.SetAttributes(attribute.String("error.msg", msg[0]))
		span.SetStatus(codes.Error, msg[0]+": "+err.Error())
		return
	}
	span.SetStatus(codes.Error, err.Error())
}

// NewTracer HTTP 请求链路追踪中间件
// 功能：
//  1. 为每个 HTTP 请求创建 Span
//  2. 记录请求参数和响应内容
//  3. 自动捕获错误并记录到 Jaeger
//
// 使用方式：在路由组中注册为中间件
//
//	group.Middleware(jaeger.NewTracer)
func NewTracer(r *ghttp.Request) {
	// 创建 Span（名称取自 controller 方法的 summary 标签）
	ctx, span := gtrace.NewSpan(r.Context(), r.GetServeHandler().GetMetaTag("summary"))
	r.SetCtx(ctx)
	defer span.End()
	// 记录请求参数
	span.SetAttributes(attribute.String("request", getParams(r)))
	// 执行后续中间件和 handler
	r.Middleware.Next()

	// 清理响应字符串，确保 UTF-8 有效（处理二进制数据如 ZIP 文件）
	response := r.Response.BufferString()
	cleanResponse := strings.ToValidUTF8(response, "")

	// 如果响应太大（如文件下载），只记录前 1000 字符
	if len(cleanResponse) > 1000 {
		cleanResponse = cleanResponse[:1000] + "... (truncated)"
	}

	span.SetAttributes(attribute.String("response", cleanResponse))
	span.SetAttributes(attribute.Int("http.status_code", r.Response.Status))

	if err := r.GetError(); err != nil {
		RecordError(ctx, err)
		return
	}
	if r.Response.Status >= 500 {
		span.SetAttributes(attribute.Bool("error", true))
		span.SetStatus(codes.Error, "http status "+strconv.Itoa(r.Response.Status))
	}
}

// getParams 提取请求参数（用于 Jaeger 记录）
func getParams(r *ghttp.Request) string {
	params := map[string]interface{}{}
	if r.Method == "POST" {
		json.Unmarshal(r.GetBody(), &params) //获取raw传参
	}
	if r.Method == "GET" {
		r.Request.ParseForm()
		form := r.Form
		for k, v := range form {
			if vl, e := strconv.Atoi(v[0]); e == nil {
				params[k] = vl
			} else {
				params[k] = v[0]
			}
		}
	}
	rp, _ := json.Marshal(&params)
	return string(rp)
}