sqlstats: move sql stats writing to be async

This commit removes the recording of sql stats at the end of statement and
transaction execution. The per-node SQLStats container has been added as a
sink to the async sql stats ingester. Instead of having sslocal.StatsCollector
record sql stats at the end of execution, it will now record the stats
asynchronously as part of event processing in sslocal.SQLStatsIngester's
ingest routine.

Epic: none
Closes: cockroachdb#141024

Release note: None

Author: xinhaoz

Diff for: pkg/sql/conn_executor.go

@@ -1230,7 +1230,6 @@ func (s *Server) newConnExecutor(
 		s.sqlStatsIngester,
 		ex.phaseTimes,
 		s.localSqlStats.GetCounters(),
-		underOuterTxn,
 		s.cfg.SQLStatsTestingKnobs,
 	)
 	ex.dataMutatorIterator.onApplicationNameChange = func(newName string) {
@@ -4079,17 +4078,7 @@ func (ex *connExecutor) txnStateTransitionsApplyWrapper(
 				return advanceInfo{}, err
 			}
 			if advInfo.txnEvent.eventType == txnUpgradeToExplicit {
-				// TODO (xinhaoz): This is a temporary hook until
-				// https://github.com/cockroachdb/cockroach/issues/141024
-				// is resolved. The reason this exists is because we
-				// need to recompute the statement fingerprint id for
-				// statements currently in the stats collector, which
-				// were computed once already for insights. There is an
-				// acknowledgement that this means the fingerprint id
-				// given to insights for upgraded transactions are
-				// currently incorrect.
 				ex.extraTxnState.txnFinishClosure.implicit = false
-				ex.statsCollector.UpgradeToExplicitTransaction()
 			}
 		}
 	case txnStart:

Diff for: pkg/sql/conn_executor_exec.go

@@ -4181,11 +4181,6 @@ func (ex *connExecutor) onTxnFinish(ctx context.Context, ev txnEvent, txnErr err
 			}
 		}
 
-		discardedStats := ex.statsCollector.EndTransaction(ctx, transactionFingerprintID)
-		if discardedStats > 0 {
-			ex.server.ServerMetrics.StatsMetrics.DiscardedStatsCount.Inc(discardedStats)
-		}
-
 		if ex.server.cfg.TestingKnobs.BeforeTxnStatsRecorded != nil {
 			ex.server.cfg.TestingKnobs.BeforeTxnStatsRecorded(
 				ex.sessionData(),
@@ -4330,6 +4325,8 @@ func (ex *connExecutor) recordTransactionFinish(
 	txnRetryLat := ex.phaseTimes.GetTransactionRetryLatency()
 	commitLat := ex.phaseTimes.GetCommitLatency()
 
+	isInternaleExec := ex.executorType == executorTypeInternal
+
 	recordedTxnStats := &sqlstats.RecordedTxnStats{
 		FingerprintID:           transactionFingerprintID,
 		SessionID:               ex.planner.extendedEvalCtx.SessionID,
@@ -4357,14 +4354,15 @@ func (ex *connExecutor) recordTransactionFinish(
 		// TODO(107318): add qos
 		// TODO(107318): add asoftime or ishistorical
 		// TODO(107318): add readonly
-		TxnErr:         txnErr,
-		Application:    ex.applicationName.Load().(string),
-		UserNormalized: ex.sessionData().User().Normalized(),
+		TxnErr:           txnErr,
+		Application:      ex.applicationName.Load().(string),
+		UserNormalized:   ex.sessionData().User().Normalized(),
+		InternalExecutor: isInternaleExec,
 	}
 
 	if ex.server.cfg.TestingKnobs.OnRecordTxnFinish != nil {
 		ex.server.cfg.TestingKnobs.OnRecordTxnFinish(
-			ex.executorType == executorTypeInternal, ex.phaseTimes, ex.planner.stmt.SQL, recordedTxnStats,
+			isInternaleExec, ex.phaseTimes, ex.planner.stmt.SQL, recordedTxnStats,
 		)
 	}
 

Diff for: pkg/sql/executor_statement_metrics.go

@@ -202,8 +202,9 @@ func (ex *connExecutor) recordStatementSummary(
 		ExecStats:            queryLevelStats,
 		// TODO(mgartner): Use a slice of struct{uint64, uint64} instead of
 		// converting to strings.
-		Indexes:  planner.instrumentation.indexesUsed.Strings(),
-		Database: planner.SessionData().Database,
+		Indexes:       planner.instrumentation.indexesUsed.Strings(),
+		Database:      planner.SessionData().Database,
+		UnderOuterTxn: ex.extraTxnState.underOuterTxn,
 	}
 
 	err := ex.statsCollector.RecordStatement(ctx, recordedStmtStats)

Diff for: pkg/sql/sqlstats/insights/provider.go

@@ -35,5 +35,8 @@ func (p *Provider) ObserveTransaction(
 	transactionStats *sqlstats.RecordedTxnStats,
 	statements []*sqlstats.RecordedStmtStats,
 ) {
+	if transactionStats == nil || transactionStats.InternalExecutor {
+		return
+	}
 	p.registry.observeTransaction(transactionStats, statements)
 }

Diff for: pkg/sql/sqlstats/sslocal/sql_stats.go

@@ -47,6 +47,8 @@ type SQLStats struct {
 	knobs *sqlstats.TestingKnobs
 }
 
+var _ SQLStatsSink = &SQLStats{}
+
 func newSQLStats(
 	st *cluster.Settings,
 	uniqueStmtFingerprintLimit *settings.IntSetting,
@@ -177,3 +179,33 @@ func (s *SQLStats) MaybeDumpStatsToLog(
 func (s *SQLStats) GetClusterSettings() *cluster.Settings {
 	return s.st
 }
+
+// ObserveTransaction implements the sslocal.SQLStatsSink interface.
+func (s *SQLStats) ObserveTransaction(
+	ctx context.Context,
+	transaction *sqlstats.RecordedTxnStats,
+	statements []*sqlstats.RecordedStmtStats,
+) {
+	if transaction != nil {
+		// Write statements.
+		appStats := s.GetApplicationStats(transaction.Application)
+		for _, stmt := range statements {
+			err := appStats.RecordStatement(ctx, stmt)
+			if err != nil {
+				// todo
+			}
+		}
+		err := appStats.RecordTransaction(ctx, transaction)
+		if err != nil {
+			// todo
+		}
+	} else {
+		for _, stmt := range statements {
+			appStats := s.GetApplicationStats(stmt.App)
+			err := appStats.RecordStatement(ctx, stmt)
+			if err != nil {
+				// todo
+			}
+		}
+	}
+}

Diff for: pkg/sql/sqlstats/sslocal/sql_stats_ingestor.go

@@ -11,6 +11,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/clusterunique"
 	"github.com/cockroachdb/cockroach/pkg/sql/contention/contentionutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
@@ -22,14 +23,15 @@ import (
 const defaultFlushInterval = time.Millisecond * 500
 
 type SQLStatsSink interface {
-	// ObserveTransaction is called by the ingester to pass along a transaction and its statementsBySessionID.
+	// ObserveTransaction is called by the ingester to pass along a transaction event (possibly nil) and its
+	// statementsBySessionID.
 	// Note that the sink should transform the transaction and statementsBySessionID into the appropriate format
 	// as these objects will be returned to the pool.
 	ObserveTransaction(ctx context.Context, transaction *sqlstats.RecordedTxnStats, statements []*sqlstats.RecordedStmtStats)
 }
 
-// SQLStatsIngester amortizes the locking cost of writing to
-// the sql stats container concurrently from multiple goroutines.
+// SQLStatsIngester collects and buffers statements and transactions per session
+// before passing them to the sinks when we receive complete information for a transaction.
 // Built around contentionutils.ConcurrentBufferGuard.
 type SQLStatsIngester struct {
 	guard struct {
@@ -66,6 +68,9 @@ type eventBufChPayload struct {
 type statementBuf []*sqlstats.RecordedStmtStats
 
 func (b *statementBuf) release() {
+	for i, n := 0, len(*b); i < n; i++ {
+		(*b)[i] = nil
+	}
 	*b = (*b)[:0]
 	statementsBufPool.Put(b)
 }
@@ -76,7 +81,7 @@ var statementsBufPool = sync.Pool{
 	},
 }
 
-// SQLStatsIngester buffers the "events" it sees (via ObserveStatement
+// SQLStatsIngester buffers the "events" it sees (via IngestStatement
 // and IngestTransaction) and passes them along to the underlying registry
 // once its buffer is full. (Or once a timeout has passed, for low-traffic
 // clusters and tests.)
@@ -184,8 +189,13 @@ func (i *SQLStatsIngester) ingest(ctx context.Context, events *eventBuffer) {
 		}
 		if e.statement != nil {
 			i.processStatement(e.statement)
+			// When under an outer transaction, we don't have a txn to associate
+			// the stmts with so we can send immediately to the sinks.
+			if e.statement.UnderOuterTxn {
+				i.flushBuffer(ctx, e.statement.SessionID, nil)
+			}
 		} else if e.transaction != nil {
-			i.flushBuffer(ctx, e.transaction)
+			i.flushBuffer(ctx, e.transaction.SessionID, e.transaction)
 		} else {
 			i.clearSession(e.sessionID)
 		}
@@ -293,11 +303,10 @@ func (i *SQLStatsIngester) processStatement(statement *sqlstats.RecordedStmtStat
 }
 
 // flushBuffer sends the buffered statementsBySessionID and provided transaction
-// to the registered sinks.
+// to the registered sinks. The transaction may be nil.
 func (i *SQLStatsIngester) flushBuffer(
-	ctx context.Context, transaction *sqlstats.RecordedTxnStats,
+	ctx context.Context, sessionID clusterunique.ID, transaction *sqlstats.RecordedTxnStats,
 ) {
-	sessionID := transaction.SessionID
 	statements, ok := func() (*statementBuf, bool) {
 		statements, ok := i.statementsBySessionID[sessionID]
 		if !ok {
@@ -315,9 +324,19 @@ func (i *SQLStatsIngester) flushBuffer(
 		return
 	}
 
-	// Set the transaction fingerprint ID for each statement.
-	for _, s := range *statements {
-		s.TransactionFingerprintID = transaction.FingerprintID
+	if transaction != nil {
+		// Set the transaction fingerprint ID for each statement.
+		// These values are only known at the time of the transaction.
+		for _, s := range *statements {
+			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)
+		}
 	}
 
 	for _, sink := range i.sinks {

Diff for: pkg/sql/sqlstats/sslocal/sql_stats_test.go

@@ -455,9 +455,7 @@ func TestExplicitTxnFingerprintAccounting(t *testing.T) {
 		nil, /* knobs */
 	)
 
-	// TODO(xinhaoz): We'll come back and add the sql stats sink once we
-	// enable the SQL stats ingestion for sql stats.
-	ingester := sslocal.NewSQLStatsIngester(nil /* testing knobs */)
+	ingester := sslocal.NewSQLStatsIngester(nil /* knobs */, sqlStats)
 
 	appStats := sqlStats.GetApplicationStats("" /* appName */)
 	statsCollector := sslocal.NewStatsCollector(
@@ -466,16 +464,14 @@ func TestExplicitTxnFingerprintAccounting(t *testing.T) {
 		ingester,
 		sessionphase.NewTimes(),
 		sqlStats.GetCounters(),
-		false, /* underOuterTxn */
-		nil,   /* knobs */
+		nil, /* knobs */
 	)
 
 	recordStats := func(testCase *tc) {
 		var txnFingerprintID appstatspb.TransactionFingerprintID
 		txnFingerprintIDHash := util.MakeFNV64()
 		statsCollector.StartTransaction()
 		defer func() {
-			statsCollector.EndTransaction(ctx, txnFingerprintID)
 			require.NoError(t,
 				statsCollector.RecordTransaction(ctx, &sqlstats.RecordedTxnStats{
 					FingerprintID:  txnFingerprintID,
@@ -586,8 +582,7 @@ func TestAssociatingStmtStatsWithTxnFingerprint(t *testing.T) {
 			ingester,
 			sessionphase.NewTimes(),
 			sqlStats.GetCounters(),
-			false, /* underOuterTxn */
-			nil,   /* knobs */
+			nil, /* knobs */
 		)
 
 		for _, txn := range simulatedTxns {
@@ -604,7 +599,6 @@ func TestAssociatingStmtStatsWithTxnFingerprint(t *testing.T) {
 			}
 
 			transactionFingerprintID := appstatspb.TransactionFingerprintID(txnFingerprintIDHash.Sum())
-			statsCollector.EndTransaction(ctx, transactionFingerprintID)
 			err := statsCollector.RecordTransaction(ctx, &sqlstats.RecordedTxnStats{
 				FingerprintID:  transactionFingerprintID,
 				UserNormalized: username.RootUser,