github.com/weedge/lib@v0.0.0-20230424045628-a36dcc1d90e4/log/tapper/trace_ctx.go (about)

     1  package tapper
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"strconv"
     8  	"sync"
     9  
    10  	hack "github.com/weedge/lib/strings"
    11  
    12  	"github.com/gin-gonic/gin"
    13  	jsoniter "github.com/json-iterator/go"
    14  	"go.uber.org/atomic"
    15  )
    16  
    17  const (
    18  	TRACECTX = "TRACECTX"
    19  	TRACELOG = "TRACELOG"
    20  )
    21  
    22  type Trace struct {
    23  	data map[string]interface{}
    24  	mu   sync.Mutex
    25  }
    26  
    27  func (trace *Trace) Marshal() string {
    28  	trace.mu.Lock()
    29  	defer trace.mu.Unlock()
    30  
    31  	b, _ := jsoniter.MarshalToString(trace.data)
    32  	return b
    33  }
    34  func (trace *Trace) Push(args ...interface{}) *Trace {
    35  	trace.mu.Lock()
    36  	defer trace.mu.Unlock()
    37  
    38  	for i := 0; i < len(args)-1; i += 2 {
    39  		switch args[i+1].(type) {
    40  		case string, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
    41  			trace.data[fmt.Sprintf("%v", args[i])] = args[i+1]
    42  		default:
    43  			if b, e := jsoniter.Marshal(args[i+1]); e == nil {
    44  				trace.data[fmt.Sprintf("%v", args[i])] = hack.String(b)
    45  			}
    46  		}
    47  	}
    48  	return trace
    49  }
    50  
    51  func GetTraceFromContext(ctx context.Context) (context.Context, *Trace) {
    52  	switch c := ctx.(type) {
    53  	case *gin.Context:
    54  		if trace, ok := c.Get(TRACECTX); ok && trace != nil {
    55  			if trace, ok := trace.(*Trace); ok && trace != nil {
    56  				return ctx, trace
    57  			}
    58  		}
    59  
    60  		trace := &Trace{data: map[string]interface{}{}}
    61  		c.Set(TRACECTX, trace)
    62  		return ctx, trace
    63  
    64  	case nil:
    65  	default:
    66  	}
    67  
    68  	trace := &Trace{data: map[string]interface{}{}}
    69  	return ctx, trace
    70  }
    71  
    72  // 参数args: [key value]...
    73  func PushTrace(ctx context.Context, args ...interface{}) {
    74  	ctx, trace := GetTraceFromContext(ctx)
    75  	trace.Push(args...)
    76  }
    77  
    78  type TraceLog struct {
    79  	SpanNum atomic.Int64 //用于计算本地spanId自增拼接
    80  	LogId   string
    81  	SpanId  string
    82  	UniqId  string
    83  
    84  	UserIP  string
    85  	Product string
    86  
    87  	Caller string
    88  	Refer  string
    89  	Path   string // 当前请求的地址，用作请求下游时，设置成refer
    90  
    91  	MqTransId string //mq tranId 比如mq push模式 push uri地址参数中会有带上事务id
    92  }
    93  
    94  func (tl *TraceLog) FormatTraceString() string {
    95  	spidSuffixNum := tl.SpanNum.Inc()
    96  	bf := bytes.Buffer{}
    97  	bf.WriteString(" [logId:")
    98  	bf.WriteString(tl.LogId)
    99  	bf.WriteString("] [module:")
   100  	bf.WriteString(tl.Caller)
   101  	bf.WriteString("] [spanId:")
   102  	spanId := ""
   103  	if len(tl.SpanId) > 0 {
   104  		spanId = tl.SpanId + "." + strconv.FormatInt(spidSuffixNum, 10)
   105  	} else {
   106  		spanId = strconv.FormatInt(spidSuffixNum, 10)
   107  	}
   108  	bf.WriteString(spanId)
   109  	if len(tl.MqTransId) > 0 {
   110  		bf.WriteString("] [mq_transId:")
   111  		bf.WriteString(tl.MqTransId)
   112  	}
   113  	bf.WriteString("]")
   114  	return bf.String()
   115  }
   116  
   117  func (tl *TraceLog) GetCurrentSpanId() string {
   118  	if len(tl.SpanId) > 0 {
   119  		return tl.SpanId + "." + strconv.FormatInt(tl.SpanNum.Load(), 10)
   120  	}
   121  	return strconv.FormatInt(tl.SpanNum.Load(), 10)
   122  }
   123  
   124  var EmptyTraceLog = &TraceLog{}
   125  
   126  // 从context中获取span信息集合
   127  func GetTraceLogFromContext(ctx context.Context) (context.Context, *TraceLog) {
   128  	if span := ctx.Value(TRACELOG); span != nil {
   129  		traceLog, _ := span.(*TraceLog)
   130  		return ctx, traceLog
   131  	}
   132  
   133  	logId := GenLogId()
   134  	traceLog := &TraceLog{
   135  		LogId:  logId,
   136  		UniqId: "",
   137  		Caller: "",
   138  		Refer:  "",
   139  	}
   140  	ctx = context.WithValue(ctx, TRACELOG, traceLog)
   141  	return ctx, traceLog
   142  }
   143  
   144  // 从gin context中获取span信息集合
   145  func GetTraceLogFromGinContext(ctx context.Context) (*TraceLog, bool) {
   146  	var ok bool
   147  	var c *gin.Context
   148  
   149  	if c, ok = ctx.(*gin.Context); !ok {
   150  		return EmptyTraceLog, false
   151  	}
   152  
   153  	var traceLog *TraceLog
   154  	if value, exists := c.Get(TRACELOG); exists && value != nil {
   155  		if traceLog, ok = value.(*TraceLog); traceLog != nil && ok {
   156  			return traceLog, true
   157  		}
   158  	}
   159  
   160  	return TraceLogger.SetTraceLogFromGinHeader(c), true
   161  }
   162