Merge #152886

152886: sql: Add txn instrumentation helper r=kyle-a-wong a=kyle-a-wong

Adds a new txnInstrumentationHelper struct for managing
the collection of transaction diagnostic bundles. This
struct has APIs for determining if diagnostics collection
should be started and should continue, to add a statement
bundle to the transaction diagnostics, and to finalize
a transaction diagnostics collection.

Epic: [CRDB-53541](https://cockroachlabs.atlassian.net/browse/CRDB-53541)
Release note: None

Co-authored-by: Kyle Wong <[email protected]>

Author: craig[bot]

pkg/sql/BUILD.bazel

@@ -276,6 +276,7 @@ go_library(
         "truncate.go",
         "two_phase_commit.go",
         "txn_fingerprint_id_cache.go",
+        "txn_instrumentation.go",
         "txn_state.go",
         "type_change.go",
         "unary.go",
@@ -760,6 +761,7 @@ go_test(
         "tenant_test.go",
         "trace_test.go",
         "txn_fingerprint_id_cache_test.go",
+        "txn_instrumentation_test.go",
         "txn_restart_test.go",
         "txn_state_test.go",
         "type_change_test.go",

pkg/sql/conn_executor.go

@@ -3477,6 +3477,12 @@ func isCommit(stmt tree.Statement) bool {
 	return ok
 }
 
+// isRollback returns true if stmt is a "ROLLBACK" statement.
+func isRollback(stmt tree.Statement) bool {
+	_, ok := stmt.(*tree.RollbackTransaction)
+	return ok
+}
+
 var retryableMinTimestampBoundUnsatisfiableError = errors.Newf(
 	"retryable MinTimestampBoundUnsatisfiableError",
 )

pkg/sql/stmtdiagnostics/statement_diagnostics.go

@@ -591,6 +591,41 @@ func (r *Registry) InsertStatementDiagnostics(
 	return diagID, err
 }
 
+func (r *Registry) insertBundleChunks(
+	ctx context.Context, bundle []byte, description string, txn isql.Txn,
+) (*tree.DArray, error) {
+	bundleChunksVal := tree.NewDArray(types.Int)
+	bundleToUpload := bundle
+	for len(bundleToUpload) > 0 {
+		chunkSize := int(bundleChunkSize.Get(&r.st.SV))
+		chunk := bundleToUpload
+		if len(chunk) > chunkSize {
+			chunk = chunk[:chunkSize]
+		}
+		bundleToUpload = bundleToUpload[len(chunk):]
+
+		// Insert the chunk into system.statement_bundle_chunks.
+		row, err := txn.QueryRowEx(
+			ctx, "stmt-bundle-chunks-insert", txn.KV(),
+			sessiondata.NodeUserSessionDataOverride,
+			"INSERT INTO system.statement_bundle_chunks(description, data) VALUES ($1, $2) RETURNING id",
+			description,
+			tree.NewDBytes(tree.DBytes(chunk)),
+		)
+		if err != nil {
+			return nil, err
+		}
+		if row == nil {
+			return nil, errors.New("failed to check statement bundle chunk")
+		}
+		chunkID := row[0].(*tree.DInt)
+		if err := bundleChunksVal.Append(chunkID); err != nil {
+			return nil, err
+		}
+	}
+	return bundleChunksVal, nil
+}
+
 func (r *Registry) innerInsertStatementDiagnostics(
 	ctx context.Context, diagnostic StmtDiagnostic, txn isql.Txn,
 ) (CollectedInstanceID, error) {
@@ -621,34 +656,9 @@ func (r *Registry) innerInsertStatementDiagnostics(
 		errorVal = tree.NewDString(diagnostic.collectionErr.Error())
 	}
 
-	bundleChunksVal := tree.NewDArray(types.Int)
-	bundleToUpload := diagnostic.bundle
-	for len(bundleToUpload) > 0 {
-		chunkSize := int(bundleChunkSize.Get(&r.st.SV))
-		chunk := bundleToUpload
-		if len(chunk) > chunkSize {
-			chunk = chunk[:chunkSize]
-		}
-		bundleToUpload = bundleToUpload[len(chunk):]
-
-		// Insert the chunk into system.statement_bundle_chunks.
-		row, err := txn.QueryRowEx(
-			ctx, "stmt-bundle-chunks-insert", txn.KV(),
-			sessiondata.NodeUserSessionDataOverride,
-			"INSERT INTO system.statement_bundle_chunks(description, data) VALUES ($1, $2) RETURNING id",
-			"statement diagnostics bundle",
-			tree.NewDBytes(tree.DBytes(chunk)),
-		)
-		if err != nil {
-			return diagID, err
-		}
-		if row == nil {
-			return diagID, errors.New("failed to check statement bundle chunk")
-		}
-		chunkID := row[0].(*tree.DInt)
-		if err := bundleChunksVal.Append(chunkID); err != nil {
-			return diagID, err
-		}
+	bundleChunksVal, err := r.insertBundleChunks(ctx, diagnostic.bundle, "statement diagnostics bundle", txn)
+	if err != nil {
+		return diagID, err
 	}
 
 	collectionTime := timeutil.Now()

pkg/sql/stmtdiagnostics/txn_diagnostics.go

@@ -23,16 +23,40 @@ import (
 // bundle should be collected.
 type TxnRequest struct {
 	txnFingerprintId    uint64
-	stmtFingerprintsId  []uint64
+	stmtFingerprintIds  []uint64
 	redacted            bool
 	username            string
 	expiresAt           time.Time
 	minExecutionLatency time.Duration
 	samplingProbability float64
 }
 
+func NewTxnRequest(
+	txnFingerprintId uint64,
+	stmtFingerprintIds []uint64,
+	redacted bool,
+	username string,
+	expiresAt time.Time,
+	minExecutionLatency time.Duration,
+	samplingProbability float64,
+) TxnRequest {
+	return TxnRequest{
+		txnFingerprintId:    txnFingerprintId,
+		stmtFingerprintIds:  stmtFingerprintIds,
+		redacted:            redacted,
+		username:            username,
+		expiresAt:           expiresAt,
+		minExecutionLatency: minExecutionLatency,
+		samplingProbability: samplingProbability,
+	}
+}
+
+func (t *TxnRequest) TxnFingerprintId() uint64 {
+	return t.txnFingerprintId
+}
+
 func (t *TxnRequest) StmtFingerprintIds() []uint64 {
-	return t.stmtFingerprintsId
+	return t.stmtFingerprintIds
 }
 
 func (t *TxnRequest) IsRedacted() bool {
@@ -56,10 +80,11 @@ func (t *TxnRequest) isConditional() bool {
 // as a transaction diagnostic bundle
 type TxnDiagnostic struct {
 	stmtDiagnostics []StmtDiagnostic
+	bundle          []byte
 }
 
-func NewTxnDiagnostic(stmtDiagnostics []StmtDiagnostic) TxnDiagnostic {
-	return TxnDiagnostic{stmtDiagnostics: stmtDiagnostics}
+func NewTxnDiagnostic(stmtDiagnostics []StmtDiagnostic, bundle []byte) TxnDiagnostic {
+	return TxnDiagnostic{stmtDiagnostics: stmtDiagnostics, bundle: bundle}
 }
 
 // TxnRegistry maintains a view on the transactions on which a diagnostic
@@ -118,7 +143,7 @@ func (r *TxnRegistry) ShouldStartTxnDiagnostic(
 	}
 
 	for id, f := range r.mu.requests {
-		if len(f.stmtFingerprintsId) > 0 && f.stmtFingerprintsId[0] == stmtFingerprintId {
+		if len(f.stmtFingerprintIds) > 0 && f.stmtFingerprintIds[0] == stmtFingerprintId {
 			if f.isExpired(timeutil.Now()) {
 				delete(r.mu.requests, id)
 				return false, 0, req
@@ -219,6 +244,13 @@ func (r *TxnRegistry) InsertTxnDiagnostic(
 	var txnDiagnosticId CollectedInstanceID
 	err := r.db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
 		txn.KV().SetDebugName("txn-diag-insert-bundle")
+
+		_, err := r.StmtRegistry.insertBundleChunks(ctx, diagnostic.bundle, "transaction diagnostics bundle", txn)
+		// Insert the transaction diagnostic bundle
+		if err != nil {
+			return err
+		}
+		// Insert all the statement diagnostics
 		stmtDiagnostics := tree.NewDArray(types.Int)
 		for _, sd := range diagnostic.stmtDiagnostics {
 			id, err := r.StmtRegistry.innerInsertStatementDiagnostics(ctx, sd, txn)
@@ -261,7 +293,7 @@ func (r *TxnRegistry) addTxnRequestInternalLocked(
 	}
 	request := TxnRequest{
 		txnFingerprintId:    txnFingerprintId,
-		stmtFingerprintsId:  stmtFingerprintsId,
+		stmtFingerprintIds:  stmtFingerprintsId,
 		redacted:            redacted,
 		username:            username,
 		expiresAt:           expiresAt,