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

   1 package opentracing
   2
   3 import (
   4         "context"
   5         "encoding/base64"
   6         "strings"
   7
   8         "github.com/opentracing/opentracing-go"
   9         "github.com/opentracing/opentracing-go/ext"
  10         "google.golang.org/grpc/metadata"
  11
  12         "github.com/go-kit/kit/log"
  13 )
  14
  15 // ContextToGRPC returns a grpc RequestFunc that injects an OpenTracing Span
  16 // found in `ctx` into the grpc Metadata. If no such Span can be found, the
  17 // RequestFunc is a noop.
  18 func ContextToGRPC(tracer opentracing.Tracer, logger log.Logger) func(ctx context.Context, md *metadata.MD) context.Context {
  19         return func(ctx context.Context, md *metadata.MD) context.Context {
  20                 if span := opentracing.SpanFromContext(ctx); span != nil {
  21                         // There's nothing we can do with an error here.
  22                         if err := tracer.Inject(span.Context(), opentracing.TextMap, metadataReaderWriter{md}); err != nil {
  23                                 logger.Log("err", err)
  24                         }
  25                 }
  26                 return ctx
  27         }
  28 }
  29
  30 // GRPCToContext returns a grpc RequestFunc that tries to join with an
  31 // OpenTracing trace found in `req` and starts a new Span called
  32 // `operationName` accordingly. If no trace could be found in `req`, the Span
  33 // will be a trace root. The Span is incorporated in the returned Context and
  34 // can be retrieved with opentracing.SpanFromContext(ctx).
  35 func GRPCToContext(tracer opentracing.Tracer, operationName string, logger log.Logger) func(ctx context.Context, md metadata.MD) context.Context {
  36         return func(ctx context.Context, md metadata.MD) context.Context {
  37                 var span opentracing.Span
  38                 wireContext, err := tracer.Extract(opentracing.TextMap, metadataReaderWriter{&md})
  39                 if err != nil && err != opentracing.ErrSpanContextNotFound {
  40                         logger.Log("err", err)
  41                 }
  42                 span = tracer.StartSpan(operationName, ext.RPCServerOption(wireContext))
  43                 return opentracing.ContextWithSpan(ctx, span)
  44         }
  45 }
  46
  47 // A type that conforms to opentracing.TextMapReader and
  48 // opentracing.TextMapWriter.
  49 type metadataReaderWriter struct {
  50         *metadata.MD
  51 }
  52
  53 func (w metadataReaderWriter) Set(key, val string) {
  54         key = strings.ToLower(key)
  55         if strings.HasSuffix(key, "-bin") {
  56                 val = string(base64.StdEncoding.EncodeToString([]byte(val)))
  57         }
  58         (*w.MD)[key] = append((*w.MD)[key], val)
  59 }
  60
  61 func (w metadataReaderWriter) ForeachKey(handler func(key, val string) error) error {
  62         for k, vals := range *w.MD {
  63                 for _, v := range vals {
  64                         if err := handler(k, v); err != nil {
  65                                 return err
  66                         }
  67                 }
  68         }
  69         return nil
  70 }