TestScope: don't prune from registry when closed

scope.go

@@ -101,6 +101,7 @@ type scope struct {
 	done        chan struct{}
 	wg          sync.WaitGroup
 	root        bool
+	testScope   bool
 }
 
 // ScopeOptions is a set of options to construct a scope.
@@ -114,6 +115,7 @@ type ScopeOptions struct {
 	SanitizeOptions *SanitizeOptions
 	MetricsOption   InternalMetricOption
 
+	testScope          bool
 	registryShardCount uint
 }
 
@@ -132,7 +134,11 @@ func NewTestScope(
 	prefix string,
 	tags map[string]string,
 ) TestScope {
-	return newRootScope(ScopeOptions{Prefix: prefix, Tags: tags}, 0)
+	return newRootScope(ScopeOptions{
+		Prefix:    prefix,
+		Tags:      tags,
+		testScope: true,
+	}, 0)
 }
 
 func newRootScope(opts ScopeOptions, interval time.Duration) *scope {
@@ -177,6 +183,7 @@ func newRootScope(opts ScopeOptions, interval time.Duration) *scope {
 		separator:       sanitizer.Name(opts.Separator),
 		timers:          make(map[string]*timer),
 		root:            true,
+		testScope:       opts.testScope,
 	}
 
 	// NB(r): Take a copy of the tags on creation

scope_registry.go

@@ -159,10 +159,10 @@ func (r *scopeRegistry) Subscope(parent *scope, prefix string, tags map[string]s
 
 	s, ok := r.lockedLookup(subscopeBucket, unsanitizedKey)
 	if ok {
-		// If this subscope isn't closed, return it. Otherwise, report it
-		// immediately and delete it so that a new (functional) scope can be
-		// returned instead.
-		if !s.closed.Load() {
+		// If this subscope isn't closed or is a test scope, return it.
+		// Otherwise, report it immediately and delete it so that a new
+		// (functional) scope can be returned instead.
+		if !s.closed.Load() || s.testScope {
 			subscopeBucket.mu.RUnlock()
 			return s
 		}
@@ -229,6 +229,7 @@ func (r *scopeRegistry) Subscope(parent *scope, prefix string, tags map[string]s
 		timers:          make(map[string]*timer),
 		bucketCache:     parent.bucketCache,
 		done:            make(chan struct{}),
+		testScope:       parent.testScope,
 	}
 	subscopeBucket.s[sanitizedKey] = subscope
 	if _, ok := r.lockedLookup(subscopeBucket, unsanitizedKey); !ok {

scope_registry_external_test.go

@@ -33,6 +33,35 @@ import (
 	"github.com/uber-go/tally/tallymock"
 )
 
+func TestTestScopesNotPruned(t *testing.T) {
+	var (
+		root     = tally.NewTestScope("", nil)
+		subscope = root.SubScope("foo")
+		counter  = subscope.Counter("bar")
+	)
+
+	counter.Inc(123)
+
+	closer, ok := subscope.(io.Closer)
+	require.True(t, ok)
+	require.NoError(t, closer.Close())
+
+	subscope = root.SubScope("foo")
+	counter = subscope.Counter("bar")
+	counter.Inc(123)
+
+	snapshot := root.Snapshot()
+	require.Len(t, snapshot.Counters(), 1)
+	require.Len(t, snapshot.Gauges(), 0)
+	require.Len(t, snapshot.Timers(), 0)
+	require.Len(t, snapshot.Histograms(), 0)
+
+	val, ok := snapshot.Counters()["foo.bar+"]
+	require.True(t, ok)
+	require.Equal(t, "foo.bar", val.Name())
+	require.EqualValues(t, 246, val.Value())
+}
+
 func TestNoDefunctSubscopes(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()