grpc.go

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

     1  package zipkin
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  
     7  	zipkin "github.com/openzipkin/zipkin-go"
     8  	"github.com/openzipkin/zipkin-go/model"
     9  	"github.com/openzipkin/zipkin-go/propagation/b3"
    10  	"google.golang.org/grpc/metadata"
    11  	"google.golang.org/grpc/status"
    12  
    13  	kitgrpc "github.com/go-kit/kit/transport/grpc"
    14  	"github.com/go-kit/log"
    15  )
    16  
    17  // GRPCClientTrace enables native Zipkin tracing of a Go kit gRPC transport
    18  // Client.
    19  //
    20  // Go kit creates gRPC transport clients per remote endpoint. This middleware
    21  // can be set-up individually by adding the endpoint name for each of the Go kit
    22  // transport clients using the Name() TracerOption.
    23  // If wanting to use the gRPC FullMethod (/service/method) as Span name you can
    24  // create a global client tracer omitting the Name() TracerOption, which you can
    25  // then feed to each Go kit gRPC transport client.
    26  // If instrumenting a client to an external (not on your platform) service, you
    27  // will probably want to disallow propagation of SpanContext using the
    28  // AllowPropagation TracerOption and setting it to false.
    29  func GRPCClientTrace(tracer *zipkin.Tracer, options ...TracerOption) kitgrpc.ClientOption {
    30  	config := tracerOptions{
    31  		tags:      make(map[string]string),
    32  		name:      "",
    33  		logger:    log.NewNopLogger(),
    34  		propagate: true,
    35  	}
    36  
    37  	for _, option := range options {
    38  		option(&config)
    39  	}
    40  
    41  	clientBefore := kitgrpc.ClientBefore(
    42  		func(ctx context.Context, md *metadata.MD) context.Context {
    43  			var (
    44  				spanContext model.SpanContext
    45  				name        string
    46  			)
    47  
    48  			if config.name != "" {
    49  				name = config.name
    50  			} else {
    51  				name = ctx.Value(kitgrpc.ContextKeyRequestMethod).(string)
    52  			}
    53  
    54  			if parent := zipkin.SpanFromContext(ctx); parent != nil {
    55  				spanContext = parent.Context()
    56  			}
    57  
    58  			span := tracer.StartSpan(
    59  				name,
    60  				zipkin.Kind(model.Client),
    61  				zipkin.Tags(config.tags),
    62  				zipkin.Parent(spanContext),
    63  				zipkin.FlushOnFinish(false),
    64  			)
    65  
    66  			if config.propagate {
    67  				if err := b3.InjectGRPC(md)(span.Context()); err != nil {
    68  					config.logger.Log("err", err)
    69  				}
    70  			}
    71  
    72  			return zipkin.NewContext(ctx, span)
    73  		},
    74  	)
    75  
    76  	clientAfter := kitgrpc.ClientAfter(
    77  		func(ctx context.Context, _ metadata.MD, _ metadata.MD) context.Context {
    78  			if span := zipkin.SpanFromContext(ctx); span != nil {
    79  				span.Finish()
    80  			}
    81  
    82  			return ctx
    83  		},
    84  	)
    85  
    86  	clientFinalizer := kitgrpc.ClientFinalizer(
    87  		func(ctx context.Context, err error) {
    88  			if span := zipkin.SpanFromContext(ctx); span != nil {
    89  				if err != nil {
    90  					zipkin.TagError.Set(span, err.Error())
    91  				}
    92  				// calling span.Finish() a second time is a noop, if we didn't get to
    93  				// ClientAfter we can at least time the early bail out by calling it
    94  				// here.
    95  				span.Finish()
    96  				// send span to the Reporter
    97  				span.Flush()
    98  			}
    99  		},
   100  	)
   101  
   102  	return func(c *kitgrpc.Client) {
   103  		clientBefore(c)
   104  		clientAfter(c)
   105  		clientFinalizer(c)
   106  	}
   107  
   108  }
   109  
   110  // GRPCServerTrace enables native Zipkin tracing of a Go kit gRPC transport
   111  // Server.
   112  //
   113  // Go kit creates gRPC transport servers per gRPC method. This middleware can be
   114  // set-up individually by adding the method name for each of the Go kit method
   115  // servers using the Name() TracerOption.
   116  // If wanting to use the gRPC FullMethod (/service/method) as Span name you can
   117  // create a global server tracer omitting the Name() TracerOption, which you can
   118  // then feed to each Go kit method server. For this to work you will need to
   119  // wire the Go kit gRPC Interceptor too.
   120  // If instrumenting a service to external (not on your platform) clients, you
   121  // will probably want to disallow propagation of a client SpanContext using
   122  // the AllowPropagation TracerOption and setting it to false.
   123  func GRPCServerTrace(tracer *zipkin.Tracer, options ...TracerOption) kitgrpc.ServerOption {
   124  	config := tracerOptions{
   125  		tags:      make(map[string]string),
   126  		name:      "",
   127  		logger:    log.NewNopLogger(),
   128  		propagate: true,
   129  	}
   130  
   131  	for _, option := range options {
   132  		option(&config)
   133  	}
   134  
   135  	serverBefore := kitgrpc.ServerBefore(
   136  		func(ctx context.Context, md metadata.MD) context.Context {
   137  			var (
   138  				spanContext model.SpanContext
   139  				name        string
   140  				tags        = make(map[string]string)
   141  			)
   142  
   143  			rpcMethod, ok := ctx.Value(kitgrpc.ContextKeyRequestMethod).(string)
   144  			if !ok {
   145  				config.logger.Log("err", "unable to retrieve method name: missing gRPC interceptor hook")
   146  			} else {
   147  				tags["grpc.method"] = rpcMethod
   148  			}
   149  
   150  			if config.name != "" {
   151  				name = config.name
   152  			} else {
   153  				name = rpcMethod
   154  			}
   155  
   156  			if config.propagate {
   157  				spanContext = tracer.Extract(b3.ExtractGRPC(&md))
   158  				if spanContext.Err != nil {
   159  					config.logger.Log("err", spanContext.Err)
   160  				}
   161  			}
   162  
   163  			span := tracer.StartSpan(
   164  				name,
   165  				zipkin.Kind(model.Server),
   166  				zipkin.Tags(config.tags),
   167  				zipkin.Tags(tags),
   168  				zipkin.Parent(spanContext),
   169  				zipkin.FlushOnFinish(false),
   170  			)
   171  
   172  			return zipkin.NewContext(ctx, span)
   173  		},
   174  	)
   175  
   176  	serverAfter := kitgrpc.ServerAfter(
   177  		func(ctx context.Context, _ *metadata.MD, _ *metadata.MD) context.Context {
   178  			if span := zipkin.SpanFromContext(ctx); span != nil {
   179  				span.Finish()
   180  			}
   181  
   182  			return ctx
   183  		},
   184  	)
   185  
   186  	serverFinalizer := kitgrpc.ServerFinalizer(
   187  		func(ctx context.Context, err error) {
   188  			if span := zipkin.SpanFromContext(ctx); span != nil {
   189  				if err != nil {
   190  					if status, ok := status.FromError(err); ok {
   191  						statusCode := strconv.FormatUint(uint64(status.Code()), 10)
   192  						zipkin.TagGRPCStatusCode.Set(span, statusCode)
   193  						zipkin.TagError.Set(span, status.Message())
   194  					} else {
   195  						zipkin.TagError.Set(span, err.Error())
   196  					}
   197  				}
   198  
   199  				// calling span.Finish() a second time is a noop, if we didn't get to
   200  				// ServerAfter we can at least time the early bail out by calling it
   201  				// here.
   202  				span.Finish()
   203  				// send span to the Reporter
   204  				span.Flush()
   205  			}
   206  		},
   207  	)
   208  
   209  	return func(s *kitgrpc.Server) {
   210  		serverBefore(s)
   211  		serverAfter(s)
   212  		serverFinalizer(s)
   213  	}
   214  }
   215
View as plain text