sql: add txn diagnostics collection to instrumentation helper

When calling the setup method on instrumentationhelper (ih), ih will
now check if transaction diagnostics should should be collected. Transaction
diagnostics collection uses a prefix based matching system to determine
whether or not transaction diagnostics collection should start, continue,
or stop. This works by creating a list of statement fingerprint ids, that make
up a transaction fingerprint, and comparing the currently executing statement
fingerprint id against this list. If at any point the expected statement fingerprint
id doesn't match, diagnostics collection is short circuited and no longer
collected.

This change also causes transaction diagnostics collection to take precedence
over statement transaction diagnostics, meaning if a transaction diagnostics request
exists that contains a statement that also has a statement diagnostics request,
the transaction diagnostic request gets priority over collecting the bundle.

When a transaction diagnostics request is being fulfilled, statement bundles for
every statement in the transaction are built and collected in memory. Only after
a transaction is finalized and it is determined that the request is fulfilled are all
the statement bundles and the transaction trace stored in the corresponding
system tables.

Resolves: CRDB-54321
Epic: CRDB-53541
Release note: None

Author: kyle-a-wong

pkg/sql/conn_executor.go

@@ -1171,6 +1171,7 @@ func (s *Server) newConnExecutor(
 			connCtx:                      ctx,
 			testingForceRealTracingSpans: s.cfg.TestingKnobs.ForceRealTracingSpans,
 			execType:                     executorType,
+			txnInstrumentationHelper:     txnInstrumentationHelper{TxnDiagnosticsRecorder: s.cfg.TxnDiagnosticsRecorder},
 		},
 		transitionCtx: transitionCtx{
 			db:           s.cfg.DB,
@@ -1195,6 +1196,7 @@ func (s *Server) newConnExecutor(
 		executorType:              executorType,
 		hasCreatedTemporarySchema: false,
 		stmtDiagnosticsRecorder:   s.cfg.StmtDiagnosticsRecorder,
+		txnDiagnosticsRecorder:    s.cfg.TxnDiagnosticsRecorder,
 		indexUsageStats:           s.indexUsageStats,
 		txnIDCacheWriter:          s.txnIDCache,
 		totalActiveTimeStopWatch:  timeutil.NewStopWatch(),
@@ -1871,6 +1873,7 @@ type connExecutor struct {
 	// stmtDiagnosticsRecorder is used to track which queries need to have
 	// information collected.
 	stmtDiagnosticsRecorder *stmtdiagnostics.Registry
+	txnDiagnosticsRecorder  *stmtdiagnostics.TxnRegistry
 
 	// indexUsageStats is used to track index usage stats.
 	indexUsageStats *idxusage.LocalIndexUsageStats

pkg/sql/conn_executor_exec.go

@@ -565,13 +565,8 @@ func (ex *connExecutor) execStmtInOpenState(
 
 	// This goroutine is the only one that can modify txnState.mu.priority and
 	// txnState.mu.autoRetryCounter, so we don't need to get a mutex here.
-	ctx = ih.Setup(
-		ctx, ex.server.cfg, ex.statsCollector, p, ex.stmtDiagnosticsRecorder,
-		&stmt, os.ImplicitTxn.Get(),
-		ex.state.mu.priority,
-		ex.extraTxnState.shouldCollectTxnExecutionStats,
-		ex.state.mu.autoRetryCounter,
-	)
+	ctx = ih.Setup(ctx, ex, p, &stmt, os.ImplicitTxn.Get(),
+		ex.state.mu.priority, ex.state.mu.autoRetryCounter)
 
 	// Note that here we always unconditionally defer a function that takes care
 	// of finishing the instrumentation helper. This is needed since in order to
@@ -580,9 +575,7 @@ func (ex *connExecutor) execStmtInOpenState(
 	defer func() {
 		if ih.needFinish {
 			retErr = ih.Finish(
-				ex.server.cfg,
-				ex.statsCollector,
-				&ex.extraTxnState.accumulatedStats,
+				ex,
 				ih.collectExecStats,
 				p,
 				ast,
@@ -1462,13 +1455,8 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 	// This goroutine is the only one that can modify txnState.mu.priority and
 	// txnState.mu.autoRetryCounter, so we don't need to get a mutex here.
 	if !portal.isPausable() || portal.pauseInfo.execStmtInOpenState.ihWrapper == nil {
-		ctx = ih.Setup(
-			ctx, ex.server.cfg, ex.statsCollector, p, ex.stmtDiagnosticsRecorder,
-			&vars.stmt, os.ImplicitTxn.Get(),
-			ex.state.mu.priority,
-			ex.extraTxnState.shouldCollectTxnExecutionStats,
-			ex.state.mu.autoRetryCounter,
-		)
+		ctx = ih.Setup(ctx, ex, p, &vars.stmt, os.ImplicitTxn.Get(),
+			ex.state.mu.priority, ex.state.mu.autoRetryCounter)
 	} else {
 		ctx = portal.pauseInfo.execStmtInOpenState.ihWrapper.ctx
 	}
@@ -1509,9 +1497,7 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 		}
 		if ihToFinish.needFinish {
 			retErr = ihToFinish.Finish(
-				ex.server.cfg,
-				ex.statsCollector,
-				&ex.extraTxnState.accumulatedStats,
+				ex,
 				ihToFinish.collectExecStats,
 				p,
 				vars.ast,

pkg/sql/conn_executor_internal_test.go

@@ -39,6 +39,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/stretchr/testify/require"
 )
@@ -293,6 +294,7 @@ func startConnExecutor(
 	// This pool should never be Stop()ed because, if the test is failing, memory
 	// is not properly released.
 	collectionFactory := descs.NewBareBonesCollectionFactory(st, keys.SystemSQLCodec)
+	registry := stmtdiagnostics.NewRegistry(nil, st)
 	cfg := &ExecutorConfig{
 		AmbientCtx: ambientCtx,
 		Settings:   st,
@@ -339,7 +341,8 @@ func startConnExecutor(
 		),
 		QueryCache:              querycache.New(0),
 		TestingKnobs:            ExecutorTestingKnobs{},
-		StmtDiagnosticsRecorder: stmtdiagnostics.NewRegistry(nil, st),
+		StmtDiagnosticsRecorder: registry,
+		TxnDiagnosticsRecorder:  stmtdiagnostics.NewTxnRegistry(nil, st, registry, timeutil.DefaultTimeSource{}),
 		HistogramWindowInterval: base.DefaultHistogramWindowInterval(),
 		CollectionFactory:       collectionFactory,
 		LicenseEnforcer:         license.NewEnforcer(nil),

pkg/sql/instrumentation.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
@@ -34,7 +35,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessionphase"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/sslocal"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/stmtdiagnostics"
 	"github.com/cockroachdb/cockroach/pkg/util"
@@ -46,6 +46,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 	"github.com/cockroachdb/errors"
+	"go.opentelemetry.io/otel/attribute"
 )
 
 var collectTxnStatsSampleRate = settings.RegisterFloatSetting(
@@ -77,6 +78,8 @@ type instrumentationHelper struct {
 	// Query fingerprint (anonymized statement).
 	fingerprint string
 
+	fingerprintId appstatspb.StmtFingerprintID
+
 	// Transaction information.
 	implicitTxn bool
 	txnPriority roachpb.UserPriority
@@ -410,16 +413,18 @@ func (ih *instrumentationHelper) finalizeSetup(ctx context.Context, cfg *Executo
 // which case Finish() is a no-op).
 func (ih *instrumentationHelper) Setup(
 	ctx context.Context,
-	cfg *ExecutorConfig,
-	statsCollector *sslocal.StatsCollector,
+	ex *connExecutor,
 	p *planner,
-	stmtDiagnosticsRecorder *stmtdiagnostics.Registry,
 	stmt *Statement,
 	implicitTxn bool,
 	txnPriority roachpb.UserPriority,
-	collectTxnExecStats bool,
 	retryCount int32,
 ) (newCtx context.Context) {
+	cfg := ex.server.cfg
+	statsCollector := ex.statsCollector
+	stmtDiagnosticsRecorder := ex.stmtDiagnosticsRecorder
+	collectTxnExecStats := ex.extraTxnState.shouldCollectTxnExecutionStats
+
 	ih.fingerprint = stmt.StmtNoConstants
 	ih.implicitTxn = implicitTxn
 	ih.txnPriority = txnPriority
@@ -431,6 +436,12 @@ func (ih *instrumentationHelper) Setup(
 	ih.isTenant = execinfra.IncludeRUEstimateInExplainAnalyze.Get(cfg.SV()) && cfg.DistSQLSrv != nil &&
 		cfg.DistSQLSrv.TenantCostController != nil
 	ih.topLevelStats = topLevelQueryStats{}
+	stmtFingerprintId := appstatspb.ConstructStatementFingerprintID(
+		stmt.StmtNoConstants, implicitTxn, p.SessionData().Database)
+	ih.fingerprintId = stmtFingerprintId
+	ih.stmtDiagnosticsRecorder = stmtDiagnosticsRecorder
+	ih.withStatementTrace = cfg.TestingKnobs.WithStatementTrace
+	defer func() { ih.finalizeSetup(newCtx, cfg) }()
 
 	switch ih.outputMode {
 	case explainAnalyzeDebugOutput:
@@ -445,17 +456,28 @@ func (ih *instrumentationHelper) Setup(
 		ih.discardRows = true
 
 	default:
-		ih.collectBundle, ih.diagRequestID, ih.diagRequest =
-			stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, stmt.StmtNoConstants, "" /* planGist */)
-		// IsRedacted will be false when ih.collectBundle is false.
-		ih.explainFlags.RedactValues = ih.explainFlags.RedactValues || ih.diagRequest.IsRedacted()
+		// Handle transaction-level diagnostics
+		var collectingDiagnostics bool
+		if ctx, collectingDiagnostics = ih.handleTransactionDiagnostics(
+			ctx,
+			&ex.state,
+			ex.transitionCtx.tracer,
+			stmt,
+			stmtFingerprintId,
+		); collectingDiagnostics {
+			// If collectingDiagnostics is true, the instrumentationHelper should
+			// already be set up for tracing, so we can return.
+			return ctx
+		} else {
+			// If no transaction diagnostics are in progress, check for statement-level
+			// diagnostics
+			ih.collectBundle, ih.diagRequestID, ih.diagRequest =
+				stmtDiagnosticsRecorder.ShouldCollectDiagnostics(ctx, stmt.StmtNoConstants, "" /* planGist */)
+			// IsRedacted will be false when ih.collectBundle is false.
+			ih.explainFlags.RedactValues = ih.explainFlags.RedactValues || ih.diagRequest.IsRedacted()
+		}
 	}
 
-	ih.stmtDiagnosticsRecorder = stmtDiagnosticsRecorder
-	ih.withStatementTrace = cfg.TestingKnobs.WithStatementTrace
-
-	defer func() { ih.finalizeSetup(newCtx, cfg) }()
-
 	if sp := tracing.SpanFromContext(ctx); sp != nil {
 		if sp.IsVerbose() {
 			// If verbose tracing was enabled at a higher level, stats
@@ -477,7 +499,7 @@ func (ih *instrumentationHelper) Setup(
 	} else {
 		collectTxnExecStats = func() bool {
 			if stmt.AST.StatementType() == tree.TypeTCL {
-				// We don't collect stats for  statements so there's no need
+				// We don't collect stats for statements, so there's no need
 				//to trace them.
 				return false
 			}
@@ -602,9 +624,7 @@ func (ih *instrumentationHelper) setupWithPlanGist(
 }
 
 func (ih *instrumentationHelper) Finish(
-	cfg *ExecutorConfig,
-	statsCollector *sslocal.StatsCollector,
-	txnStats *execstats.QueryLevelStats,
+	ex *connExecutor,
 	collectExecStats bool,
 	p *planner,
 	ast tree.Statement,
@@ -614,6 +634,10 @@ func (ih *instrumentationHelper) Finish(
 	retErr error,
 ) error {
 	ctx := ih.origCtx
+	cfg := ex.server.cfg
+	statsCollector := ex.statsCollector
+	txnStats := &ex.extraTxnState.accumulatedStats
+	txnHelper := &ex.state.txnInstrumentationHelper
 	if _, ok := ih.Tracing(); !ok {
 		return retErr
 	}
@@ -711,17 +735,28 @@ func (ih *instrumentationHelper) Finish(
 				stmtRawSQL, &p.curPlan, planString, trace, placeholders, res.ErrAllowReleased(),
 				payloadErr, retErr, &p.extendedEvalCtx.Settings.SV, ih.inFlightTraceCollector,
 			)
-			// Include all non-critical errors as warnings. Note that these
-			// error strings might contain PII, but the warnings are only shown
-			// to the current user and aren't included into the bundle.
-			warnings = append(warnings, bundle.errorStrings...)
-			bundle.insert(
-				bundleCtx, ih.fingerprint, ast, cfg.StmtDiagnosticsRecorder, ih.diagRequestID, ih.diagRequest,
-			)
+
+			if !txnHelper.DiagnosticsInProgress() {
+				// Include all non-critical errors as warnings. Note that these
+				// error strings might contain PII, but the warnings are only shown
+				// to the current user and aren't included into the bundle.
+				warnings = append(warnings, bundle.errorStrings...)
+				bundle.insert(
+					bundleCtx, ih.fingerprint, ast, cfg.StmtDiagnosticsRecorder, ih.diagRequestID, ih.diagRequest,
+				)
+			}
 			telemetry.Inc(sqltelemetry.StatementDiagnosticsCollectedCounter)
 		}
 	}
 
+	if txnHelper.DiagnosticsInProgress() {
+		// If we're collecting a transaction bundle, add the statement to the
+		// current txn diagnostic bundle. These will be persisted at the end
+		// of the transaction if the transaction matches the request.
+		// NB: It is safe for bundle to be empty here.
+		txnHelper.AddStatementBundle(ctx, ast, uint64(ih.fingerprintId), ih.fingerprint, bundle)
+	}
+
 	// If there was a communication error already, no point in setting any
 	// results.
 	if retErr != nil {
@@ -1214,3 +1249,31 @@ func (ih *instrumentationHelper) SetIndexRecommendations(
 		reset,
 	)
 }
+
+// handleTransactionDiagnostics manages transaction-level diagnostics
+// collection. If transaction diagnostics are being collected, the
+// instrumentationHelper is set up to collect the bundle for the statement.
+func (ih *instrumentationHelper) handleTransactionDiagnostics(
+	ctx context.Context,
+	txnState *txnState,
+	tracer *tracing.Tracer,
+	stmt *Statement,
+	stmtFingerprintId appstatspb.StmtFingerprintID,
+) (newCtx context.Context, collectingDiagnostics bool) {
+	newCtx, collectingDiagnostics = txnState.shouldCollectTxnDiagnostics(ctx, uint64(stmtFingerprintId), stmt, tracer)
+
+	if collectingDiagnostics {
+		var stmtSp *tracing.Span
+		newCtx, stmtSp = tracing.ChildSpan(newCtx, "txn-diag-sql-query")
+		stmtSp.SetTag("statement", attribute.StringValue(stmt.SQL))
+		ih.collectBundle = true
+		ih.diagRequestID = 0
+		ih.diagRequest = stmtdiagnostics.Request{}
+		ih.explainFlags.RedactValues = txnState.txnInstrumentationHelper.ShouldRedact()
+		ih.shouldFinishSpan = true
+		ih.needFinish = true
+		ih.sp = stmtSp
+	}
+
+	return newCtx, collectingDiagnostics
+}

pkg/sql/stmtdiagnostics/BUILD.bazel

@@ -28,7 +28,7 @@ go_library(
 
 go_test(
     name = "stmtdiagnostics_test",
-    size = "medium",
+    size = "large",
     srcs = [
         "main_test.go",
         "statement_diagnostics_helpers_test.go",
@@ -37,6 +37,7 @@ go_test(
         "txn_diagnostics_test.go",
     ],
     embed = [":stmtdiagnostics"],
+    exec_properties = {"test.Pool": "large"},
     tags = ["no-remote-exec"],
     deps = [
         "//pkg/base",
@@ -53,6 +54,7 @@ go_test(
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/isql",
         "//pkg/sql/sqlerrors",
+        "//pkg/sql/sqlstats",
         "//pkg/testutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",