jsonrpc.go

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

     1  package opencensus
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  
     7  	"go.opencensus.io/plugin/ochttp"
     8  	"go.opencensus.io/plugin/ochttp/propagation/b3"
     9  	"go.opencensus.io/trace"
    10  
    11  	kithttp "github.com/go-kit/kit/transport/http"
    12  	jsonrpc "github.com/go-kit/kit/transport/http/jsonrpc"
    13  )
    14  
    15  // JSONRPCClientTrace enables OpenCensus tracing of a Go kit JSONRPC transport client.
    16  func JSONRPCClientTrace(options ...TracerOption) jsonrpc.ClientOption {
    17  	cfg := TracerOptions{}
    18  
    19  	for _, option := range options {
    20  		option(&cfg)
    21  	}
    22  
    23  	if !cfg.Public && cfg.HTTPPropagate == nil {
    24  		cfg.HTTPPropagate = &b3.HTTPFormat{}
    25  	}
    26  
    27  	clientBefore := jsonrpc.ClientBefore(
    28  		func(ctx context.Context, req *http.Request) context.Context {
    29  			var name string
    30  
    31  			if cfg.Name != "" {
    32  				name = cfg.Name
    33  			} else {
    34  				// OpenCensus states Path being default naming for a client span
    35  				name = ctx.Value(jsonrpc.ContextKeyRequestMethod).(string)
    36  			}
    37  
    38  			ctx, span := trace.StartSpan(
    39  				ctx,
    40  				name,
    41  				trace.WithSampler(cfg.Sampler),
    42  				trace.WithSpanKind(trace.SpanKindClient),
    43  			)
    44  
    45  			span.AddAttributes(
    46  				trace.StringAttribute(ochttp.HostAttribute, req.URL.Host),
    47  				trace.StringAttribute(ochttp.MethodAttribute, req.Method),
    48  				trace.StringAttribute(ochttp.PathAttribute, req.URL.Path),
    49  				trace.StringAttribute(ochttp.UserAgentAttribute, req.UserAgent()),
    50  			)
    51  
    52  			if !cfg.Public {
    53  				cfg.HTTPPropagate.SpanContextToRequest(span.SpanContext(), req)
    54  			}
    55  
    56  			return ctx
    57  		},
    58  	)
    59  
    60  	clientAfter := jsonrpc.ClientAfter(
    61  		func(ctx context.Context, res *http.Response) context.Context {
    62  			if span := trace.FromContext(ctx); span != nil {
    63  				span.SetStatus(ochttp.TraceStatus(res.StatusCode, http.StatusText(res.StatusCode)))
    64  				span.AddAttributes(
    65  					trace.Int64Attribute(ochttp.StatusCodeAttribute, int64(res.StatusCode)),
    66  				)
    67  			}
    68  			return ctx
    69  		},
    70  	)
    71  
    72  	clientFinalizer := jsonrpc.ClientFinalizer(
    73  		func(ctx context.Context, err error) {
    74  			if span := trace.FromContext(ctx); span != nil {
    75  				if err != nil {
    76  					span.SetStatus(trace.Status{
    77  						Code:    trace.StatusCodeUnknown,
    78  						Message: err.Error(),
    79  					})
    80  				}
    81  				span.End()
    82  			}
    83  		},
    84  	)
    85  
    86  	return func(c *jsonrpc.Client) {
    87  		clientBefore(c)
    88  		clientAfter(c)
    89  		clientFinalizer(c)
    90  	}
    91  }
    92  
    93  // JSONRPCServerTrace enables OpenCensus tracing of a Go kit JSONRPC transport server.
    94  func JSONRPCServerTrace(options ...TracerOption) jsonrpc.ServerOption {
    95  	cfg := TracerOptions{}
    96  
    97  	for _, option := range options {
    98  		option(&cfg)
    99  	}
   100  
   101  	if !cfg.Public && cfg.HTTPPropagate == nil {
   102  		cfg.HTTPPropagate = &b3.HTTPFormat{}
   103  	}
   104  
   105  	serverBeforeCodec := jsonrpc.ServerBeforeCodec(
   106  		func(ctx context.Context, httpReq *http.Request, req jsonrpc.Request) context.Context {
   107  			var (
   108  				spanContext trace.SpanContext
   109  				span        *trace.Span
   110  				name        string
   111  				ok          bool
   112  			)
   113  
   114  			if cfg.Name != "" {
   115  				name = cfg.Name
   116  			} else {
   117  				name = ctx.Value(jsonrpc.ContextKeyRequestMethod).(string)
   118  				if name == "" {
   119  					// we can't find the rpc method. probably the
   120  					// unaryInterceptor was not wired up.
   121  					name = "unknown jsonrpc method"
   122  				}
   123  			}
   124  
   125  			spanContext, ok = cfg.HTTPPropagate.SpanContextFromRequest(httpReq)
   126  			if ok && !cfg.Public {
   127  				ctx, span = trace.StartSpanWithRemoteParent(
   128  					ctx,
   129  					name,
   130  					spanContext,
   131  					trace.WithSpanKind(trace.SpanKindServer),
   132  					trace.WithSampler(cfg.Sampler),
   133  				)
   134  			} else {
   135  				ctx, span = trace.StartSpan(
   136  					ctx,
   137  					name,
   138  					trace.WithSpanKind(trace.SpanKindServer),
   139  					trace.WithSampler(cfg.Sampler),
   140  				)
   141  				if ok {
   142  					span.AddLink(trace.Link{
   143  						TraceID:    spanContext.TraceID,
   144  						SpanID:     spanContext.SpanID,
   145  						Type:       trace.LinkTypeChild,
   146  						Attributes: nil,
   147  					})
   148  				}
   149  			}
   150  
   151  			span.AddAttributes(
   152  				trace.StringAttribute(ochttp.MethodAttribute, httpReq.Method),
   153  				trace.StringAttribute(ochttp.PathAttribute, httpReq.URL.Path),
   154  			)
   155  
   156  			return ctx
   157  		},
   158  	)
   159  
   160  	serverFinalizer := jsonrpc.ServerFinalizer(
   161  		func(ctx context.Context, code int, r *http.Request) {
   162  			if span := trace.FromContext(ctx); span != nil {
   163  				span.SetStatus(ochttp.TraceStatus(code, http.StatusText(code)))
   164  
   165  				if rs, ok := ctx.Value(kithttp.ContextKeyResponseSize).(int64); ok {
   166  					span.AddAttributes(
   167  						trace.Int64Attribute("http.response_size", rs),
   168  					)
   169  				}
   170  
   171  				span.End()
   172  			}
   173  		},
   174  	)
   175  
   176  	return func(s *jsonrpc.Server) {
   177  		serverBeforeCodec(s)
   178  		serverFinalizer(s)
   179  	}
   180  }
   181
View as plain text