From c8d3af79b2924c5027a99e0d74ab53aa688c583e Mon Sep 17 00:00:00 2001
From: Tobias Grieger <tobias.b.grieger@gmail.com>
Date: Tue, 14 Jan 2025 16:53:33 +0100
Subject: [PATCH] testutils,kvserver: add StartExecTrace and adopt in
 TestPromoteNonVoterInAddVoter

Touches https://github.com/cockroachdb/cockroach/issues/134383.
Epic: None
Release note: None
---
 pkg/kv/kvserver/replicate_queue_test.go |  5 +-
 pkg/testutils/BUILD.bazel               |  2 +
 pkg/testutils/exectrace.go              | 71 +++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 pkg/testutils/exectrace.go

diff --git a/pkg/kv/kvserver/replicate_queue_test.go b/pkg/kv/kvserver/replicate_queue_test.go
index 5b221d7cd463..893466b802ee 100644
--- a/pkg/kv/kvserver/replicate_queue_test.go
+++ b/pkg/kv/kvserver/replicate_queue_test.go
@@ -2140,7 +2140,8 @@ func iterateOverAllStores(
 // the range log where the added replica type is a LEARNER.
 func TestPromoteNonVoterInAddVoter(t *testing.T) {
 	defer leaktest.AfterTest(t)()
-	defer log.Scope(t).Close(t)
+	scope := log.Scope(t)
+	defer scope.Close(t)
 
 	// This test is slow under stress/race and can time out when upreplicating /
 	// rebalancing to ensure all stores have the same range count initially, due
@@ -2149,6 +2150,8 @@ func TestPromoteNonVoterInAddVoter(t *testing.T) {
 	skip.UnderDeadlock(t)
 	skip.UnderRace(t)
 
+	defer testutils.StartExecTrace(t, scope.GetDirectory()).Finish(t)
+
 	ctx := context.Background()
 
 	// Create 7 stores: 3 in Region 1, 2 in Region 2, and 2 in Region 3.
diff --git a/pkg/testutils/BUILD.bazel b/pkg/testutils/BUILD.bazel
index 5fb1beba8f2e..f0fcdb14104c 100644
--- a/pkg/testutils/BUILD.bazel
+++ b/pkg/testutils/BUILD.bazel
@@ -6,6 +6,7 @@ go_library(
         "backup.go",
         "dir.go",
         "error.go",
+        "exectrace.go",
         "files.go",
         "hook.go",
         "keys.go",
@@ -33,6 +34,7 @@ go_library(
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_dataexmachina_dev_side_eye_go//sideeye",
+        "@com_github_petermattis_goid//:goid",
         "@com_github_stretchr_testify//require",
     ],
 )
diff --git a/pkg/testutils/exectrace.go b/pkg/testutils/exectrace.go
new file mode 100644
index 000000000000..c9c4e60dfad5
--- /dev/null
+++ b/pkg/testutils/exectrace.go
@@ -0,0 +1,71 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package testutils
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime/trace"
+
+	"github.com/petermattis/goid"
+)
+
+type ActiveExecTrace struct {
+	name string
+	file *os.File
+	reg  *trace.Region
+}
+
+// Finish stops the ongoing execution trace, if there is one, and closes the
+// file. It must be called only once.
+func (a *ActiveExecTrace) Finish(t interface {
+	Failed() bool
+	Logf(string, ...interface{})
+}) {
+	if a == nil {
+		return
+	}
+	a.reg.End()
+	trace.Stop()
+	_ = a.file.Close()
+	if !t.Failed() {
+		_ = os.Remove(a.file.Name())
+	} else {
+		t.Logf("execution trace written to %s", a.file.Name())
+	}
+}
+
+// StartExecTrace starts a Go execution trace and returns a handle that allows
+// stopping it. If a trace cannot be started, this is logged and nil is returned.
+// It is valid to stop a nil ActiveExecTrace.
+//
+// This helper is intended to instrument tests for which an execution trace is
+// desired on the next failure.
+func StartExecTrace(
+	t interface {
+		Name() string
+		Logf(string, ...interface{})
+	}, dir string,
+) *ActiveExecTrace {
+	path := filepath.Join(dir, fmt.Sprintf("exectrace_goid_%d.bin", goid.Get()))
+	f, err := os.Create(path)
+	if err != nil {
+		t.Logf("could not create file for execution trace: %s", err)
+		return nil
+	}
+	if err := trace.Start(f); err != nil {
+		_ = f.Close()
+		t.Logf("could not start execution trace: %s", err)
+		return nil
+	}
+	return &ActiveExecTrace{
+		name: t.Name(),
+		file: f,
+		reg:  trace.StartRegion(context.Background(), t.Name()),
+	}
+}