roachtest: add profile collection step at the end of tpcc

Similar to what we do in sysbench, this commit performs a small 75s run
at the end of tpcc roachtest runs, and collects the different profiles.

This will be useful to update our PGO process to use profiles taken from
tpcc runs instead of sysbench.

References: #147400

Release note: None

Author: iskettaneh

pkg/cmd/roachtest/tests/tpcc.go

@@ -42,6 +42,7 @@ import (
 	"github.com/cockroachdb/ttycolor"
 	"github.com/cockroachdb/version"
 	"github.com/codahale/hdrhistogram"
+	"github.com/google/pprof/profile"
 	"github.com/lib/pq"
 	promapi "github.com/prometheus/client_golang/api"
 	promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
@@ -469,6 +470,98 @@ func runTPCC(
 	}
 	m.Wait()
 
+	// Collect profiles by running a short workload
+	t.Status("running 75 second workload to collect profiles")
+
+	// Capture existing workload log files before collecting profiles. This will
+	// allow us to clean up any new workload log files that are created during
+	// the profile collection run.
+	logFileNamePattern := "run_*_n*_cockroach-workload-*.log"
+	existingWorkloadLogs := make(map[string]struct{})
+	if matches, _ := filepath.Glob(filepath.Join(t.ArtifactsDir(),
+		logFileNamePattern)); matches != nil {
+		for _, logPath := range matches {
+			existingWorkloadLogs[logPath] = struct{}{}
+		}
+	}
+
+	profilesDir := filepath.Join(t.ArtifactsDir(), "1.perf", "profiles")
+	if err := os.MkdirAll(profilesDir, 0755); err != nil {
+		t.L().Errorf("Failed to create profiles directory %s: %v", profilesDir, err)
+	} else {
+		// Start a short TPCC test in order to collect the profiles from an
+		// active cluster.
+		profileM := t.NewErrorGroup(task.WithContext(ctx))
+		profileM.Go(
+			func(ctx context.Context, l *logger.Logger) error {
+				// Run workload for 75 seconds
+				profileDuration := 75 * time.Second
+
+				fileName := roachtestutil.GetBenchmarkMetricsFileName(t)
+				histogramsPath := fmt.Sprintf("%s/profile_%s", t.PerfArtifactsDir(), fileName)
+
+				cmd := roachtestutil.NewCommand("%s workload run %s",
+					test.DefaultCockroachPath, opts.getWorkloadCmd()).
+					MaybeFlag(opts.DB != "", "db", opts.DB).
+					Flag("warehouses", opts.Warehouses).
+					MaybeFlag(!opts.DisableHistogram, "histograms", histogramsPath).
+					Flag("ramp", 0*time.Second). // No ramp for profile collection
+					Flag("duration", profileDuration).
+					Arg("%s", opts.ExtraRunArgs).
+					Arg("%s", pgURLs[0])
+				return c.RunE(ctx, option.WithNodes(c.WorkloadNode()), cmd.String())
+			},
+		)
+
+		// Wait for 30 seconds to give a chance to the workload to start, and then
+		// collect CPU, mutex, allocs profiles.
+		t.L().Printf("waiting 30 seconds for workload to ramp up before collecting profiles")
+		time.Sleep(30 * time.Second)
+
+		collectionDuration := 30 * time.Second
+		t.L().Printf("starting profile collection from %d nodes for %s",
+			len(c.CRDBNodes()), collectionDuration)
+
+		// Collect the profiles.
+		profiles := map[string][]*profile.Profile{"cpu": {}, "allocs": {}, "mutex": {}}
+		for typ := range profiles {
+			typ := typ // Capture for goroutine
+			profileM.Go(
+				func(ctx context.Context, l *logger.Logger) error {
+					var err error
+					profiles[typ], err = roachtestutil.GetProfile(ctx, t, c, typ,
+						collectionDuration, c.CRDBNodes())
+					return err
+				},
+			)
+		}
+
+		// If there is a problem executing the workload or there is a problem
+		// collecting the profiles we need to clean up the directory and log the error.
+		if err := profileM.WaitE(); err != nil {
+			t.L().Errorf("failed to collect profiles: %v", err)
+			_ = os.RemoveAll(profilesDir)
+		} else {
+			// At this point we know that the workload has not crashed, and we have
+			// collected all the individual profiles. We can now merge and export them.
+			if err := mergeAndExportTPCCProfiles(t, c, collectionDuration, profiles, profilesDir); err != nil {
+				t.L().Errorf("failed to merge and export profiles: %v", err)
+				_ = os.RemoveAll(profilesDir)
+			}
+		}
+
+		// Clean up the profile collection workload log file to reduce clutter.
+		// Delete any workload-r logs that were created after the main test run.
+		if matches, _ := filepath.Glob(filepath.Join(t.ArtifactsDir(), logFileNamePattern)); matches != nil {
+			for _, logPath := range matches {
+				if _, ok := existingWorkloadLogs[logPath]; !ok {
+					// This log was created by the profile collection run, remove it
+					_ = os.Remove(logPath)
+				}
+			}
+		}
+	}
+
 	if !opts.SkipPostRunCheck {
 		cmd := roachtestutil.NewCommand("%s workload check %s", test.DefaultCockroachPath, opts.getWorkloadCmd()).
 			MaybeFlag(opts.DB != "", "db", opts.DB).
@@ -488,6 +581,50 @@ func runTPCC(
 	}
 }
 
+// mergeAndExportTPCCProfiles accepts a map of individual profiles of each
+// node of different types (cpu, allocs, mutex), and exports them to the
+// specified directory. Also, it merges them and exports the merged profiles
+// to the same directory.
+func mergeAndExportTPCCProfiles(
+	t test.Test,
+	c cluster.Cluster,
+	duration time.Duration,
+	profiles map[string][]*profile.Profile,
+	profilesDir string,
+) error {
+	// Merge the profiles.
+	mergedProfiles := map[string]*profile.Profile{"cpu": {}, "allocs": {}, "mutex": {}}
+	for typ := range mergedProfiles {
+		var err error
+		if mergedProfiles[typ], err = profile.Merge(profiles[typ]); err != nil {
+			return errors.Wrapf(err, "failed to merge profiles type: %s", typ)
+		}
+	}
+
+	// Export the merged profiles.
+	for typ := range mergedProfiles {
+		filename := fmt.Sprintf("merged.%s.pb.gz", typ)
+		if err := roachtestutil.ExportProfile(mergedProfiles[typ], profilesDir, filename); err != nil {
+			return errors.Wrapf(err, "failed to export merged profiles: %s", typ)
+		}
+		t.L().Printf("successfully exported merged profile: %s of type %s", filename, typ)
+	}
+
+	// Export the individual profiles as well.
+	numNodes := len(c.CRDBNodes())
+	for i := range numNodes {
+		for typ := range profiles {
+			filename := fmt.Sprintf("n%d.%s%s.pb.gz", i+1, typ, duration)
+			if err := roachtestutil.ExportProfile(profiles[typ][i], profilesDir, filename); err != nil {
+				return errors.Wrapf(err, "failed to export individual profile type: %s for node %d", typ, i+1)
+			}
+			t.L().Printf("successfully exported individual profile %s of type %s for node %d", filename, typ, i+1)
+		}
+	}
+
+	return nil
+}
+
 // tpccSupportedWarehouses returns our claim for the maximum number of tpcc
 // warehouses we support for a given hardware configuration.
 //