feat(#90): adds support for finished span handler.

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() {

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

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

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