roachtest: add utility to get one ts

This commit updates `setupStatCollector` in `allocation_bench` to be
more general, so it can be reused by other roachtests. This commit also
adds a new utility function to the `clusterstats` pkg, `GetOneSeries`,
which enables conveniently grabbing the only series, if applicable,
from the `clusterstats.CollectInterval` time series response. This
is useful when just one series is expected, such as in aggregation
queries.

Epic: none

Release note: None

pkg/cmd/roachtest/clusterstats/helpers.go

@@ -26,6 +26,29 @@ import (
 	promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 )
 
+// GetOneSeries accepts the output of StatCollector.CollectInterval and returns
+// either 1 or no StatSeries. 1 StatSeries will be returned if there is exactly
+// one metric label and for that label, exactly one value. If either of these
+// conditions are not met, nil is returned instead. This type of time series
+// output is common when aggregating over every label value e.g.
+// `sum(sql_query_count)`, which will return the sum of every nodes sql query
+// count.
+func GetOneSeries(taggedTimeSeries map[string]map[string]StatSeries) StatSeries {
+	if len(taggedTimeSeries) != 1 {
+		return nil
+	}
+
+	for _, tagged := range taggedTimeSeries {
+		if len(tagged) != 1 {
+			return nil
+		}
+		for _, series := range tagged {
+			return series
+		}
+	}
+	panic("Unreachable")
+}
+
 // SetupCollectorPromClient instantiates a prometheus client for the given
 // configuration. It assumes that a prometheus instance has been created for
 // the configuration already.

pkg/cmd/roachtest/tests/allocation_bench.go

@@ -300,19 +300,21 @@ func setupAllocationBench(
 		c.Start(ctx, t.L(), startOpts, install.MakeClusterSettings(), c.Node(i))
 	}
 
-	return setupStatCollector(ctx, t, c, spec)
+	return setupStatCollector(ctx, t, c, c.Node(c.Spec().NodeCount), c.Range(1, spec.nodes))
 }
 
+// setupStatCollector is a helper function which installs prometheus and
+// grafana on the promNode, scraping the list of clusterNodes. Using the
+// prometheus client, the function creates and returns a new
+// clusterstats.StatCollector and cleanup function, which should be called
+// before finishing the test.
 func setupStatCollector(
-	ctx context.Context, t test.Test, c cluster.Cluster, spec allocationBenchSpec,
+	ctx context.Context, t test.Test, c cluster.Cluster, promNode, clusterNodes option.NodeListOption,
 ) (clusterstats.StatCollector, func(context.Context)) {
-	t.Status("setting up prometheus and grafana")
-
 	// Setup the prometheus instance and client.
-	clusNodes := c.Range(1, spec.nodes)
-	promNode := c.Node(c.Spec().NodeCount)
+	t.Status("setting up prometheus and grafana")
 	cfg := (&prometheus.Config{}).
-		WithCluster(clusNodes.InstallNodes()).
+		WithCluster(clusterNodes.InstallNodes()).
 		WithPrometheusNode(promNode.InstallNodes()[0])
 
 	err := c.StartGrafana(ctx, t.L(), cfg)