go/stats: improve performance of safeJoinLabels (#16953)

Signed-off-by: Matt Robenolt <[email protected]>

go/stats/export.go

@@ -384,19 +384,76 @@ func IsDimensionCombined(name string) bool {
 // them apart later. The function also replaces specific label values with "all"
 // if a dimenstion is marked as true in combinedLabels.
 func safeJoinLabels(labels []string, combinedLabels []bool) string {
-	sanitizedLabels := make([]string, len(labels))
+	// fast path that potentially requires 0 allocations
+	switch len(labels) {
+	case 0:
+		return ""
+	case 1:
+		if combinedLabels == nil || !combinedLabels[0] {
+			return safeLabel(labels[0])
+		}
+		return StatsAllStr
+	}
+
+	var b strings.Builder
+	size := len(labels) - 1 // number of separators
 	for idx, label := range labels {
 		if combinedLabels != nil && combinedLabels[idx] {
-			sanitizedLabels[idx] = StatsAllStr
+			size += len(StatsAllStr)
+		} else {
+			size += len(label)
+		}
+	}
+	b.Grow(size)
+
+	for idx, label := range labels {
+		if idx > 0 {
+			b.WriteByte('.')
+		}
+		if combinedLabels != nil && combinedLabels[idx] {
+			b.WriteString(StatsAllStr)
+		} else {
+			appendSafeLabel(&b, label)
+		}
+	}
+	return b.String()
+}
+
+// appendSafeLabel is a more efficient version equivalent
+// to strings.ReplaceAll(label, ".", "_"), but appends into
+// a strings.Builder.
+func appendSafeLabel(b *strings.Builder, label string) {
+	// first quickly check if there are any periods to be replaced
+	found := false
+	for i := 0; i < len(label); i++ {
+		if label[i] == '.' {
+			found = true
+			break
+		}
+	}
+	// if there are none, we can just write the label as-is into the
+	// Builder.
+	if !found {
+		b.WriteString(label)
+		return
+	}
+
+	for i := 0; i < len(label); i++ {
+		if label[i] == '.' {
+			b.WriteByte('_')
 		} else {
-			sanitizedLabels[idx] = safeLabel(label)
+			b.WriteByte(label[i])
 		}
 	}
-	return strings.Join(sanitizedLabels, ".")
 }
 
 func safeLabel(label string) string {
-	return strings.Replace(label, ".", "_", -1)
+	// XXX: strings.ReplaceAll is optimal in the case where '.' does not
+	// exist in the label name, and will return the string as-is without
+	// allocations. So if we are working with a single label, it's preferrable
+	// over appendSafeLabel, since appendSafeLabel is required to allocate
+	// into a strings.Builder.
+	return strings.ReplaceAll(label, ".", "_")
 }
 
 func isVarDropped(name string) bool {

go/stats/export_test.go

@@ -19,6 +19,7 @@ package stats
 import (
 	"expvar"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -189,3 +190,73 @@ func TestStringMapWithMultiLabels(t *testing.T) {
 
 	require.Equal(t, c.ValueLabel(), "ccc")
 }
+
+func TestSafeJoinLabels(t *testing.T) {
+	cases := []struct {
+		labels   []string
+		combined []bool
+		want     string
+	}{
+		{
+			labels: []string{},
+			want:   "",
+		},
+		{
+			labels: []string{"foo"},
+			want:   "foo",
+		},
+		{
+			labels: []string{"foo.bar"},
+			want:   "foo_bar",
+		},
+		{
+			labels:   []string{"foo"},
+			combined: []bool{true},
+			want:     "all",
+		},
+		{
+			labels: []string{"foo", "bar"},
+			want:   "foo.bar",
+		},
+		{
+			labels:   []string{"foo", "bar"},
+			combined: []bool{true, false},
+			want:     "all.bar",
+		},
+		{
+			labels:   []string{"foo", "bar"},
+			combined: []bool{true, true},
+			want:     "all.all",
+		},
+		{
+			labels:   []string{"foo", "bar"},
+			combined: []bool{false, true},
+			want:     "foo.all",
+		},
+		{
+			labels: []string{"foo.bar", "bar.baz"},
+			want:   "foo_bar.bar_baz",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(strings.Join(tc.labels, ","), func(t *testing.T) {
+			require.Equal(t, tc.want, safeJoinLabels(tc.labels, tc.combined))
+		})
+	}
+}
+func BenchmarkSafeJoinLabels(b *testing.B) {
+	labels := [5]string{"foo:bar", "foo.bar", "c1a", "testing", "testing.a.b"}
+	combined := [5]bool{true, true, true, true, true}
+	b.Run("no combined", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = safeJoinLabels(labels[:], nil)
+		}
+	})
+	b.Run("combined", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = safeJoinLabels(labels[:], combined[:])
+		}
+	})
+}