stmtdiagnostics: fix upload in face of transaction retry

pkg/sql/explain_bundle_test.go

@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -34,6 +35,41 @@ import (
 	"github.com/lib/pq"
 )
 
+func TestExplainAnalyzeDebugWithTxnRetries(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	retryFilter, verifyRetryHit := testutils.TestingRequestFilterRetryTxnWithPrefix(t, "stmt-diag-", 1)
+	srv, godb, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Insecure: true,
+		Knobs: base.TestingKnobs{
+			Store: &kvserver.StoreTestingKnobs{
+				TestingRequestFilter: retryFilter,
+			},
+		},
+	})
+	defer srv.Stopper().Stop(ctx)
+	r := sqlutils.MakeSQLRunner(godb)
+	r.Exec(t, `CREATE TABLE abc (a INT PRIMARY KEY, b INT, c INT UNIQUE);
+CREATE SCHEMA s;
+CREATE TABLE s.a (a INT PRIMARY KEY);`)
+
+	base := "statement.sql trace.json trace.txt trace-jaeger.json env.sql"
+	plans := "schema.sql opt.txt opt-v.txt opt-vv.txt plan.txt"
+
+	// Set a small chunk size to test splitting into chunks. The bundle files are
+	// on the order of 10KB.
+	r.Exec(t, "SET CLUSTER SETTING sql.stmt_diagnostics.bundle_chunk_size = '2000'")
+
+	rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM abc WHERE c=1")
+	checkBundle(
+		t, fmt.Sprint(rows), "public.abc", nil, false, /* expectErrors */
+		base, plans, "stats-defaultdb.public.abc.sql", "distsql.html vec.txt vec-v.txt",
+	)
+	verifyRetryHit()
+}
+
 func TestExplainAnalyzeDebug(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/sql/stmtdiagnostics/statement_diagnostics.go

@@ -314,6 +314,7 @@ func (r *Registry) insertRequestInternal(
 	var reqID RequestID
 	var expiresAt time.Time
 	err := r.db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		txn.KV().SetDebugName("stmt-diag-insert-request")
 		// Check if there's already a pending request for this fingerprint.
 		row, err := txn.QueryRowEx(ctx, "stmt-diag-check-pending", txn.KV(),
 			sessiondata.RootUserSessionDataOverride,
@@ -530,7 +531,9 @@ func (r *Registry) InsertStatementDiagnostics(
 		ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second) // nolint:context
 		defer cancel()
 	}
+
 	err := r.db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+		txn.KV().SetDebugName("stmt-diag-insert-bundle")
 		if requestID != 0 {
 			row, err := txn.QueryRowEx(ctx, "stmt-diag-check-completed", txn.KV(),
 				sessiondata.RootUserSessionDataOverride,
@@ -558,13 +561,14 @@ func (r *Registry) InsertStatementDiagnostics(
 		}
 
 		bundleChunksVal := tree.NewDArray(types.Int)
-		for len(bundle) > 0 {
+		bundleToUpload := bundle
+		for len(bundleToUpload) > 0 {
 			chunkSize := int(bundleChunkSize.Get(&r.st.SV))
-			chunk := bundle
+			chunk := bundleToUpload
 			if len(chunk) > chunkSize {
 				chunk = chunk[:chunkSize]
 			}
-			bundle = bundle[len(chunk):]
+			bundleToUpload = bundleToUpload[len(chunk):]
 
 			// Insert the chunk into system.statement_bundle_chunks.
 			row, err := txn.QueryRowEx(