gocommon/pkg/trace/gin_trace.go

package trace

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"time"

	"git.shuncheng.lu/bigthing/gocommon/pkg/internal/util"

	"git.shuncheng.lu/bigthing/gocommon/pkg/net/engines"
	"github.com/SkyAPM/go2sky"
	"github.com/SkyAPM/go2sky/propagation"
	v3 "github.com/SkyAPM/go2sky/reporter/grpc/language-agent"
)

const (
	ginRequestContextKey = "_gin_request_context_key"
	ginLocalSpan         = "_gin_local_span_key"
)

var (
	tracerNotFoundError = errors.New("tracer not found err")
	hostName            = util.CurrentServerHostName()
)

/**
ctx拿到的是 request 中 绑定的ctx
*/

/**
这里设计的主要目的：
1、gin的ctx和context.Context无法转换
2、所以所有的span信息都存入在 requestContext中
3、查找原则，如果发现ginRequestContextKey，则返回，没有发现则返回 ctx，那么链路会丢失
*/
func getGinRequestContext(ctx context.Context) context.Context {
	result, isOk := ctx.Value(ginRequestContextKey).(context.Context)
	if isOk {
		return result
	}
	return ctx
}

/**
设置 ctx
*/
func setGinRequestContext(ctx *engines.Context, reqCtx context.Context) {
	ctx.Set(ginRequestContextKey, reqCtx)
}

/**
设置LocalSpan 上下文
*/
func setGinLocalSpan(ctx context.Context, span go2sky.Span) context.Context {
	if span == nil || ctx == nil {
		return ctx
	}
	return context.WithValue(ctx, ginLocalSpan, span)
}

/**
存在 request的ctx中 span
*/
func GetGinLocalSpan(ctx context.Context) go2sky.Span {
	ctx = getGinRequestContext(ctx)
	value, isOk := ctx.Value(ginLocalSpan).(go2sky.Span)
	if isOk {
		return value
	}
	return nil
}

func GinLogInfo(ctx context.Context, log string) {
	if ctx == nil || log == "" {
		return
	}
	span := GetGinLocalSpan(ctx)
	if span != nil {
		span.Log(time.Now(), GetCallLine(3), log)
	}
}

// 这里只会记录error
func GinLogError(ctx context.Context, log string) {
	if ctx == nil || log == "" {
		return
	}
	span := GetGinLocalSpan(ctx)
	if span != nil {
		span.Error(time.Now(), GetCallLine(3), log)
	}
}

/**
这里需要注意，默认生成的一个空的trace_id，所以需要判断，是否未空trace_id
*/
func GinGetTraceId(ctx *engines.Context) string {
	traceId := go2sky.TraceID(ctx.Request.Context())
	if traceId == go2sky.EmptyTraceID {
		return ""
	}
	return traceId
}

func GetTraceId(ctx context.Context) string {
	traceId := go2sky.TraceID(ctx)
	if traceId == go2sky.EmptyTraceID {
		return ""
	}
	return traceId
}

/**
 */
func GinMiddleware(engine *engines.Engine, tracer *go2sky.Tracer) engines.HandlerFunc {
	if engine == nil || tracer == nil {
		return func(c *engines.Context) {
			c.Next()
		}
	}
	return func(c *engines.Context) {
		operationName := "/" + c.Request.Method + c.Request.URL.Path
		// 创建一个 entry span，这里需要解析header
		span, ctx, err := tracer.CreateEntrySpan(c.Request.Context(), operationName, func() (string, error) {
			return c.Request.Header.Get(propagation.Header), nil
		})

		// 如果异常则直接继续
		if span == nil || err != nil {
			c.Next()
			return
		}
		span.SetComponent(HttpServerComponentID)
		span.SetSpanLayer(v3.SpanLayer_Http)

		// 记录日志
		startTime := time.Now()
		defer func() {
			span.Tag("http.status", strconv.Itoa(c.Writer.Status()))
			span.Tag("spend", util.TimeToSeconds(time.Now().Sub(startTime).Seconds()))
			if c.Errors != nil && len(c.Errors) > 0 {
				span.Error(time.Now(), "Error", fmt.Sprintf("%v", c.Errors))
			}
			span.End()
		}()
		span.Tag("server.host", hostName)
		span.Tag("client.ip", c.ClientIP())
		span.Tag("http.method", c.Request.Method)
		span.Tag("http.url", c.Request.URL.String())

		// 继续创建
		localSpan, ctx, err := tracer.CreateLocalSpan(ctx)
		if localSpan != nil && err == nil {
			localSpan.SetOperationName(c.HandlerName())
			defer localSpan.End()
			// 将localSpan绑定到上下文
			ctx = setGinLocalSpan(ctx, localSpan)
		}

		// 绑定request
		c.Request = c.Request.WithContext(ctx)
		setGinRequestContext(c, c.Request.Context())

		// 继续调用
		c.Next()
	}
}

func GinNewExitSpan(ctx context.Context, opName, peerName string, handlerHeader propagation.Injector) (go2sky.Span, error) {
	tracer := GetSkyWalkingTracer()
	if tracer == nil {
		return nil, tracerNotFoundError
	}
	return tracer.CreateExitSpan(getGinRequestContext(ctx), opName, peerName, handlerHeader)
}

func GinNewLocalSpan(ctx context.Context) (go2sky.Span, context.Context, error) {
	tracer := GetSkyWalkingTracer()
	if tracer == nil {
		return nil, ctx, tracerNotFoundError
	}
	return tracer.CreateLocalSpan(getGinRequestContext(ctx))
}

func GinNewEntrySpan(ctx context.Context, opName string, handlerHeader propagation.Extractor) (go2sky.Span, context.Context, error) {
	tracer := GetSkyWalkingTracer()
	if tracer == nil {
		return nil, nil, tracerNotFoundError
	}
	return tracer.CreateEntrySpan(getGinRequestContext(ctx), opName, handlerHeader)
}