Storage: Fix max_chunks_per_query functionality

pkg/storage/stores/composite_store.go

@@ -2,6 +2,7 @@ package stores
 
 import (
 	"context"
+	"fmt"
 	"sort"
 
 	"github.com/prometheus/common/model"
@@ -10,6 +11,7 @@ import (
 	"github.com/grafana/loki/pkg/logproto"
 	"github.com/grafana/loki/pkg/storage/chunk"
 	"github.com/grafana/loki/pkg/storage/chunk/fetcher"
+	"github.com/grafana/loki/pkg/storage/errors"
 	"github.com/grafana/loki/pkg/storage/stores/index"
 	"github.com/grafana/loki/pkg/storage/stores/index/seriesvolume"
 	"github.com/grafana/loki/pkg/storage/stores/index/stats"
@@ -171,6 +173,14 @@ func (c CompositeStore) GetChunks(ctx context.Context, userID string, from, thro
 		fetchers = append(fetchers, fetcher...)
 		return nil
 	})
+
+	// Protect ourselves against OOMing.
+	maxChunksPerQuery := c.limits.MaxChunksPerQuery(userID)
+	if maxChunksPerQuery > 0 && len(chunkIDs) > maxChunksPerQuery {
+		err := errors.QueryError(fmt.Sprintf("Query %v fetched too many chunks (%d > %d)", predicate, len(chunkIDs), maxChunksPerQuery))
+		return nil, nil, err
+	}
+
 	return chunkIDs, fetchers, err
 }
 

pkg/storage/stores/composite_store_entry.go

@@ -25,7 +25,7 @@ import (
 )
 
 type StoreLimits interface {
-	MaxChunksPerQueryFromStore(string) int
+	MaxChunksPerQuery(string) int
 	MaxQueryLength(context.Context, string) time.Duration
 }
 

pkg/storage/stores/composite_store_test.go

@@ -5,12 +5,14 @@ import (
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/pkg/errors"
 
 	"github.com/grafana/loki/pkg/logproto"
 
 	"github.com/grafana/dskit/test"
+	"github.com/grafana/dskit/user"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/stretchr/testify/require"
@@ -301,6 +303,83 @@ func TestCompositeStore_GetChunkFetcher(t *testing.T) {
 	}
 }
 
+type mockStoreGetChunks struct {
+	mockStoreGetChunkFetcher
+	chunkCount int
+}
+
+func (m mockStoreGetChunks) GetChunks(_ context.Context, _ string, _, _ model.Time, _ ...*labels.Matcher) ([][]chunk.Chunk, []*fetcher.Fetcher, error) {
+	return make([][]chunk.Chunk, m.chunkCount), []*fetcher.Fetcher{m.chunkFetcher}, nil
+}
+
+type mockStoreLimits struct {
+	MaxChunksPerQuery int
+}
+
+func (ml mockStoreLimits) MaxChunksPerQuery(string) int {
+	return ml.MaxChunksPerQuery
+}
+
+func (ml mockStoreLimits) MaxQueryLength(context.Context, string) time.Duration {
+	// mock
+	return time.Duration(1000 * 1000 * 1000 * 60 * 60 * 24)
+}
+
+var (
+	ctx = user.InjectOrgID(context.Background(), "fake")
+)
+
+func TestCompositeStore_GetChunks_MaxChunksPerQuery(t *testing.T) {
+	cs := []CompositeStore{
+		{
+			limits: mockStoreLimits{10},
+			stores: []compositeStoreEntry{
+				{
+					model.TimeFromUnix(10),
+					mockStoreGetChunks{mockStoreGetChunkFetcher{mockStore(0), &fetcher.Fetcher{}}, 5},
+				},
+			},
+		},
+		{
+			limits: mockStoreLimits{10},
+			stores: []compositeStoreEntry{
+				{
+					model.TimeFromUnix(10),
+					mockStoreGetChunks{mockStoreGetChunkFetcher{mockStore(2), &fetcher.Fetcher{}}, 100},
+				},
+			},
+		},
+	}
+
+	for _, tc := range []struct {
+		name                string
+		tmf                 model.Time
+		tmt                 model.Time
+		cs                  CompositeStore
+		expectedChunkLength int
+	}{
+		{
+			name:                "store with 5 chunks",
+			tmf:                 model.TimeFromUnix(10),
+			tmt:                 model.TimeFromUnix(20),
+			cs:                  cs[0],
+			expectedChunkLength: 5,
+		},
+		{
+			name:                "store with 100 chunks",
+			tmf:                 model.TimeFromUnix(10),
+			tmt:                 model.TimeFromUnix(20),
+			cs:                  cs[1],
+			expectedChunkLength: 0,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			actualChunks, _, _ := tc.cs.GetChunks(ctx, "fake", tc.tmf, tc.tmt)
+			require.Equal(t, tc.expectedChunkLength, len(actualChunks))
+		})
+	}
+}
+
 type mockStoreVolume struct {
 	mockStore
 	value *logproto.VolumeResponse