[mock tracer] Keep track of unfinished spans (#239)

* Add started spans to mock tracer

* Add an assertion that all started spans are finished

* Remove spans from start when they are finished

* Rename started span to unfinished and remove assertion function

* Check both trace ID and span ID when trying to remove span

mocktracer/mockspan.go

@@ -178,7 +178,7 @@ func (s *MockSpan) Finish() {
 	s.Lock()
 	s.FinishTime = time.Now()
 	s.Unlock()
-	s.tracer.recordSpan(s)
+	s.tracer.recordFinishedSpan(s)
 }
 
 // FinishWithOptions belongs to the Span interface
@@ -205,7 +205,7 @@ func (s *MockSpan) FinishWithOptions(opts opentracing.FinishOptions) {
 		}
 	}
 
-	s.tracer.recordSpan(s)
+	s.tracer.recordFinishedSpan(s)
 }
 
 // String allows printing span for debugging

mocktracer/mocktracer.go

@@ -11,6 +11,7 @@ import (
 func New() *MockTracer {
 	t := &MockTracer{
 		finishedSpans: []*MockSpan{},
+		startedSpans:  []*MockSpan{},
 		injectors:     make(map[interface{}]Injector),
 		extractors:    make(map[interface{}]Extractor),
 	}
@@ -34,10 +35,22 @@ func New() *MockTracer {
 type MockTracer struct {
 	sync.RWMutex
 	finishedSpans []*MockSpan
+	startedSpans  []*MockSpan
 	injectors     map[interface{}]Injector
 	extractors    map[interface{}]Extractor
 }
 
+// UnfinishedSpans returns all spans that have been started and not finished since the
+// MockTracer was constructed or since the last call to its Reset() method.
+func (t *MockTracer) UnfinishedSpans() []*MockSpan {
+	t.RLock()
+	defer t.RUnlock()
+
+	spans := make([]*MockSpan, len(t.startedSpans))
+	copy(spans, t.startedSpans)
+	return spans
+}
+
 // FinishedSpans returns all spans that have been Finish()'ed since the
 // MockTracer was constructed or since the last call to its Reset() method.
 func (t *MockTracer) FinishedSpans() []*MockSpan {
@@ -54,6 +67,7 @@ func (t *MockTracer) FinishedSpans() []*MockSpan {
 func (t *MockTracer) Reset() {
 	t.Lock()
 	defer t.Unlock()
+	t.startedSpans = []*MockSpan{}
 	t.finishedSpans = []*MockSpan{}
 }
 
@@ -63,7 +77,10 @@ func (t *MockTracer) StartSpan(operationName string, opts ...opentracing.StartSp
 	for _, o := range opts {
 		o.Apply(&sso)
 	}
-	return newMockSpan(t, operationName, sso)
+
+	span := newMockSpan(t, operationName, sso)
+	t.recordStartedSpan(span)
+	return span
 }
 
 // RegisterInjector registers injector for given format
@@ -98,8 +115,22 @@ func (t *MockTracer) Extract(format interface{}, carrier interface{}) (opentraci
 	return extractor.Extract(carrier)
 }
 
-func (t *MockTracer) recordSpan(span *MockSpan) {
+func (t *MockTracer) recordStartedSpan(span *MockSpan) {
+	t.Lock()
+	defer t.Unlock()
+	t.startedSpans = append(t.startedSpans, span)
+}
+
+func (t *MockTracer) recordFinishedSpan(span *MockSpan) {
 	t.Lock()
 	defer t.Unlock()
 	t.finishedSpans = append(t.finishedSpans, span)
+
+	for i := range t.startedSpans {
+		if t.startedSpans[i].SpanContext.SpanID == span.SpanContext.SpanID &&
+			t.startedSpans[i].SpanContext.TraceID == span.SpanContext.TraceID {
+			t.startedSpans = append(t.startedSpans[:i], t.startedSpans[i+1:]...)
+			return
+		}
+	}
 }

mocktracer/mocktracer_test.go

@@ -25,11 +25,13 @@ func TestMockTracer_StartSpan(t *testing.T) {
 		"", opentracing.ChildOf(span1.Context()))
 	span2.Finish()
 	span1.Finish()
-	spans := tracer.FinishedSpans()
-	assert.Equal(t, 2, len(spans))
+	finishedSpans := tracer.FinishedSpans()
+	unfinishedSpans := tracer.UnfinishedSpans()
+	assert.Equal(t, 2, len(finishedSpans))
+	assert.Equal(t, 0, len(unfinishedSpans))
 
-	parent := spans[1]
-	child := spans[0]
+	parent := finishedSpans[1]
+	child := finishedSpans[0]
 	assert.Equal(t, map[string]interface{}{"x": "y"}, parent.Tags())
 	assert.Equal(t, child.ParentID, parent.Context().(MockSpanContext).SpanID)
 }
@@ -83,14 +85,18 @@ func TestMockTracer_FinishedSpans_and_Reset(t *testing.T) {
 	tracer := New()
 	span := tracer.StartSpan("x")
 	span.SetTag("x", "y")
+	assert.Len(t, tracer.UnfinishedSpans(), 1)
 	span.Finish()
 	spans := tracer.FinishedSpans()
-	assert.Equal(t, 1, len(spans))
+	assert.Len(t, spans, 1)
+	assert.Len(t, tracer.UnfinishedSpans(), 0)
 	assert.Equal(t, map[string]interface{}{"x": "y"}, spans[0].Tags())
 
+	tracer.StartSpan("z")
 	tracer.Reset()
 	spans = tracer.FinishedSpans()
-	assert.Equal(t, 0, len(spans))
+	assert.Len(t, spans, 0)
+	assert.Len(t, tracer.UnfinishedSpans(), 0)
 }
 
 func zeroOutTimestamps(recs []MockLogRecord) {