tracing/zipkin: use idiomatic next endpoint.Endpoint
Peter Bourgon
8 years ago
| 53 | 53 | // and submits the span to the collector. If no span is found in the context, |
| 54 | 54 | // a new span is generated and inserted. |
| 55 | 55 | func AnnotateServer(newSpan NewSpanFunc, c Collector) endpoint.Middleware { |
| 56 | return func(e endpoint.Endpoint) endpoint.Endpoint { | |
| 56 | return func(next endpoint.Endpoint) endpoint.Endpoint { | |
| 57 | 57 | return func(ctx context.Context, request interface{}) (interface{}, error) { |
| 58 | 58 | span, ok := fromContext(ctx) |
| 59 | 59 | if !ok { |
| 62 | 62 | } |
| 63 | 63 | span.Annotate(ServerReceive) |
| 64 | 64 | defer func() { span.Annotate(ServerSend); c.Collect(span) }() |
| 65 | return e(ctx, request) | |
| 65 | return next(ctx, request) | |
| 66 | 66 | } |
| 67 | 67 | } |
| 68 | 68 | } |
| 73 | 73 | // collector. If no span is found in the context, a new span is generated and |
| 74 | 74 | // inserted. |
| 75 | 75 | func AnnotateClient(newSpan NewSpanFunc, c Collector) endpoint.Middleware { |
| 76 | return func(e endpoint.Endpoint) endpoint.Endpoint { | |
| 76 | return func(next endpoint.Endpoint) endpoint.Endpoint { | |
| 77 | 77 | return func(ctx context.Context, request interface{}) (interface{}, error) { |
| 78 | 78 | var clientSpan *Span |
| 79 | 79 | parentSpan, ok := fromContext(ctx) |
| 86 | 86 | defer func() { ctx = context.WithValue(ctx, SpanContextKey, parentSpan) }() // reset |
| 87 | 87 | clientSpan.Annotate(ClientSend) |
| 88 | 88 | defer func() { clientSpan.Annotate(ClientReceive); c.Collect(clientSpan) }() |
| 89 | return e(ctx, request) | |
| 89 | return next(ctx, request) | |
| 90 | 90 | } |
| 91 | 91 | } |
| 92 | 92 | } |