sqlstats: reuse the temporary statement stats container

This commit cleans up the confusing swapping and uses of the
2 fields on the StatsCollector, `ApplicationStats` and
`flushTarget`.

The original purpose of `flushTarget` seems to have been for
explicit transactions, where we would need to use a temp container
to record statement stats and flush them to the app container at
the end of transaction execution in order to correlate the txn
fingerprint id. The assumption then was that implicit transactions
have 1 statement, thus we'd only need to use a flush target for
explicit transactions - for implicit txns we could get away with
using the `ApplicationStats` field directly.

Before:
If implict txn:
- do nothing, stats will be recorded to `ApplicationStats` directly
If explicit txn:
- set `flushTarget` as `ApplicationStats`, allocate temp container
for `ApplicationStats` where the txn's statements will be recorded
- on txn finish, flush `ApplicationStats` to `flushTarget`, set
`ApplicationStats` to `flushTarget`.

Eventually we corrected that assumption and are now separating the
current transacion's statements and the flushTarget for every execution.
There's no need to perform the little dance above, and we're likely not
managing our statement stats container as efficiently as we could be.

Now:
- `FlushTarget` is always defined for a stats collector. It
represents the current application's sql stats
- `ApplicationStats` is renamed to `activeTxnStatementStats` to reflect
that it is use to store the current transaction's statement stats
- Instead of allocating a new container for the current transaction
each time, we'll clear and reuse `activeTxnStatementStats`

Fixes: cockroachdb#94650

pkg/sql/sqlstats/sslocal/sslocal_stats_collector.go

@@ -28,7 +28,8 @@ import (
 // StatsCollector is used to collect statistics for transactions and
 // statements for the entire lifetime of a session.
 type StatsCollector struct {
-	// currentTransactionStatements is the current transaction's statement statistics.
+
+	// ApplicationStats contains the current transaction's statement statistics.
 	// They will be flushed to flushTarget when the transaction is done.
 	ApplicationStats sqlstats.ApplicationStats
 
@@ -56,8 +57,14 @@ type StatsCollector struct {
 	// statement insights.
 	sendInsights bool
 
+	// flushTarget is the sql stats container for the current application.
+	// This is the target where the statement stats are flushed to upon
+	// transaction completion. Note that these are the global stats for the
+	// application.
 	flushTarget sqlstats.ApplicationStats
 
+	// uniqueServerCounts is a pointer to the statement and transaction
+	// fingerprint counters tracked per server.
 	uniqueServerCounts *ssmemstorage.SQLStatsAtomicCounters
 
 	st    *cluster.Settings
@@ -74,7 +81,8 @@ func NewStatsCollector(
 	knobs *sqlstats.TestingKnobs,
 ) *StatsCollector {
 	return &StatsCollector{
-		ApplicationStats:   appStats,
+		flushTarget:        appStats,
+		ApplicationStats:   appStats.NewApplicationStatsWithInheritedOptions(),
 		insightsWriter:     insights,
 		phaseTimes:         phaseTime.Clone(),
 		uniqueServerCounts: uniqueServerCounts,
@@ -119,8 +127,6 @@ func (s *StatsCollector) Reset(appStats sqlstats.ApplicationStats, phaseTime *se
 // The current application stats are reset for the new transaction.
 func (s *StatsCollector) StartTransaction() {
 	s.sendInsights = s.shouldObserveInsights()
-	s.flushTarget = s.ApplicationStats
-	s.ApplicationStats = s.flushTarget.NewApplicationStatsWithInheritedOptions()
 }
 
 // EndTransaction informs the StatsCollector that the current txn has
@@ -144,19 +150,12 @@ func (s *StatsCollector) EndTransaction(
 		ctx, s.ApplicationStats, transactionFingerprintID,
 	)
 
-	discardedStats += s.flushTarget.MergeApplicationTransactionStats(
-		ctx,
-		s.ApplicationStats,
-	)
-
 	// Avoid taking locks if no stats are discarded.
 	if discardedStats > 0 {
 		s.flushTarget.MaybeLogDiscardMessage(ctx)
 	}
 
-	s.ApplicationStats.Free(ctx)
-	s.ApplicationStats = s.flushTarget
-	s.flushTarget = nil
+	s.ApplicationStats.Clear(ctx)
 }
 
 // ShouldSample returns two booleans, the first one indicates whether we
@@ -166,17 +165,8 @@ func (s *StatsCollector) EndTransaction(
 func (s *StatsCollector) ShouldSample(
 	fingerprint string, implicitTxn bool, database string,
 ) (previouslySampled bool, savePlanForStats bool) {
-	sampledInFlushTarget := false
-	savePlanForStatsInFlushTarget := true
-
-	if s.flushTarget != nil {
-		sampledInFlushTarget, savePlanForStatsInFlushTarget = s.flushTarget.ShouldSample(fingerprint, implicitTxn, database)
-	}
 
-	sampledInAppStats, savePlanForStatsInAppStats := s.ApplicationStats.ShouldSample(fingerprint, implicitTxn, database)
-	previouslySampled = sampledInFlushTarget || sampledInAppStats
-	savePlanForStats = savePlanForStatsInFlushTarget && savePlanForStatsInAppStats
-	return previouslySampled, savePlanForStats
+	return s.flushTarget.ShouldSample(fingerprint, implicitTxn, database)
 }
 
 // UpgradeImplicitTxn informs the StatsCollector that the current txn has been
@@ -330,14 +320,15 @@ func (s *StatsCollector) StatementsContainerFull() bool {
 func (s *StatsCollector) RecordStatement(
 	ctx context.Context, key appstatspb.StatementStatisticsKey, value sqlstats.RecordedStmtStats,
 ) (appstatspb.StmtFingerprintID, error) {
-	return s.ApplicationStats.RecordStatement(ctx, key, value)
+	return s.currentTransactionStatementStats.RecordStatement(ctx, key, value)
 }
 
 // RecordTransaction records the statistics of a transaction.
+// Transaction stats are always recorded directly on the flushTarget.
 func (s *StatsCollector) RecordTransaction(
 	ctx context.Context, key appstatspb.TransactionFingerprintID, value sqlstats.RecordedTxnStats,
 ) error {
-	return s.ApplicationStats.RecordTransaction(ctx, key, value)
+	return s.flushTarget.RecordTransaction(ctx, key, value)
 }
 
 func (s *StatsCollector) RecordStatementExecStats(

pkg/sql/sqlstats/ssmemstorage/ss_mem_storage.go

@@ -361,12 +361,6 @@ func (t *txnStats) sizeUnsafeLocked() int64 {
 	return txnStatsShallowSize + stmtFingerprintIDsSize + dataSize
 }
 
-func (t *txnStats) mergeStats(stats *appstatspb.TransactionStatistics) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	t.mu.data.Add(stats)
-}
-
 // stmtStats holds per-statement statistics.
 type stmtStats struct {
 	// ID is the statementFingerprintID constructed using the stmtKey fields.

pkg/sql/sqlstats/ssprovider.go

@@ -103,6 +103,9 @@ type ApplicationStats interface {
 	// Free frees the current ApplicationStats and zeros out the memory counts
 	// and fingerprint counts.
 	Free(context.Context)
+
+	// Clear is like Free but also prepares the container for reuse.
+	Clear(context.Context)
 }
 
 // IteratorOptions provides the ability to the caller to change how it iterates