context.go raw

   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  }
  97