Add endpoint name middleware

endpoint/endpoint.go

@@ -38,3 +38,30 @@ 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)
+		}
+	}
+}
+
+// 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 (
+	// ContextKeyEndpointName is populated in the context by EndpointNameMiddleware.
+	ContextKeyEndpointName contextKey = iota
+)

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, _ = endpoint.EndpointName(ctx)
+
+		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)
+	}
+}

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 := endpoint.EndpointName(ctx)
+				if ok && endpointName != "" {
+					name = endpointName
+				}
+			}
+
 			ctx, span := trace.StartSpan(ctx, name)
 			if len(cfg.Attributes) > 0 {
 				span.AddAttributes(cfg.Attributes...)

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)
+	}
 }

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 := endpoint.EndpointName(ctx)
+				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(

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)
+	}
+}

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 := endpoint.EndpointName(ctx)
+				if ok && endpointName != "" {
+					name = endpointName
+				}
+			}
+
 			var sc model.SpanContext
 			if parentSpan := zipkin.SpanFromContext(ctx); parentSpan != nil {
 				sc = parentSpan.Context()

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)
+	}
+}