github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/query/api/v1/handler/prometheus/remote/write_test.go (about)

     1  // Copyright (c) 2018 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package remote
    22  
    23  import (
    24  	"bytes"
    25  	"context"
    26  	"crypto/rand"
    27  	"errors"
    28  	"fmt"
    29  	"io/ioutil"
    30  	"math"
    31  	"net/http"
    32  	"net/http/httptest"
    33  	"strings"
    34  	"testing"
    35  	"time"
    36  
    37  	"github.com/m3db/m3/src/cmd/services/m3coordinator/ingest"
    38  	"github.com/m3db/m3/src/cmd/services/m3query/config"
    39  	"github.com/m3db/m3/src/dbnode/generated/proto/annotation"
    40  	"github.com/m3db/m3/src/metrics/policy"
    41  	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/handleroptions"
    42  	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/remote/test"
    43  	"github.com/m3db/m3/src/query/api/v1/options"
    44  	"github.com/m3db/m3/src/query/generated/proto/prompb"
    45  	"github.com/m3db/m3/src/query/models"
    46  	"github.com/m3db/m3/src/query/storage/m3/storagemetadata"
    47  	xclock "github.com/m3db/m3/src/x/clock"
    48  	xerrors "github.com/m3db/m3/src/x/errors"
    49  	"github.com/m3db/m3/src/x/headers"
    50  	"github.com/m3db/m3/src/x/instrument"
    51  	xtest "github.com/m3db/m3/src/x/test"
    52  
    53  	"github.com/golang/mock/gomock"
    54  	"github.com/stretchr/testify/assert"
    55  	"github.com/stretchr/testify/require"
    56  	"github.com/uber-go/tally"
    57  )
    58  
    59  func makeOptions(ds ingest.DownsamplerAndWriter) options.HandlerOptions {
    60  	return options.EmptyHandlerOptions().
    61  		SetNowFn(time.Now).
    62  		SetDownsamplerAndWriter(ds).
    63  		SetTagOptions(models.NewTagOptions()).
    64  		SetConfig(config.Configuration{
    65  			WriteForwarding: config.WriteForwardingConfiguration{
    66  				PromRemoteWrite: handleroptions.PromWriteHandlerForwardingOptions{},
    67  			},
    68  		}).
    69  		SetStoreMetricsType(true)
    70  }
    71  
    72  func TestPromWriteParsing(t *testing.T) {
    73  	ctrl := xtest.NewController(t)
    74  	defer ctrl.Finish()
    75  
    76  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
    77  	handlerOpts := makeOptions(mockDownsamplerAndWriter)
    78  	handler, err := NewPromWriteHandler(handlerOpts)
    79  	require.NoError(t, err)
    80  
    81  	promReq := test.GeneratePromWriteRequest()
    82  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
    83  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
    84  
    85  	r, err := handler.(*PromWriteHandler).parseRequest(req)
    86  	require.Nil(t, err, "unable to parse request")
    87  	require.Equal(t, len(r.Request.Timeseries), 2)
    88  	require.Equal(t, ingest.WriteOptions{}, r.Options)
    89  }
    90  
    91  func TestPromWrite(t *testing.T) {
    92  	ctrl := xtest.NewController(t)
    93  	defer ctrl.Finish()
    94  
    95  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
    96  	mockDownsamplerAndWriter.
    97  		EXPECT().
    98  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any())
    99  
   100  	opts := makeOptions(mockDownsamplerAndWriter)
   101  	handler, err := NewPromWriteHandler(opts)
   102  	require.NoError(t, err)
   103  
   104  	promReq := test.GeneratePromWriteRequest()
   105  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   106  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   107  
   108  	writer := httptest.NewRecorder()
   109  	handler.ServeHTTP(writer, req)
   110  	resp := writer.Result()
   111  	require.Equal(t, http.StatusOK, resp.StatusCode)
   112  }
   113  
   114  func TestPromWriteError(t *testing.T) {
   115  	ctrl := xtest.NewController(t)
   116  	defer ctrl.Finish()
   117  
   118  	multiErr := xerrors.NewMultiError().Add(errors.New("an error"))
   119  	batchErr := ingest.BatchError(multiErr)
   120  
   121  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   122  	mockDownsamplerAndWriter.EXPECT().
   123  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any()).
   124  		Return(batchErr)
   125  
   126  	opts := makeOptions(mockDownsamplerAndWriter)
   127  	handler, err := NewPromWriteHandler(opts)
   128  	require.NoError(t, err)
   129  
   130  	promReq := test.GeneratePromWriteRequest()
   131  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   132  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   133  	require.NoError(t, err)
   134  
   135  	writer := httptest.NewRecorder()
   136  	handler.ServeHTTP(writer, req)
   137  	resp := writer.Result()
   138  	require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
   139  
   140  	body, err := ioutil.ReadAll(resp.Body)
   141  	require.NoError(t, err)
   142  	require.True(t, bytes.Contains(body, []byte(batchErr.Error())))
   143  }
   144  
   145  func TestWriteErrorMetricCount(t *testing.T) {
   146  	ctrl := xtest.NewController(t)
   147  	defer ctrl.Finish()
   148  
   149  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   150  
   151  	scope := tally.NewTestScope("",
   152  		map[string]string{"test": "error-metric-test"})
   153  
   154  	iopts := instrument.NewOptions().SetMetricsScope(scope)
   155  	opts := makeOptions(mockDownsamplerAndWriter).SetInstrumentOpts(iopts)
   156  	handler, err := NewPromWriteHandler(opts)
   157  	require.NoError(t, err)
   158  
   159  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, nil)
   160  	handler.ServeHTTP(httptest.NewRecorder(), req)
   161  
   162  	foundMetric := xclock.WaitUntil(func() bool {
   163  		found, ok := scope.Snapshot().Counters()["write.errors+code=4XX,handler=remote-write,test=error-metric-test"]
   164  		return ok && found.Value() == 1
   165  	}, 5*time.Second)
   166  	require.True(t, foundMetric)
   167  }
   168  
   169  func TestWriteDatapointDelayMetric(t *testing.T) {
   170  	ctrl := xtest.NewController(t)
   171  	defer ctrl.Finish()
   172  
   173  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   174  	mockDownsamplerAndWriter.
   175  		EXPECT().
   176  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any())
   177  
   178  	scope := tally.NewTestScope("",
   179  		map[string]string{"test": "delay-metric-test"})
   180  
   181  	iopts := instrument.NewOptions().SetMetricsScope(scope)
   182  	opts := makeOptions(mockDownsamplerAndWriter).SetInstrumentOpts(iopts)
   183  	handler, err := NewPromWriteHandler(opts)
   184  	require.NoError(t, err)
   185  
   186  	writeHandler, ok := handler.(*PromWriteHandler)
   187  	require.True(t, ok)
   188  
   189  	buckets := writeHandler.metrics.ingestLatencyBuckets
   190  
   191  	// NB(r): Bucket length is tested just to sanity check how many buckets we are creating
   192  	require.Equal(t, 80, len(buckets.AsDurations()))
   193  
   194  	// NB(r): Bucket values are tested to sanity check they look right
   195  	expected := "[0s 100ms 200ms 300ms 400ms 500ms 600ms 700ms 800ms 900ms 1s 1.5s 2s 2.5s 3s 3.5s 4s 4.5s 5s 5.5s 6s 6.5s 7s 7.5s 8s 8.5s 9s 9.5s 10s 15s 20s 25s 30s 35s 40s 45s 50s 55s 1m0s 5m0s 10m0s 15m0s 20m0s 25m0s 30m0s 35m0s 40m0s 45m0s 50m0s 55m0s 1h0m0s 1h30m0s 2h0m0s 2h30m0s 3h0m0s 3h30m0s 4h0m0s 4h30m0s 5h0m0s 5h30m0s 6h0m0s 6h30m0s 7h0m0s 8h0m0s 9h0m0s 10h0m0s 11h0m0s 12h0m0s 13h0m0s 14h0m0s 15h0m0s 16h0m0s 17h0m0s 18h0m0s 19h0m0s 20h0m0s 21h0m0s 22h0m0s 23h0m0s 24h0m0s]"
   196  	actual := fmt.Sprintf("%v", buckets.AsDurations())
   197  	require.Equal(t, expected, actual)
   198  
   199  	// Ensure buckets increasing in order
   200  	lastValue := time.Duration(math.MinInt64)
   201  	for _, value := range buckets.AsDurations() {
   202  		require.True(t, value > lastValue,
   203  			fmt.Sprintf("%s must be greater than last bucket value %s", value, lastValue))
   204  		lastValue = value
   205  	}
   206  
   207  	promReq := test.GeneratePromWriteRequest()
   208  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   209  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   210  	handler.ServeHTTP(httptest.NewRecorder(), req)
   211  
   212  	foundMetric := xclock.WaitUntil(func() bool {
   213  		values, found := scope.Snapshot().Histograms()["ingest.latency+handler=remote-write,test=delay-metric-test"]
   214  		if !found {
   215  			return false
   216  		}
   217  		for _, valuesInBucket := range values.Durations() {
   218  			if valuesInBucket > 0 {
   219  				return true
   220  			}
   221  		}
   222  		return false
   223  	}, 5*time.Second)
   224  	require.True(t, foundMetric)
   225  }
   226  
   227  func TestPromWriteUnaggregatedMetricsWithHeader(t *testing.T) {
   228  	ctrl := xtest.NewController(t)
   229  	defer ctrl.Finish()
   230  
   231  	expectedIngestWriteOptions := ingest.WriteOptions{
   232  		DownsampleOverride:     true,
   233  		DownsampleMappingRules: nil,
   234  		WriteOverride:          false,
   235  		WriteStoragePolicies:   nil,
   236  	}
   237  
   238  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   239  	mockDownsamplerAndWriter.
   240  		EXPECT().
   241  		WriteBatch(gomock.Any(), gomock.Any(), expectedIngestWriteOptions)
   242  
   243  	opts := makeOptions(mockDownsamplerAndWriter)
   244  	handler, err := NewPromWriteHandler(opts)
   245  	require.NoError(t, err)
   246  
   247  	promReq := test.GeneratePromWriteRequest()
   248  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   249  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   250  	req.Header.Add(headers.MetricsTypeHeader,
   251  		storagemetadata.UnaggregatedMetricsType.String())
   252  
   253  	writer := httptest.NewRecorder()
   254  	handler.ServeHTTP(writer, req)
   255  	resp := writer.Result()
   256  	require.Equal(t, http.StatusOK, resp.StatusCode)
   257  }
   258  
   259  func TestPromWriteAggregatedMetricsWithHeader(t *testing.T) {
   260  	ctrl := xtest.NewController(t)
   261  	defer ctrl.Finish()
   262  
   263  	expectedIngestWriteOptions := ingest.WriteOptions{
   264  		DownsampleOverride:     true,
   265  		DownsampleMappingRules: nil,
   266  		WriteOverride:          true,
   267  		WriteStoragePolicies: policy.StoragePolicies{
   268  			policy.MustParseStoragePolicy("1m:21d"),
   269  		},
   270  	}
   271  
   272  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   273  	mockDownsamplerAndWriter.
   274  		EXPECT().
   275  		WriteBatch(gomock.Any(), gomock.Any(), expectedIngestWriteOptions)
   276  
   277  	opts := makeOptions(mockDownsamplerAndWriter)
   278  	writeHandler, err := NewPromWriteHandler(opts)
   279  	require.NoError(t, err)
   280  
   281  	promReq := test.GeneratePromWriteRequest()
   282  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   283  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   284  	req.Header.Add(headers.MetricsTypeHeader,
   285  		storagemetadata.AggregatedMetricsType.String())
   286  	req.Header.Add(headers.MetricsStoragePolicyHeader,
   287  		"1m:21d")
   288  
   289  	writer := httptest.NewRecorder()
   290  	writeHandler.ServeHTTP(writer, req)
   291  	resp := writer.Result()
   292  	require.Equal(t, http.StatusOK, resp.StatusCode)
   293  }
   294  
   295  func TestPromWriteOpenMetricsTypes(t *testing.T) {
   296  	ctrl := xtest.NewController(t)
   297  	defer ctrl.Finish()
   298  
   299  	var capturedIter ingest.DownsampleAndWriteIter
   300  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   301  	mockDownsamplerAndWriter.
   302  		EXPECT().
   303  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any()).
   304  		Do(func(_ context.Context, iter ingest.DownsampleAndWriteIter, _ ingest.WriteOptions) ingest.BatchError {
   305  			capturedIter = iter
   306  			return nil
   307  		})
   308  
   309  	opts := makeOptions(mockDownsamplerAndWriter)
   310  
   311  	promReq := &prompb.WriteRequest{
   312  		Timeseries: []prompb.TimeSeries{
   313  			{Type: prompb.MetricType_UNKNOWN},
   314  			{Type: prompb.MetricType_COUNTER},
   315  			{Type: prompb.MetricType_GAUGE},
   316  			{Type: prompb.MetricType_GAUGE},
   317  			{Type: prompb.MetricType_SUMMARY},
   318  			{Type: prompb.MetricType_HISTOGRAM},
   319  			{Type: prompb.MetricType_GAUGE_HISTOGRAM},
   320  			{Type: prompb.MetricType_INFO},
   321  			{Type: prompb.MetricType_STATESET},
   322  			{},
   323  		},
   324  	}
   325  
   326  	executeWriteRequest(t, opts, promReq)
   327  
   328  	firstValue := verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_UNKNOWN, false)
   329  	secondValue := verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_COUNTER, true)
   330  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_GAUGE, false)
   331  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_GAUGE, false)
   332  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_SUMMARY, false)
   333  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_HISTOGRAM, true)
   334  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_GAUGE_HISTOGRAM, false)
   335  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_INFO, false)
   336  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_STATESET, false)
   337  	verifyIterValueAnnotation(t, capturedIter, annotation.OpenMetricsFamilyType_UNKNOWN, false)
   338  
   339  	require.False(t, capturedIter.Next())
   340  	require.NoError(t, capturedIter.Error())
   341  
   342  	assert.Nil(t, firstValue.Annotation, "first annotation invalidation")
   343  
   344  	secondAnnotationPayload := unmarshalAnnotation(t, secondValue.Annotation)
   345  	assert.Equal(t, annotation.Payload{
   346  		OpenMetricsFamilyType:        annotation.OpenMetricsFamilyType_COUNTER,
   347  		OpenMetricsHandleValueResets: true,
   348  	}, secondAnnotationPayload, "second annotation invalidated")
   349  }
   350  
   351  func TestPromWriteGraphiteMetricsTypes(t *testing.T) {
   352  	ctrl := xtest.NewController(t)
   353  	defer ctrl.Finish()
   354  
   355  	var capturedIter ingest.DownsampleAndWriteIter
   356  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   357  	mockDownsamplerAndWriter.
   358  		EXPECT().
   359  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any()).
   360  		Do(func(_ context.Context, iter ingest.DownsampleAndWriteIter, _ ingest.WriteOptions) ingest.BatchError {
   361  			capturedIter = iter
   362  			return nil
   363  		})
   364  
   365  	opts := makeOptions(mockDownsamplerAndWriter)
   366  
   367  	promReq := &prompb.WriteRequest{
   368  		Timeseries: []prompb.TimeSeries{
   369  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_TIMER},
   370  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_COUNTER},
   371  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_GAUGE},
   372  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_GAUGE},
   373  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_TIMER},
   374  			{Source: prompb.Source_GRAPHITE, M3Type: prompb.M3Type_M3_COUNTER},
   375  		},
   376  	}
   377  
   378  	executeWriteRequest(t, opts, promReq)
   379  
   380  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_TIMER)
   381  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_COUNTER)
   382  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_GAUGE)
   383  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_GAUGE)
   384  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_TIMER)
   385  	verifyIterValueAnnotationGraphite(t, capturedIter, annotation.GraphiteType_GRAPHITE_COUNTER)
   386  
   387  	require.False(t, capturedIter.Next())
   388  	require.NoError(t, capturedIter.Error())
   389  }
   390  
   391  func TestPromWriteDisabledMetricsTypes(t *testing.T) {
   392  	ctrl := xtest.NewController(t)
   393  	defer ctrl.Finish()
   394  
   395  	var capturedIter ingest.DownsampleAndWriteIter
   396  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   397  	mockDownsamplerAndWriter.
   398  		EXPECT().
   399  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any()).
   400  		Do(func(_ context.Context, iter ingest.DownsampleAndWriteIter, _ ingest.WriteOptions) ingest.BatchError {
   401  			capturedIter = iter
   402  			return nil
   403  		})
   404  
   405  	opts := makeOptions(mockDownsamplerAndWriter).SetStoreMetricsType(false)
   406  
   407  	promReq := &prompb.WriteRequest{
   408  		Timeseries: []prompb.TimeSeries{
   409  			{Type: prompb.MetricType_COUNTER},
   410  			{},
   411  		},
   412  	}
   413  
   414  	executeWriteRequest(t, opts, promReq)
   415  
   416  	verifyIterValueNoAnnotation(t, capturedIter)
   417  	verifyIterValueNoAnnotation(t, capturedIter)
   418  
   419  	require.False(t, capturedIter.Next())
   420  	require.NoError(t, capturedIter.Error())
   421  }
   422  
   423  func TestPromWriteLiteralIsTooLongError(t *testing.T) {
   424  	ctrl := xtest.NewController(t)
   425  	defer ctrl.Finish()
   426  
   427  	opts := makeOptions(ingest.NewMockDownsamplerAndWriter(ctrl))
   428  	handler, err := NewPromWriteHandler(opts)
   429  	require.NoError(t, err)
   430  
   431  	veryLongLiteral := strings.Repeat("x", int(opts.TagOptions().MaxTagLiteralLength())+1)
   432  	promReq := &prompb.WriteRequest{
   433  		Timeseries: []prompb.TimeSeries{
   434  			{
   435  				Labels: []prompb.Label{
   436  					{Name: []byte("name1"), Value: []byte("value1")},
   437  					{Name: []byte("name2"), Value: []byte(veryLongLiteral)},
   438  				},
   439  			},
   440  		},
   441  	}
   442  
   443  	for i := 0; i < maxLiteralIsTooLongLogCount*2; i++ {
   444  		promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   445  		req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   446  		writer := httptest.NewRecorder()
   447  		handler.ServeHTTP(writer, req)
   448  		resp := writer.Result()
   449  		require.Equal(t, http.StatusBadRequest, resp.StatusCode)
   450  		require.NoError(t, resp.Body.Close())
   451  	}
   452  }
   453  
   454  func TestPromWriteForwardWithShadow(t *testing.T) {
   455  	for _, tt := range []struct {
   456  		percent         float64
   457  		numSeries       int
   458  		allowedVariance float64
   459  	}{
   460  		{0, 10000, 0},
   461  		{0.25, 10000, 0.05},
   462  		{0.5, 10000, 0.05},
   463  		{0.75, 10000, 0.05},
   464  		{1, 10000, 0},
   465  	} {
   466  		for _, h := range []string{"", "murmur3", "xxhash"} {
   467  			h := h
   468  			t.Run(fmt.Sprintf("hash='%s', params=%+v", h, tt), func(t *testing.T) {
   469  				testPromWriteForwardWithShadow(t, testPromWriteForwardWithShadowOptions{
   470  					hash:                         h,
   471  					numSeries:                    tt.numSeries,
   472  					percent:                      tt.percent,
   473  					expectedFwded:                int(float64(tt.numSeries) * tt.percent),
   474  					expectedFwdedAllowedVariance: tt.allowedVariance,
   475  				})
   476  			})
   477  		}
   478  	}
   479  }
   480  
   481  type testPromWriteForwardWithShadowOptions struct {
   482  	numSeries                    int
   483  	percent                      float64
   484  	hash                         string
   485  	expectedFwded                int
   486  	expectedFwdedAllowedVariance float64
   487  }
   488  
   489  func testPromWriteForwardWithShadow(
   490  	t *testing.T,
   491  	testOpts testPromWriteForwardWithShadowOptions,
   492  ) {
   493  	ctrl := xtest.NewController(t)
   494  	defer ctrl.Finish()
   495  
   496  	// Create forwarding receiver.
   497  	forwardRecvReqCh := make(chan *prompb.WriteRequest, 1)
   498  	forwardRecvSvr := httptest.NewServer(
   499  		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   500  			forwardRecvReqCh <- test.ReadPromWriteRequestBody(t, r.Body)
   501  			w.WriteHeader(http.StatusOK)
   502  		}))
   503  	defer forwardRecvSvr.Close()
   504  
   505  	target := handleroptions.PromWriteHandlerForwardTargetOptions{
   506  		URL:     forwardRecvSvr.URL,
   507  		Method:  http.MethodPost,
   508  		NoRetry: true,
   509  		Shadow: &handleroptions.PromWriteHandlerForwardTargetShadowOptions{
   510  			Percent: testOpts.percent,
   511  			Hash:    testOpts.hash,
   512  		},
   513  	}
   514  
   515  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   516  	mockDownsamplerAndWriter.
   517  		EXPECT().
   518  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any())
   519  
   520  	// Setup opts and modify config.
   521  	opts := makeOptions(mockDownsamplerAndWriter)
   522  
   523  	cfg := opts.Config()
   524  	cfg.WriteForwarding.PromRemoteWrite.Targets = append(cfg.WriteForwarding.PromRemoteWrite.Targets, target)
   525  
   526  	opts = opts.SetConfig(cfg)
   527  
   528  	handler, err := NewPromWriteHandler(opts)
   529  	require.NoError(t, err)
   530  
   531  	promReq := &prompb.WriteRequest{}
   532  	for i := 0; i < testOpts.numSeries; i++ {
   533  		series := prompb.TimeSeries{
   534  			Labels: []prompb.Label{
   535  				{Name: []byte("__name__"), Value: []byte(fmt.Sprintf("name_%d", i))},
   536  			},
   537  			Samples: []prompb.Sample{
   538  				{Timestamp: time.Now().UnixMilli(), Value: 42},
   539  			},
   540  		}
   541  
   542  		// Add some labels, unsorted.
   543  		for j := 0; j < 5; j++ {
   544  			label := prompb.Label{Name: make([]byte, 16), Value: make([]byte, 16)}
   545  			_, err = rand.Reader.Read(label.Name)
   546  			require.NoError(t, err)
   547  			_, err = rand.Reader.Read(label.Value)
   548  			require.NoError(t, err)
   549  			// Add to start or end.
   550  			if j%2 == 0 {
   551  				series.Labels = append(series.Labels, label)
   552  			} else {
   553  				series.Labels = append([]prompb.Label{label}, series.Labels...)
   554  			}
   555  		}
   556  
   557  		promReq.Timeseries = append(promReq.Timeseries, series)
   558  	}
   559  
   560  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   561  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   562  	writer := httptest.NewRecorder()
   563  	handler.ServeHTTP(writer, req)
   564  	resp := writer.Result()
   565  	require.Equal(t, http.StatusOK, resp.StatusCode)
   566  	require.NoError(t, resp.Body.Close())
   567  
   568  	select {
   569  	case fwdReq := <-forwardRecvReqCh:
   570  		if testOpts.expectedFwdedAllowedVariance > 0 {
   571  			assert.InEpsilon(t, testOpts.expectedFwded, len(fwdReq.Timeseries),
   572  				testOpts.expectedFwdedAllowedVariance,
   573  				fmt.Sprintf("expected=%v, actual=%v, allowed_variance=%v",
   574  					testOpts.expectedFwded, len(fwdReq.Timeseries),
   575  					testOpts.expectedFwdedAllowedVariance))
   576  		} else {
   577  			assert.Equal(t, testOpts.expectedFwded, len(fwdReq.Timeseries),
   578  				fmt.Sprintf("expected=%v, actual=%v",
   579  					testOpts.expectedFwded, len(fwdReq.Timeseries)))
   580  		}
   581  	case <-time.After(10 * time.Second):
   582  		require.FailNow(t, "timeout waiting for fwd request")
   583  	}
   584  }
   585  
   586  func BenchmarkWriteDatapoints(b *testing.B) {
   587  	ctrl := xtest.NewController(b)
   588  	defer ctrl.Finish()
   589  
   590  	mockDownsamplerAndWriter := ingest.NewMockDownsamplerAndWriter(ctrl)
   591  	mockDownsamplerAndWriter.
   592  		EXPECT().
   593  		WriteBatch(gomock.Any(), gomock.Any(), gomock.Any()).
   594  		AnyTimes()
   595  
   596  	opts := makeOptions(mockDownsamplerAndWriter)
   597  	handler, err := NewPromWriteHandler(opts)
   598  	require.NoError(b, err)
   599  
   600  	promReq := test.GeneratePromWriteRequest()
   601  	promReqBody := test.GeneratePromWriteRequestBodyBytes(b, promReq)
   602  	promReqBodyReader := bytes.NewReader(nil)
   603  
   604  	for i := 0; i < b.N; i++ {
   605  		promReqBodyReader.Reset(promReqBody)
   606  		req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBodyReader)
   607  		handler.ServeHTTP(httptest.NewRecorder(), req)
   608  	}
   609  }
   610  
   611  func verifyIterValueAnnotation(
   612  	t *testing.T,
   613  	iter ingest.DownsampleAndWriteIter,
   614  	expectedMetricType annotation.OpenMetricsFamilyType,
   615  	expectedHandleValueResets bool,
   616  ) ingest.IterValue {
   617  	require.True(t, iter.Next())
   618  	value := iter.Current()
   619  
   620  	expectedPayload := annotation.Payload{
   621  		SourceFormat:                 annotation.SourceFormat_OPEN_METRICS,
   622  		OpenMetricsFamilyType:        expectedMetricType,
   623  		OpenMetricsHandleValueResets: expectedHandleValueResets,
   624  	}
   625  	assert.Equal(t, expectedPayload, unmarshalAnnotation(t, value.Annotation))
   626  
   627  	return value
   628  }
   629  
   630  func verifyIterValueAnnotationGraphite(
   631  	t *testing.T,
   632  	iter ingest.DownsampleAndWriteIter,
   633  	expectedMetricType annotation.GraphiteType,
   634  ) {
   635  	require.True(t, iter.Next())
   636  	value := iter.Current()
   637  
   638  	expectedPayload := annotation.Payload{
   639  		SourceFormat: annotation.SourceFormat_GRAPHITE,
   640  		GraphiteType: expectedMetricType,
   641  	}
   642  	assert.Equal(t, expectedPayload, unmarshalAnnotation(t, value.Annotation))
   643  }
   644  
   645  func verifyIterValueNoAnnotation(t *testing.T, iter ingest.DownsampleAndWriteIter) {
   646  	require.True(t, iter.Next())
   647  	value := iter.Current()
   648  	assert.Nil(t, value.Annotation)
   649  }
   650  
   651  func unmarshalAnnotation(t *testing.T, annot []byte) annotation.Payload {
   652  	payload := annotation.Payload{}
   653  	require.NoError(t, payload.Unmarshal(annot))
   654  	return payload
   655  }
   656  
   657  func executeWriteRequest(t *testing.T, handlerOpts options.HandlerOptions, promReq *prompb.WriteRequest) {
   658  	handler, err := NewPromWriteHandler(handlerOpts)
   659  	require.NoError(t, err)
   660  
   661  	promReqBody := test.GeneratePromWriteRequestBody(t, promReq)
   662  	req := httptest.NewRequest(PromWriteHTTPMethod, PromWriteURL, promReqBody)
   663  
   664  	writer := httptest.NewRecorder()
   665  	handler.ServeHTTP(writer, req)
   666  	resp := writer.Result()
   667  	require.Equal(t, http.StatusOK, resp.StatusCode)
   668  }