Repository: prometheus/prometheus
File: storage/remote/write_handler_test.go (Go)

Code before:
func BenchmarkRemoteWriteHandler(b *testing.B) {
	testCases := []struct {
		name        string
		payloadFunc func() ([]byte, error)
		protoFormat config.RemoteWriteProtoMsg
	}{
		{
			name: "V1 Write",
			payloadFunc: func() ([]byte, error) {
				num := strings.Repeat(strconv.Itoa(b.N), 16)
				buf, _, _, err := buildWriteRequest(nil, []prompb.TimeSeries{{
					Labels: []prompb.Label{
						{Name: "__name__", Value: "test_metric"},
						{Name: "test_label_name_" + num, Value: "abcdefg'hijlmn234!@#$%^&*()_+~`\"{}[],./<>?hello0123hiOlá你好Dzieńdobry9Zd8ra765v4stvuyte" + num},
					},
					Histograms: []prompb.Histogram{prompb.FromIntHistogram(0, &testHistogram)},
				}}, nil, nil, nil, nil, "snappy")
				return buf, err
			},
			protoFormat: config.RemoteWriteProtoMsgV1,
		},
		{
			name: "V2 Write",
			payloadFunc: func() ([]byte, error) {
				num := strings.Repeat(strconv.Itoa(b.N), 16)
				labels := []string{"__name__", "test_metric", "test_label_name_" + num, "abcdefg'hijlmn234!@#$%^&*()_+~`\"{}[],./<>?hello0123hiOlá你好Dzieńdobry9Zd8ra765v4stvuyte" + num}
				buf, _, _, err := buildV2WriteRequest(promslog.NewNopLogger(), []writev2.TimeSeries{{
					LabelsRefs: []uint32{0, 1, 2, 3},
					Histograms: []writev2.Histogram{writev2.FromIntHistogram(0, &testHistogram)},
				}}, labels,
					nil, nil, nil, "snappy")
				return buf, err
			},
			protoFormat: config.RemoteWriteProtoMsgV2,
		},
	}
	for _, tc := range testCases {
		b.Run(tc.name, func(b *testing.B) {
			appendable := &mockAppendable{}
			handler := NewWriteHandler(promslog.NewNopLogger(), nil, appendable, []config.RemoteWriteProtoMsg{tc.protoFormat}, false)
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				b.StopTimer()
				buf, err := tc.payloadFunc()
				require.NoError(b, err)
				req, err := http.NewRequest("", "", bytes.NewReader(buf))
				require.NoError(b, err)
				b.StartTimer()

				recorder := httptest.NewRecorder()
				handler.ServeHTTP(recorder, req)

Review comment (suggestion):
Sorry for the delayed review here.

Can we make one last change? It would be nice if we could define the labels strings once for both v1 and v2, and recreate the required proto from that.

I think we can just define this same slice of strings at the beginning of the test and reuse it here, and then also define a 2nd var up at the top that uses `labels.FromStrings(slice)` and `prompb.FromLabels` to turn them into prompb labels for v1.

That way if we ever run into a situation where we need to change the contents of the labels to test some case we've missed the labels content only has to change in one place.

Code after:
	labelStrings := []string{
		"__name__", "test_metric",
		"test_label_name", "abcdefg'hijlmn234!@#$%^&*()_+~`\"{}[],./<>?hello0123hiOlá你好Dzieńdobry9Zd8ra765v4stvuyte",
	}
	v1Labels := prompb.FromLabels(labels.FromStrings(labelStrings...), nil)

	testCases := []struct {
		name        string
		payloadFunc func() ([]byte, error)
		protoFormat config.RemoteWriteProtoMsg
	}{
		{
			name: "V1 Write",
			payloadFunc: func() ([]byte, error) {
				buf, _, _, err := buildWriteRequest(nil, []prompb.TimeSeries{{
					Labels:     v1Labels,
					Histograms: []prompb.Histogram{prompb.FromIntHistogram(0, &testHistogram)},
				}}, nil, nil, nil, nil, "snappy")
				return buf, err
			},
			protoFormat: config.RemoteWriteProtoMsgV1,
		},
		{
			name: "V2 Write",
			payloadFunc: func() ([]byte, error) {
				buf, _, _, err := buildV2WriteRequest(promslog.NewNopLogger(), []writev2.TimeSeries{{
					LabelsRefs: []uint32{0, 1, 2, 3},
					Histograms: []writev2.Histogram{writev2.FromIntHistogram(0, &testHistogram)},
				}}, labelStrings,
					nil, nil, nil, "snappy")
				return buf, err
			},
			protoFormat: config.RemoteWriteProtoMsgV2,
		},
	}
	for _, tc := range testCases {
		b.Run(tc.name, func(b *testing.B) {
			appendable := &mockAppendable{}
			handler := NewWriteHandler(promslog.NewNopLogger(), nil, appendable, []config.RemoteWriteProtoMsg{tc.protoFormat}, false)
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				b.StopTimer()
				buf, err := tc.payloadFunc()
				require.NoError(b, err)
				req, err := http.NewRequest("", "", bytes.NewReader(buf))
				require.NoError(b, err)
				b.StartTimer()

				recorder := httptest.NewRecorder()
				handler.ServeHTTP(recorder, req)
			}