diff --git a/tracing/opencensus/endpoint.go b/tracing/opencensus/endpoint.go index 70c4e87e6..f211a1767 100644 --- a/tracing/opencensus/endpoint.go +++ b/tracing/opencensus/endpoint.go @@ -31,12 +31,8 @@ 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) { - var attributes []trace.Attribute - - if cfg.GetSpanDetails != nil { - var newName string - - newName, attributes = cfg.GetSpanDetails(ctx, name) + if cfg.GetName != nil { + newName := cfg.GetName(ctx, name) if newName != "" { name = newName @@ -47,11 +43,15 @@ func TraceEndpoint(name string, options ...EndpointOption) endpoint.Middleware { if len(cfg.Attributes) > 0 { span.AddAttributes(cfg.Attributes...) } - if len(attributes) > 0 { - span.AddAttributes(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 { diff --git a/tracing/opencensus/endpoint_options.go b/tracing/opencensus/endpoint_options.go index a91dfb561..88fc52b0f 100644 --- a/tracing/opencensus/endpoint_options.go +++ b/tracing/opencensus/endpoint_options.go @@ -16,12 +16,15 @@ type EndpointOptions struct { // creation by our Endpoint middleware. Attributes []trace.Attribute - // GetSpanDetails holds the function to use for generating the span name + // GetName holds the function used for generating the span name // based on the current name and from the information found in the incoming Request. - // It can also return additional attributes for the span. // - // A returned empty name defaults to the name that the middleware was initialized with. - GetSpanDetails func(ctx context.Context, name string) (string, []trace.Attribute) + // If the returned name is empty, the existing name for the endpoint is kept. + GetName func(ctx context.Context, name string) string + + // GetAttributes holds the function used for extracting additional attributes + // from the information found in the incoming Request. + GetAttributes func(ctx context.Context) []trace.Attribute } // EndpointOption allows for functional options to our OpenCensus endpoint @@ -52,9 +55,16 @@ func WithIgnoreBusinessError(val bool) EndpointOption { } } -// WithSpanDetails extracts details from the request context (like span name and additional attributes). -func WithSpanDetails(fn func(ctx context.Context, name string) (string, []trace.Attribute)) EndpointOption { +// 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.GetSpanDetails = fn + o.GetAttributes = fn } } diff --git a/tracing/opencensus/endpoint_test.go b/tracing/opencensus/endpoint_test.go index 813d6eee6..3a63a54be 100644 --- a/tracing/opencensus/endpoint_test.go +++ b/tracing/opencensus/endpoint_test.go @@ -88,8 +88,11 @@ func TestTraceEndpoint(t *testing.T) { } mw = opencensus.TraceEndpoint( "", - opencensus.WithSpanDetails(func(ctx context.Context, name string) (string, []trace.Attribute) { - return span6, span6Attrs + 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)