Add a span name getter to opencensus endpoint options

tracing/opencensus/endpoint.go

@@ -31,12 +31,25 @@ 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 {
+ attributes := cfg.GetAttributes(ctx)
+ if len(attributes) > 0 {
+ span.AddAttributes(attributes...)
+ }
+ }
+
  defer func() {
  if err != nil {
  if lberr, ok := err.(lb.RetryError); ok {

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 attaches additional attributes to the span
+ // based on the information found in the context.
+ 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
+ }
+}

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