From f8f74599feeecc98f2b2e83554c37c4b938dc83e Mon Sep 17 00:00:00 2001
From: Thomas Hardy <thardy@cockroachlabs.com>
Date: Wed, 29 Jun 2022 10:28:50 -0400
Subject: [PATCH] sql: fix and rename sql stats session transaction received
 time

Resolves: #82894

Due to a change from #76792, implicit transactions can start before
`SessionQueryReceived` session phase time is set by the sqlstats system.
In turn, this caused the `SessionTransactionReceived` (now renamed as
`SessionTransactionStarted`) session phase time to be recorded
incorrectly, causing extremely large transactions times on the UI. This
change fixes this mistake by setting the actual transaction start time
as the `SessionTransactionStarted` session phase time, instead of
`SessionQueryReceived`.

Release note (bug fix): The `SessionTransactionReceived` session phase
time is no longer recorded incorrectly, fixing large transaction times
from appearing on the UI, also renamed to `SessionTransactionStarted`.
---
 pkg/sql/conn_executor_exec.go              |  9 ++--
 pkg/sql/exec_util.go                       |  5 ++
 pkg/sql/sessionphase/session_phase.go      |  8 +--
 pkg/sql/sqlstats/sslocal/BUILD.bazel       |  1 +
 pkg/sql/sqlstats/sslocal/sql_stats_test.go | 61 ++++++++++++++++++++++
 5 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
index 986162b7be74..5a90906e9794 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/pkg/sql/conn_executor_exec.go
@@ -2182,10 +2182,7 @@ func (ex *connExecutor) recordTransactionStart(txnID uuid.UUID) {
 	ex.state.mu.RUnlock()
 	implicit := ex.implicitTxn()
 
-	// Transaction received time is the time at which the statement that prompted
-	// the creation of this transaction was received.
-	ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionTransactionReceived,
-		ex.phaseTimes.GetSessionPhaseTime(sessionphase.SessionQueryReceived))
+	ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionTransactionStarted, txnStart)
 	ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionFirstStartExecTransaction, timeutil.Now())
 	ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionMostRecentStartExecTransaction,
 		ex.phaseTimes.GetSessionPhaseTime(sessionphase.SessionFirstStartExecTransaction))
@@ -2281,6 +2278,10 @@ func (ex *connExecutor) recordTransactionFinish(
 		BytesRead:               ex.extraTxnState.bytesRead,
 	}
 
+	if ex.server.cfg.TestingKnobs.OnRecordTxnFinish != nil {
+		ex.server.cfg.TestingKnobs.OnRecordTxnFinish(ex.executorType == executorTypeInternal, ex.phaseTimes, ex.planner.stmt.SQL)
+	}
+
 	return ex.statsCollector.RecordTransaction(
 		ctx,
 		transactionFingerprintID,
diff --git a/pkg/sql/exec_util.go b/pkg/sql/exec_util.go
index 485fa7a0306a..fc519e953a0c 100644
--- a/pkg/sql/exec_util.go
+++ b/pkg/sql/exec_util.go
@@ -86,6 +86,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessioninit"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessionphase"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
@@ -1462,6 +1463,10 @@ type ExecutorTestingKnobs struct {
 	// AfterBackupCheckpoint if set will be called after a BACKUP-CHECKPOINT
 	// is written.
 	AfterBackupCheckpoint func()
+
+	// OnRecordTxnFinish, if set, will be called as we record a transaction
+	// finishing.
+	OnRecordTxnFinish func(isInternal bool, phaseTimes *sessionphase.Times, stmt string)
 }
 
 // PGWireTestingKnobs contains knobs for the pgwire module.
diff --git a/pkg/sql/sessionphase/session_phase.go b/pkg/sql/sessionphase/session_phase.go
index daf5d78e3920..0218003d164a 100644
--- a/pkg/sql/sessionphase/session_phase.go
+++ b/pkg/sql/sessionphase/session_phase.go
@@ -59,9 +59,9 @@ const (
 	// have no execution, like SHOW TRANSACTION STATUS.
 	SessionQueryServiced
 
-	// SessionTransactionReceived is the SessionPhase when a transaction is
-	// received.
-	SessionTransactionReceived
+	// SessionTransactionStarted is the SessionPhase when a transaction is
+	// started.
+	SessionTransactionStarted
 
 	// SessionFirstStartExecTransaction is the SessionPhase when a transaction
 	// is started for the first time.
@@ -197,7 +197,7 @@ func (t *Times) GetTransactionRetryLatency() time.Duration {
 // GetTransactionServiceLatency returns the total time to service the
 // transaction.
 func (t *Times) GetTransactionServiceLatency() time.Duration {
-	return t.times[SessionEndExecTransaction].Sub(t.times[SessionTransactionReceived])
+	return t.times[SessionEndExecTransaction].Sub(t.times[SessionTransactionStarted])
 }
 
 // GetCommitLatency returns the total time spent for the transaction to
diff --git a/pkg/sql/sqlstats/sslocal/BUILD.bazel b/pkg/sql/sqlstats/sslocal/BUILD.bazel
index 7e8144f61697..7d3fffa5b144 100644
--- a/pkg/sql/sqlstats/sslocal/BUILD.bazel
+++ b/pkg/sql/sqlstats/sslocal/BUILD.bazel
@@ -69,6 +69,7 @@ go_test(
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "//pkg/util/mon",
+        "@com_github_jackc_pgx_v4//:pgx",
         "@com_github_stretchr_testify//require",
     ],
 )
diff --git a/pkg/sql/sqlstats/sslocal/sql_stats_test.go b/pkg/sql/sqlstats/sslocal/sql_stats_test.go
index cf6fd11d8b89..0617e8582d63 100644
--- a/pkg/sql/sqlstats/sslocal/sql_stats_test.go
+++ b/pkg/sql/sqlstats/sslocal/sql_stats_test.go
@@ -13,7 +13,9 @@ package sslocal_test
 import (
 	"context"
 	"math"
+	"net/url"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -37,6 +39,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/jackc/pgx/v4"
 	"github.com/stretchr/testify/require"
 )
 
@@ -672,3 +675,61 @@ func TestUnprivilegedUserReset(t *testing.T) {
 
 	require.Contains(t, err.Error(), "requires admin privilege")
 }
+
+func TestTransactionServiceLatencyOnExtendedProtocol(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	testData := []*struct {
+		query        string
+		placeholders []interface{}
+		phaseTimes   *sessionphase.Times
+	}{
+		{
+			query:        "SELECT $1::INT8",
+			placeholders: []interface{}{1},
+			phaseTimes:   nil,
+		},
+	}
+
+	waitTxnFinish := make(chan struct{})
+	currentTestCaseIdx := 0
+	const latencyThreshold = time.Second * 5
+
+	params, _ := tests.CreateTestServerParams()
+	params.Knobs.SQLExecutor = &sql.ExecutorTestingKnobs{
+		OnRecordTxnFinish: func(isInternal bool, phaseTimes *sessionphase.Times, stmt string) {
+			if !isInternal && testData[currentTestCaseIdx].query == stmt {
+				testData[currentTestCaseIdx].phaseTimes = phaseTimes.Clone()
+				go func() {
+					waitTxnFinish <- struct{}{}
+				}()
+			}
+		},
+	}
+	s, _, _ := serverutils.StartServer(t, params)
+	defer s.Stopper().Stop(ctx)
+
+	pgURL, cleanupGoDB := sqlutils.PGUrl(
+		t, s.ServingSQLAddr(), "StartServer", url.User(username.RootUser))
+	defer cleanupGoDB()
+	c, err := pgx.Connect(ctx, pgURL.String())
+	require.NoError(t, err, "error connecting with pg url")
+
+	for currentTestCaseIdx < len(testData) {
+		tc := testData[currentTestCaseIdx]
+		// Make extended protocol query
+		_ = c.QueryRow(ctx, tc.query, tc.placeholders...)
+		require.NoError(t, err, "error scanning row")
+		<-waitTxnFinish
+
+		// Ensure test case phase times are populated by query txn.
+		require.True(t, tc.phaseTimes != nil)
+		// Ensure SessionTransactionStarted variable is populated.
+		require.True(t, !tc.phaseTimes.GetSessionPhaseTime(sessionphase.SessionTransactionStarted).IsZero())
+		// Ensure compute transaction service latency is within a reasonable threshold.
+		require.True(t, tc.phaseTimes.GetTransactionServiceLatency() < latencyThreshold)
+		currentTestCaseIdx++
+	}
+}