src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

context.go (3148B)

      1 package tracing
      2 
      3 import "context"
      4 
      5 type (
      6 	operationTracerKey struct{}
      7 	spanLineageKey     struct{}
      8 )
      9 
     10 // GetSpan returns the active trace Span on the context.
     11 //
     12 // The boolean in the return indicates whether a Span was actually in the
     13 // context, but a no-op implementation will be returned if not, so callers
     14 // can generally disregard the boolean unless they wish to explicitly confirm
     15 // presence/absence of a Span.
     16 func GetSpan(ctx context.Context) (Span, bool) {
     17 	lineage := getLineage(ctx)
     18 	if len(lineage) == 0 {
     19 		return nopSpan{}, false
     20 	}
     21 
     22 	return lineage[len(lineage)-1], true
     23 }
     24 
     25 // WithSpan sets the active trace Span on the context.
     26 func WithSpan(parent context.Context, span Span) context.Context {
     27 	lineage := getLineage(parent)
     28 	if len(lineage) == 0 {
     29 		return context.WithValue(parent, spanLineageKey{}, []Span{span})
     30 	}
     31 
     32 	lineage = append(lineage, span)
     33 	return context.WithValue(parent, spanLineageKey{}, lineage)
     34 }
     35 
     36 // PopSpan pops the current Span off the context, setting the active Span on
     37 // the returned Context back to its parent and returning the REMOVED one.
     38 //
     39 // PopSpan on a context with no active Span will return a no-op instance.
     40 //
     41 // This is mostly necessary for the runtime to manage base trace spans due to
     42 // the wrapped-function nature of the middleware stack. End-users of Smithy
     43 // clients SHOULD NOT generally be using this API.
     44 func PopSpan(parent context.Context) (context.Context, Span) {
     45 	lineage := getLineage(parent)
     46 	if len(lineage) == 0 {
     47 		return parent, nopSpan{}
     48 	}
     49 
     50 	span := lineage[len(lineage)-1]
     51 	lineage = lineage[:len(lineage)-1]
     52 	return context.WithValue(parent, spanLineageKey{}, lineage), span
     53 }
     54 
     55 func getLineage(ctx context.Context) []Span {
     56 	v := ctx.Value(spanLineageKey{})
     57 	if v == nil {
     58 		return nil
     59 	}
     60 
     61 	return v.([]Span)
     62 }
     63 
     64 // GetOperationTracer returns the embedded operation-scoped Tracer on a
     65 // Context.
     66 //
     67 // The boolean in the return indicates whether a Tracer was actually in the
     68 // context, but a no-op implementation will be returned if not, so callers
     69 // can generally disregard the boolean unless they wish to explicitly confirm
     70 // presence/absence of a Tracer.
     71 func GetOperationTracer(ctx context.Context) (Tracer, bool) {
     72 	v := ctx.Value(operationTracerKey{})
     73 	if v == nil {
     74 		return nopTracer{}, false
     75 	}
     76 
     77 	return v.(Tracer), true
     78 }
     79 
     80 // WithOperationTracer returns a child Context embedding the given Tracer.
     81 //
     82 // The runtime will use this embed a scoped tracer for client operations,
     83 // Smithy/SDK client callers DO NOT need to do this explicitly.
     84 func WithOperationTracer(parent context.Context, tracer Tracer) context.Context {
     85 	return context.WithValue(parent, operationTracerKey{}, tracer)
     86 }
     87 
     88 // StartSpan is a convenience API for creating tracing Spans from a Context.
     89 //
     90 // StartSpan uses the operation-scoped Tracer, previously stored using
     91 // [WithOperationTracer], to start the Span. If a Tracer has not been embedded
     92 // the returned Span will be a no-op implementation.
     93 func StartSpan(ctx context.Context, name string, opts ...SpanOption) (context.Context, Span) {
     94 	tracer, _ := GetOperationTracer(ctx)
     95 	return tracer.StartSpan(ctx, name, opts...)
     96 }