From db186479100ec87dddf8fe47b1ef87b8df8bb507 Mon Sep 17 00:00:00 2001
From: Matt Way <mway@users.noreply.github.com>
Date: Sat, 10 Jun 2023 20:19:23 -0400
Subject: [PATCH] TestScope: don't prune from registry when closed (#222)

---
 scope.go                        | 26 +++++++++++++++++---------
 scope_registry.go               |  9 +++++----
 scope_registry_external_test.go | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/scope.go b/scope.go
index 9a6ee67..d63bd51 100644
--- a/scope.go
+++ b/scope.go
@@ -92,20 +92,23 @@ type scope struct {
 	done        chan struct{}
 	wg          sync.WaitGroup
 	root        bool
+	testScope   bool
 }
 
 // ScopeOptions is a set of options to construct a scope.
 type ScopeOptions struct {
-	Tags                   map[string]string
-	Prefix                 string
-	Reporter               StatsReporter
-	CachedReporter         CachedStatsReporter
-	Separator              string
-	DefaultBuckets         Buckets
-	SanitizeOptions        *SanitizeOptions
-	registryShardCount     uint
+	Tags            map[string]string
+	Prefix          string
+	Reporter        StatsReporter
+	CachedReporter  CachedStatsReporter
+	Separator       string
+	DefaultBuckets  Buckets
+	SanitizeOptions *SanitizeOptions
 	OmitCardinalityMetrics bool
 	CardinalityMetricsTags map[string]string
+
+	testScope          bool
+	registryShardCount uint
 }
 
 // NewRootScope creates a new root Scope with a set of options and
@@ -129,7 +132,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 {
@@ -174,6 +181,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
diff --git a/scope_registry.go b/scope_registry.go
index fe26723..d1481da 100644
--- a/scope_registry.go
+++ b/scope_registry.go
@@ -188,10 +188,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
 		}
@@ -258,6 +258,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 {
diff --git a/scope_registry_external_test.go b/scope_registry_external_test.go
index e6dac92..5df241d 100644
--- a/scope_registry_external_test.go
+++ b/scope_registry_external_test.go
@@ -33,6 +33,38 @@ import (
 	"github.com/uber-go/tally/v4/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)
+
+	var (
+		snapshot = root.Snapshot()
+		counters = snapshot.Counters()
+	)
+	require.Len(t, counters, 1)
+	require.Len(t, snapshot.Gauges(), 0)
+	require.Len(t, snapshot.Timers(), 0)
+	require.Len(t, snapshot.Histograms(), 0)
+
+	val, ok := 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()