Merge #158527

158527: sqlstats: flush ingester buffer if above max statements per txn r=kyle-a-wong a=kyle-a-wong

There have been a couple incidents recently where workloads running
multi day / million statement txns cause sql stats to have unbounded
memory growth. This is happening because SQL Stats are buffered by
sessionID until a transaction commits, which causes the buffer to
flush and the statement stats to be recorded in the SQL stats
subsystem. When there is a long running transaction with many
statements, we buffer is never flushed and continues to grow.

To fix this, we have introduced a new "force flush" event that
forces the SQL stats ingester to flush sql stats in the current
session, if a certain threshold is met. This threshold is currently
set to the value of a new cluster setting:
`sql.metrics.transaction_details.max_statement_stats`.

If this threshold is met, the stats are automatically flushed. The side effect
of this is that these stats will not have an associated transaction
fingerprint id. This is because the transaction fingerprint id
isn't finalized until a transaction is complete and we don't know
the transaction fingerprint id at the time of this forced flush.
There is also no way to "backfill" or set the transaction
fingerprint id once the statement stat has been recorded to the
SQL Stats subsystem.

This commit also includes a change to the
`SQLStatsIngester.flushBuffer` method to recording transaction
stats when there are no statements stats in the buffer. This is
the case when the buffer is being force flushed in the scenario
mentioned above.

Fixes: #158800
Epic: None
Release note (bug fix): Addresses a bug with unbounded memory
growth when long running transactions are contain many statements.
In this case, the SQL Stats system will automatically flush these
before the transaction commits. The side effect of this is that
these statement statistics will no longer have an associated
transaction fingerprint id.

Co-authored-by: Kyle Wong <[email protected]>

Author: craig[bot]

pkg/sql/executor_statement_metrics.go

@@ -187,8 +187,8 @@ func (ex *connExecutor) recordStatementSummary(
 	ex.metrics.EngineMetrics.StatementBytesRead.Inc(stats.bytesRead)
 	ex.metrics.EngineMetrics.StatementIndexRowsWritten.Inc(stats.indexRowsWritten)
 	ex.metrics.EngineMetrics.StatementIndexBytesWritten.Inc(stats.indexBytesWritten)
-
 	if ex.statsCollector.EnabledForTransaction() {
+		maxRecordedStats := sqlstats.TxnStatsNumStmtFingerprintStatsToRecord.Get(&ex.server.cfg.Settings.SV)
 		b := sqlstats.NewRecordedStatementStatsBuilder(
 			stmtFingerprintID,
 			planner.SessionData().Database,
@@ -228,6 +228,19 @@ func (ex *connExecutor) recordStatementSummary(
 		}
 
 		ex.statsCollector.RecordStatement(ctx, b.Build())
+		if int64(ex.state.mu.stmtCount)%(maxRecordedStats) == 0 {
+			// If the statement count is equal to the configured limit,
+			// force flush the SQL stats ingestion buffer.
+			// Note that these statements that are force flushed cannot be associated
+			// with the transaction in SQL Stats since this association can only be
+			// done once a transaction has been committed. As a result, these
+			// recorded stats will be given a static transaction fingerprint ID to
+			// indicate that they are force flushed. Since these force flushes occur
+			// in batches, any leftover statement stats will be flushed when the
+			// transaction is committed, causing them to be associated with the
+			// correct transaction fingerprint.
+			ex.statsCollector.Flush(ex.planner.extendedEvalCtx.SessionID)
+		}
 	}
 
 	// Record statement execution statistics if span is recorded and no error was

pkg/sql/sqlstats/cluster_settings.go

@@ -29,6 +29,23 @@ var TxnStatsNumStmtFingerprintIDsToRecord = settings.RegisterIntSetting(
 	settings.PositiveInt,
 )
 
+// TxnStatsNumStmtFingerprintStatsToRecord limits the number of recorded
+// statement statistics that may be associated with transaction statistics for
+// a single transaction. If the number of statements executed exceeds this
+// value, SQL Stats will force the ingester to flush the buffered stats. These
+// stats will not be associated with the current transaction. SQL Stats will
+// continue to buffer statements until this limit is reached again. In the case
+// that it isn't reached again, the buffered statements will be flushed when
+// the transaction is committed, and they will be associated with the
+// transaction.
+var TxnStatsNumStmtFingerprintStatsToRecord = settings.RegisterIntSetting(
+	settings.ApplicationLevel,
+	"sql.metrics.transaction_details.max_statement_stats",
+	"max number of statement statistics that may be associated with transaction statistics",
+	100_000,
+	settings.PositiveInt,
+)
+
 // TxnStatsEnable determines whether to collect per-application transaction
 // statistics.
 // TODO(117690): Unify StmtStatsEnable and TxnStatsEnable into a single cluster setting.

pkg/sql/sqlstats/sqlstats_test.go

@@ -122,6 +122,7 @@ func GetTxnStats(t *testing.T, conn *sqlutils.SQLRunner, appName string, dbName
 SELECT 
 	encode(fingerprint_id, 'hex') AS txn_fingerprint_id,
 	metadata->'stmtFingerprintIDs' AS statement_fingerprint_ids,
+  statistics -> 'statistics'->> 'cnt' AS count,
 	(statistics -> 'statistics' -> 'commitLat' -> 'mean')::FLOAT > 0 AS commit_lat_not_zero,
 	(statistics -> 'statistics' -> 'svcLat' -> 'mean')::FLOAT > 0 AS svc_lat_not_zero
 FROM crdb_internal.transaction_statistics ts

pkg/sql/sqlstats/sslocal/sql_stats_ingestor.go

@@ -7,6 +7,7 @@ package sslocal
 
 import (
 	"context"
+	"math"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -25,6 +26,8 @@ import (
 // will go before flushing its contents to the registry.
 const defaultFlushInterval = time.Millisecond * 500
 
+const forceFlushTransactionFingerprintId = appstatspb.TransactionFingerprintID(math.MaxUint64)
+
 // Metrics holds running measurements of various ingester-related runtime stats.
 type Metrics struct {
 	// NumProcessed tracks how many items have been processed.
@@ -149,6 +152,7 @@ type event struct {
 	sessionID   clusterunique.ID
 	transaction *sqlstats.RecordedTxnStats
 	statement   *sqlstats.RecordedStmtStats
+	forceFlush  bool
 }
 
 type BufferOpt func(i *SQLStatsIngester)
@@ -250,14 +254,18 @@ func (i *SQLStatsIngester) ingest(ctx context.Context, events *eventBuffer) {
 			// the stmts with and the statement's txn id will be nil. In that case
 			// we can send immediately to the sinks.
 			if e.statement.UnderOuterTxn {
-				i.flushBuffer(ctx, e.statement.SessionID, nil)
+				i.flushBuffer(ctx, e.statement.SessionID, nil, 0)
 			}
 		} else if e.transaction != nil {
-			i.flushBuffer(ctx, e.transaction.SessionID, e.transaction)
+			i.flushBuffer(ctx, e.transaction.SessionID, e.transaction, e.transaction.FingerprintID)
 			i.metrics.NumProcessed.Inc(1)
 			i.metrics.QueueSize.Dec(1)
 		} else {
-			i.clearSession(e.sessionID)
+			if e.forceFlush {
+				i.flushBuffer(ctx, e.sessionID, nil, forceFlushTransactionFingerprintId)
+			} else {
+				i.clearSession(e.sessionID)
+			}
 			i.metrics.NumProcessed.Inc(1)
 			i.metrics.QueueSize.Dec(1)
 		}
@@ -323,6 +331,18 @@ func (i *SQLStatsIngester) ClearSession(sessionID clusterunique.ID) {
 	})
 }
 
+// FlushBuffer sends a signal to the underlying registry to flush any cached
+// data associated with the given sessionID. This is an async operation.
+func (i *SQLStatsIngester) FlushBuffer(sessionID clusterunique.ID) {
+	i.guard.AtomicWrite(func(writerIdx int64) {
+		i.guard.eventBuffer[writerIdx] = event{
+			sessionID:  sessionID,
+			forceFlush: true,
+		}
+		i.metrics.QueueSize.Inc(1)
+	})
+}
+
 func NewSQLStatsIngester(
 	st *cluster.Settings, knobs *sqlstats.TestingKnobs, metrics Metrics, sinks ...SQLStatsSink,
 ) *SQLStatsIngester {
@@ -397,45 +417,39 @@ func (i *SQLStatsIngester) processStatement(statement *sqlstats.RecordedStmtStat
 // flushBuffer sends the buffered statementsBySessionID and provided transaction
 // to the registered sinks. The transaction may be nil.
 func (i *SQLStatsIngester) flushBuffer(
-	ctx context.Context, sessionID clusterunique.ID, transaction *sqlstats.RecordedTxnStats,
+	ctx context.Context,
+	sessionID clusterunique.ID,
+	transaction *sqlstats.RecordedTxnStats,
+	transactionFingerprintID appstatspb.TransactionFingerprintID,
 ) {
-	statements, ok := func() (*statementBuf, bool) {
-		statements, ok := i.statementsBySessionID[sessionID]
-		if !ok {
-			return nil, false
-		}
+	var statements statementBuf
+	if sessionStatements, ok := i.statementsBySessionID[sessionID]; ok {
+		defer sessionStatements.release()
+		statements = *sessionStatements
 		delete(i.statementsBySessionID, sessionID)
-		return statements, true
-	}()
-	if !ok {
+	} else if transaction == nil {
+		// No statements or transactions to flush, return early
 		return
 	}
-	defer statements.release()
 
-	if len(*statements) == 0 {
-		return
-	}
-
-	if transaction != nil {
-		// Here we'll set the transaction fingerprint ID for each statement if the
-		// below cluster setting is enabled.
-		shouldAssociateWithTxn := AssociateStmtWithTxnFingerprint.Get(&i.settings.SV)
-		// These values are only known at the time of the transaction.
-		for _, s := range *statements {
-			if shouldAssociateWithTxn {
-				s.TransactionFingerprintID = transaction.FingerprintID
-			}
-			if s.ImplicitTxn == transaction.ImplicitTxn {
-				continue
-			}
-			// We need to recompute the fingerprint ID.
-			s.ImplicitTxn = transaction.ImplicitTxn
-			s.FingerprintID = appstatspb.ConstructStatementFingerprintID(
-				s.Query, s.ImplicitTxn, s.Database)
+	// Here we'll set the transaction fingerprint ID for each statement if the
+	// below cluster setting is enabled.
+	shouldAssociateWithTxn := AssociateStmtWithTxnFingerprint.Get(&i.settings.SV)
+	// These values are only known at the time of the transaction.
+	for _, s := range statements {
+		if shouldAssociateWithTxn {
+			s.TransactionFingerprintID = transactionFingerprintID
+		}
+		if transaction == nil || s.ImplicitTxn == transaction.ImplicitTxn {
+			continue
 		}
+		// We need to recompute the fingerprint ID.
+		s.ImplicitTxn = transaction.ImplicitTxn
+		s.FingerprintID = appstatspb.ConstructStatementFingerprintID(
+			s.Query, s.ImplicitTxn, s.Database)
 	}
 
 	for _, sink := range i.sinks {
-		sink.ObserveTransaction(ctx, transaction, *statements)
+		sink.ObserveTransaction(ctx, transaction, statements)
 	}
 }

pkg/sql/sqlstats/sslocal/sslocal_stats_collector.go

@@ -164,6 +164,14 @@ func (s *StatsCollector) RecordTransaction(_ctx context.Context, value *sqlstats
 	s.statsIngester.RecordTransaction(value)
 }
 
+func (s *StatsCollector) Flush(sessionID clusterunique.ID) {
+	if !s.sendStats {
+		return
+	}
+
+	s.statsIngester.FlushBuffer(sessionID)
+}
+
 func (s *StatsCollector) EnabledForTransaction() bool {
 	return s.sendStats
 }

pkg/sql/sqlstats/testdata/max_txn_statements

@@ -0,0 +1,73 @@
+# In this test, we verify that statement statistics are force flushed based on
+# the `sql.metrics.transaction_details.max_statement_ids` cluster setting. If
+# the expected criteria is met, statement statistics for a session will be
+# flushed from the buffer, regardless of whether the transaction has ended.
+# When statement statistics are force flushed, there will be no transaction
+# fingerprint ID associated with the flushed statement statistics, since this
+# requires a transaction to have ended to compute.
+
+# Setup baseline to test against with default settings
+exec-sql db=max_txn_stmts_test app-name=exec
+BEGIN;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+END;
+----
+
+exec-sql db=max_txn_stmts_test app-name=exec
+BEGIN;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+END;
+----
+
+show-stats db=max_txn_stmts_test app-name=exec
+----
+{"count": "5", "fingerprint_id": "bdb329db6493db17", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT _", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT _", "svc_lat_not_zero": true, "transaction_fingerprint_id": "023097855475259c"}
+{"count": "3", "fingerprint_id": "bdb329db6493db17", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT _", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT _", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3b5be2cb288f18aa"}
+
+show-txn-stats db=max_txn_stmts_test app-name=exec
+----
+{"commit_lat_not_zero": true, "count": "1", "statement_fingerprint_ids": ["bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17"], "svc_lat_not_zero": true, "txn_fingerprint_id": "023097855475259c"}
+{"commit_lat_not_zero": true, "count": "1", "statement_fingerprint_ids": ["bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17"], "svc_lat_not_zero": true, "txn_fingerprint_id": "3b5be2cb288f18aa"}
+
+exec-sql db=max_txn_stmts_test app-name=setup-limit
+SET CLUSTER SETTING sql.metrics.transaction_details.max_statement_stats=4;
+----
+
+# This is over the configured settings but below the threshold for force flush,
+# so it will still contribute to the same SQL Stats rows as above.
+exec-sql db=max_txn_stmts_test app-name=exec
+BEGIN;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+END;
+----
+
+# This is over the threshold for force flush, so it results in a new row in
+# SQL stats with a blank transaction fingerprint ID.
+exec-sql db=max_txn_stmts_test app-name=exec
+BEGIN;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+SELECT 1;
+END;
+----
+
+show-stats db=max_txn_stmts_test app-name=exec
+----
+{"count": "6", "fingerprint_id": "bdb329db6493db17", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT _", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT _", "svc_lat_not_zero": true, "transaction_fingerprint_id": "023097855475259c"}
+{"count": "6", "fingerprint_id": "bdb329db6493db17", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT _", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT _", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3b5be2cb288f18aa"}
+{"count": "4", "fingerprint_id": "bdb329db6493db17", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT _", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT _", "svc_lat_not_zero": true, "transaction_fingerprint_id": "ffffffffffffffff"}
+
+show-txn-stats db=max_txn_stmts_test app-name=exec
+----
+{"commit_lat_not_zero": true, "count": "2", "statement_fingerprint_ids": ["bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17"], "svc_lat_not_zero": true, "txn_fingerprint_id": "023097855475259c"}
+{"commit_lat_not_zero": true, "count": "2", "statement_fingerprint_ids": ["bdb329db6493db17", "bdb329db6493db17", "bdb329db6493db17"], "svc_lat_not_zero": true, "txn_fingerprint_id": "3b5be2cb288f18aa"}

pkg/sql/sqlstats/testdata/query

@@ -35,4 +35,4 @@ show-stats db=random_db app-name=transactions
 
 show-txn-stats db=random_db app-name=transactions
 ----
-{"commit_lat_not_zero": true, "statement_fingerprint_ids": ["5fe99858726d3688", "5fe99858726d3688"], "svc_lat_not_zero": true, "txn_fingerprint_id": "78d7c1c32632f05d"}
+{"commit_lat_not_zero": true, "count": "2", "statement_fingerprint_ids": ["5fe99858726d3688", "5fe99858726d3688"], "svc_lat_not_zero": true, "txn_fingerprint_id": "78d7c1c32632f05d"}