sql: collect exec stats for a statement if we see it for the first time

Previously, we would decide whether we collect the execution stats for
a statement based on `sql.txn_stats.sample_rate` that will determine the
stats collection on a txn basis. This commit additionally begins
collecting the stats on a per stmt basis if it is the first time we see
the stmt. We are careful not to populate txn stats for explicit txns if
the exec stats collection wasn't enabled on a txn level so that we don't
provide an incomplete picture.

Release note (sql change): CockroachDB will now collect execution
stats for all statements when seen for the first time. In order to
disable this behavior, `sql.txn_stats.sample_rate` cluster setting needs
to be set to 0 (this will disable all execution stats collection).

pkg/sql/app_stats.go

@@ -481,18 +481,15 @@ func (a *appStats) recordTransaction(
 	}
 }
 
-// shouldSaveLogicalPlanDescription returns whether we should save this as a
-// sample logical plan for its corresponding fingerprint. We use
-// `logicalPlanCollectionPeriod` to assess how frequently to sample logical
-// plans.
-func (a *appStats) shouldSaveLogicalPlanDescription(anonymizedStmt string, implicitTxn bool) bool {
+// shouldSaveLogicalPlanDescription returns whether we should save the sample
+// logical plan for a fingerprint (represented implicitly by the corresponding
+// stmtStats object). stats is nil if it is the first time we see the
+// fingerprint. We use `logicalPlanCollectionPeriod` to assess how frequently to
+// sample logical plans.
+func (a *appStats) shouldSaveLogicalPlanDescription(stats *stmtStats) bool {
 	if !sampleLogicalPlans.Get(&a.st.SV) {
 		return false
 	}
-	// We don't know yet if we will hit an error, so we assume we don't. The worst
-	// that can happen is that for statements that always error out, we will
-	// always save the tree plan.
-	stats, _ := a.getStatsForStmt(anonymizedStmt, implicitTxn, nil /* error */, false /* createIfNonexistent */)
 	if stats == nil {
 		// Save logical plan the first time we see new statement fingerprint.
 		return true

pkg/sql/conn_executor.go

@@ -1093,9 +1093,9 @@ type connExecutor struct {
 
 		schemaChangerState SchemaChangerState
 
-		// shouldCollectExecutionStats specifies whether the statements in this
-		// transaction should collect execution stats.
-		shouldCollectExecutionStats bool
+		// shouldCollectTxnExecutionStats specifies whether the statements in
+		// this transaction should collect execution stats.
+		shouldCollectTxnExecutionStats bool
 		// accumulatedStats are the accumulated stats of all statements.
 		accumulatedStats execstats.QueryLevelStats
 		// rowsRead and bytesRead are separate from QueryLevelStats because they are

pkg/sql/conn_executor_exec.go

@@ -390,7 +390,7 @@ func (ex *connExecutor) execStmtInOpenState(
 	var needFinish bool
 	ctx, needFinish = ih.Setup(
 		ctx, ex.server.cfg, ex.appStats, p, ex.stmtDiagnosticsRecorder,
-		stmt.AnonymizedStr, os.ImplicitTxn.Get(), ex.extraTxnState.shouldCollectExecutionStats,
+		stmt.AnonymizedStr, os.ImplicitTxn.Get(), ex.extraTxnState.shouldCollectTxnExecutionStats,
 	)
 	if needFinish {
 		sql := stmt.SQL
@@ -399,6 +399,7 @@ func (ex *connExecutor) execStmtInOpenState(
 				ex.server.cfg,
 				ex.appStats,
 				&ex.extraTxnState.accumulatedStats,
+				ex.extraTxnState.shouldCollectTxnExecutionStats,
 				ex.statsCollector,
 				p,
 				ast,
@@ -1492,12 +1493,12 @@ func (ex *connExecutor) recordTransactionStart() (onTxnFinish func(txnEvent), on
 	ex.extraTxnState.transactionStatementsHash = util.MakeFNV64()
 	ex.extraTxnState.transactionStatementIDs = nil
 	ex.extraTxnState.numRows = 0
-	ex.extraTxnState.shouldCollectExecutionStats = false
+	ex.extraTxnState.shouldCollectTxnExecutionStats = false
 	ex.extraTxnState.accumulatedStats = execstats.QueryLevelStats{}
 	ex.extraTxnState.rowsRead = 0
 	ex.extraTxnState.bytesRead = 0
-	if execStatsSampleRate := collectTxnStatsSampleRate.Get(&ex.server.GetExecutorConfig().Settings.SV); execStatsSampleRate > 0 {
-		ex.extraTxnState.shouldCollectExecutionStats = execStatsSampleRate > ex.rng.Float64()
+	if txnExecStatsSampleRate := collectTxnStatsSampleRate.Get(&ex.server.GetExecutorConfig().Settings.SV); txnExecStatsSampleRate > 0 {
+		ex.extraTxnState.shouldCollectTxnExecutionStats = txnExecStatsSampleRate > ex.rng.Float64()
 	}
 
 	ex.metrics.EngineMetrics.SQLTxnsOpen.Inc(1)
@@ -1511,7 +1512,7 @@ func (ex *connExecutor) recordTransactionStart() (onTxnFinish func(txnEvent), on
 		ex.extraTxnState.transactionStatementIDs = nil
 		ex.extraTxnState.transactionStatementsHash = util.MakeFNV64()
 		ex.extraTxnState.numRows = 0
-		// accumulatedStats are cleared, but shouldCollectExecutionStats is
+		// accumulatedStats are cleared, but shouldCollectTxnExecutionStats is
 		// unchanged.
 		ex.extraTxnState.accumulatedStats = execstats.QueryLevelStats{}
 		ex.extraTxnState.rowsRead = 0
@@ -1541,7 +1542,7 @@ func (ex *connExecutor) recordTransaction(ev txnEvent, implicit bool, txnStart t
 		txnRetryLat,
 		commitLat,
 		ex.extraTxnState.numRows,
-		ex.extraTxnState.shouldCollectExecutionStats,
+		ex.extraTxnState.shouldCollectTxnExecutionStats,
 		ex.extraTxnState.accumulatedStats,
 		ex.extraTxnState.rowsRead,
 		ex.extraTxnState.bytesRead,

pkg/sql/instrumentation.go

@@ -137,7 +137,7 @@ func (ih *instrumentationHelper) Setup(
 	stmtDiagnosticsRecorder *stmtdiagnostics.Registry,
 	fingerprint string,
 	implicitTxn bool,
-	collectExecStats bool,
+	collectTxnExecStats bool,
 ) (newCtx context.Context, needFinish bool) {
 	ih.fingerprint = fingerprint
 	ih.implicitTxn = implicitTxn
@@ -162,7 +162,12 @@ func (ih *instrumentationHelper) Setup(
 
 	ih.withStatementTrace = cfg.TestingKnobs.WithStatementTrace
 
-	ih.savePlanForStats = appStats.shouldSaveLogicalPlanDescription(fingerprint, implicitTxn)
+	// We don't know yet if we will hit an error, so we assume we don't. The
+	// worst that can happen is that for statements that always error out, we
+	// will always save the tree plan.
+	stats, _ := appStats.getStatsForStmt(fingerprint, implicitTxn, nil /* error */, false /* createIfNonexistent */)
+
+	ih.savePlanForStats = appStats.shouldSaveLogicalPlanDescription(stats)
 
 	if sp := tracing.SpanFromContext(ctx); sp != nil {
 		if sp.IsVerbose() {
@@ -178,7 +183,15 @@ func (ih *instrumentationHelper) Setup(
 		}
 	}
 
-	ih.collectExecStats = collectExecStats
+	ih.collectExecStats = collectTxnExecStats
+	if !collectTxnExecStats && stats == nil {
+		// We don't collect the execution stats for statements in this txn, but
+		// this is the first time we see this statement ever, so we'll collect
+		// its execution stats anyway (unless the user disabled txn stats
+		// collection entirely).
+		statsCollectionDisabled := collectTxnStatsSampleRate.Get(&cfg.Settings.SV) == 0
+		ih.collectExecStats = !statsCollectionDisabled
+	}
 
 	if !ih.collectBundle && ih.withStatementTrace == nil && ih.outputMode == unmodifiedOutput {
 		if ih.collectExecStats {
@@ -203,6 +216,7 @@ func (ih *instrumentationHelper) Finish(
 	cfg *ExecutorConfig,
 	appStats *appStats,
 	txnStats *execstats.QueryLevelStats,
+	collectTxnExecStats bool,
 	statsCollector *sqlStatsCollector,
 	p *planner,
 	ast tree.Statement,
@@ -250,7 +264,16 @@ func (ih *instrumentationHelper) Finish(
 		stmtStats, _ := appStats.getStatsForStmt(ih.fingerprint, ih.implicitTxn, retErr, false)
 		if stmtStats != nil {
 			stmtStats.recordExecStats(queryLevelStats)
-			txnStats.Accumulate(queryLevelStats)
+			if collectTxnExecStats || ih.implicitTxn {
+				// Only accumulate into the txn stats if we're collecting
+				// execution stats for all statements in the txn or the txn is
+				// implicit (indicating that it consists of a single statement).
+				//
+				// The goal is that we don't want to provide an incomplete
+				// picture for explicit txns for which we didn't decide to
+				// collect the execution stats on a txn level.
+				txnStats.Accumulate(queryLevelStats)
+			}
 		}
 	}