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 }