chore: add support for fetching chunks with sizing info from the index to optimize GetShards calls (#19221)

Signed-off-by: Sandeep Sukhani <[email protected]>
Co-authored-by: Christian Haudum <[email protected]>

Author: sandeepsukhani

pkg/indexgateway/gateway.go

@@ -418,7 +418,7 @@ func (g *Gateway) GetShards(request *logproto.ShardsRequest, server logproto.Ind
 		return err
 	}
 
-	forSeries, ok := g.indexQuerier.HasForSeries(request.From, request.Through)
+	ok := g.indexQuerier.HasChunkSizingInfo(request.From, request.Through)
 	if !ok {
 		sp.AddEvent("index does not support forSeries", trace.WithAttributes(
 			attribute.String("action", "falling back to indexQuerier.GetShards impl"),
@@ -438,7 +438,7 @@ func (g *Gateway) GetShards(request *logproto.ShardsRequest, server logproto.Ind
 		return server.Send(shards)
 	}
 
-	return g.boundedShards(ctx, request, server, instanceID, p, forSeries)
+	return g.boundedShards(ctx, request, server, instanceID, p)
 }
 
 // boundedShards handles bounded shard requests, optionally returning precomputed chunks.
@@ -448,7 +448,6 @@ func (g *Gateway) boundedShards(
 	server logproto.IndexGateway_GetShardsServer,
 	instanceID string,
 	p chunk.Predicate,
-	forSeries sharding.ForSeries,
 ) error {
 	// TODO(owen-d): instead of using GetChunks which buffers _all_ the chunks
 	// (expensive when looking at the full fingerprint space), we should
@@ -467,27 +466,16 @@ func (g *Gateway) boundedShards(
 	defer sp.End()
 
 	// 1) for all bounds, get chunk refs
-	grps, _, err := g.indexQuerier.GetChunks(ctx, instanceID, req.From, req.Through, p, nil)
+	refs, err := g.indexQuerier.GetChunkRefsWithSizingInfo(ctx, instanceID, req.From, req.Through, p)
 	if err != nil {
 		return err
 	}
 
-	var ct int
-	for _, g := range grps {
-		ct += len(g)
-	}
+	ct := len(refs)
 
 	sp.AddEvent("queried local index", trace.WithAttributes(
 		attribute.Int("index_chunks_resolved", ct),
 	))
-	// TODO(owen-d): pool
-	refs := make([]*logproto.ChunkRef, 0, ct)
-
-	for _, cs := range grps {
-		for j := range cs {
-			refs = append(refs, &cs[j].ChunkRef)
-		}
-	}
 
 	filtered := refs
 
@@ -523,7 +511,7 @@ func (g *Gateway) boundedShards(
 		}
 
 	} else {
-		shards, chunkGrps, err := accumulateChunksToShards(ctx, instanceID, forSeries, req, p, filtered)
+		shards, chunkGrps, err := accumulateChunksToShards(req, refs)
 		if err != nil {
 			return err
 		}
@@ -597,71 +585,14 @@ func ExtractShardRequestMatchersAndAST(query string) (chunk.Predicate, error) {
 	}), nil
 }
 
-// TODO(owen-d): consider extending index impl to support returning chunkrefs _with_ sizing info
-// TODO(owen-d): perf, this is expensive :(
 func accumulateChunksToShards(
-	ctx context.Context,
-	user string,
-	forSeries sharding.ForSeries,
 	req *logproto.ShardsRequest,
-	p chunk.Predicate,
-	filtered []*logproto.ChunkRef,
+	filtered []logproto.ChunkRefWithSizingInfo,
 ) ([]logproto.Shard, []logproto.ChunkRefGroup, error) {
 	// map for looking up post-filtered chunks in O(n) while iterating the index again for sizing info
-	filteredM := make(map[model.Fingerprint][]refWithSizingInfo, 1024)
+	filteredM := make(map[model.Fingerprint][]logproto.ChunkRefWithSizingInfo, 1024)
 	for _, ref := range filtered {
-		x := refWithSizingInfo{ref: ref}
-		filteredM[model.Fingerprint(ref.Fingerprint)] = append(filteredM[model.Fingerprint(ref.Fingerprint)], x)
-	}
-
-	var mtx sync.Mutex
-
-	if err := forSeries.ForSeries(
-		ctx,
-		user,
-		v1.NewBounds(filtered[0].FingerprintModel(), filtered[len(filtered)-1].FingerprintModel()),
-		req.From, req.Through,
-		func(l labels.Labels, fp model.Fingerprint, chks []tsdb_index.ChunkMeta) (stop bool) {
-			mtx.Lock()
-			defer mtx.Unlock()
-
-			// check if this is a fingerprint we need
-			if _, ok := filteredM[fp]; !ok {
-				return false
-			}
-
-			filteredChks := filteredM[fp]
-			var j int
-
-		outer:
-			for i := range filteredChks {
-				for j < len(chks) {
-					switch filteredChks[i].Cmp(chks[j]) {
-					case iter.Less:
-						// this chunk is not in the queried index, continue checking other chunks
-						continue outer
-					case iter.Greater:
-						// next chunk in index but didn't pass filter; continue
-						j++
-						continue
-					case iter.Eq:
-						// a match; set the sizing info
-						filteredChks[i].KB = chks[j].KB
-						filteredChks[i].Entries = chks[j].Entries
-						j++
-						continue outer
-					}
-				}
-
-				// we've finished this index's chunks; no need to keep checking filtered chunks
-				break
-			}
-
-			return false
-		},
-		p.Matchers...,
-	); err != nil {
-		return nil, nil, err
+		filteredM[model.Fingerprint(ref.Fingerprint)] = append(filteredM[model.Fingerprint(ref.Fingerprint)], ref)
 	}
 
 	collectedSeries := sharding.SizedFPs(sharding.SizedFPsPool.Get(len(filteredM)))
@@ -689,13 +620,22 @@ func accumulateChunksToShards(
 			return filtered[i].Fingerprint > uint64(s.Bounds.Max)
 		})
 		chkGrps = append(chkGrps, logproto.ChunkRefGroup{
-			Refs: filtered[from:through],
+			Refs: refsWithSizingInfoToRefs(filtered[from:through]),
 		})
 	}
 
 	return shards, chkGrps, nil
 }
 
+func refsWithSizingInfoToRefs(refsWithSizingInfo []logproto.ChunkRefWithSizingInfo) []*logproto.ChunkRef {
+	refs := make([]*logproto.ChunkRef, 0, len(refsWithSizingInfo))
+	for _, refWithSizingInfo := range refsWithSizingInfo {
+		refs = append(refs, &refWithSizingInfo.ChunkRef)
+	}
+
+	return refs
+}
+
 type refWithSizingInfo struct {
 	ref     *logproto.ChunkRef
 	KB      uint32

pkg/indexgateway/gateway_test.go

@@ -16,7 +16,6 @@ import (
 
 	v2 "github.com/grafana/loki/v3/pkg/iter/v2"
 	"github.com/grafana/loki/v3/pkg/logproto"
-	"github.com/grafana/loki/v3/pkg/storage/chunk"
 	"github.com/grafana/loki/v3/pkg/storage/config"
 	"github.com/grafana/loki/v3/pkg/storage/stores/series/index"
 	tsdb_index "github.com/grafana/loki/v3/pkg/storage/stores/shipper/indexshipper/tsdb/index"
@@ -426,135 +425,47 @@ func TestRefWithSizingInfo(t *testing.T) {
 // TODO(owen-d): more testing for specific cases
 func TestAccumulateChunksToShards(t *testing.T) {
 	// only check eq by checksum for convenience -- we're not testing the comparison function here
-	mkRef := func(fp model.Fingerprint, checksum uint32) *logproto.ChunkRef {
-		return &logproto.ChunkRef{
+	mkRef := func(fp model.Fingerprint, checksum uint32) logproto.ChunkRef {
+		return logproto.ChunkRef{
 			Fingerprint: uint64(fp),
 			Checksum:    checksum,
 		}
 	}
 
-	sized := func(ref *logproto.ChunkRef, kb, entries uint32) refWithSizingInfo {
-		return refWithSizingInfo{
-			ref:     ref,
-			KB:      kb,
-			Entries: entries,
+	sized := func(ref logproto.ChunkRef, kb, entries uint32) logproto.ChunkRefWithSizingInfo {
+		return logproto.ChunkRefWithSizingInfo{
+			ChunkRef: ref,
+			KB:       kb,
+			Entries:  entries,
 		}
 
 	}
 
-	fsImpl := func(series [][]refWithSizingInfo) sharding.ForSeriesFunc {
-		return sharding.ForSeriesFunc(
-			func(
-				_ context.Context,
-				_ string,
-				_ tsdb_index.FingerprintFilter,
-				_, _ model.Time,
-				fn func(
-					_ labels.Labels,
-					fp model.Fingerprint,
-					chks []tsdb_index.ChunkMeta,
-				) (stop bool), _ ...*labels.Matcher) error {
-
-				for _, s := range series {
-					chks := []tsdb_index.ChunkMeta{}
-					for _, r := range s {
-						chks = append(chks, tsdb_index.ChunkMeta{
-							Checksum: r.ref.Checksum,
-							KB:       r.KB,
-							Entries:  r.Entries,
-						})
-					}
-
-					if stop := fn(labels.EmptyLabels(), s[0].ref.FingerprintModel(), chks); stop {
-						return nil
-					}
-				}
-				return nil
-			},
-		)
-	}
-
-	filtered := []*logproto.ChunkRef{
+	filtered := []logproto.ChunkRefWithSizingInfo{
 		// shard 0
-		mkRef(1, 0),
-		mkRef(1, 1),
-		mkRef(1, 2),
+		sized(mkRef(1, 0), 100, 1),
+		sized(mkRef(1, 1), 100, 1),
+		sized(mkRef(1, 2), 100, 1),
 
 		// shard 1
-		mkRef(2, 10),
-		mkRef(2, 20),
-		mkRef(2, 30),
+		sized(mkRef(2, 10), 100, 1),
+		sized(mkRef(2, 20), 100, 1),
+		sized(mkRef(2, 30), 100, 1),
 
 		// shard 2 split across multiple series
-		mkRef(3, 10),
-		mkRef(4, 10),
-		mkRef(4, 20),
+		sized(mkRef(3, 10), 50, 1),
+		sized(mkRef(4, 10), 30, 1),
+		sized(mkRef(4, 20), 30, 1),
 
 		// last shard contains leftovers + skip a few fps in between
-		mkRef(7, 10),
+		sized(mkRef(7, 10), 25, 1),
 	}
 
-	series := [][]refWithSizingInfo{
-		{
-			// first series creates one shard since a shard can't contain partial series.
-			// no chunks were filtered out
-			sized(mkRef(1, 0), 100, 1),
-			sized(mkRef(1, 1), 100, 1),
-			sized(mkRef(1, 2), 100, 1),
-		},
-		{
-			// second shard also contains one series, but this series has chunks filtered out.
-			sized(mkRef(2, 0), 100, 1),  // filtered out
-			sized(mkRef(2, 10), 100, 1), // included
-			sized(mkRef(2, 11), 100, 1), // filtered out
-			sized(mkRef(2, 20), 100, 1), // included
-			sized(mkRef(2, 21), 100, 1), // filtered out
-			sized(mkRef(2, 30), 100, 1), // included
-			sized(mkRef(2, 31), 100, 1), // filtered out
-		},
-
-		// third shard contains multiple series.
-		// combined they have 110kb, which is above the target of 100kb
-		// but closer than leaving the second series out which would create
-		// a shard with 50kb
-		{
-			// first series, 50kb
-			sized(mkRef(3, 10), 50, 1), // 50kb
-			sized(mkRef(3, 11), 50, 1), // 50kb, not included
-		},
-		{
-			// second series
-			sized(mkRef(4, 10), 30, 1), // 30kb
-			sized(mkRef(4, 11), 30, 1), // 30kb, not included
-			sized(mkRef(4, 20), 30, 1), // 30kb
-		},
-
-		// Fourth shard contains a single series with 25kb,
-		// but iterates over non-included fp(s) before it
-		{
-			// register a series in the index which is not included in the filtered list
-			sized(mkRef(6, 10), 100, 1), // not included
-			sized(mkRef(6, 11), 100, 1), // not included
-		},
-		{
-			// last shard contains leftovers
-			sized(mkRef(7, 10), 25, 1),
-			sized(mkRef(7, 11), 100, 1), // not included
-		},
-	}
-
-	shards, grps, err := accumulateChunksToShards(
-		context.Background(),
-		"",
-		fsImpl(series),
-		&logproto.ShardsRequest{
-			TargetBytesPerShard: 100 << 10,
-		},
-		chunk.NewPredicate(nil, nil), // we're not checking matcher injection here
-		filtered,
-	)
+	shards, grps, err := accumulateChunksToShards(&logproto.ShardsRequest{
+		TargetBytesPerShard: 100 << 10,
+	}, filtered)
 
-	expectedChks := [][]*logproto.ChunkRef{
+	expectedChks := [][]logproto.ChunkRefWithSizingInfo{
 		filtered[0:3],
 		filtered[3:6],
 		filtered[6:9],
@@ -604,7 +515,7 @@ func TestAccumulateChunksToShards(t *testing.T) {
 	for i := range shards {
 		require.Equal(t, exp[i], shards[i], "invalid shard at index %d", i)
 		for j := range grps[i].Refs {
-			require.Equal(t, expectedChks[i][j], grps[i].Refs[j], "invalid chunk in grp %d at index %d", i, j)
+			require.Equal(t, &expectedChks[i][j].ChunkRef, grps[i].Refs[j], "invalid chunk in grp %d at index %d", i, j)
 		}
 	}
 	require.Equal(t, len(exp), len(shards))

pkg/ingester/flush_test.go

@@ -500,6 +500,14 @@ func (s *testStore) HasForSeries(_, _ model.Time) (sharding.ForSeries, bool) {
 	return nil, false
 }
 
+func (s *testStore) HasChunkSizingInfo(_, _ model.Time) bool {
+	return false
+}
+
+func (s *testStore) GetChunkRefsWithSizingInfo(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate) ([]logproto.ChunkRefWithSizingInfo, error) {
+	return nil, nil
+}
+
 func (s *testStore) GetSchemaConfigs() []config.PeriodConfig {
 	return defaultPeriodConfigs
 }

pkg/ingester/ingester_test.go

@@ -515,6 +515,14 @@ func (s *mockStore) HasForSeries(_, _ model.Time) (sharding.ForSeries, bool) {
 	return nil, false
 }
 
+func (s *mockStore) HasChunkSizingInfo(_, _ model.Time) bool {
+	return false
+}
+
+func (s *mockStore) GetChunkRefsWithSizingInfo(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate) ([]logproto.ChunkRefWithSizingInfo, error) {
+	return nil, nil
+}
+
 func (s *mockStore) Volume(_ context.Context, _ string, _, _ model.Time, limit int32, _ []string, _ string, _ ...*labels.Matcher) (*logproto.VolumeResponse, error) {
 	return &logproto.VolumeResponse{
 		Volumes: []logproto.Volume{

pkg/logproto/types.go

@@ -7,3 +7,27 @@ import (
 func (c *ChunkRef) FingerprintModel() model.Fingerprint {
 	return model.Fingerprint(c.Fingerprint)
 }
+
+type ChunkRefWithSizingInfo struct {
+	ChunkRef
+	KB      uint32
+	Entries uint32
+}
+
+// Less Compares chunks by (Fp, From, Through, checksum)
+// Assumes User is equivalent
+func (c *ChunkRef) Less(x ChunkRef) bool {
+	if c.Fingerprint != x.Fingerprint {
+		return c.Fingerprint < x.Fingerprint
+	}
+
+	if c.From != x.From {
+		return c.From < x.From
+	}
+
+	if c.Through != x.Through {
+		return c.Through < x.Through
+	}
+
+	return c.Checksum < x.Checksum
+}