fix: Correctly propagate index stats to metrics.go log line (#14941)

Signed-off-by: Christian Haudum <christian.haudum@gmail.com>

Author: chaudum

pkg/logqlmodel/stats/context.go

@@ -108,6 +108,13 @@ func (c *Context) Index() Index {
 	return c.index
 }
 
+// Merge index stats from multiple respones in a concurrency-safe manner
+func (c *Context) MergeIndex(i Index) {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	c.index.Merge(i)
+}
+
 // Caches returns the cache statistics accumulated so far.
 func (c *Context) Caches() Caches {
 	return Caches{

pkg/storage/stores/series/series_index_gateway_store.go

@@ -60,8 +60,7 @@ func (c *IndexGatewayClientStore) GetChunkRefs(ctx context.Context, _ string, fr
 	}
 
 	statsCtx := statscontext.FromContext(ctx)
-	statsCtx.AddIndexTotalChunkRefs(response.Stats.TotalChunks)
-	statsCtx.AddIndexPostFilterChunkRefs(response.Stats.PostFilterChunks)
+	statsCtx.MergeIndex(response.Stats)
 
 	return result, nil
 }
@@ -131,23 +130,13 @@ func (c *IndexGatewayClientStore) Volume(ctx context.Context, _ string, from, th
 	})
 }
 
-func (c *IndexGatewayClientStore) GetShards(
-	ctx context.Context,
-	_ string,
-	from, through model.Time,
-	targetBytesPerShard uint64,
-	predicate chunk.Predicate,
-) (*logproto.ShardsResponse, error) {
-	resp, err := c.client.GetShards(ctx, &logproto.ShardsRequest{
+func (c *IndexGatewayClientStore) GetShards(ctx context.Context, _ string, from, through model.Time, targetBytesPerShard uint64, predicate chunk.Predicate) (*logproto.ShardsResponse, error) {
+	return c.client.GetShards(ctx, &logproto.ShardsRequest{
 		From:                from,
 		Through:             through,
 		Query:               predicate.Plan().AST.String(),
 		TargetBytesPerShard: targetBytesPerShard,
 	})
-	if err != nil {
-		return nil, err
-	}
-	return resp, nil
 }
 
 func (c *IndexGatewayClientStore) SetChunkFilterer(_ chunk.RequestChunkFilterer) {

pkg/storage/stores/series/series_index_gateway_store_test.go

@@ -9,18 +9,52 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/v3/pkg/logproto"
+	"github.com/grafana/loki/v3/pkg/logqlmodel/stats"
+	"github.com/grafana/loki/v3/pkg/storage/chunk"
 )
 
-type fakeClient struct {
+type mockClient struct {
 	GatewayClient
 }
 
-func (fakeClient) GetSeries(_ context.Context, _ *logproto.GetSeriesRequest) (*logproto.GetSeriesResponse, error) {
+func (mockClient) GetSeries(_ context.Context, _ *logproto.GetSeriesRequest) (*logproto.GetSeriesResponse, error) {
 	return &logproto.GetSeriesResponse{}, nil
 }
 
-func Test_IndexGatewayClient(t *testing.T) {
-	idx := NewIndexGatewayClientStore(fakeClient{}, log.NewNopLogger())
-	_, err := idx.GetSeries(context.Background(), "foo", model.Earliest, model.Latest)
+func (mockClient) GetChunkRef(_ context.Context, _ *logproto.GetChunkRefRequest) (*logproto.GetChunkRefResponse, error) {
+	return &logproto.GetChunkRefResponse{
+		Refs: []*logproto.ChunkRef{},
+		Stats: stats.Index{
+			TotalChunks:      1000,
+			PostFilterChunks: 10,
+			ShardsDuration:   0,
+			UsedBloomFilters: true,
+		},
+	}, nil
+}
+
+func Test_IndexGatewayClientStore_GetSeries(t *testing.T) {
+	idx := NewIndexGatewayClientStore(&mockClient{}, log.NewNopLogger())
+	_, err := idx.GetSeries(context.Background(), "tenant", model.Earliest, model.Latest)
 	require.NoError(t, err)
 }
+
+func Test_IndexGatewayClientStore_GetChunkRefs(t *testing.T) {
+	idx := NewIndexGatewayClientStore(&mockClient{}, log.NewNopLogger())
+
+	t.Run("stats context is merged correctly", func(t *testing.T) {
+		ctx := context.Background()
+		_, ctx = stats.NewContext(ctx)
+
+		_, err := idx.GetChunkRefs(ctx, "tenant", model.Earliest, model.Latest, chunk.NewPredicate(nil, nil))
+		require.NoError(t, err)
+
+		_, err = idx.GetChunkRefs(ctx, "tenant", model.Earliest, model.Latest, chunk.NewPredicate(nil, nil))
+		require.NoError(t, err)
+
+		statsCtx := stats.FromContext(ctx)
+		require.True(t, statsCtx.Index().UsedBloomFilters)
+		require.Equal(t, int64(2000), statsCtx.Index().TotalChunks)
+		require.Equal(t, int64(20), statsCtx.Index().PostFilterChunks)
+	})
+}