tracing: add registry of unfinished local root Spans

pkg/util/tracing/alloc_test.go

@@ -60,7 +60,7 @@ func BenchmarkTracer_StartSpanCtx(b *testing.B) {
 			for i := 0; i < b.N; i++ {
 				newCtx, sp := tr.StartSpanCtx(ctx, "benching", tc.opts...)
 				_ = newCtx
-				_ = sp
+				sp.Finish() // clean up
 			}
 		})
 	}

pkg/util/tracing/span.go

@@ -436,6 +436,9 @@ func (s *Span) Finish() {
 	if s.netTr != nil {
 		s.netTr.Finish()
 	}
+	s.tracer.activeSpans.Lock()
+	delete(s.tracer.activeSpans.m, s)
+	s.tracer.activeSpans.Unlock()
 }
 
 // Meta returns the information which needs to be propagated across

pkg/util/tracing/tracer.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/logtags"
 	opentracing "github.com/opentracing/opentracing-go"
 	"golang.org/x/net/trace"
@@ -140,13 +141,33 @@ type Tracer struct {
 
 	// Pointer to shadowTracer, if using one.
 	shadowTracer unsafe.Pointer
+
+	// activeSpans is a map that references all non-Finish'ed local root spans
+	// (i.e. those for which no WithLocalParent(<non-nil>) option was supplied).
+	//
+	// In normal operation, a local root Span is inserted on creation and
+	// removed on .Finish().
+	//
+	// The map can be introspected by `Tracer.VisitSpans`.
+	activeSpans struct {
+		// NB: it might be tempting to use a sync.Map here, but
+		// this incurs an allocation per Span (sync.Map does
+		// not use a sync.Pool for its internal *entry type).
+		//
+		// The bare map approach is essentially allocation-free once the map
+		// has grown to accommodate the usual number of active local root spans,
+		// and the critical sections of the mutex are very small.
+		syncutil.Mutex
+		m map[*Span]struct{}
+	}
 }
 
 // NewTracer creates a Tracer. It initially tries to run with minimal overhead
 // and collects essentially nothing; use Configure() to enable various tracing
 // backends.
 func NewTracer() *Tracer {
 	t := &Tracer{}
+	t.activeSpans.m = map[*Span]struct{}{}
 	t.noopSpan = &Span{tracer: t}
 	return t
 }
@@ -397,16 +418,26 @@ func (t *Tracer) startSpanGeneric(
 	// spans contained in Span.
 	//
 	// NB: this could be optimized.
-	if opts.Parent != nil && !opts.Parent.isNoop() {
-		opts.Parent.crdb.mu.Lock()
-		m := opts.Parent.crdb.mu.Baggage
-		for k, v := range m {
-			s.SetBaggageItem(k, v)
+	if opts.Parent != nil {
+		if !opts.Parent.isNoop() {
+			opts.Parent.crdb.mu.Lock()
+			m := opts.Parent.crdb.mu.Baggage
+			for k, v := range m {
+				s.SetBaggageItem(k, v)
+			}
+			opts.Parent.crdb.mu.Unlock()
 		}
-		opts.Parent.crdb.mu.Unlock()
-	} else if opts.RemoteParent != nil {
-		for k, v := range opts.RemoteParent.Baggage {
-			s.SetBaggageItem(k, v)
+	} else {
+		// Local root span - put it into the registry of active local root spans.
+		// `Span.Finish` takes care of deleting it again.
+		t.activeSpans.Lock()
+		t.activeSpans.m[s] = struct{}{}
+		t.activeSpans.Unlock()
+
+		if opts.RemoteParent != nil {
+			for k, v := range opts.RemoteParent.Baggage {
+				s.SetBaggageItem(k, v)
+			}
 		}
 	}
 
@@ -577,6 +608,20 @@ func (t *Tracer) Extract(format interface{}, carrier interface{}) (*SpanMeta, er
 	}, nil
 }
 
+// VisitSpans invokes the visitor with all active Spans.
+func (t *Tracer) VisitSpans(visitor func(*Span)) {
+	t.activeSpans.Lock()
+	sl := make([]*Span, 0, len(t.activeSpans.m))
+	for sp := range t.activeSpans.m {
+		sl = append(sl, sp)
+	}
+	t.activeSpans.Unlock()
+
+	for _, sp := range sl {
+		visitor(sp)
+	}
+}
+
 // ForkCtxSpan checks if ctx has a Span open; if it does, it creates a new Span
 // that "follows from" the original Span. This allows the resulting context to be
 // used in an async task that might outlive the original operation.

pkg/util/tracing/tracer_test.go

@@ -11,6 +11,7 @@
 package tracing
 
 import (
+	"sort"
 	"testing"
 
 	"github.com/cockroachdb/logtags"
@@ -336,3 +337,50 @@ func TestLightstepContext(t *testing.T) {
 	require.Equal(t, exp, s2.Meta().Baggage)
 	require.Equal(t, exp, shadowBaggage)
 }
+
+func getSortedActiveSpanOps(tr *Tracer) []string {
+	var sl []string
+	tr.VisitSpans(func(sp *Span) {
+		for _, rec := range sp.GetRecording() {
+			sl = append(sl, rec.Operation)
+		}
+	})
+	sort.Strings(sl)
+	return sl
+}
+
+// TestTracer_VisitSpans verifies that in-flight local root Spans
+// are tracked by the Tracer, and that Finish'ed Spans are not.
+func TestTracer_VisitSpans(t *testing.T) {
+	tr1 := NewTracer()
+	tr2 := NewTracer()
+
+	root := tr1.StartSpan("root", WithForceRealSpan())
+	root.SetVerbose(true)
+	child := tr1.StartSpan("root.child", WithParentAndAutoCollection(root))
+	require.Len(t, tr1.activeSpans.m, 1)
+
+	childChild := tr2.StartSpan("root.child.remotechild", WithParentAndManualCollection(child.Meta()))
+	childChildFinished := tr2.StartSpan("root.child.remotechilddone", WithParentAndManualCollection(child.Meta()))
+	require.Len(t, tr2.activeSpans.m, 2)
+
+	require.NoError(t, child.ImportRemoteSpans(childChildFinished.GetRecording()))
+
+	childChildFinished.Finish()
+	require.Len(t, tr2.activeSpans.m, 1)
+
+	// Even though only `root` is tracked by tr1, we also reach
+	// root.child and (via ImportRemoteSpans) the remote child.
+	require.Equal(t, []string{"root", "root.child", "root.child.remotechilddone"}, getSortedActiveSpanOps(tr1))
+	require.Equal(t, []string{"root.child.remotechild"}, getSortedActiveSpanOps(tr2))
+
+	childChild.Finish()
+	child.Finish()
+	root.Finish()
+
+	// Nothing is tracked any more.
+	require.Len(t, getSortedActiveSpanOps(tr1), 0)
+	require.Len(t, getSortedActiveSpanOps(tr2), 0)
+	require.Len(t, tr1.activeSpans.m, 0)
+	require.Len(t, tr2.activeSpans.m, 0)
+}