store: add ability to limit max num of samples / concurrent queries

cmd/thanos/store.go

@@ -36,6 +36,12 @@ func registerStore(m map[string]setupFunc, app *kingpin.Application, name string
 	chunkPoolSize := cmd.Flag("chunk-pool-size", "Maximum size of concurrently allocatable bytes for chunks.").
 		Default("2GB").Bytes()
 
+	maxSampleCount := cmd.Flag("grpc-sample-limit",
+		"Maximum amount of samples returned via a single Series call. 0 means no limit. NOTE: may overestimate the number of samples that would be needed to respond to a query.").
+		Default("50000000").Uint()
+
+	maxConcurrent := cmd.Flag("grpc-concurrent-limit", "Maximum number of concurrent Series calls. 0 means no limit.").Default("20").Int()
+
 	objStoreConfig := regCommonObjStoreFlags(cmd, "", true)
 
 	syncInterval := cmd.Flag("sync-block-duration", "Repeat interval for syncing the blocks between local and remote view.").
@@ -63,6 +69,8 @@ func registerStore(m map[string]setupFunc, app *kingpin.Application, name string
 			peer,
 			uint64(*indexCacheSize),
 			uint64(*chunkPoolSize),
+			uint64(*maxSampleCount),
+			int(*maxConcurrent),
 			name,
 			debugLogging,
 			*syncInterval,
@@ -87,6 +95,8 @@ func runStore(
 	peer cluster.Peer,
 	indexCacheSizeBytes uint64,
 	chunkPoolSizeBytes uint64,
+	maxSampleCount uint64,
+	maxConcurrent int,
 	component string,
 	verbose bool,
 	syncInterval time.Duration,
@@ -117,6 +127,8 @@ func runStore(
 			dataDir,
 			indexCacheSizeBytes,
 			chunkPoolSizeBytes,
+			maxSampleCount,
+			maxConcurrent,
 			verbose,
 			blockSyncConcurrency,
 		)

docs/components/store.md

@@ -104,6 +104,14 @@ Flags:
       --index-cache-size=250MB   Maximum size of items held in the index cache.
       --chunk-pool-size=2GB      Maximum size of concurrently allocatable bytes
                                  for chunks.
+      --grpc-sample-limit=50000000  
+                                 Maximum amount of samples returned via a single
+                                 Series call. 0 means no limit. NOTE: may
+                                 overestimate the number of samples that would
+                                 be needed to respond to a query.
+      --grpc-concurrent-limit=20  
+                                 Maximum number of concurrent Series calls. 0
+                                 means no limit.
       --objstore.config-file=<bucket.config-yaml-path>  
                                  Path to YAML file that contains object store
                                  configuration.

pkg/store/bucket.go

@@ -42,6 +42,14 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// Approximately this is the max number of samples that we may have in any given chunk. This is needed
+// for precalculating the number of samples that we may have to retrieve and decode for any given query
+// without downloading them. Please take a look at https://github.com/prometheus/tsdb/pull/397 to know
+// where this number comes from. Long story short: TSDB is made in such a way, and it is made in such a way
+// because you barely get any improvements in compression when the number of samples is beyond this.
+// Take a look at Figure 6 in this whitepaper http://www.vldb.org/pvldb/vol8/p1816-teller.pdf.
+const maxSamplesPerChunk uint64 = 120
+
 type bucketStoreMetrics struct {
 	blocksLoaded          prometheus.Gauge
 	blockLoads            prometheus.Counter
@@ -57,6 +65,8 @@ type bucketStoreMetrics struct {
 	seriesMergeDuration   prometheus.Histogram
 	resultSeriesCount     prometheus.Summary
 	chunkSizeBytes        prometheus.Histogram
+	queriesLimited        prometheus.Counter
+	queriesLimit          prometheus.Gauge
 }
 
 func newBucketStoreMetrics(reg prometheus.Registerer) *bucketStoreMetrics {
@@ -132,6 +142,15 @@ func newBucketStoreMetrics(reg prometheus.Registerer) *bucketStoreMetrics {
 		},
 	})
 
+	m.queriesLimited = prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "thanos_bucket_store_queries_limited",
+		Help: "Number of queries that were dropped due to the sample limit.",
+	})
+	m.queriesLimit = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "thanos_bucket_store_queries_limit",
+		Help: "Number of maximum concurrent queries.",
+	})
+
 	if reg != nil {
 		reg.MustRegister(
 			m.blockLoads,
@@ -148,6 +167,8 @@ func newBucketStoreMetrics(reg prometheus.Registerer) *bucketStoreMetrics {
 			m.seriesMergeDuration,
 			m.resultSeriesCount,
 			m.chunkSizeBytes,
+			m.queriesLimited,
+			m.queriesLimit,
 		)
 	}
 	return &m
@@ -173,7 +194,12 @@ type BucketStore struct {
 	// Number of goroutines to use when syncing blocks from object storage.
 	blockSyncConcurrency int
 
-	partitioner partitioner
+	// Query gate which limits the maximum amount of concurrent queries.
+	queryGate *Gate
+
+	// Samples limiter which limits the number of samples per each Series() call.
+	samplesLimiter *Limiter
+	partitioner    partitioner
 }
 
 // NewBucketStore creates a new bucket backed store that implements the store API against
@@ -185,12 +211,19 @@ func NewBucketStore(
 	dir string,
 	indexCacheSizeBytes uint64,
 	maxChunkPoolBytes uint64,
+	maxSampleCount uint64,
+	maxConcurrent int,
 	debugLogging bool,
 	blockSyncConcurrency int,
 ) (*BucketStore, error) {
 	if logger == nil {
 		logger = log.NewNopLogger()
 	}
+
+	if maxConcurrent < 0 {
+		return nil, errors.Errorf("max concurrency value cannot be lower than 0 (got %v)", maxConcurrent)
+	}
+
 	indexCache, err := newIndexCache(reg, indexCacheSizeBytes)
 	if err != nil {
 		return nil, errors.Wrap(err, "create index cache")
@@ -202,6 +235,7 @@ func NewBucketStore(
 
 	const maxGapSize = 512 * 1024
 
+	metrics := newBucketStoreMetrics(reg)
 	s := &BucketStore{
 		logger:               logger,
 		bucket:               bucket,
@@ -212,14 +246,18 @@ func NewBucketStore(
 		blockSets:            map[uint64]*bucketBlockSet{},
 		debugLogging:         debugLogging,
 		blockSyncConcurrency: blockSyncConcurrency,
+		queryGate:            NewGate(maxConcurrent, reg),
+		samplesLimiter:       NewLimiter(maxSampleCount, &metrics.queriesLimited),
 		partitioner:          gapBasedPartitioner{maxGapSize: maxGapSize},
 	}
-	s.metrics = newBucketStoreMetrics(reg)
+	s.metrics = metrics
 
 	if err := os.MkdirAll(dir, 0777); err != nil {
 		return nil, errors.Wrap(err, "create dir")
 	}
 
+	s.metrics.queriesLimit.Set(float64(maxConcurrent))
+
 	return s, nil
 }
 
@@ -472,14 +510,15 @@ func (s *bucketSeriesSet) Err() error {
 	return s.err
 }
 
-func (s *BucketStore) blockSeries(
+func blockSeries(
 	ctx context.Context,
 	ulid ulid.ULID,
 	extLset map[string]string,
 	indexr *bucketIndexReader,
 	chunkr *bucketChunkReader,
 	matchers []labels.Matcher,
 	req *storepb.SeriesRequest,
+	samplesLimiter *Limiter,
 ) (storepb.SeriesSet, *queryStats, error) {
 	ps, err := indexr.ExpandedPostings(matchers)
 	if err != nil {
@@ -557,7 +596,7 @@ func (s *BucketStore) blockSeries(
 	}
 
 	// Preload all chunks that were marked in the previous stage.
-	if err := chunkr.preload(); err != nil {
+	if err := chunkr.preload(samplesLimiter); err != nil {
 		return nil, nil, errors.Wrap(err, "preload chunks")
 	}
 
@@ -661,10 +700,17 @@ func debugFoundBlockSetOverview(logger log.Logger, mint, maxt int64, lset labels
 }
 
 // Series implements the storepb.StoreServer interface.
-// TODO(bwplotka): It buffers all chunks in memory and only then streams to client.
-// 1. Either count chunk sizes and error out too big query.
-// 2. Stream posting -> series -> chunk all together.
 func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storepb.Store_SeriesServer) error {
+	{
+		span, _ := tracing.StartSpan(srv.Context(), "bucket_store_gate")
+		err := s.queryGate.IsMyTurn(srv.Context())
+		span.Finish()
+		if err != nil {
+			return errors.Wrapf(err, "failed to wait for turn")
+		}
+		defer s.queryGate.Done()
+	}
+
 	matchers, err := translateMatchers(req.Matchers)
 	if err != nil {
 		return status.Error(codes.InvalidArgument, err.Error())
@@ -703,13 +749,14 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storepb.Store_Serie
 			defer runutil.CloseWithLogOnErr(s.logger, chunkr, "series block")
 
 			g.Add(func() error {
-				part, pstats, err := s.blockSeries(ctx,
+				part, pstats, err := blockSeries(ctx,
 					b.meta.ULID,
 					b.meta.Thanos.Labels,
 					indexr,
 					chunkr,
 					blockMatchers,
 					req,
+					s.samplesLimiter,
 				)
 				if err != nil {
 					return errors.Wrapf(err, "fetch series for block %s", b.meta.ULID)
@@ -1585,11 +1632,22 @@ func (r *bucketChunkReader) addPreload(id uint64) error {
 }
 
 // preload all added chunk IDs. Must be called before the first call to Chunk is made.
-func (r *bucketChunkReader) preload() error {
+func (r *bucketChunkReader) preload(samplesLimiter *Limiter) error {
 	const maxChunkSize = 16000
 
 	var g run.Group
 
+	numChunks := uint64(0)
+	for _, offsets := range r.preloads {
+		for range offsets {
+			numChunks++
+		}
+	}
+	numSamples := numChunks * maxSamplesPerChunk
+	if err := samplesLimiter.Check(numSamples); err != nil {
+		return errors.Wrapf(err, "exceeded samples limit")
+	}
+
 	for seq, offsets := range r.preloads {
 		sort.Slice(offsets, func(i, j int) bool {
 			return offsets[i] < offsets[j]

pkg/store/bucket_e2e_test.go

@@ -35,7 +35,7 @@ func (s *storeSuite) Close() {
 	s.wg.Wait()
 }
 
-func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, manyParts bool) *storeSuite {
+func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, manyParts bool, maxSampleCount uint64) *storeSuite {
 	series := []labels.Labels{
 		labels.FromStrings("a", "1", "b", "1"),
 		labels.FromStrings("a", "1", "b", "2"),
@@ -87,7 +87,7 @@ func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, m
 		testutil.Ok(t, os.RemoveAll(dir2))
 	}
 
-	store, err := NewBucketStore(log.NewLogfmtLogger(os.Stderr), nil, bkt, dir, 100, 0, false, 20)
+	store, err := NewBucketStore(log.NewLogfmtLogger(os.Stderr), nil, bkt, dir, 100, 0, maxSampleCount, 20, false, 20)
 	testutil.Ok(t, err)
 
 	s.store = store
@@ -321,7 +321,7 @@ func TestBucketStore_e2e(t *testing.T) {
 		testutil.Ok(t, err)
 		defer func() { testutil.Ok(t, os.RemoveAll(dir)) }()
 
-		s := prepareStoreWithTestBlocks(t, dir, bkt, false)
+		s := prepareStoreWithTestBlocks(t, dir, bkt, false, 0)
 		defer s.Close()
 
 		testBucketStore_e2e(t, ctx, s)
@@ -350,7 +350,7 @@ func TestBucketStore_ManyParts_e2e(t *testing.T) {
 		testutil.Ok(t, err)
 		defer func() { testutil.Ok(t, os.RemoveAll(dir)) }()
 
-		s := prepareStoreWithTestBlocks(t, dir, bkt, true)
+		s := prepareStoreWithTestBlocks(t, dir, bkt, true, 0)
 		defer s.Close()
 
 		testBucketStore_e2e(t, ctx, s)

pkg/store/bucket_test.go

@@ -283,7 +283,7 @@ func TestBucketStore_Info(t *testing.T) {
 	dir, err := ioutil.TempDir("", "prometheus-test")
 	testutil.Ok(t, err)
 
-	bucketStore, err := NewBucketStore(nil, nil, nil, dir, 2e5, 2e5, false, 20)
+	bucketStore, err := NewBucketStore(nil, nil, nil, dir, 2e5, 2e5, 0, 0, false, 20)
 	testutil.Ok(t, err)
 
 	resp, err := bucketStore.Info(ctx, &storepb.InfoRequest{})

pkg/store/gate.go

@@ -0,0 +1,54 @@
+package store
+
+import (
+	"context"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/pkg/gate"
+)
+
+// Gate wraps the Prometheus gate with extra metrics.
+type Gate struct {
+	g              *gate.Gate
+	currentQueries prometheus.Gauge
+	gateTiming     prometheus.Summary
+}
+
+// NewGate returns a new gate.
+func NewGate(maxConcurrent int, reg prometheus.Registerer) *Gate {
+	g := &Gate{
+		g: gate.New(maxConcurrent),
+	}
+	g.currentQueries = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "thanos_bucket_store_queries_total",
+		Help: "Total number of currently executing queries.",
+	})
+	g.gateTiming = prometheus.NewSummary(prometheus.SummaryOpts{
+		Name: "thanos_bucket_store_gate_seconds",
+		Help: "How many seconds it took for a query to pass through the gate.",
+	})
+
+	if reg != nil {
+		reg.MustRegister(g.currentQueries, g.gateTiming)
+	}
+
+	return g
+}
+
+// IsMyTurn iniates a new query and waits until it's our turn to fulfill a query request.
+func (g *Gate) IsMyTurn(ctx context.Context) error {
+	g.currentQueries.Inc()
+	start := time.Now()
+	if err := g.g.Start(ctx); err != nil {
+		return err
+	}
+	g.gateTiming.Observe(float64(time.Now().Sub(start)))
+	return nil
+}
+
+// Done finishes a query.
+func (g *Gate) Done() {
+	g.currentQueries.Dec()
+	g.g.Done()
+}

pkg/store/limiter.go

@@ -0,0 +1,33 @@
+package store
+
+import (
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+// Limiter is a simple mechanism for checking if something has passed a certain threshold.
+type Limiter struct {
+	limit uint64
+
+	// A ptr to a counter metric which we will increase if Check() fails.
+	failedCounter *prometheus.Counter
+}
+
+// NewLimiter returns a new limiter with a specified limit. 0 disables the limit.
+func NewLimiter(limit uint64, ctr *prometheus.Counter) *Limiter {
+	return &Limiter{limit: limit, failedCounter: ctr}
+}
+
+// Check checks if the passed number exceeds the limits or not.
+func (l *Limiter) Check(num uint64) error {
+	if l.limit == 0 {
+		return nil
+	}
+	if num > l.limit {
+		if l.failedCounter != nil {
+			(*l.failedCounter).Inc()
+		}
+		return errors.Errorf("limit %v violated (got %v)", l.limit, num)
+	}
+	return nil
+}