diff --git a/span_implementation.go b/span_implementation.go index 72904a8..022fa75 100644 --- a/span_implementation.go +++ b/span_implementation.go @@ -25,9 +25,10 @@ import ( type spanImpl struct { mtx sync.RWMutex model.SpanModel - tracer *Tracer - mustCollect int32 // used as atomic bool (1 = true, 0 = false) - flushOnFinish bool + tracer *Tracer + mustCollect int32 // used as atomic bool (1 = true, 0 = false) + flushOnFinish bool + finishedSpanHandler func(*model.SpanModel) bool } func (s *spanImpl) Context() model.SpanContext { @@ -77,21 +78,83 @@ func (s *spanImpl) Tag(key, value string) { } func (s *spanImpl) Finish() { + d := time.Since(s.Timestamp) if atomic.CompareAndSwapInt32(&s.mustCollect, 1, 0) { - s.Duration = time.Since(s.Timestamp) - if s.flushOnFinish { + s.mtx.Lock() + s.Duration = d + s.mtx.Unlock() + + shouldRecord := true + if s.finishedSpanHandler != nil { + shouldRecord = s.finishedSpanHandler(&s.SpanModel) + } + + if shouldRecord && s.flushOnFinish { s.tracer.reporter.Send(s.SpanModel) } + return } + + var hasDuration bool + s.mtx.Lock() + hasDuration = s.Duration == 0 + s.mtx.Unlock() + + if hasDuration { + // it was not meant to be recorded because the CompareAndSwap + // did not happen (meaning that s.mustCollect is 0 at this moment) + // and duration is still zero value. + s.mtx.Lock() + s.Duration = d + s.mtx.Unlock() + + shouldRecord := false + if s.finishedSpanHandler != nil { + shouldRecord = s.finishedSpanHandler(&s.SpanModel) + } + + if shouldRecord && s.flushOnFinish { + s.tracer.reporter.Send(s.SpanModel) + } + } // else the span is being finished concurrently by another goroutine } func (s *spanImpl) FinishedWithDuration(d time.Duration) { if atomic.CompareAndSwapInt32(&s.mustCollect, 1, 0) { + s.mtx.Lock() s.Duration = d - if s.flushOnFinish { + s.mtx.Unlock() + + shouldRecord := true + if s.finishedSpanHandler != nil { + shouldRecord = s.finishedSpanHandler(&s.SpanModel) + } + + if shouldRecord && s.flushOnFinish { s.tracer.reporter.Send(s.SpanModel) } + return } + + var hasDuration bool + s.mtx.Lock() + hasDuration = s.Duration == 0 + s.mtx.Unlock() + + if hasDuration { + // it was not meant to be recorded because the CompareAndSwap + // did not happen (meaning that s.mustCollect is 0 at this moment) + // and duration is still zero value. + s.Duration = d + shouldRecord := false + if s.finishedSpanHandler != nil { + shouldRecord = s.finishedSpanHandler(&s.SpanModel) + } + + if shouldRecord && s.flushOnFinish { + s.tracer.reporter.Send(s.SpanModel) + } + } // else the span is being finished concurrently by another goroutine } func (s *spanImpl) Flush() { diff --git a/tracer.go b/tracer.go index 0f294cf..d4a2bac 100644 --- a/tracer.go +++ b/tracer.go @@ -37,6 +37,7 @@ type Tracer struct { noop int32 // used as atomic bool (1 = true, 0 = false) sharedSpans bool unsampledNoop bool + finishedSpanHandler func(*model.SpanModel) bool } // NewTracer returns a new Zipkin Tracer. @@ -93,8 +94,9 @@ func (t *Tracer) StartSpan(name string, options ...SpanOption) Span { Annotations: make([]model.Annotation, 0), Tags: make(map[string]string), }, - flushOnFinish: true, - tracer: t, + flushOnFinish: true, + tracer: t, + finishedSpanHandler: t.finishedSpanHandler, } // add default tracer tags to span diff --git a/tracer_options.go b/tracer_options.go index 533c5e4..8d6a3c6 100644 --- a/tracer_options.go +++ b/tracer_options.go @@ -136,3 +136,13 @@ func WithNoopTracer(tracerNoop bool) TracerOption { return nil } } + +// WithFinishedSpanHandler if set, can mutate all span data and decide if a span +// should be recorded or not. E.g. user could decide to sample all requests having +// an error tag set or duration over certain value. +func WithFinishedSpanHandler(handler func(*model.SpanModel) bool) TracerOption { + return func(o *Tracer) error { + o.finishedSpanHandler = handler + return nil + } +} diff --git a/tracer_test.go b/tracer_test.go index d90d217..64c58ac 100644 --- a/tracer_test.go +++ b/tracer_test.go @@ -21,6 +21,8 @@ import ( "testing" "time" + "github.com/openzipkin/zipkin-go/reporter/recorder" + "github.com/openzipkin/zipkin-go/idgenerator" "github.com/openzipkin/zipkin-go/model" "github.com/openzipkin/zipkin-go/reporter" @@ -808,3 +810,49 @@ func TestLocalEndpoint(t *testing.T) { t.Errorf("IPv6 endpoint want %+v, have %+v", want.IPv6, have.IPv6) } } + +func TestFinishedSpanHandlerAvoidsReporting(t *testing.T) { + rep := recorder.NewReporter() + defer rep.Close() + + tracer, _ := NewTracer( + rep, + WithNoopSpan(false), + WithSampler(AlwaysSample), + WithFinishedSpanHandler(func(s *model.SpanModel) bool { + return false + }), + ) + sp := tracer.StartSpan("test") + sp.Finish() + + if want, have := 0, len(rep.Flush()); want != have { + t.Errorf("unexpected number of spans, want: %d, have: %d", want, have) + } +} + +func TestFinishedSpanAddsTagsToSpan(t *testing.T) { + rep := recorder.NewReporter() + defer rep.Close() + + tracer, _ := NewTracer( + rep, + WithNoopSpan(false), + WithSampler(AlwaysSample), + WithFinishedSpanHandler(func(s *model.SpanModel) bool { + s.Tags["my_key"] = "my_value" + return true + }), + ) + sp := tracer.StartSpan("test") + sp.Finish() + + repSans := rep.Flush() + if want, have := 1, len(repSans); want != have { + t.Errorf("unexpected number of spans, want: %d, have: %d", want, have) + } + + if want, have := "my_value", repSans[0].Tags["my_key"]; want != have { + t.Errorf("unexpected number of spans, want: %q, have: %q", want, have) + } +}