From b3e2b6836d1286aed677c751f4a696905abb5cea Mon Sep 17 00:00:00 2001
From: Azhng <archer.xn@gmail.com>
Date: Thu, 31 Mar 2022 22:01:28 +0000
Subject: [PATCH] sql: collect contention stats for sql stats when tracing is
 enabled

Resolves #78675

Previously, contention stats was not collected for SQL Stats even when
tracing was enabled. This was caused by two bugs:
1. instrumentationHelper skipping analyzing traces when tracing is enabled
2. transaction statistics ignore traces when tracing is turned on at the
higher level.

This commit ensures that contention stats is collected when tracing is
turned on.

Release note (bug fix): Contention statistics are now being collected for
SQL Stats when tracing is enabled.
---
 pkg/server/status_test.go     | 17 +++++++++++++++++
 pkg/sql/conn_executor_exec.go |  4 ++--
 pkg/sql/instrumentation.go    | 13 +++++++------
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/pkg/server/status_test.go b/pkg/server/status_test.go
index c72abb5fed66..342954164ba1 100644
--- a/pkg/server/status_test.go
+++ b/pkg/server/status_test.go
@@ -3120,6 +3120,7 @@ func TestStatusAPIContentionEvents(t *testing.T) {
 
 	server1Conn.Exec(t, "USE test")
 	server2Conn.Exec(t, "USE test")
+	server2Conn.Exec(t, "SET application_name = 'contentionTest'")
 
 	server1Conn.Exec(t, `
 SET TRACING=on;
@@ -3164,6 +3165,22 @@ SET TRACING=off;
 
 	require.True(t, found,
 		"expect to find contention event for table %d, but found %+v", testTableID, resp)
+
+	server1Conn.CheckQueryResults(t, `
+  SELECT count(*)
+  FROM crdb_internal.statement_statistics
+  WHERE
+    (statistics -> 'execution_statistics' -> 'contentionTime' ->> 'mean')::FLOAT > 0
+    AND app_name = 'contentionTest'
+`, [][]string{{"1"}})
+
+	server1Conn.CheckQueryResults(t, `
+  SELECT count(*)
+  FROM crdb_internal.transaction_statistics
+  WHERE
+    (statistics -> 'execution_statistics' -> 'contentionTime' ->> 'mean')::FLOAT > 0
+    AND app_name = 'contentionTest'
+`, [][]string{{"1"}})
 }
 
 func TestStatusCancelSessionGatewayMetadataPropagation(t *testing.T) {
diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
index 4eba91d0524b..ece2b6275ecc 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/pkg/sql/conn_executor_exec.go
@@ -470,7 +470,7 @@ func (ex *connExecutor) execStmtInOpenState(
 				ex.server.cfg,
 				ex.statsCollector,
 				&ex.extraTxnState.accumulatedStats,
-				ex.extraTxnState.shouldCollectTxnExecutionStats,
+				ih.collectExecStats,
 				p,
 				ast,
 				sql,
@@ -2241,7 +2241,7 @@ func (ex *connExecutor) recordTransactionFinish(
 		RetryLatency:            txnRetryLat,
 		CommitLatency:           commitLat,
 		RowsAffected:            ex.extraTxnState.numRows,
-		CollectedExecStats:      ex.extraTxnState.shouldCollectTxnExecutionStats,
+		CollectedExecStats:      ex.planner.instrumentation.collectExecStats,
 		ExecStats:               ex.extraTxnState.accumulatedStats,
 		RowsRead:                ex.extraTxnState.rowsRead,
 		RowsWritten:             ex.extraTxnState.rowsWritten,
diff --git a/pkg/sql/instrumentation.go b/pkg/sql/instrumentation.go
index d71f6c36403b..409d8fea5fc7 100644
--- a/pkg/sql/instrumentation.go
+++ b/pkg/sql/instrumentation.go
@@ -196,11 +196,12 @@ func (ih *instrumentationHelper) Setup(
 			// collection is enabled so that stats are shown in the traces, but
 			// no extra work is needed by the instrumentationHelper.
 			ih.collectExecStats = true
-			// We still, however, want to finish the instrumentationHelper in
-			// case we're collecting a bundle. We also capture the span in order
-			// to fetch the trace from it, but the span won't be finished.
+			// We always want to finish the instrumentationHelper in order
+			// to record the execution statistics. Note that we capture the
+			// span in order to fetch the trace from it, but the span won't be
+			// finished.
 			ih.sp = sp
-			return ctx, ih.collectBundle
+			return ctx, true /* needFinish */
 		}
 	} else {
 		if buildutil.CrdbTestBuild {
@@ -244,7 +245,7 @@ func (ih *instrumentationHelper) Finish(
 	cfg *ExecutorConfig,
 	statsCollector sqlstats.StatsCollector,
 	txnStats *execstats.QueryLevelStats,
-	collectTxnExecStats bool,
+	collectExecStats bool,
 	p *planner,
 	ast tree.Statement,
 	stmtRawSQL string,
@@ -314,7 +315,7 @@ func (ih *instrumentationHelper) Finish(
 				log.Warningf(ctx, "unable to record statement exec stats: %s", err)
 			}
 		}
-		if collectTxnExecStats || ih.implicitTxn {
+		if collectExecStats || ih.implicitTxn {
 			txnStats.Accumulate(queryLevelStats)
 		}
 	}