Merge #133035

133035: roachtest: add code changes to benchmarks to emit openmetrics r=nameisbhaskar,DarrylWong a=sambhav-jain-16

#129221 and #132023 added changes to
exporters to emit openmetrics. This PR makes changes to the roachtests
to make use of the changes in the above PRs. This change also made some
changes to some roachtests that use neither of the above approaches

Epic: https://cockroachlabs.atlassian.net/browse/CRDB-41852

Release note: None

Co-authored-by: Sambhav Jain <sambhav.jain@cockroachlabs.com>

pkg/cmd/roachprod-microbench/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "main.go",
         "metadata.go",
         "report.go",
+        "roachprod_util.go",
         "slack.go",
         "stage.go",
     ],

pkg/cmd/roachprod-microbench/executor.go

@@ -109,7 +109,7 @@ func newExecutor(config executorConfig) (*executor, error) {
 
 	roachprodConfig.Quiet = config.quiet
 	timestamp := timeutil.Now()
-	l := util.InitLogger(filepath.Join(config.outputDir, fmt.Sprintf("roachprod-microbench-%s.log", timestamp.Format(util.TimeFormat))))
+	l := InitLogger(filepath.Join(config.outputDir, fmt.Sprintf("roachprod-microbench-%s.log", timestamp.Format(util.TimeFormat))))
 
 	excludeBenchmarks := util.GetRegexExclusionPairs(config.excludeList)
 	return &executor{
@@ -275,7 +275,7 @@ func (e *executor) executeBenchmarks() error {
 	}
 
 	// Init `roachprod` and get the number of nodes in the cluster.
-	util.InitRoachprod()
+	InitRoachprod()
 	statuses, err := roachprod.Status(context.Background(), e.log, e.cluster, "")
 	if err != nil {
 		return err

pkg/cmd/roachprod-microbench/roachprod_util.go

@@ -0,0 +1,46 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/cockroachdb/cockroach/pkg/roachprod"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+)
+
+// InitRoachprod initializes the roachprod providers by calling InitProviders.
+// This function sets up the environment for running roachprod commands.
+func InitRoachprod() {
+	_ = roachprod.InitProviders()
+}
+
+// RoachprodRun runs a command on a roachprod cluster with the given cluster name and logger.
+// It takes a list of command arguments and passes them to the roachprod command execution.
+func RoachprodRun(clusterName string, l *logger.Logger, cmdArray []string) error {
+	// Execute the roachprod command with the provided context, logger, cluster name, and options.
+	return roachprod.Run(
+		context.Background(), l, clusterName, "", "", false,
+		os.Stdout, os.Stderr, cmdArray, install.DefaultRunOptions(),
+	)
+}
+
+// InitLogger initializes and returns a logger based on the provided log file path.
+// If the logger configuration fails, the program prints an error and exits.
+func InitLogger(path string) *logger.Logger {
+	loggerCfg := logger.Config{Stdout: os.Stdout, Stderr: os.Stderr} // Create a logger config with standard output and error.
+	var loggerError error
+	l, loggerError := loggerCfg.NewLogger(path) // Create a new logger based on the configuration.
+	if loggerError != nil {
+		// If there is an error initializing the logger, print the error message and exit the program.
+		_, _ = fmt.Fprintf(os.Stderr, "unable to configure logger: %s\n", loggerError)
+		os.Exit(1)
+	}
+	return l // Return the initialized logger.
+}

pkg/cmd/roachprod-microbench/stage.go

@@ -11,7 +11,6 @@ import (
 	"path"
 	"strings"
 
-	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod-microbench/util"
 	"github.com/cockroachdb/cockroach/pkg/roachprod"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
 	"github.com/cockroachdb/errors"
@@ -22,7 +21,7 @@ import (
 func stage(cluster, archivePath, remoteDest string) (err error) {
 	ctx := context.Background()
 
-	util.InitRoachprod()
+	InitRoachprod()
 	loggerCfg := logger.Config{Stdout: os.Stdout, Stderr: os.Stderr}
 	l, err := loggerCfg.NewLogger("")
 	if err != nil {
@@ -34,22 +33,22 @@ func stage(cluster, archivePath, remoteDest string) (err error) {
 
 	defer func() {
 		// Remove the remote archive after we're done.
-		cleanUpErr := util.RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath})
+		cleanUpErr := RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath})
 		err = errors.CombineErrors(err, errors.Wrapf(cleanUpErr, "removing remote archive: %s", archiveRemotePath))
 	}()
 
 	// Remove the remote archive and destination directory if they exist.
-	if err = util.RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", archiveRemotePath}); err != nil {
 		return errors.Wrapf(err, "removing remote archive: %s", archiveRemotePath)
 	}
-	if err = util.RoachprodRun(cluster, l, []string{"rm", "-rf", remoteDest}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"rm", "-rf", remoteDest}); err != nil {
 		return errors.Wrapf(err, "removing remote destination: %s", remoteDest)
 	}
 
 	// Copy the archive to the remote machine.
 	copyFromGCS := strings.HasPrefix(archivePath, "gs://")
 	if copyFromGCS {
-		if err = util.RoachprodRun(cluster, l, []string{"gsutil", "-q", "-m", "cp", archivePath, archiveRemotePath}); err != nil {
+		if err = RoachprodRun(cluster, l, []string{"gsutil", "-q", "-m", "cp", archivePath, archiveRemotePath}); err != nil {
 			return errors.Wrapf(err, "copying archive from GCS: %s", archivePath)
 		}
 	} else {
@@ -59,10 +58,10 @@ func stage(cluster, archivePath, remoteDest string) (err error) {
 	}
 
 	// Extract the archive on the remote machine.
-	if err = util.RoachprodRun(cluster, l, []string{"mkdir", "-p", remoteDest}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"mkdir", "-p", remoteDest}); err != nil {
 		return errors.Wrapf(err, "creating remote destination: %s", remoteDest)
 	}
-	if err = util.RoachprodRun(cluster, l, []string{"tar", "-C", remoteDest, "-xzf", archiveRemotePath}); err != nil {
+	if err = RoachprodRun(cluster, l, []string{"tar", "-C", remoteDest, "-xzf", archiveRemotePath}); err != nil {
 		return errors.Wrapf(err, "extracting archive: %s", archiveRemotePath)
 	}
 

pkg/cmd/roachprod-microbench/util/BUILD.bazel

@@ -6,9 +6,6 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod-microbench/util",
     visibility = ["//visibility:public"],
     deps = [
-        "//pkg/roachprod",
-        "//pkg/roachprod/install",
-        "//pkg/roachprod/logger",
         "@com_github_spf13_cobra//:cobra",
         "@org_golang_x_exp//maps",
     ],

pkg/cmd/roachprod-microbench/util/util.go

@@ -6,16 +6,12 @@
 package util
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"regexp"
 	"sort"
 	"strings"
 
-	"github.com/cockroachdb/cockroach/pkg/roachprod"
-	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
-	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
 	"github.com/spf13/cobra"
 	"golang.org/x/exp/maps"
 )
@@ -25,13 +21,15 @@ const TimeFormat = "2006-01-02T15_04_05"
 const PackageSeparator = "→"
 
 var (
-	// invalidCharRegex matches
+	// invalidCharKeyRegex matches
 	// the first character if it is not a letter (a-z, A-Z) or an underscore (_)
 	// or
 	// any character that is not a letter (a-z, A-Z), digit (0-9), or underscore (_).
-	invalidCharKeyRegex = regexp.MustCompile(`(^[^a-zA-Z_])|([^a-zA-Z0-9_])`)
-	// invalidCharValueRegex
+	invalidCharKeyRegex   = regexp.MustCompile(`(^[^a-zA-Z_])|([^a-zA-Z0-9_])`)
 	invalidCharValueRegex = regexp.MustCompile(`[\\\n"]`)
+
+	// invalidMetricNameRegex matches any other character other than _, :, characters and digits
+	invalidMetricNameRegex = regexp.MustCompile(`[^a-zA-Z0-9_:]`)
 )
 
 // LabelMapToString converts a map of labels (key-value pairs) into a formatted string.
@@ -55,6 +53,11 @@ func SanitizeKey(input string) string {
 	return invalidCharKeyRegex.ReplaceAllString(input, "_")
 }
 
+// SanitizeMetricName replaces all invalid characters in input string with underscores
+func SanitizeMetricName(input string) string {
+	return invalidMetricNameRegex.ReplaceAllString(input, "_")
+}
+
 // SanitizeValue replaces all \,\n and " with underscores (_).
 func SanitizeValue(input string) string {
 	// Replace all characters that match as per the regex with an underscore.
@@ -120,33 +123,3 @@ func GetRegexExclusionPairs(excludeList []string) [][]*regexp.Regexp {
 	}
 	return excludeRegexes
 }
-
-// InitRoachprod initializes the roachprod providers by calling InitProviders.
-// This function sets up the environment for running roachprod commands.
-func InitRoachprod() {
-	_ = roachprod.InitProviders()
-}
-
-// RoachprodRun runs a command on a roachprod cluster with the given cluster name and logger.
-// It takes a list of command arguments and passes them to the roachprod command execution.
-func RoachprodRun(clusterName string, l *logger.Logger, cmdArray []string) error {
-	// Execute the roachprod command with the provided context, logger, cluster name, and options.
-	return roachprod.Run(
-		context.Background(), l, clusterName, "", "", false,
-		os.Stdout, os.Stderr, cmdArray, install.DefaultRunOptions(),
-	)
-}
-
-// InitLogger initializes and returns a logger based on the provided log file path.
-// If the logger configuration fails, the program prints an error and exits.
-func InitLogger(path string) *logger.Logger {
-	loggerCfg := logger.Config{Stdout: os.Stdout, Stderr: os.Stderr} // Create a logger config with standard output and error.
-	var loggerError error
-	l, loggerError := loggerCfg.NewLogger(path) // Create a new logger based on the configuration.
-	if loggerError != nil {
-		// If there is an error initializing the logger, print the error message and exit the program.
-		_, _ = fmt.Fprintf(os.Stderr, "unable to configure logger: %s\n", loggerError)
-		os.Exit(1)
-	}
-	return l // Return the initialized logger.
-}

pkg/cmd/roachtest/clusterstats/BUILD.bazel

@@ -36,6 +36,7 @@ go_library(
         "@com_github_prometheus_client_golang//api",
         "@com_github_prometheus_client_golang//api/prometheus/v1:prometheus",
         "@com_github_prometheus_common//model",
+        "@org_golang_x_exp//maps",
     ],
 )
 

pkg/cmd/roachtest/clusterstats/exporter.go

@@ -11,6 +11,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod-microbench/util"
@@ -21,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
+	"golang.org/x/exp/maps"
 )
 
 // ClusterStat represents a filtered query by the given LabelName. For example,
@@ -128,7 +130,7 @@ func (r *ClusterStatRun) serializeOpenmetricsOutRun(
 	ctx context.Context, t test.Test, c cluster.Cluster,
 ) error {
 
-	labelString := GetDefaultOpenmetricsLabelString(t, c)
+	labelString := GetOpenmetricsLabelString(t, c, nil)
 	report, err := serializeOpenmetricsReport(*r, &labelString)
 	if err != nil {
 		return errors.Wrap(err, "failed to serialize perf artifacts")
@@ -142,13 +144,13 @@ func serializeOpenmetricsReport(r ClusterStatRun, labelString *string) (*bytes.B
 
 	// Emit summary metrics from Total
 	for key, value := range r.Total {
-		buffer.WriteString(fmt.Sprintf("# TYPE %s gauge\n", util.SanitizeKey(key)))
+		buffer.WriteString(fmt.Sprintf("# TYPE %s gauge\n", util.SanitizeMetricName(key)))
 		buffer.WriteString(fmt.Sprintf("%s{%s} %f %d\n", util.SanitizeKey(key), *labelString, value, timeutil.Now().UTC().Unix()))
 	}
 
 	// Emit histogram metrics from Stats
 	for _, stat := range r.Stats {
-		buffer.WriteString(fmt.Sprintf("# TYPE %s gauge\n", util.SanitizeKey(stat.Tag)))
+		buffer.WriteString(fmt.Sprintf("# TYPE %s gauge\n", util.SanitizeMetricName(stat.Tag)))
 		for i, timestamp := range stat.Time {
 			t := timeutil.Unix(0, timestamp)
 			buffer.WriteString(
@@ -397,15 +399,49 @@ func (cs *clusterStatCollector) getStatSummary(
 	return ret, nil
 }
 
-func GetDefaultOpenmetricsLabelString(t test.Test, c cluster.Cluster) string {
-	return util.LabelMapToString(GetDefaultOpenmetricsLabelMap(t, c))
+// GetOpenmetricsLabelString creates a string that follows the openmetrics labels format
+func GetOpenmetricsLabelString(t test.Test, c cluster.Cluster, labels map[string]string) string {
+	return util.LabelMapToString(GetOpenmetricsLabelMap(t, c, labels))
 }
 
-func GetDefaultOpenmetricsLabelMap(t test.Test, c cluster.Cluster) map[string]string {
-	return map[string]string{
-		"test":  t.Name(),
+// GetOpenmetricsLabelMap creates a map of label keys and values
+// It takes roachtest parameters and create relevant labels.
+// Test name is split and each split is added as a subtype
+func GetOpenmetricsLabelMap(
+	t test.Test, c cluster.Cluster, labels map[string]string,
+) map[string]string {
+	defaultMap := map[string]string{
 		"cloud": c.Cloud().String(),
 		"owner": string(t.Spec().(*registry.TestSpec).Owner),
 		"suite": t.Spec().(*registry.TestSpec).Suites.String(),
 	}
+
+	// Since the roachtest have / delimiter for subtests
+	// Partitioning the name by '/'
+	testNameArray := strings.Split(t.Name(), "/")
+	defaultMap["test"] = testNameArray[0]
+
+	subTestIterator := 1
+	for i := 1; i < len(testNameArray); i++ {
+		testSubName := testNameArray[i]
+
+		// If the partition has '=', add the key as a label itself
+		// and the value as its value
+		if strings.Contains(testSubName, "=") {
+			testLabels := strings.Split(testSubName, "=")
+			defaultMap[testLabels[0]] = testLabels[1]
+		} else {
+
+			// Add the subtest label with the iterator since there can be nested subtests
+			testSubType := fmt.Sprintf("subtest-%d", subTestIterator)
+			subTestIterator++
+			defaultMap[testSubType] = testSubName
+		}
+	}
+
+	// If the tests has passed some custom labels, copy them to the map created above
+	if labels != nil {
+		maps.Copy(defaultMap, labels)
+	}
+	return defaultMap
 }

pkg/cmd/roachtest/roachtestutil/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/cmd/roachtest/cluster",
+        "//pkg/cmd/roachtest/clusterstats",
         "//pkg/cmd/roachtest/option",
         "//pkg/cmd/roachtest/spec",
         "//pkg/cmd/roachtest/test",
@@ -36,6 +37,7 @@ go_library(
         "//pkg/util/retry",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
+        "//pkg/workload/histogram/exporter",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],