Codebase list golang-github-go-kit-kit / b69f23a
Merge pull request #225 from basvanbeek/master Export zipkin.FromContext for annotating spans inside service implementations Peter Bourgon 8 years ago
2 changed file(s) with 75 addition(s) and 33 deletion(s). Raw diff Collapse all Expand all
6060 func AnnotateServer(newSpan NewSpanFunc, c Collector) endpoint.Middleware {
6161 return func(next endpoint.Endpoint) endpoint.Endpoint {
6262 return func(ctx context.Context, request interface{}) (interface{}, error) {
63 span, ok := fromContext(ctx)
63 span, ok := FromContext(ctx)
6464 if !ok {
6565 span = newSpan(newID(), newID(), 0)
6666 ctx = context.WithValue(ctx, SpanContextKey, span)
8181 return func(next endpoint.Endpoint) endpoint.Endpoint {
8282 return func(ctx context.Context, request interface{}) (interface{}, error) {
8383 var clientSpan *Span
84 parentSpan, ok := fromContext(ctx)
84 parentSpan, ok := FromContext(ctx)
8585 if ok {
8686 clientSpan = newSpan(parentSpan.TraceID(), newID(), parentSpan.SpanID())
8787 } else {
123123 // a child/client span.
124124 func ToRequest(newSpan NewSpanFunc) func(ctx context.Context, r *http.Request) context.Context {
125125 return func(ctx context.Context, r *http.Request) context.Context {
126 span, ok := fromContext(ctx)
126 span, ok := FromContext(ctx)
127127 if !ok {
128128 span = newSpan(newID(), newID(), 0)
129129 }
147147 // a child/client span.
148148 func ToGRPCRequest(newSpan NewSpanFunc) func(ctx context.Context, md *metadata.MD) context.Context {
149149 return func(ctx context.Context, md *metadata.MD) context.Context {
150 span, ok := fromContext(ctx)
150 span, ok := FromContext(ctx)
151151 if !ok {
152152 span = newSpan(newID(), newID(), 0)
153153 }
242242 return newSpan(traceID, spanID, parentSpanID)
243243 }
244244
245 func fromContext(ctx context.Context) (*Span, bool) {
245 // FromContext extracts an existing Zipkin span if it is stored in the provided
246 // context. If you add context.Context as the first parameter in your service
247 // methods you can annotate spans from within business logic. Typical use case
248 // is to AnnotateDuration on interaction with resources like databases.
249 func FromContext(ctx context.Context) (*Span, bool) {
246250 val := ctx.Value(SpanContextKey)
247251 if val == nil {
248252 return nil, false
6060 }
6161 }
6262
63 func TestToGRPCContext(t *testing.T) {
64 const (
65 hostport = "5.5.5.5:5555"
66 serviceName = "foo-service"
67 methodName = "foo-method"
68 traceID int64 = 12
69 spanID int64 = 34
70 parentSpanID int64 = 56
71 )
72
73 md := metadata.MD{
74 "x-b3-traceid": []string{strconv.FormatInt(traceID, 16)},
75 "x-b3-spanid": []string{strconv.FormatInt(spanID, 16)},
76 "x-b3-parentspanid": []string{strconv.FormatInt(parentSpanID, 16)},
77 }
78
79 newSpan := zipkin.MakeNewSpanFunc(hostport, serviceName, methodName)
80 toContext := zipkin.ToGRPCContext(newSpan, log.NewLogfmtLogger(ioutil.Discard))
81
82 ctx := toContext(context.Background(), &md)
83 val := ctx.Value(zipkin.SpanContextKey)
84 if val == nil {
85 t.Fatalf("%s returned no value", zipkin.SpanContextKey)
86 }
87 span, ok := val.(*zipkin.Span)
63 func TestFromContext(t *testing.T) {
64 const (
65 hostport = "5.5.5.5:5555"
66 serviceName = "foo-service"
67 methodName = "foo-method"
68 traceID int64 = 14
69 spanID int64 = 36
70 parentSpanID int64 = 58
71 )
72
73 ctx := context.WithValue(
74 context.Background(),
75 zipkin.SpanContextKey,
76 zipkin.NewSpan(hostport, serviceName, methodName, traceID, spanID, parentSpanID),
77 )
78
79 span, ok := zipkin.FromContext(ctx)
8880 if !ok {
89 t.Fatalf("%s was not a Span object", zipkin.SpanContextKey)
90 }
91
81 t.Fatalf("expected a context value in %q", zipkin.SpanContextKey)
82 }
83 if span == nil {
84 t.Fatal("expected a Zipkin span object")
85 }
9286 for want, haveFunc := range map[int64]func() int64{
9387 traceID: span.TraceID,
9488 spanID: span.SpanID,
10498 }
10599 }
106100
101 func TestToGRPCContext(t *testing.T) {
102 const (
103 hostport = "5.5.5.5:5555"
104 serviceName = "foo-service"
105 methodName = "foo-method"
106 traceID int64 = 12
107 spanID int64 = 34
108 parentSpanID int64 = 56
109 )
110
111 md := metadata.MD{
112 "x-b3-traceid": []string{strconv.FormatInt(traceID, 16)},
113 "x-b3-spanid": []string{strconv.FormatInt(spanID, 16)},
114 "x-b3-parentspanid": []string{strconv.FormatInt(parentSpanID, 16)},
115 }
116
117 newSpan := zipkin.MakeNewSpanFunc(hostport, serviceName, methodName)
118 toContext := zipkin.ToGRPCContext(newSpan, log.NewLogfmtLogger(ioutil.Discard))
119
120 ctx := toContext(context.Background(), &md)
121 val := ctx.Value(zipkin.SpanContextKey)
122 if val == nil {
123 t.Fatalf("%s returned no value", zipkin.SpanContextKey)
124 }
125 span, ok := val.(*zipkin.Span)
126 if !ok {
127 t.Fatalf("%s was not a Span object", zipkin.SpanContextKey)
128 }
129
130 for want, haveFunc := range map[int64]func() int64{
131 traceID: span.TraceID,
132 spanID: span.SpanID,
133 parentSpanID: span.ParentSpanID,
134 } {
135 if have := haveFunc(); want != have {
136 name := runtime.FuncForPC(reflect.ValueOf(haveFunc).Pointer()).Name()
137 name = strings.Split(name, "ยท")[0]
138 toks := strings.Split(name, ".")
139 name = toks[len(toks)-1]
140 t.Errorf("%s: want %d, have %d", name, want, have)
141 }
142 }
143 }
144
107145 func TestToRequest(t *testing.T) {
108146 const (
109147 hostport = "5.5.5.5:5555"