From cc938d52e0cdf4c811ab203f428fcd06f9d9a148 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A1rk=20S=C3=A1gi-Kaz=C3=A1r?=
 <sagikazarmark@users.noreply.github.com>
Date: Mon, 17 Feb 2020 20:15:32 +0100
Subject: [PATCH] Add a span name getter to opencensus endpoint options (#948)

* Add a span name getter to opencensus endpoint options

* Add GetSpanDetails option

* Split GetSpanDetails function

* Simplify getting the name

* Improve documentation

* Simplify attributes

* Update documentation

* Update tracing/opencensus/endpoint_options.go

Co-authored-by: Peter Bourgon <peterbourgon@users.noreply.github.com>
---
 tracing/opencensus/endpoint.go         | 12 +++++++++++
 tracing/opencensus/endpoint_options.go | 30 +++++++++++++++++++++++++-
 tracing/opencensus/endpoint_test.go    | 30 ++++++++++++++++++++++++--
 3 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/tracing/opencensus/endpoint.go b/tracing/opencensus/endpoint.go
index 4f678fc7d..16283a543 100644
--- a/tracing/opencensus/endpoint.go
+++ b/tracing/opencensus/endpoint.go
@@ -31,12 +31,24 @@ 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) {
+			if cfg.GetName != nil {
+				if newName := cfg.GetName(ctx, name); newName != "" {
+					name = newName
+				}
+			}
+
 			ctx, span := trace.StartSpan(ctx, name)
 			if len(cfg.Attributes) > 0 {
 				span.AddAttributes(cfg.Attributes...)
 			}
 			defer span.End()
 
+			if cfg.GetAttributes != nil {
+				if attrs := cfg.GetAttributes(ctx); len(attrs) > 0 {
+					span.AddAttributes(attrs...)
+				}
+			}
+
 			defer func() {
 				if err != nil {
 					if lberr, ok := err.(lb.RetryError); ok {
diff --git a/tracing/opencensus/endpoint_options.go b/tracing/opencensus/endpoint_options.go
index 3af145681..8eedd59b2 100644
--- a/tracing/opencensus/endpoint_options.go
+++ b/tracing/opencensus/endpoint_options.go
@@ -1,6 +1,10 @@
 package opencensus
 
-import "go.opencensus.io/trace"
+import (
+	"context"
+
+	"go.opencensus.io/trace"
+)
 
 // EndpointOptions holds the options for tracing an endpoint
 type EndpointOptions struct {
@@ -11,6 +15,16 @@ type EndpointOptions struct {
 	// Attributes holds the default attributes which will be set on span
 	// creation by our Endpoint middleware.
 	Attributes []trace.Attribute
+
+	// GetName is an optional function that can set the span name based on the existing name
+	// for the endpoint and information in the context.
+	//
+	// If the function is nil, or the returned name is empty, the existing name for the endpoint is used.
+	GetName func(ctx context.Context, name string) string
+
+	// GetAttributes is an optional function that can extract trace attributes 
+	// from the context and add them to the span.
+	GetAttributes func(ctx context.Context) []trace.Attribute
 }
 
 // EndpointOption allows for functional options to our OpenCensus endpoint
@@ -40,3 +54,17 @@ func WithIgnoreBusinessError(val bool) EndpointOption {
 		o.IgnoreBusinessError = val
 	}
 }
+
+// WithSpanName extracts additional attributes from the request context.
+func WithSpanName(fn func(ctx context.Context, name string) string) EndpointOption {
+	return func(o *EndpointOptions) {
+		o.GetName = fn
+	}
+}
+
+// WithSpanAttributes extracts additional attributes from the request context.
+func WithSpanAttributes(fn func(ctx context.Context) []trace.Attribute) EndpointOption {
+	return func(o *EndpointOptions) {
+		o.GetAttributes = fn
+	}
+}
diff --git a/tracing/opencensus/endpoint_test.go b/tracing/opencensus/endpoint_test.go
index 924560a6a..3a63a54be 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,29 @@ 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
+	span6Attrs := []trace.Attribute{
+		trace.StringAttribute("string", "value"),
+		trace.Int64Attribute("int64", 42),
+	}
+	mw = opencensus.TraceEndpoint(
+		"",
+		opencensus.WithSpanName(func(ctx context.Context, name string) string {
+			return span6
+		}),
+		opencensus.WithSpanAttributes(func(ctx context.Context) []trace.Attribute {
+			return span6Attrs
+		}),
+	)
+	mw(endpoint.Nop)(ctx, 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 +173,13 @@ 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)
+	}
+
+	if want, have := 2, len(span.Attributes); want != have {
+		t.Fatalf("incorrect attribute count, wanted %d, got %d", want, have)
+	}
 }