From 6f5b1fcddd902796d6fd8bfc6e234d6183df3159 Mon Sep 17 00:00:00 2001 From: Mark Sagi-Kazar Date: Fri, 10 Jan 2020 01:44:30 +0100 Subject: [PATCH 1/5] Add endpoint name middleware --- endpoint/endpoint.go | 19 +++++++++++++++++++ endpoint/endpoint_test.go | 28 ++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 endpoint/endpoint_test.go diff --git a/endpoint/endpoint.go b/endpoint/endpoint.go index 6e9da3679..da0861abf 100644 --- a/endpoint/endpoint.go +++ b/endpoint/endpoint.go @@ -38,3 +38,22 @@ func Chain(outer Middleware, others ...Middleware) Middleware { type Failer interface { Failed() error } + +// EndpointNameMiddleware populates the context with a common name for the endpoint. +// It can be used in subsequent endpoints in the chain to identify the operation. +func EndpointNameMiddleware(name string) Middleware { + return func(next Endpoint) Endpoint { + return func(ctx context.Context, req interface{}) (interface{}, error) { + ctx = context.WithValue(ctx, ContextKeyEndpointName, name) + + return next(ctx, req) + } + } +} + +type contextKey int + +const ( + // ContextKeyEndpointName is populated in the context by EndpointNameMiddleware. + ContextKeyEndpointName contextKey = iota +) diff --git a/endpoint/endpoint_test.go b/endpoint/endpoint_test.go new file mode 100644 index 000000000..33439d24f --- /dev/null +++ b/endpoint/endpoint_test.go @@ -0,0 +1,28 @@ +package endpoint_test + +import ( + "context" + "testing" + + "github.com/go-kit/kit/endpoint" +) + +func TestEndpointNameMiddleware(t *testing.T) { + ctx := context.Background() + + var name string + + ep := func(ctx context.Context, request interface{}) (interface{}, error) { + name = ctx.Value(endpoint.ContextKeyEndpointName).(string) + + return nil, nil + } + + mw := endpoint.EndpointNameMiddleware("go-kit/endpoint") + + mw(ep)(ctx, nil) + + if want, have := "go-kit/endpoint", name; want != have { + t.Fatalf("unexpected endpoint name, wanted %q, got %q", want, have) + } +} From 97d756b77855750b8128387e67af1b1d3dd49610 Mon Sep 17 00:00:00 2001 From: Mark Sagi-Kazar Date: Fri, 10 Jan 2020 01:52:30 +0100 Subject: [PATCH 2/5] Add endpoint name support to opencensus tracing --- tracing/opencensus/endpoint.go | 8 ++++++++ tracing/opencensus/endpoint_test.go | 14 ++++++++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/tracing/opencensus/endpoint.go b/tracing/opencensus/endpoint.go index 4f678fc7d..f1f4bc740 100644 --- a/tracing/opencensus/endpoint.go +++ b/tracing/opencensus/endpoint.go @@ -31,6 +31,14 @@ func TraceEndpoint(name string, options ...EndpointOption) endpoint.Middleware { return func(next endpoint.Endpoint) endpoint.Endpoint { return func(ctx context.Context, request interface{}) (response interface{}, err error) { + // Use endpoint name from the context if there is no operation name specified + if name == TraceEndpointDefaultName { + endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + if ok && endpointName != "" { + name = endpointName + } + } + ctx, span := trace.StartSpan(ctx, name) if len(cfg.Attributes) > 0 { span.AddAttributes(cfg.Attributes...) diff --git a/tracing/opencensus/endpoint_test.go b/tracing/opencensus/endpoint_test.go index 924560a6a..67d6f6f67 100644 --- a/tracing/opencensus/endpoint_test.go +++ b/tracing/opencensus/endpoint_test.go @@ -20,6 +20,7 @@ const ( span3 = "SPAN-3" span4 = "SPAN-4" span5 = "SPAN-5" + span6 = "SPAN-6" ) var ( @@ -76,13 +77,17 @@ func TestTraceEndpoint(t *testing.T) { mw = opencensus.TraceEndpoint(span4) mw(passEndpoint)(ctx, failedResponse{err: err3}) - // span4 + // span5 mw = opencensus.TraceEndpoint(span5, opencensus.WithIgnoreBusinessError(true)) mw(passEndpoint)(ctx, failedResponse{err: err4}) + // span6 + mw = opencensus.TraceEndpoint("") + mw(endpoint.Nop)(context.WithValue(ctx, endpoint.ContextKeyEndpointName, span6), nil) + // check span count spans := e.Flush() - if want, have := 5, len(spans); want != have { + if want, have := 6, len(spans); want != have { t.Fatalf("incorrected number of spans, wanted %d, got %d", want, have) } @@ -156,4 +161,9 @@ func TestTraceEndpoint(t *testing.T) { t.Fatalf("incorrect attribute count, wanted %d, got %d", want, have) } + // test span 6 + span = spans[5] + if want, have := span6, span.Name; want != have { + t.Errorf("incorrect span name, wanted %q, got %q", want, have) + } } From 9b7382764e2fde5dfbfea1d4eb80f6cd03474da3 Mon Sep 17 00:00:00 2001 From: Mark Sagi-Kazar Date: Fri, 10 Jan 2020 01:58:44 +0100 Subject: [PATCH 3/5] Add endpoint name support to opentracing --- tracing/opentracing/endpoint.go | 16 +++++++++ tracing/opentracing/endpoint_test.go | 49 ++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/tracing/opentracing/endpoint.go b/tracing/opentracing/endpoint.go index 0482e9c0d..07889a939 100644 --- a/tracing/opentracing/endpoint.go +++ b/tracing/opentracing/endpoint.go @@ -17,6 +17,14 @@ import ( func TraceServer(tracer opentracing.Tracer, operationName string) endpoint.Middleware { return func(next endpoint.Endpoint) endpoint.Endpoint { return func(ctx context.Context, request interface{}) (interface{}, error) { + // Use endpoint name from the context if there is no operation name specified + if operationName == "" { + endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + if ok && endpointName != "" { + operationName = endpointName + } + } + serverSpan := opentracing.SpanFromContext(ctx) if serverSpan == nil { // All we can do is create a new root span. @@ -37,6 +45,14 @@ func TraceServer(tracer opentracing.Tracer, operationName string) endpoint.Middl func TraceClient(tracer opentracing.Tracer, operationName string) endpoint.Middleware { return func(next endpoint.Endpoint) endpoint.Endpoint { return func(ctx context.Context, request interface{}) (interface{}, error) { + // Use endpoint name from the context if there is no operation name specified + if operationName == "" { + endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + if ok && endpointName != "" { + operationName = endpointName + } + } + var clientSpan opentracing.Span if parentSpan := opentracing.SpanFromContext(ctx); parentSpan != nil { clientSpan = tracer.StartSpan( diff --git a/tracing/opentracing/endpoint_test.go b/tracing/opentracing/endpoint_test.go index a26f7bab6..d312bed9e 100644 --- a/tracing/opentracing/endpoint_test.go +++ b/tracing/opentracing/endpoint_test.go @@ -62,6 +62,30 @@ func TestTraceServerNoContextSpan(t *testing.T) { } } +func TestTraceServerEndpointName(t *testing.T) { + tracer := mocktracer.New() + + // Initialize the ctx with a nameless Span. + contextSpan := tracer.StartSpan("").(*mocktracer.MockSpan) + ctx := opentracing.ContextWithSpan(context.Background(), contextSpan) + + tracedEndpoint := kitot.TraceServer(tracer, "")(endpoint.Nop) + if _, err := tracedEndpoint(context.WithValue(ctx, endpoint.ContextKeyEndpointName, "testOp"), struct{}{}); err != nil { + t.Fatal(err) + } + + finishedSpans := tracer.FinishedSpans() + if want, have := 1, len(finishedSpans); want != have { + t.Fatalf("Want %v span(s), found %v", want, have) + } + + // Test that the op name is updated + endpointSpan := finishedSpans[0] + if want, have := "testOp", endpointSpan.OperationName; want != have { + t.Fatalf("Want %q, have %q", want, have) + } +} + func TestTraceClient(t *testing.T) { tracer := mocktracer.New() @@ -115,3 +139,28 @@ func TestTraceClientNoContextSpan(t *testing.T) { t.Fatalf("Want %q, have %q", want, have) } } + +func TestTraceClientEndpointName(t *testing.T) { + tracer := mocktracer.New() + + // Initialize the ctx with a parent Span. + parentSpan := tracer.StartSpan("parent").(*mocktracer.MockSpan) + defer parentSpan.Finish() + ctx := opentracing.ContextWithSpan(context.Background(), parentSpan) + + tracedEndpoint := kitot.TraceClient(tracer, "")(endpoint.Nop) + if _, err := tracedEndpoint(context.WithValue(ctx, endpoint.ContextKeyEndpointName, "testOp"), struct{}{}); err != nil { + t.Fatal(err) + } + + // tracedEndpoint created a new Span. + finishedSpans := tracer.FinishedSpans() + if want, have := 1, len(finishedSpans); want != have { + t.Fatalf("Want %v span(s), found %v", want, have) + } + + endpointSpan := finishedSpans[0] + if want, have := "testOp", endpointSpan.OperationName; want != have { + t.Fatalf("Want %q, have %q", want, have) + } +} From c92784b1e185e95f32ac03da5dc97b1f24ea42ab Mon Sep 17 00:00:00 2001 From: Mark Sagi-Kazar Date: Fri, 10 Jan 2020 02:01:32 +0100 Subject: [PATCH 4/5] Add endpoint name support to zipkin tracing --- tracing/zipkin/endpoint.go | 8 ++++++++ tracing/zipkin/endpoint_test.go | 17 +++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/tracing/zipkin/endpoint.go b/tracing/zipkin/endpoint.go index e004bf24f..6b48ee977 100644 --- a/tracing/zipkin/endpoint.go +++ b/tracing/zipkin/endpoint.go @@ -16,6 +16,14 @@ import ( func TraceEndpoint(tracer *zipkin.Tracer, name string) endpoint.Middleware { return func(next endpoint.Endpoint) endpoint.Endpoint { return func(ctx context.Context, request interface{}) (interface{}, error) { + // Use endpoint name from the context if there is no operation name specified + if name == "" { + endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + if ok && endpointName != "" { + name = endpointName + } + } + var sc model.SpanContext if parentSpan := zipkin.SpanFromContext(ctx); parentSpan != nil { sc = parentSpan.Context() diff --git a/tracing/zipkin/endpoint_test.go b/tracing/zipkin/endpoint_test.go index 7ed8aac8a..a7dd5e5e1 100644 --- a/tracing/zipkin/endpoint_test.go +++ b/tracing/zipkin/endpoint_test.go @@ -29,3 +29,20 @@ func TestTraceEndpoint(t *testing.T) { t.Fatalf("incorrect span name, wanted %s, got %s", want, have) } } + +func TestTraceEndpointName(t *testing.T) { + rec := recorder.NewReporter() + tr, _ := zipkin.NewTracer(rec) + mw := zipkinkit.TraceEndpoint(tr, "") + mw(endpoint.Nop)(context.WithValue(context.Background(), endpoint.ContextKeyEndpointName, spanName), nil) + + spans := rec.Flush() + + if want, have := 1, len(spans); want != have { + t.Fatalf("incorrect number of spans, wanted %d, got %d", want, have) + } + + if want, have := spanName, spans[0].Name; want != have { + t.Fatalf("incorrect span name, wanted %s, got %s", want, have) + } +} \ No newline at end of file From 155c0f15ec1d35539fbd90118f4a79eae964339e Mon Sep 17 00:00:00 2001 From: Mark Sagi-Kazar Date: Fri, 10 Jan 2020 12:19:16 +0100 Subject: [PATCH 5/5] Add getter for endpoint name context value --- endpoint/endpoint.go | 8 ++++++++ endpoint/endpoint_test.go | 2 +- tracing/opencensus/endpoint.go | 2 +- tracing/opentracing/endpoint.go | 2 +- tracing/zipkin/endpoint.go | 2 +- 5 files changed, 12 insertions(+), 4 deletions(-) diff --git a/endpoint/endpoint.go b/endpoint/endpoint.go index da0861abf..2018e42c1 100644 --- a/endpoint/endpoint.go +++ b/endpoint/endpoint.go @@ -51,6 +51,14 @@ func EndpointNameMiddleware(name string) Middleware { } } +// EndpointName fetches the endpoint name from the context (if any). +// If an endpoint name is not found or it isn't string, the second return argument is false. +func EndpointName(ctx context.Context) (string, bool) { + name, ok := ctx.Value(ContextKeyEndpointName).(string) + + return name, ok +} + type contextKey int const ( diff --git a/endpoint/endpoint_test.go b/endpoint/endpoint_test.go index 33439d24f..989fac4eb 100644 --- a/endpoint/endpoint_test.go +++ b/endpoint/endpoint_test.go @@ -13,7 +13,7 @@ func TestEndpointNameMiddleware(t *testing.T) { var name string ep := func(ctx context.Context, request interface{}) (interface{}, error) { - name = ctx.Value(endpoint.ContextKeyEndpointName).(string) + name, _ = endpoint.EndpointName(ctx) return nil, nil } diff --git a/tracing/opencensus/endpoint.go b/tracing/opencensus/endpoint.go index f1f4bc740..ff51031ab 100644 --- a/tracing/opencensus/endpoint.go +++ b/tracing/opencensus/endpoint.go @@ -33,7 +33,7 @@ func TraceEndpoint(name string, options ...EndpointOption) endpoint.Middleware { return func(ctx context.Context, request interface{}) (response interface{}, err error) { // Use endpoint name from the context if there is no operation name specified if name == TraceEndpointDefaultName { - endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + endpointName, ok := endpoint.EndpointName(ctx) if ok && endpointName != "" { name = endpointName } diff --git a/tracing/opentracing/endpoint.go b/tracing/opentracing/endpoint.go index 07889a939..4ac21ce04 100644 --- a/tracing/opentracing/endpoint.go +++ b/tracing/opentracing/endpoint.go @@ -19,7 +19,7 @@ func TraceServer(tracer opentracing.Tracer, operationName string) endpoint.Middl return func(ctx context.Context, request interface{}) (interface{}, error) { // Use endpoint name from the context if there is no operation name specified if operationName == "" { - endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + endpointName, ok := endpoint.EndpointName(ctx) if ok && endpointName != "" { operationName = endpointName } diff --git a/tracing/zipkin/endpoint.go b/tracing/zipkin/endpoint.go index 6b48ee977..92f54c97e 100644 --- a/tracing/zipkin/endpoint.go +++ b/tracing/zipkin/endpoint.go @@ -18,7 +18,7 @@ func TraceEndpoint(tracer *zipkin.Tracer, name string) endpoint.Middleware { return func(ctx context.Context, request interface{}) (interface{}, error) { // Use endpoint name from the context if there is no operation name specified if name == "" { - endpointName, ok := ctx.Value(endpoint.ContextKeyEndpointName).(string) + endpointName, ok := endpoint.EndpointName(ctx) if ok && endpointName != "" { name = endpointName }