grpc.go

Documentation: github.com/go-kit/kit/tracing/opencensus

     1  package opencensus
     2  
     3  import (
     4  	"context"
     5  
     6  	"go.opencensus.io/trace"
     7  	"go.opencensus.io/trace/propagation"
     8  	"google.golang.org/grpc/codes"
     9  	"google.golang.org/grpc/metadata"
    10  	"google.golang.org/grpc/status"
    11  
    12  	kitgrpc "github.com/go-kit/kit/transport/grpc"
    13  )
    14  
    15  const propagationKey = "grpc-trace-bin"
    16  
    17  // GRPCClientTrace enables OpenCensus tracing of a Go kit gRPC transport client.
    18  func GRPCClientTrace(options ...TracerOption) kitgrpc.ClientOption {
    19  	cfg := TracerOptions{}
    20  
    21  	for _, option := range options {
    22  		option(&cfg)
    23  	}
    24  
    25  	clientBefore := kitgrpc.ClientBefore(
    26  		func(ctx context.Context, md *metadata.MD) context.Context {
    27  			var name string
    28  
    29  			if cfg.Name != "" {
    30  				name = cfg.Name
    31  			} else {
    32  				name = ctx.Value(kitgrpc.ContextKeyRequestMethod).(string)
    33  			}
    34  
    35  			ctx, span := trace.StartSpan(
    36  				ctx,
    37  				name,
    38  				trace.WithSampler(cfg.Sampler),
    39  				trace.WithSpanKind(trace.SpanKindClient),
    40  			)
    41  
    42  			if !cfg.Public {
    43  				traceContextBinary := string(propagation.Binary(span.SpanContext()))
    44  				(*md)[propagationKey] = append((*md)[propagationKey], traceContextBinary)
    45  			}
    46  
    47  			return ctx
    48  		},
    49  	)
    50  
    51  	clientFinalizer := kitgrpc.ClientFinalizer(
    52  		func(ctx context.Context, err error) {
    53  			if span := trace.FromContext(ctx); span != nil {
    54  				if s, ok := status.FromError(err); ok {
    55  					span.SetStatus(trace.Status{Code: int32(s.Code()), Message: s.Message()})
    56  				} else {
    57  					span.SetStatus(trace.Status{Code: int32(codes.Unknown), Message: err.Error()})
    58  				}
    59  				span.End()
    60  			}
    61  		},
    62  	)
    63  
    64  	return func(c *kitgrpc.Client) {
    65  		clientBefore(c)
    66  		clientFinalizer(c)
    67  	}
    68  }
    69  
    70  // GRPCServerTrace enables OpenCensus tracing of a Go kit gRPC transport server.
    71  func GRPCServerTrace(options ...TracerOption) kitgrpc.ServerOption {
    72  	cfg := TracerOptions{}
    73  
    74  	for _, option := range options {
    75  		option(&cfg)
    76  	}
    77  
    78  	serverBefore := kitgrpc.ServerBefore(
    79  		func(ctx context.Context, md metadata.MD) context.Context {
    80  			var name string
    81  
    82  			if cfg.Name != "" {
    83  				name = cfg.Name
    84  			} else {
    85  				name, _ = ctx.Value(kitgrpc.ContextKeyRequestMethod).(string)
    86  				if name == "" {
    87  					// we can't find the gRPC method. probably the
    88  					// unaryInterceptor was not wired up.
    89  					name = "unknown grpc method"
    90  				}
    91  			}
    92  
    93  			var (
    94  				parentContext trace.SpanContext
    95  				traceContext  = md[propagationKey]
    96  				ok            bool
    97  			)
    98  
    99  			if len(traceContext) > 0 {
   100  				traceContextBinary := []byte(traceContext[0])
   101  				parentContext, ok = propagation.FromBinary(traceContextBinary)
   102  				if ok && !cfg.Public {
   103  					ctx, _ = trace.StartSpanWithRemoteParent(
   104  						ctx,
   105  						name,
   106  						parentContext,
   107  						trace.WithSpanKind(trace.SpanKindServer),
   108  						trace.WithSampler(cfg.Sampler),
   109  					)
   110  					return ctx
   111  				}
   112  			}
   113  			ctx, span := trace.StartSpan(
   114  				ctx,
   115  				name,
   116  				trace.WithSpanKind(trace.SpanKindServer),
   117  				trace.WithSampler(cfg.Sampler),
   118  			)
   119  			if ok {
   120  				span.AddLink(
   121  					trace.Link{
   122  						TraceID: parentContext.TraceID,
   123  						SpanID:  parentContext.SpanID,
   124  						Type:    trace.LinkTypeChild,
   125  					},
   126  				)
   127  			}
   128  			return ctx
   129  		},
   130  	)
   131  
   132  	serverFinalizer := kitgrpc.ServerFinalizer(
   133  		func(ctx context.Context, err error) {
   134  			if span := trace.FromContext(ctx); span != nil {
   135  				if s, ok := status.FromError(err); ok {
   136  					span.SetStatus(trace.Status{Code: int32(s.Code()), Message: s.Message()})
   137  				} else {
   138  					span.SetStatus(trace.Status{Code: int32(codes.Internal), Message: err.Error()})
   139  				}
   140  				span.End()
   141  			}
   142  		},
   143  	)
   144  
   145  	return func(s *kitgrpc.Server) {
   146  		serverBefore(s)
   147  		serverFinalizer(s)
   148  	}
   149  }
   150
View as plain text