github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/limiter/query_limiter.go (about)

     1  package limiter
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"github.com/prometheus/common/model"
     9  	"go.uber.org/atomic"
    10  
    11  	"github.com/grafana/loki/pkg/ingester/client"
    12  	"github.com/grafana/loki/pkg/logproto"
    13  )
    14  
    15  type queryLimiterCtxKey struct{}
    16  
    17  var (
    18  	ctxKey                    = &queryLimiterCtxKey{}
    19  	ErrMaxSeriesHit           = "the query hit the max number of series limit (limit: %d series)"
    20  	ErrMaxChunkBytesHit       = "the query hit the aggregated chunks size limit (limit: %d bytes)"
    21  	ErrMaxChunksPerQueryLimit = "the query hit the max number of chunks limit (limit: %d chunks)"
    22  )
    23  
    24  type QueryLimiter struct {
    25  	uniqueSeriesMx sync.Mutex
    26  	uniqueSeries   map[model.Fingerprint]struct{}
    27  
    28  	chunkBytesCount atomic.Int64
    29  	chunkCount      atomic.Int64
    30  
    31  	maxSeriesPerQuery     int
    32  	maxChunkBytesPerQuery int
    33  	maxChunksPerQuery     int
    34  }
    35  
    36  // NewQueryLimiter makes a new per-query limiter. Each query limiter
    37  // is configured using the `maxSeriesPerQuery` limit.
    38  func NewQueryLimiter(maxSeriesPerQuery, maxChunkBytesPerQuery int, maxChunksPerQuery int) *QueryLimiter {
    39  	return &QueryLimiter{
    40  		uniqueSeriesMx: sync.Mutex{},
    41  		uniqueSeries:   map[model.Fingerprint]struct{}{},
    42  
    43  		maxSeriesPerQuery:     maxSeriesPerQuery,
    44  		maxChunkBytesPerQuery: maxChunkBytesPerQuery,
    45  		maxChunksPerQuery:     maxChunksPerQuery,
    46  	}
    47  }
    48  
    49  func AddQueryLimiterToContext(ctx context.Context, limiter *QueryLimiter) context.Context {
    50  	return context.WithValue(ctx, ctxKey, limiter)
    51  }
    52  
    53  // QueryLimiterFromContextWithFallback returns a QueryLimiter from the current context.
    54  // If there is not a QueryLimiter on the context it will return a new no-op limiter.
    55  func QueryLimiterFromContextWithFallback(ctx context.Context) *QueryLimiter {
    56  	ql, ok := ctx.Value(ctxKey).(*QueryLimiter)
    57  	if !ok {
    58  		// If there's no limiter return a new unlimited limiter as a fallback
    59  		ql = NewQueryLimiter(0, 0, 0)
    60  	}
    61  	return ql
    62  }
    63  
    64  // AddSeries adds the input series and returns an error if the limit is reached.
    65  func (ql *QueryLimiter) AddSeries(seriesLabels []logproto.LabelAdapter) error {
    66  	// If the max series is unlimited just return without managing map
    67  	if ql.maxSeriesPerQuery == 0 {
    68  		return nil
    69  	}
    70  	fingerprint := client.FastFingerprint(seriesLabels)
    71  
    72  	ql.uniqueSeriesMx.Lock()
    73  	defer ql.uniqueSeriesMx.Unlock()
    74  
    75  	ql.uniqueSeries[fingerprint] = struct{}{}
    76  	if len(ql.uniqueSeries) > ql.maxSeriesPerQuery {
    77  		// Format error with max limit
    78  		return fmt.Errorf(ErrMaxSeriesHit, ql.maxSeriesPerQuery)
    79  	}
    80  	return nil
    81  }
    82  
    83  // AddChunkBytes adds the input chunk size in bytes and returns an error if the limit is reached.
    84  func (ql *QueryLimiter) AddChunkBytes(chunkSizeInBytes int) error {
    85  	if ql.maxChunkBytesPerQuery == 0 {
    86  		return nil
    87  	}
    88  	if ql.chunkBytesCount.Add(int64(chunkSizeInBytes)) > int64(ql.maxChunkBytesPerQuery) {
    89  		return fmt.Errorf(ErrMaxChunkBytesHit, ql.maxChunkBytesPerQuery)
    90  	}
    91  	return nil
    92  }
    93  
    94  func (ql *QueryLimiter) AddChunks(count int) error {
    95  	if ql.maxChunksPerQuery == 0 {
    96  		return nil
    97  	}
    98  
    99  	if ql.chunkCount.Add(int64(count)) > int64(ql.maxChunksPerQuery) {
   100  		return fmt.Errorf(fmt.Sprintf(ErrMaxChunksPerQueryLimit, ql.maxChunksPerQuery))
   101  	}
   102  	return nil
   103  }