Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions tracing/opencensus/endpoint.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
30 changes: 29 additions & 1 deletion tracing/opencensus/endpoint_options.go
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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
}
}
30 changes: 28 additions & 2 deletions tracing/opencensus/endpoint_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ const (
span3 = "SPAN-3"
span4 = "SPAN-4"
span5 = "SPAN-5"
span6 = "SPAN-6"
)

var (
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -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)
}
}