vendor/github.com/go-kit/kit/tracing/opentracing/endpoint.go

   1 package opentracing
   2
   3 import (
   4         "context"
   5
   6         "github.com/opentracing/opentracing-go"
   7         otext "github.com/opentracing/opentracing-go/ext"
   8
   9         "github.com/go-kit/kit/endpoint"
  10 )
  11
  12 // TraceServer returns a Middleware that wraps the `next` Endpoint in an
  13 // OpenTracing Span called `operationName`.
  14 //
  15 // If `ctx` already has a Span, it is re-used and the operation name is
  16 // overwritten. If `ctx` does not yet have a Span, one is created here.
  17 func TraceServer(tracer opentracing.Tracer, operationName string) endpoint.Middleware {
  18         return func(next endpoint.Endpoint) endpoint.Endpoint {
  19                 return func(ctx context.Context, request interface{}) (interface{}, error) {
  20                         serverSpan := opentracing.SpanFromContext(ctx)
  21                         if serverSpan == nil {
  22                                 // All we can do is create a new root span.
  23                                 serverSpan = tracer.StartSpan(operationName)
  24                         } else {
  25                                 serverSpan.SetOperationName(operationName)
  26                         }
  27                         defer serverSpan.Finish()
  28                         otext.SpanKindRPCServer.Set(serverSpan)
  29                         ctx = opentracing.ContextWithSpan(ctx, serverSpan)
  30                         return next(ctx, request)
  31                 }
  32         }
  33 }
  34
  35 // TraceClient returns a Middleware that wraps the `next` Endpoint in an
  36 // OpenTracing Span called `operationName`.
  37 func TraceClient(tracer opentracing.Tracer, operationName string) endpoint.Middleware {
  38         return func(next endpoint.Endpoint) endpoint.Endpoint {
  39                 return func(ctx context.Context, request interface{}) (interface{}, error) {
  40                         var clientSpan opentracing.Span
  41                         if parentSpan := opentracing.SpanFromContext(ctx); parentSpan != nil {
  42                                 clientSpan = tracer.StartSpan(
  43                                         operationName,
  44                                         opentracing.ChildOf(parentSpan.Context()),
  45                                 )
  46                         } else {
  47                                 clientSpan = tracer.StartSpan(operationName)
  48                         }
  49                         defer clientSpan.Finish()
  50                         otext.SpanKindRPCClient.Set(clientSpan)
  51                         ctx = opentracing.ContextWithSpan(ctx, clientSpan)
  52                         return next(ctx, request)
  53                 }
  54         }
  55 }