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