Adds chunk filter hook for ingesters. (#3603)

* Adds chunk filter hook for ingesters.

Follow up on #3569 this is the missing piece.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

* That interface is not required to change.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

pkg/ingester/checkpoint.go

@@ -2,6 +2,7 @@ package ingester
 
 import (
 	"bytes"
+	"context"
 	fmt "fmt"
 	"io/ioutil"
 	"os"
@@ -206,7 +207,7 @@ func newStreamsIterator(ing ingesterInstances) *streamIterator {
 		inst.streamsMtx.RLock()
 		streams := make([]*stream, 0, len(inst.streams))
 		inst.streamsMtx.RUnlock()
-		_ = inst.forAllStreams(func(s *stream) error {
+		_ = inst.forAllStreams(context.Background(), func(s *stream) error {
 			streams = append(streams, s)
 			return nil
 		})

pkg/ingester/checkpoint_test.go

@@ -242,7 +242,6 @@ func TestUnflushedChunks(t *testing.T) {
 }
 
 func TestIngesterWALBackpressureSegments(t *testing.T) {
-
 	walDir, err := ioutil.TempDir(os.TempDir(), "loki-wal")
 	require.Nil(t, err)
 	defer os.RemoveAll(walDir)
@@ -287,7 +286,6 @@ func TestIngesterWALBackpressureSegments(t *testing.T) {
 }
 
 func TestIngesterWALBackpressureCheckpoint(t *testing.T) {
-
 	walDir, err := ioutil.TempDir(os.TempDir(), "loki-wal")
 	require.Nil(t, err)
 	defer os.RemoveAll(walDir)
@@ -353,7 +351,6 @@ func expectCheckpoint(t *testing.T, walDir string, shouldExist bool, max time.Du
 			return
 		}
 	}
-
 }
 
 // mkPush makes approximately totalSize bytes of log lines across min(500, totalSize) streams
@@ -456,7 +453,7 @@ func Test_SeriesIterator(t *testing.T) {
 	limiter := NewLimiter(limits, &ringCountMock{count: 1}, 1)
 
 	for i := 0; i < 3; i++ {
-		inst := newInstance(defaultConfig(), fmt.Sprintf("%d", i), limiter, runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, nil)
+		inst := newInstance(defaultConfig(), fmt.Sprintf("%d", i), limiter, runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, nil, nil)
 		require.NoError(t, inst.Push(context.Background(), &logproto.PushRequest{Streams: []logproto.Stream{stream1}}))
 		require.NoError(t, inst.Push(context.Background(), &logproto.PushRequest{Streams: []logproto.Stream{stream2}}))
 		instances = append(instances, inst)
@@ -506,7 +503,7 @@ func Benchmark_SeriesIterator(b *testing.B) {
 	limiter := NewLimiter(limits, &ringCountMock{count: 1}, 1)
 
 	for i := range instances {
-		inst := newInstance(defaultConfig(), fmt.Sprintf("instance %d", i), limiter, nil, noopWAL{}, NilMetrics, nil)
+		inst := newInstance(defaultConfig(), fmt.Sprintf("instance %d", i), limiter, nil, noopWAL{}, NilMetrics, nil, nil)
 
 		require.NoError(b,
 			inst.Push(context.Background(), &logproto.PushRequest{

pkg/ingester/flush_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/grafana/loki/pkg/logql"
 	"github.com/grafana/loki/pkg/logql/log"
+	"github.com/grafana/loki/pkg/storage"
 	"github.com/grafana/loki/pkg/util/runtime"
 
 	"github.com/cortexproject/cortex/pkg/chunk"
@@ -340,6 +341,8 @@ func (s *testStore) GetSchemaConfigs() []chunk.PeriodConfig {
 
 func (s *testStore) Stop() {}
 
+func (s *testStore) SetChunkFilterer(_ storage.RequestChunkFilterer) {}
+
 func pushTestSamples(t *testing.T, ing logproto.PusherServer) map[string][]logproto.Stream {
 	userIDs := []string{"1", "2", "3"}
 

pkg/ingester/ingester.go

@@ -158,6 +158,8 @@ type Ingester struct {
 	metrics *ingesterMetrics
 
 	wal WAL
+
+	chunkFilter storage.RequestChunkFilterer
 }
 
 // ChunkStore is the interface we need to store chunks.
@@ -220,6 +222,10 @@ func New(cfg Config, clientConfig client.Config, store ChunkStore, limits *valid
 	return i, nil
 }
 
+func (i *Ingester) SetChunkFilterer(chunkFilter storage.RequestChunkFilterer) {
+	i.chunkFilter = chunkFilter
+}
+
 func (i *Ingester) starting(ctx context.Context) error {
 	if i.cfg.WAL.Enabled {
 		// Ignore retain period during wal replay.
@@ -404,7 +410,7 @@ func (i *Ingester) getOrCreateInstance(instanceID string) *instance {
 	defer i.instancesMtx.Unlock()
 	inst, ok = i.instances[instanceID]
 	if !ok {
-		inst = newInstance(&i.cfg, instanceID, i.limiter, i.tenantConfigs, i.wal, i.metrics, i.flushOnShutdownSwitch)
+		inst = newInstance(&i.cfg, instanceID, i.limiter, i.tenantConfigs, i.wal, i.metrics, i.flushOnShutdownSwitch, i.chunkFilter)
 		i.instances[instanceID] = inst
 	}
 	return inst
@@ -677,7 +683,7 @@ func (i *Ingester) Tail(req *logproto.TailRequest, queryServer logproto.Querier_
 		return err
 	}
 
-	if err := instance.addNewTailer(tailer); err != nil {
+	if err := instance.addNewTailer(queryServer.Context(), tailer); err != nil {
 		return err
 	}
 	tailer.loop()

pkg/ingester/ingester_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/grafana/loki/pkg/iter"
 	"github.com/grafana/loki/pkg/logproto"
 	"github.com/grafana/loki/pkg/logql"
+	"github.com/grafana/loki/pkg/storage"
 	"github.com/grafana/loki/pkg/util/runtime"
 	"github.com/grafana/loki/pkg/util/validation"
 )
@@ -284,6 +285,9 @@ func (s *mockStore) GetSchemaConfigs() []chunk.PeriodConfig {
 	return nil
 }
 
+func (s *mockStore) SetChunkFilterer(_ storage.RequestChunkFilterer) {
+}
+
 type mockQuerierServer struct {
 	ctx   context.Context
 	resps []*logproto.QueryResponse
@@ -448,7 +452,6 @@ func TestIngester_boltdbShipperMaxLookBack(t *testing.T) {
 }
 
 func TestValidate(t *testing.T) {
-
 	for i, tc := range []struct {
 		in       Config
 		err      bool

pkg/ingester/instance.go

@@ -28,6 +28,7 @@ import (
 	"github.com/grafana/loki/pkg/logproto"
 	"github.com/grafana/loki/pkg/logql"
 	"github.com/grafana/loki/pkg/logql/stats"
+	"github.com/grafana/loki/pkg/storage"
 	"github.com/grafana/loki/pkg/util/runtime"
 	"github.com/grafana/loki/pkg/util/validation"
 )
@@ -89,9 +90,11 @@ type instance struct {
 	flushOnShutdownSwitch *OnceSwitch
 
 	metrics *ingesterMetrics
+
+	chunkFilter storage.RequestChunkFilterer
 }
 
-func newInstance(cfg *Config, instanceID string, limiter *Limiter, configs *runtime.TenantConfigs, wal WAL, metrics *ingesterMetrics, flushOnShutdownSwitch *OnceSwitch) *instance {
+func newInstance(cfg *Config, instanceID string, limiter *Limiter, configs *runtime.TenantConfigs, wal WAL, metrics *ingesterMetrics, flushOnShutdownSwitch *OnceSwitch, chunkFilter storage.RequestChunkFilterer) *instance {
 	i := &instance{
 		cfg:         cfg,
 		streams:     map[string]*stream{},
@@ -110,6 +113,8 @@ func newInstance(cfg *Config, instanceID string, limiter *Limiter, configs *runt
 		wal:                   wal,
 		metrics:               metrics,
 		flushOnShutdownSwitch: flushOnShutdownSwitch,
+
+		chunkFilter: chunkFilter,
 	}
 	i.mapper = newFPMapper(i.getLabelsFromFingerprint)
 	return i
@@ -295,7 +300,9 @@ func (i *instance) Query(ctx context.Context, req logql.SelectLogParams) ([]iter
 
 	ingStats := stats.GetIngesterData(ctx)
 	var iters []iter.EntryIterator
+
 	err = i.forMatchingStreams(
+		ctx,
 		expr.Matchers(),
 		func(stream *stream) error {
 			iter, err := stream.Iterator(ctx, ingStats, req.Start, req.End, req.Direction, pipeline.ForStream(stream.labels))
@@ -326,6 +333,7 @@ func (i *instance) QuerySample(ctx context.Context, req logql.SelectSampleParams
 	ingStats := stats.GetIngesterData(ctx)
 	var iters []iter.SampleIterator
 	err = i.forMatchingStreams(
+		ctx,
 		expr.Selector().Matchers(),
 		func(stream *stream) error {
 			iter, err := stream.SampleIterator(ctx, ingStats, req.Start, req.End, extractor.ForStream(stream.labels))
@@ -363,7 +371,7 @@ func (i *instance) Label(_ context.Context, req *logproto.LabelRequest) (*logpro
 	}, nil
 }
 
-func (i *instance) Series(_ context.Context, req *logproto.SeriesRequest) (*logproto.SeriesResponse, error) {
+func (i *instance) Series(ctx context.Context, req *logproto.SeriesRequest) (*logproto.SeriesResponse, error) {
 	groups, err := loghttp.Match(req.GetGroups())
 	if err != nil {
 		return nil, err
@@ -374,7 +382,7 @@ func (i *instance) Series(_ context.Context, req *logproto.SeriesRequest) (*logp
 	// If no matchers were supplied we include all streams.
 	if len(groups) == 0 {
 		series = make([]logproto.SeriesIdentifier, 0, len(i.streams))
-		err = i.forAllStreams(func(stream *stream) error {
+		err = i.forAllStreams(ctx, func(stream *stream) error {
 			// consider the stream only if it overlaps the request time range
 			if shouldConsiderStream(stream, req) {
 				series = append(series, logproto.SeriesIdentifier{
@@ -389,7 +397,7 @@ func (i *instance) Series(_ context.Context, req *logproto.SeriesRequest) (*logp
 	} else {
 		dedupedSeries := make(map[uint64]logproto.SeriesIdentifier)
 		for _, matchers := range groups {
-			err = i.forMatchingStreams(matchers, func(stream *stream) error {
+			err = i.forMatchingStreams(ctx, matchers, func(stream *stream) error {
 				// consider the stream only if it overlaps the request time range
 				if shouldConsiderStream(stream, req) {
 					// exit early when this stream was added by an earlier group
@@ -426,11 +434,18 @@ func (i *instance) numStreams() int {
 
 // forAllStreams will execute a function for all streams in the instance.
 // It uses a function in order to enable generic stream access without accidentally leaking streams under the mutex.
-func (i *instance) forAllStreams(fn func(*stream) error) error {
+func (i *instance) forAllStreams(ctx context.Context, fn func(*stream) error) error {
 	i.streamsMtx.RLock()
 	defer i.streamsMtx.RUnlock()
+	var chunkFilter storage.ChunkFilterer
+	if i.chunkFilter != nil {
+		chunkFilter = i.chunkFilter.ForRequest(ctx)
+	}
 
 	for _, stream := range i.streams {
+		if chunkFilter != nil && chunkFilter.ShouldFilter(stream.labels) {
+			continue
+		}
 		err := fn(stream)
 		if err != nil {
 			return err
@@ -442,6 +457,7 @@ func (i *instance) forAllStreams(fn func(*stream) error) error {
 // forMatchingStreams will execute a function for each stream that satisfies a set of requirements (time range, matchers, etc).
 // It uses a function in order to enable generic stream access without accidentally leaking streams under the mutex.
 func (i *instance) forMatchingStreams(
+	ctx context.Context,
 	matchers []*labels.Matcher,
 	fn func(*stream) error,
 ) error {
@@ -450,7 +466,10 @@ func (i *instance) forMatchingStreams(
 
 	filters, matchers := cutil.SplitFiltersAndMatchers(matchers)
 	ids := i.index.Lookup(matchers)
-
+	var chunkFilter storage.ChunkFilterer
+	if i.chunkFilter != nil {
+		chunkFilter = i.chunkFilter.ForRequest(ctx)
+	}
 outer:
 	for _, streamID := range ids {
 		stream, ok := i.streamsByFP[streamID]
@@ -462,7 +481,9 @@ outer:
 				continue outer
 			}
 		}
-
+		if chunkFilter != nil && chunkFilter.ShouldFilter(stream.labels) {
+			continue
+		}
 		err := fn(stream)
 		if err != nil {
 			return err
@@ -471,8 +492,8 @@ outer:
 	return nil
 }
 
-func (i *instance) addNewTailer(t *tailer) error {
-	if err := i.forMatchingStreams(t.matchers, func(s *stream) error {
+func (i *instance) addNewTailer(ctx context.Context, t *tailer) error {
+	if err := i.forMatchingStreams(ctx, t.matchers, func(s *stream) error {
 		s.addTailer(t)
 		return nil
 	}); err != nil {
@@ -494,8 +515,15 @@ func (i *instance) addTailersToNewStream(stream *stream) {
 		if t.isClosed() {
 			continue
 		}
+		var chunkFilter storage.ChunkFilterer
+		if i.chunkFilter != nil {
+			chunkFilter = i.chunkFilter.ForRequest(t.conn.Context())
+		}
 
 		if isMatching(stream.labels, t.matchers) {
+			if chunkFilter != nil && chunkFilter.ShouldFilter(stream.labels) {
+				continue
+			}
 			stream.addTailer(t)
 		}
 	}