Girish Kumar | 2ed051b | 2020-07-28 16:35:25 +0000 | [diff] [blame] | 1 | // Copyright 2017 Michal Witkowski. All Rights Reserved. |
| 2 | // See LICENSE for licensing terms. |
| 3 | |
| 4 | package grpc_opentracing |
| 5 | |
| 6 | import ( |
| 7 | "context" |
| 8 | |
| 9 | "github.com/opentracing/opentracing-go" |
| 10 | ) |
| 11 | |
| 12 | var ( |
| 13 | defaultOptions = &options{ |
| 14 | filterOutFunc: nil, |
| 15 | tracer: nil, |
| 16 | } |
| 17 | ) |
| 18 | |
| 19 | // FilterFunc allows users to provide a function that filters out certain methods from being traced. |
| 20 | // |
| 21 | // If it returns false, the given request will not be traced. |
| 22 | type FilterFunc func(ctx context.Context, fullMethodName string) bool |
| 23 | |
David K. Bainbridge | e05cf0c | 2021-08-19 03:16:50 +0000 | [diff] [blame] | 24 | // UnaryRequestHandlerFunc is a custom request handler |
| 25 | type UnaryRequestHandlerFunc func(span opentracing.Span, req interface{}) |
| 26 | |
| 27 | // OpNameFunc is a func that allows custom operation names instead of the gRPC method. |
| 28 | type OpNameFunc func(method string) string |
| 29 | |
Girish Kumar | 2ed051b | 2020-07-28 16:35:25 +0000 | [diff] [blame] | 30 | type options struct { |
David K. Bainbridge | e05cf0c | 2021-08-19 03:16:50 +0000 | [diff] [blame] | 31 | filterOutFunc FilterFunc |
| 32 | tracer opentracing.Tracer |
| 33 | traceHeaderName string |
| 34 | unaryRequestHandlerFunc UnaryRequestHandlerFunc |
| 35 | opNameFunc OpNameFunc |
Girish Kumar | 2ed051b | 2020-07-28 16:35:25 +0000 | [diff] [blame] | 36 | } |
| 37 | |
| 38 | func evaluateOptions(opts []Option) *options { |
| 39 | optCopy := &options{} |
| 40 | *optCopy = *defaultOptions |
| 41 | for _, o := range opts { |
| 42 | o(optCopy) |
| 43 | } |
| 44 | if optCopy.tracer == nil { |
| 45 | optCopy.tracer = opentracing.GlobalTracer() |
| 46 | } |
David K. Bainbridge | e05cf0c | 2021-08-19 03:16:50 +0000 | [diff] [blame] | 47 | if optCopy.traceHeaderName == "" { |
| 48 | optCopy.traceHeaderName = "uber-trace-id" |
| 49 | } |
Girish Kumar | 2ed051b | 2020-07-28 16:35:25 +0000 | [diff] [blame] | 50 | return optCopy |
| 51 | } |
| 52 | |
| 53 | type Option func(*options) |
| 54 | |
| 55 | // WithFilterFunc customizes the function used for deciding whether a given call is traced or not. |
| 56 | func WithFilterFunc(f FilterFunc) Option { |
| 57 | return func(o *options) { |
| 58 | o.filterOutFunc = f |
| 59 | } |
| 60 | } |
| 61 | |
David K. Bainbridge | e05cf0c | 2021-08-19 03:16:50 +0000 | [diff] [blame] | 62 | // WithTraceHeaderName customizes the trace header name where trace metadata passed with requests. |
| 63 | // Default one is `uber-trace-id` |
| 64 | func WithTraceHeaderName(name string) Option { |
| 65 | return func(o *options) { |
| 66 | o.traceHeaderName = name |
| 67 | } |
| 68 | } |
| 69 | |
Girish Kumar | 2ed051b | 2020-07-28 16:35:25 +0000 | [diff] [blame] | 70 | // WithTracer sets a custom tracer to be used for this middleware, otherwise the opentracing.GlobalTracer is used. |
| 71 | func WithTracer(tracer opentracing.Tracer) Option { |
| 72 | return func(o *options) { |
| 73 | o.tracer = tracer |
| 74 | } |
| 75 | } |
David K. Bainbridge | e05cf0c | 2021-08-19 03:16:50 +0000 | [diff] [blame] | 76 | |
| 77 | // WithUnaryRequestHandlerFunc sets a custom handler for the request |
| 78 | func WithUnaryRequestHandlerFunc(f UnaryRequestHandlerFunc) Option { |
| 79 | return func(o *options) { |
| 80 | o.unaryRequestHandlerFunc = f |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | // WithOpName customizes the trace Operation name |
| 85 | func WithOpName(f OpNameFunc) Option { |
| 86 | return func(o *options) { |
| 87 | o.opNameFunc = f |
| 88 | } |
| 89 | } |