6 "github.com/opentracing/opentracing-go"
7 otext "github.com/opentracing/opentracing-go/ext"
9 "github.com/go-kit/kit/endpoint"
12 // TraceServer returns a Middleware that wraps the `next` Endpoint in an
13 // OpenTracing Span called `operationName`.
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)
25 serverSpan.SetOperationName(operationName)
27 defer serverSpan.Finish()
28 otext.SpanKindRPCServer.Set(serverSpan)
29 ctx = opentracing.ContextWithSpan(ctx, serverSpan)
30 return next(ctx, request)
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(
44 opentracing.ChildOf(parentSpan.Context()),
47 clientSpan = tracer.StartSpan(operationName)
49 defer clientSpan.Finish()
50 otext.SpanKindRPCClient.Set(clientSpan)
51 ctx = opentracing.ContextWithSpan(ctx, clientSpan)
52 return next(ctx, request)