Add limiter (#81)

* example limiter

Signed-off-by: yeya24 <[email protected]>

* address comments and add tests

Signed-off-by: yeya24 <[email protected]>

---------

Signed-off-by: yeya24 <[email protected]>

Author: yeya24

queryable/parquet_queryable.go

@@ -34,11 +34,19 @@ import (
 type ShardsFinderFunction func(ctx context.Context, mint, maxt int64) ([]storage.ParquetShard, error)
 
 type queryableOpts struct {
-	concurrency int
+	concurrency                int
+	rowCountLimitFunc          search.QuotaLimitFunc
+	chunkBytesLimitFunc        search.QuotaLimitFunc
+	dataBytesLimitFunc         search.QuotaLimitFunc
+	materializedSeriesCallback search.MaterializedSeriesFunc
 }
 
 var DefaultQueryableOpts = queryableOpts{
-	concurrency: runtime.GOMAXPROCS(0),
+	concurrency:                runtime.GOMAXPROCS(0),
+	rowCountLimitFunc:          search.NoopQuotaLimitFunc,
+	chunkBytesLimitFunc:        search.NoopQuotaLimitFunc,
+	dataBytesLimitFunc:         search.NoopQuotaLimitFunc,
+	materializedSeriesCallback: search.NoopMaterializedSeriesFunc,
 }
 
 type QueryableOpts func(*queryableOpts)
@@ -50,6 +58,35 @@ func WithConcurrency(concurrency int) QueryableOpts {
 	}
 }
 
+// WithRowCountLimitFunc sets a callback function to get limit for matched row count.
+func WithRowCountLimitFunc(fn search.QuotaLimitFunc) QueryableOpts {
+	return func(opts *queryableOpts) {
+		opts.rowCountLimitFunc = fn
+	}
+}
+
+// WithChunkBytesLimitFunc sets a callback function to get limit for chunk column page bytes fetched.
+func WithChunkBytesLimitFunc(fn search.QuotaLimitFunc) QueryableOpts {
+	return func(opts *queryableOpts) {
+		opts.chunkBytesLimitFunc = fn
+	}
+}
+
+// WithDataBytesLimitFunc sets a callback function to get limit for data (including label and chunk)
+// column page bytes fetched.
+func WithDataBytesLimitFunc(fn search.QuotaLimitFunc) QueryableOpts {
+	return func(opts *queryableOpts) {
+		opts.dataBytesLimitFunc = fn
+	}
+}
+
+// WithMaterializedSeriesCallback sets a callback function to process the materialized series.
+func WithMaterializedSeriesCallback(fn search.MaterializedSeriesFunc) QueryableOpts {
+	return func(opts *queryableOpts) {
+		opts.materializedSeriesCallback = fn
+	}
+}
+
 type parquetQueryable struct {
 	shardsFinder ShardsFinderFunction
 	d            *schema.PrometheusParquetChunksDecoder
@@ -191,8 +228,11 @@ func (p parquetQuerier) queryableShards(ctx context.Context, mint, maxt int64) (
 		return nil, err
 	}
 	qBlocks := make([]*queryableShard, len(shards))
+	rowCountQuota := search.NewQuota(p.opts.rowCountLimitFunc(ctx))
+	chunkBytesQuota := search.NewQuota(p.opts.chunkBytesLimitFunc(ctx))
+	dataBytesQuota := search.NewQuota(p.opts.dataBytesLimitFunc(ctx))
 	for i, shard := range shards {
-		qb, err := newQueryableShard(p.opts, shard, p.d)
+		qb, err := newQueryableShard(p.opts, shard, p.d, rowCountQuota, chunkBytesQuota, dataBytesQuota)
 		if err != nil {
 			return nil, err
 		}
@@ -207,12 +247,12 @@ type queryableShard struct {
 	concurrency int
 }
 
-func newQueryableShard(opts *queryableOpts, block storage.ParquetShard, d *schema.PrometheusParquetChunksDecoder) (*queryableShard, error) {
+func newQueryableShard(opts *queryableOpts, block storage.ParquetShard, d *schema.PrometheusParquetChunksDecoder, rowCountQuota *search.Quota, chunkBytesQuota *search.Quota, dataBytesQuota *search.Quota) (*queryableShard, error) {
 	s, err := block.TSDBSchema()
 	if err != nil {
 		return nil, err
 	}
-	m, err := search.NewMaterializer(s, d, block, opts.concurrency)
+	m, err := search.NewMaterializer(s, d, block, opts.concurrency, rowCountQuota, chunkBytesQuota, dataBytesQuota, opts.materializedSeriesCallback)
 	if err != nil {
 		return nil, err
 	}

queryable/parquet_queryable_test.go

@@ -35,6 +35,7 @@ import (
 
 	"github.com/prometheus-community/parquet-common/convert"
 	"github.com/prometheus-community/parquet-common/schema"
+	"github.com/prometheus-community/parquet-common/search"
 	"github.com/prometheus-community/parquet-common/storage"
 	"github.com/prometheus-community/parquet-common/util"
 )
@@ -269,6 +270,126 @@ func TestQueryable(t *testing.T) {
 					require.Equal(t, expectedLabelValues, lValues)
 				})
 			})
+
+			t.Run("RowCountQuota", func(t *testing.T) {
+				// Test with limited row count quota
+				limitedRowQuota := func(ctx context.Context) int64 {
+					return 10 // Only allow 10 rows
+				}
+				queryable, err := createQueryable(shard, WithRowCountLimitFunc(limitedRowQuota))
+				require.NoError(t, err)
+				querier, err := queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				// Try to query more rows than quota allows
+				matchers := []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "unique", "unique_0")}
+				ss := querier.Select(ctx, true, nil, matchers...)
+
+				// This should fail due to row count quota
+				for ss.Next() {
+					_ = ss.At()
+				}
+				require.Error(t, ss.Err())
+				require.Contains(t, ss.Err().Error(), "would fetch too many rows")
+				require.True(t, search.IsResourceExhausted(ss.Err()))
+
+				// Test with sufficient quota
+				sufficientRowQuota := func(ctx context.Context) int64 {
+					return 1000 // Allow 1000 rows
+				}
+				queryable, err = createQueryable(shard, WithRowCountLimitFunc(sufficientRowQuota))
+				require.NoError(t, err)
+				querier, err = queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				ss = querier.Select(ctx, true, nil, matchers...)
+				var series []prom_storage.Series
+				for ss.Next() {
+					series = append(series, ss.At())
+				}
+				require.NoError(t, ss.Err())
+				require.NotEmpty(t, series)
+			})
+
+			t.Run("ChunkBytesQuota", func(t *testing.T) {
+				// Test with limited chunk bytes quota
+				limitedChunkQuota := func(ctx context.Context) int64 {
+					return 100 // Only allow 100 bytes
+				}
+				queryable, err := createQueryable(shard, WithChunkBytesLimitFunc(limitedChunkQuota))
+				require.NoError(t, err)
+				querier, err := queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				// Try to query chunks that exceed the quota
+				matchers := []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "unique", "unique_0")}
+				ss := querier.Select(ctx, true, nil, matchers...)
+
+				// This should fail due to chunk bytes quota
+				for ss.Next() {
+					_ = ss.At()
+				}
+				require.Error(t, ss.Err())
+				require.Contains(t, ss.Err().Error(), "would fetch too many chunk bytes")
+				require.True(t, search.IsResourceExhausted(ss.Err()))
+
+				// Test with sufficient quota
+				sufficientChunkQuota := func(ctx context.Context) int64 {
+					return 1000000 // Allow 1MB
+				}
+				queryable, err = createQueryable(shard, WithChunkBytesLimitFunc(sufficientChunkQuota))
+				require.NoError(t, err)
+				querier, err = queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				ss = querier.Select(ctx, true, nil, matchers...)
+				var series []prom_storage.Series
+				for ss.Next() {
+					series = append(series, ss.At())
+				}
+				require.NoError(t, ss.Err())
+				require.NotEmpty(t, series)
+			})
+
+			t.Run("DataBytesQuota", func(t *testing.T) {
+				// Test with limited data bytes quota
+				limitedDataQuota := func(ctx context.Context) int64 {
+					return 100 // Only allow 100 bytes
+				}
+				queryable, err := createQueryable(shard, WithDataBytesLimitFunc(limitedDataQuota))
+				require.NoError(t, err)
+				querier, err := queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				// Try to query data that exceeds the quota
+				matchers := []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "unique", "unique_0")}
+				ss := querier.Select(ctx, true, nil, matchers...)
+
+				// This should fail due to data bytes quota
+				for ss.Next() {
+					_ = ss.At()
+				}
+				require.Error(t, ss.Err())
+				require.Contains(t, ss.Err().Error(), "would fetch too many data bytes")
+				require.True(t, search.IsResourceExhausted(ss.Err()))
+
+				// Test with sufficient quota
+				sufficientDataQuota := func(ctx context.Context) int64 {
+					return 1000000 // Allow 1MB
+				}
+				queryable, err = createQueryable(shard, WithDataBytesLimitFunc(sufficientDataQuota))
+				require.NoError(t, err)
+				querier, err = queryable.Querier(data.MinTime, data.MaxTime)
+				require.NoError(t, err)
+
+				ss = querier.Select(ctx, true, nil, matchers...)
+				var series []prom_storage.Series
+				for ss.Next() {
+					series = append(series, ss.At())
+				}
+				require.NoError(t, ss.Err())
+				require.NotEmpty(t, series)
+			})
 		})
 	}
 }
@@ -338,11 +459,11 @@ func queryWithQueryable(t *testing.T, mint, maxt int64, shard storage.ParquetSha
 	return found
 }
 
-func createQueryable(shard storage.ParquetShard) (prom_storage.Queryable, error) {
+func createQueryable(shard storage.ParquetShard, opts ...QueryableOpts) (prom_storage.Queryable, error) {
 	d := schema.NewPrometheusParquetChunksDecoder(chunkenc.NewPool())
 	return NewParquetQueryable(d, func(ctx context.Context, mint, maxt int64) ([]storage.ParquetShard, error) {
 		return []storage.ParquetShard{shard}, nil
-	})
+	}, opts...)
 }
 
 var benchmarkCases = []struct {

search/limits.go

@@ -0,0 +1,82 @@
+// Copyright The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+// This package is a modified copy from
+// https://github.com/thanos-io/thanos-parquet-gateway/blob/cfc1279f605d1c629c4afe8b1e2a340e8b15ecdc/internal/limits/limit.go.
+
+package search
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+)
+
+type resourceExhausted struct {
+	used int64
+}
+
+func (re *resourceExhausted) Error() string {
+	return fmt.Sprintf("resource exhausted (used %d)", re.used)
+}
+
+// IsResourceExhausted checks if the error is a resource exhausted error.
+func IsResourceExhausted(err error) bool {
+	var re *resourceExhausted
+	return errors.As(err, &re)
+}
+
+// Quota is a limiter for a resource.
+type Quota struct {
+	mu sync.Mutex
+	q  int64
+	u  int64
+}
+
+// NewQuota creates a new quota with the given limit.
+func NewQuota(n int64) *Quota {
+	return &Quota{q: n, u: n}
+}
+
+// UnlimitedQuota creates a new quota with no limit.
+func UnlimitedQuota() *Quota {
+	return NewQuota(0)
+}
+
+func (q *Quota) Reserve(n int64) error {
+	if q.q == 0 {
+		return nil
+	}
+
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if q.u-n < 0 {
+		return &resourceExhausted{used: q.q}
+	}
+	q.u -= n
+	return nil
+}
+
+// QuotaLimitFunc is a function that returns the limit value.
+type QuotaLimitFunc func(ctx context.Context) int64
+
+// NoopQuotaLimitFunc returns 0 which means no limit.
+func NoopQuotaLimitFunc(ctx context.Context) int64 {
+	return 0
+}