gopkgsv2/trace/trace.go

// Copyright 2020 getensh.com. All rights reserved.
// Use of this source code is governed by getensh.com.

/**
 * @Author: mac
 * @Description:
 * @Date: 2020/3/25 9:46
 */

package trace

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/uber/jaeger-client-go"
	"github.com/uber/jaeger-client-go/config"
	"google.golang.org/grpc/metadata"
	"io"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"

	"github.com/opentracing-contrib/go-stdlib/nethttp"
	"github.com/opentracing/opentracing-go"
)

type JaegerConf struct {
	Open          bool
	AgentHostPort string
	SamplerType   string
	SamplerParam  float64
	LogSpans      bool
}

/* url path 过滤名单，加入的url path不会进行调用链追踪*/
// 不确定全路径，需要正则的加到此处
var urlPathFiltersRegexp = []string{
	"/swagger/*",
	"/*.cfg",
}

// 可以确定全路径的加到此处
var urlPathFiltersMap = map[string]bool{
	"/":            true,
	"/metrics":     true,
	"/favicon.ico": true,
	"/version":     true,
}

// 增加要跳过的全路径
func FilterPathAdd(filters ...string) {
	for _, path := range filters {
		urlPathFiltersMap[path] = true
	}
}

// 增加要跳过的路径正则
func FilterPathRegexpAdd(filters ...string) {
	urlPathFiltersRegexp = append(urlPathFiltersRegexp, filters...)
}

type MDReaderWriter struct {
	metadata.MD
}

// ForeachKey 实现 ForeachKey of opentracing.TextMapReader
func (c MDReaderWriter) ForeachKey(handler func(key, val string) error) error {
	for k, vs := range c.MD {
		for _, v := range vs {
			if err := handler(k, v); err != nil {
				return err
			}
		}
	}
	return nil
}

// Set 实现 Set() of opentracing.TextMapWriter
func (c MDReaderWriter) Set(key, val string) {
	key = strings.ToLower(key)
	c.MD[key] = append(c.MD[key], val)
}

// NewJaegerTracer jaeger初始化
func NewJaegerTracer(service string, jaegerConf JaegerConf) (opentracing.Tracer, io.Closer) {
	cfg := &config.Configuration{
		Sampler: &config.SamplerConfig{
			Type:  jaegerConf.SamplerType,
			Param: jaegerConf.SamplerParam,
		},
		Reporter: &config.ReporterConfig{
			LogSpans:           jaegerConf.LogSpans, // 关闭打印span上报日志
			LocalAgentHostPort: jaegerConf.AgentHostPort,
		},
	}
	tracer, closer, err := cfg.New(service, config.Logger(jaeger.StdLogger))
	if err != nil {
		log.Fatalf("ERROR: cannot init tracer: %v\n", err)
	}
	opentracing.SetGlobalTracer(tracer)
	return tracer, closer
}

// JaegerPluginsForGin gin jaeger 中间件，在每次http调用前拦截，进行trace注入
func JaegerPluginsForGin() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 如果不需要链路的，直接跳过
		urlPath := c.Request.URL.Path
		if _, ok := urlPathFiltersMap[urlPath]; ok {
			c.Next()
			return
		}
		for _, pattern := range urlPathFiltersRegexp {
			if mathed, _ := regexp.MatchString(pattern, urlPath); mathed {
				c.Next()
				return
			}
		}

		var parentSpan opentracing.Span

		tracer := opentracing.GlobalTracer()

		spCtx, err := tracer.Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(c.Request.Header))
		opName := "HTTP " + c.Request.Method + " " + urlPath
		if err != nil {
			parentSpan = tracer.StartSpan(opName)
			defer parentSpan.Finish()
		} else {
			parentSpan = opentracing.StartSpan(
				opName,
				opentracing.ChildOf(spCtx),
				opentracing.Tag{Key: string(ext.Component), Value: "HTTP"},
				ext.SpanKindRPCServer,
			)
			defer parentSpan.Finish()
		}

		ctx := opentracing.ContextWithSpan(c.Request.Context(), parentSpan)

		// 设置 request ID,可用于日志记录, 使用ctx.Value("RequestId")可以取出
		if sc, ok := parentSpan.Context().(jaeger.SpanContext); ok {
			traceId := sc.TraceID().String()
			ctx = context.WithValue(ctx, "RequestId", traceId)
			c.Header("X-Request-Id", traceId)
		}

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

		c.Next()
	}
}

// JaegerPluginsForBeego beego jaeger中间件，在每次http调用前拦截，进行trace注入
func JaegerPluginsForBeego(componentName string) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		opts := []nethttp.MWOption{
			nethttp.MWComponentName(componentName),
			nethttp.OperationNameFunc(func(r *http.Request) string {
				return "HTTP " + r.Method + " " + r.URL.Path
			}),
			nethttp.MWURLTagFunc(func(u *url.URL) string {
				return u.String()
			}),
			nethttp.MWSpanFilter(func(r *http.Request) bool {
				// 如果不需要链路的，直接跳过
				if _, ok := urlPathFiltersMap[r.URL.Path]; ok {
					return false
				}
				for _, pattern := range urlPathFiltersRegexp {
					if mathed, _ := regexp.MatchString(pattern, r.URL.Path); mathed {
						return false
					}
				}
				return true
			}),
			nethttp.MWSpanObserver(func(span opentracing.Span, r *http.Request) {
				ctx := opentracing.ContextWithSpan(r.Context(), span)
				// 设置 request ID,可用于日志记录, 使用ctx.Value("RequestId")可以取出
				if sc, ok := span.Context().(jaeger.SpanContext); ok {
					traceId := sc.TraceID().String()
					ctx = context.WithValue(ctx, "RequestId", traceId)
					r.Header.Set("X-Request-Id", traceId)
				}
			}),
		}

		return nethttp.Middleware(opentracing.GlobalTracer(), h, opts...)
	}
}

// 继承 ctx，放置断链
func CtxFromParent(parentCtx context.Context) context.Context {
	span := opentracing.SpanFromContext(parentCtx)
	return opentracing.ContextWithSpan(context.Background(), span)
}

// 根据父级ctx生成子级span和ctx,同服务调用时使用(如mysql、redis或者服务内方法)
func StartChildSpan(parentCtx context.Context, operationName string) (opentracing.Span, context.Context) {
	return opentracing.StartSpanFromContext(parentCtx, operationName)
}