github.com/thanos-io/thanos@v0.32.5/pkg/store/limiter_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package store
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/efficientgo/core/testutil"
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"github.com/prometheus/client_golang/prometheus/promauto"
    14  	prom_testutil "github.com/prometheus/client_golang/prometheus/testutil"
    15  	"github.com/prometheus/prometheus/model/labels"
    16  
    17  	"github.com/thanos-io/thanos/pkg/store/storepb"
    18  )
    19  
    20  func TestLimiter(t *testing.T) {
    21  	c := promauto.With(nil).NewCounter(prometheus.CounterOpts{})
    22  	l := NewLimiter(10, c)
    23  
    24  	testutil.Ok(t, l.Reserve(5))
    25  	testutil.Equals(t, float64(0), prom_testutil.ToFloat64(c))
    26  
    27  	testutil.Ok(t, l.Reserve(5))
    28  	testutil.Equals(t, float64(0), prom_testutil.ToFloat64(c))
    29  
    30  	testutil.NotOk(t, l.Reserve(1))
    31  	testutil.Equals(t, float64(1), prom_testutil.ToFloat64(c))
    32  
    33  	testutil.NotOk(t, l.Reserve(2))
    34  	testutil.Equals(t, float64(1), prom_testutil.ToFloat64(c))
    35  }
    36  
    37  func TestRateLimitedServer(t *testing.T) {
    38  	numSamples := 60
    39  	series := []*storepb.SeriesResponse{
    40  		storeSeriesResponse(t, labels.FromStrings("series", "1"), makeSamples(numSamples)),
    41  		storeSeriesResponse(t, labels.FromStrings("series", "2"), makeSamples(numSamples)),
    42  		storeSeriesResponse(t, labels.FromStrings("series", "3"), makeSamples(numSamples)),
    43  	}
    44  	tests := []struct {
    45  		name   string
    46  		limits SeriesSelectLimits
    47  		series []*storepb.SeriesResponse
    48  		err    string
    49  	}{
    50  		{
    51  			name: "no limits",
    52  			limits: SeriesSelectLimits{
    53  				SeriesPerRequest:  0,
    54  				SamplesPerRequest: 0,
    55  			},
    56  			series: series,
    57  		},
    58  		{
    59  			name: "series bellow limit",
    60  			limits: SeriesSelectLimits{
    61  				SeriesPerRequest:  3,
    62  				SamplesPerRequest: 0,
    63  			},
    64  			series: series,
    65  		},
    66  		{
    67  			name: "series over limit",
    68  			limits: SeriesSelectLimits{
    69  				SeriesPerRequest:  2,
    70  				SamplesPerRequest: 0,
    71  			},
    72  			series: series,
    73  			err:    "failed to send series: limit 2 violated (got 3)",
    74  		},
    75  		{
    76  			name: "chunks bellow limit",
    77  			limits: SeriesSelectLimits{
    78  				SeriesPerRequest:  0,
    79  				SamplesPerRequest: uint64(3 * numSamples * MaxSamplesPerChunk),
    80  			},
    81  			series: series,
    82  		},
    83  		{
    84  			name: "chunks over limit",
    85  			limits: SeriesSelectLimits{
    86  				SeriesPerRequest:  0,
    87  				SamplesPerRequest: 50,
    88  			},
    89  			series: series,
    90  			err:    "failed to send samples: limit 50 violated (got 120)",
    91  		},
    92  	}
    93  	for _, test := range tests {
    94  		t.Run(test.name, func(t *testing.T) {
    95  			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
    96  			defer cancel()
    97  
    98  			store := NewLimitedStoreServer(newStoreServerStub(test.series), prometheus.NewRegistry(), test.limits)
    99  			seriesServer := storepb.NewInProcessStream(ctx, 10)
   100  			err := store.Series(&storepb.SeriesRequest{}, seriesServer)
   101  			if test.err == "" {
   102  				testutil.Ok(t, err)
   103  			} else {
   104  				testutil.NotOk(t, err)
   105  				testutil.Assert(t, test.err == err.Error(), "want %s, got %s", test.err, err.Error())
   106  			}
   107  		})
   108  	}
   109  }
   110  
   111  func makeSamples(numSamples int) []sample {
   112  	samples := make([]sample, numSamples)
   113  	for i := range samples {
   114  		samples[i] = sample{t: int64(i), v: float64(i)}
   115  	}
   116  	return samples
   117  }
   118  
   119  type testStoreServer struct {
   120  	storepb.StoreServer
   121  	responses []*storepb.SeriesResponse
   122  }
   123  
   124  func newStoreServerStub(responses []*storepb.SeriesResponse) *testStoreServer {
   125  	return &testStoreServer{responses: responses}
   126  }
   127  
   128  func (m *testStoreServer) Series(_ *storepb.SeriesRequest, server storepb.Store_SeriesServer) error {
   129  	for _, r := range m.responses {
   130  		if err := server.Send(r); err != nil {
   131  			return err
   132  		}
   133  	}
   134  	return nil
   135  }