topsql: fix memory leak in stmtstats (#34558)

ref #34502, ref #34525

util/topsql/stmtstats/aggregator.go

@@ -68,28 +68,25 @@ func (m *aggregator) run() {
 		case <-m.ctx.Done():
 			return
 		case <-tick.C:
-			m.aggregate(state.TopSQLEnabled())
+			m.aggregate()
 		}
 	}
 }
 
 // aggregate data from all associated StatementStats.
 // If StatementStats has been closed, collect will remove it from the map.
-func (m *aggregator) aggregate(take bool) {
+func (m *aggregator) aggregate() {
 	total := StatementStatsMap{}
 	m.statsSet.Range(func(statsR, _ interface{}) bool {
 		stats := statsR.(*StatementStats)
 		if stats.Finished() {
 			m.unregister(stats)
-			total.Merge(stats.Take())
-			return true
-		}
-		if take {
-			total.Merge(stats.Take())
 		}
+		total.Merge(stats.Take())
 		return true
 	})
-	if len(total) > 0 {
+	// If TopSQL is not enabled, just drop them.
+	if len(total) > 0 && state.TopSQLEnabled() {
 		m.collectors.Range(func(c, _ interface{}) bool {
 			c.(Collector).CollectStmtStatsMap(total)
 			return true

util/topsql/stmtstats/aggregator_test.go

@@ -16,7 +16,6 @@ package stmtstats
 
 import (
 	"math/rand"
-	"sync"
 	"testing"
 	"time"
 
@@ -51,6 +50,8 @@ func Test_RegisterUnregisterCollector(t *testing.T) {
 }
 
 func Test_aggregator_register_collect(t *testing.T) {
+	state.EnableTopSQL()
+	defer state.DisableTopSQL()
 	a := newAggregator()
 	stats := &StatementStats{
 		data:     StatementStatsMap{},
@@ -63,7 +64,7 @@ func Test_aggregator_register_collect(t *testing.T) {
 	a.registerCollector(newMockCollector(func(data StatementStatsMap) {
 		total.Merge(data)
 	}))
-	a.aggregate(true)
+	a.aggregate()
 	assert.NotEmpty(t, total)
 	assert.Equal(t, uint64(1), total[SQLPlanDigest{SQLDigest: "SQL-1"}].ExecCount)
 	assert.Equal(t, uint64(time.Millisecond.Nanoseconds()), total[SQLPlanDigest{SQLDigest: "SQL-1"}].SumDurationNs)
@@ -90,37 +91,36 @@ func Test_aggregator_run_close(t *testing.T) {
 }
 
 func TestAggregatorDisableAggregate(t *testing.T) {
-	var mu sync.Mutex
 	total := StatementStatsMap{}
 	a := newAggregator()
 	a.registerCollector(newMockCollector(func(data StatementStatsMap) {
-		mu.Lock()
 		total.Merge(data)
-		mu.Unlock()
 	}))
 
-	a.start()
-
+	state.DisableTopSQL()
 	stats := &StatementStats{
 		data: StatementStatsMap{
-			SQLPlanDigest{SQLDigest: BinaryDigest("")}: &StatementStatsItem{},
+			SQLPlanDigest{SQLDigest: ""}: &StatementStatsItem{},
 		},
 		finished: atomic.NewBool(false),
 	}
-	state.DisableTopSQL()
 	a.register(stats)
-	time.Sleep(1500 * time.Millisecond)
-	mu.Lock()
-	require.Empty(t, total)
-	mu.Unlock()
+	a.aggregate()
+	require.Empty(t, stats.data) // a.aggregate() will take all data even if TopSQL is not enabled.
+	require.Empty(t, total)      // But just drop them.
+
 	state.EnableTopSQL()
-	time.Sleep(1500 * time.Millisecond)
-	mu.Lock()
+	stats = &StatementStats{
+		data: StatementStatsMap{
+			SQLPlanDigest{SQLDigest: ""}: &StatementStatsItem{},
+		},
+		finished: atomic.NewBool(false),
+	}
+	a.register(stats)
+	a.aggregate()
+	require.Empty(t, stats.data)
 	require.Len(t, total, 1)
-	mu.Unlock()
 	state.DisableTopSQL()
-
-	a.close()
 }
 
 type mockCollector struct {