github.com/palisadeinc/bor@v0.0.0-20230615125219-ab7196213d15/common/tracing/context.go (about)

     1  package tracing
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"go.opentelemetry.io/otel"
     8  	"go.opentelemetry.io/otel/attribute"
     9  	"go.opentelemetry.io/otel/trace"
    10  )
    11  
    12  type tracerKey struct{}
    13  
    14  type Option func(context.Context, trace.Span)
    15  
    16  func WithTracer(ctx context.Context, tr trace.Tracer) context.Context {
    17  	return context.WithValue(ctx, tracerKey{}, tr)
    18  }
    19  
    20  func FromContext(ctx context.Context) trace.Tracer {
    21  	tr, _ := ctx.Value(tracerKey{}).(trace.Tracer)
    22  
    23  	return tr
    24  }
    25  
    26  func StartSpan(ctx context.Context, snapName string) (context.Context, trace.Span) {
    27  	tr := FromContext(ctx)
    28  
    29  	if tr == nil {
    30  		return ctx, nil
    31  	}
    32  
    33  	ctx, span := tr.Start(ctx, snapName)
    34  	ctx = WithTracer(ctx, tr)
    35  
    36  	return ctx, span
    37  }
    38  
    39  func EndSpan(span trace.Span) {
    40  	if span != nil {
    41  		span.End()
    42  	}
    43  }
    44  
    45  func Trace(ctx context.Context, spanName string) (context.Context, trace.Span) {
    46  	tr := FromContext(ctx)
    47  
    48  	if tr == nil {
    49  		return ctx, nil
    50  	}
    51  
    52  	return tr.Start(ctx, spanName)
    53  }
    54  
    55  func Exec(ctx context.Context, instrumentationName, spanName string, opts ...Option) {
    56  	var span trace.Span
    57  
    58  	tr := FromContext(ctx)
    59  
    60  	if tr == nil && len(instrumentationName) != 0 {
    61  		tr = otel.GetTracerProvider().Tracer(instrumentationName)
    62  		ctx = WithTracer(ctx, tr)
    63  	}
    64  
    65  	if tr != nil {
    66  		ctx, span = tr.Start(ctx, spanName)
    67  	}
    68  
    69  	for _, optFn := range opts {
    70  		optFn(ctx, span)
    71  	}
    72  
    73  	if tr != nil {
    74  		span.End()
    75  	}
    76  }
    77  
    78  func WithTime(fn func(context.Context, trace.Span)) Option {
    79  	return func(ctx context.Context, span trace.Span) {
    80  		ElapsedTime(ctx, span, "elapsed", fn)
    81  	}
    82  }
    83  
    84  func ElapsedTime(ctx context.Context, span trace.Span, msg string, fn func(context.Context, trace.Span)) {
    85  	var now time.Time
    86  
    87  	if span != nil {
    88  		now = time.Now()
    89  	}
    90  
    91  	fn(ctx, span)
    92  
    93  	if span != nil {
    94  		span.SetAttributes(attribute.Int(msg, int(time.Since(now).Microseconds())))
    95  	}
    96  }
    97  
    98  func SetAttributes(span trace.Span, kvs ...attribute.KeyValue) {
    99  	if span != nil {
   100  		span.SetAttributes(kvs...)
   101  	}
   102  }