Merge cockroachdb#138041

138041: sql: do not consider TCL stmts in sql stats tracing r=xinhaoz a=xinhaoz

For SQL stats we trace every new fingerprint in an application
container in order to populate statistics that are only available
via tracing for each fingerprint. This sampling logic worked by
tracking the SQL fingerprint strings encountered by each application
but did not factor in fingerprints that are not tracked by SQL stats,
such as TCL statements (BEGIN, COMMIT, ROLLBACK), were being treated
as new fingerprints on each execution which resulted in tracing being
turned on. This commit ensures that we don't consider TCL statements
in this sampling strategy. TestSampledStatsCollectionOnNewFingerprint
is also fixed as it was assuming that previously seen statements in
new transactions would trigger this sampling behaviour when it should
not - this was being erronneously triggered by `BEGIN`.

```
name                                   old time/op    new time/op    delta
Sysbench/SQL/3node/oltp_read_write-10    5.61ms ± 3%    5.69ms ± 4%    ~     (p=0.167 n=8+9)

name                                   old alloc/op   new alloc/op   delta
Sysbench/SQL/3node/oltp_read_write-10    2.15MB ± 2%    2.12MB ± 2%  -1.24%  (p=0.029 n=10+10)

name                                   old allocs/op  new allocs/op  delta
Sysbench/SQL/3node/oltp_read_write-10     10.7k ± 1%     10.5k ± 0%  -1.69%  (p=0.000 n=10+9)
```

Epic: none
Part of: cockroachdb#133307

Release note: None

Co-authored-by: Xin Hao Zhang <[email protected]>

pkg/sql/conn_executor_exec.go

@@ -508,7 +508,7 @@ func (ex *connExecutor) execStmtInOpenState(
 
 	ctx = ih.Setup(
 		ctx, ex.server.cfg, ex.statsCollector, p, ex.stmtDiagnosticsRecorder,
-		stmt.StmtNoConstants, os.ImplicitTxn.Get(),
+		&stmt, os.ImplicitTxn.Get(),
 		// This goroutine is the only one that can modify
 		// txnState.mu.priority, so we don't need to get a mutex here.
 		ex.state.mu.priority,
@@ -1462,7 +1462,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 	if !portal.isPausable() || portal.pauseInfo.execStmtInOpenState.ihWrapper == nil {
 		ctx = ih.Setup(
 			ctx, ex.server.cfg, ex.statsCollector, p, ex.stmtDiagnosticsRecorder,
-			vars.stmt.StmtNoConstants, os.ImplicitTxn.Get(),
+			&vars.stmt, os.ImplicitTxn.Get(),
 			// This goroutine is the only one that can modify
 			// txnState.mu.priority, so we don't need to get a mutex here.
 			ex.state.mu.priority,

pkg/sql/instrumentation.go

@@ -418,12 +418,12 @@ func (ih *instrumentationHelper) Setup(
 	statsCollector *sslocal.StatsCollector,
 	p *planner,
 	stmtDiagnosticsRecorder *stmtdiagnostics.Registry,
-	fingerprint string,
+	stmt *Statement,
 	implicitTxn bool,
 	txnPriority roachpb.UserPriority,
 	collectTxnExecStats bool,
 ) (newCtx context.Context) {
-	ih.fingerprint = fingerprint
+	ih.fingerprint = stmt.StmtNoConstants
 	ih.implicitTxn = implicitTxn
 	ih.txnPriority = txnPriority
 	ih.codec = cfg.Codec
@@ -447,7 +447,7 @@ func (ih *instrumentationHelper) Setup(
 
 	default:
 		ih.collectBundle, ih.diagRequestID, ih.diagRequest =
-			stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, fingerprint, "" /* planGist */)
+			stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, stmt.StmtNoConstants, "" /* planGist */)
 		// IsRedacted will be false when ih.collectBundle is false.
 		ih.explainFlags.RedactValues = ih.explainFlags.RedactValues || ih.diagRequest.IsRedacted()
 	}
@@ -456,7 +456,7 @@ func (ih *instrumentationHelper) Setup(
 	ih.withStatementTrace = cfg.TestingKnobs.WithStatementTrace
 
 	var previouslySampled bool
-	previouslySampled, ih.savePlanForStats = statsCollector.ShouldSample(fingerprint, implicitTxn, p.SessionData().Database)
+	previouslySampled, ih.savePlanForStats = statsCollector.ShouldSample(stmt.StmtNoConstants, implicitTxn, p.SessionData().Database)
 
 	defer func() { ih.finalizeSetup(newCtx, cfg) }()
 
@@ -481,6 +481,12 @@ func (ih *instrumentationHelper) Setup(
 	}
 
 	shouldSampleFirstEncounter := func() bool {
+		if stmt.AST.StatementType() == tree.TypeTCL {
+			// We don't collect stats for TCL statements so
+			// there's no need to trace them.
+			return false
+		}
+
 		// If this is the first time we see this statement in the current stats
 		// container, we'll collect its execution stats anyway (unless the user
 		// disabled txn or stmt stats collection entirely).

pkg/sql/instrumentation_test.go

@@ -241,7 +241,8 @@ func TestSampledStatsCollectionOnNewFingerprint(t *testing.T) {
 			"SELECT 1, 2, 3",
 			"CREATE TABLE IF NOT EXISTS foo (x INT)",
 			"SELECT * FROM foo",
-			// An explicit txn results in the queries inside being recorded as 'new' statements.
+			// Since the sampling key does not include the txn fingerprint, no
+			// statements in this txn should be sampled.
 			"BEGIN; SELECT 1; COMMIT;",
 		}
 
@@ -251,10 +252,11 @@ func TestSampledStatsCollectionOnNewFingerprint(t *testing.T) {
 
 		require.Equal(t, len(queries), len(collectedTxnStats))
 
-		// We should have collected stats for each of the queries.
-		for i := range collectedTxnStats {
+		// We should have collected stats for each of the queries except the last.
+		for i := range collectedTxnStats[:len(queries)-1] {
 			require.True(t, collectedTxnStats[i].CollectedExecStats)
 		}
+		require.False(t, collectedTxnStats[len(queries)-1].CollectedExecStats)
 	})
 
 	collectedTxnStats = nil