feat(blooms): bloom integration in query planning

pkg/logproto/indexgateway.proto

@@ -2,6 +2,7 @@ syntax = "proto3";
 
 package indexgatewaypb;
 
+import "gogoproto/gogo.proto";
 import "pkg/logproto/logproto.proto";
 
 option go_package = "github.com/grafana/loki/pkg/logproto";
@@ -25,4 +26,44 @@ service IndexGateway {
   // Note: this MUST be the same as the variant defined in
   // logproto.proto on the Querier service.
   rpc GetVolume(logproto.VolumeRequest) returns (logproto.VolumeResponse) {}
+
+  // GetShards is an optimized implemented shard-planning implementation
+  // on the index gateway and not on the ingester.
+  rpc GetShards(ShardsRequest) returns (stream ShardsResponse);
+}
+
+message ShardsRequest {
+  int64 from = 1 [
+    (gogoproto.customtype) = "github.com/prometheus/common/model.Time",
+    (gogoproto.nullable) = false
+  ];
+  int64 through = 2 [
+    (gogoproto.customtype) = "github.com/prometheus/common/model.Time",
+    (gogoproto.nullable) = false
+  ];
+  string matchers = 3;
+  uint64 target_bytes_per_shard = 4;
+  repeated FPBounds shards = 5;
+}
+
+message ShardsResponse {
+  repeated Shard shards = 1;
+}
+
+message Shard {
+  FPBounds bounds = 1 [
+    (gogoproto.nullable) = false
+  ];
+  logproto.IndexStatsResponse stats = 2;
+}
+
+// FPBounds is identical to the definition in `pkg/storage/bloom/v1/bounds.FingerprintBounds`
+// which ensures we can cast between them without allocations.
+message FPBounds {
+  uint64 min = 1 [
+    (gogoproto.casttype) = "github.com/prometheus/common/model.Fingerprint"
+  ];
+  uint64 max = 2 [
+    (gogoproto.casttype) = "github.com/prometheus/common/model.Fingerprint"
+  ];
 }

pkg/storage/bloom/v1/bounds.go

@@ -4,12 +4,13 @@ import (
 	"fmt"
 	"hash"
 	"strings"
+	"unsafe"
 
+	"github.com/grafana/loki/pkg/logproto"
+	"github.com/grafana/loki/pkg/util/encoding"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
 	"golang.org/x/exp/slices"
-
-	"github.com/grafana/loki/pkg/util/encoding"
 )
 
 type BoundsCheck uint8
@@ -20,6 +21,24 @@ const (
 	After
 )
 
+type FingerprintBounds struct {
+	Min, Max model.Fingerprint
+}
+
+// Proto compat
+// compiler check ensuring equal repr of underlying types
+var _ FingerprintBounds = FingerprintBounds(logproto.FPBounds{})
+
+func BoundsFromProto(pb logproto.FPBounds) FingerprintBounds {
+	return FingerprintBounds(pb)
+}
+
+// Unsafe cast to avoid allocation. This _requires_ that the underlying types are the same
+// which is checked by the compiler above
+func MultiBoundsFromProto(pb []logproto.FPBounds) MultiFingerprintBounds {
+	return MultiFingerprintBounds(*(*MultiFingerprintBounds)(unsafe.Pointer(&pb)))
+}
+
 // ParseBoundsFromAddr parses a fingerprint bounds from a string
 func ParseBoundsFromAddr(s string) (FingerprintBounds, error) {
 	parts := strings.Split(s, "-")
@@ -40,10 +59,6 @@ func ParseBoundsFromParts(a, b string) (FingerprintBounds, error) {
 	return NewBounds(minFingerprint, maxFingerprint), nil
 }
 
-type FingerprintBounds struct {
-	Min, Max model.Fingerprint
-}
-
 func NewBounds(min, max model.Fingerprint) FingerprintBounds {
 	return FingerprintBounds{Min: min, Max: max}
 }

pkg/storage/bloom/v1/bounds_test.go

@@ -3,10 +3,30 @@ package v1
 import (
 	"testing"
 
+	"github.com/grafana/loki/pkg/logproto"
 	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestBoundsFromProto(t *testing.T) {
+	bounds := BoundsFromProto(logproto.FPBounds{
+		Min: 10,
+		Max: 2000,
+	})
+	assert.Equal(t, NewBounds(10, 2000), bounds)
+}
+
+func TestMultiBoundsFromProto(t *testing.T) {
+	bounds := MultiBoundsFromProto([]logproto.FPBounds{
+		{Min: 10, Max: 2000},
+		{Min: 2001, Max: 4000},
+	})
+	assert.Equal(t, MultiFingerprintBounds{
+		NewBounds(10, 2000),
+		NewBounds(2001, 4000),
+	}, bounds)
+}
+
 func Test_ParseFingerprint(t *testing.T) {
 	t.Parallel()
 	fp, err := model.ParseFingerprint("7d0")

pkg/storage/stores/composite_store.go

@@ -206,6 +206,41 @@ func (c CompositeStore) Volume(ctx context.Context, userID string, from, through
 	return res, err
 }
 
+func (c CompositeStore) GetShards(
+	ctx context.Context,
+	userID string,
+	from, through model.Time,
+	targetBytesPerShard uint64,
+	matchers ...*labels.Matcher,
+) ([]*logproto.Shard, error) {
+	// TODO(owen-d): improve. Since shards aren't easily merge-able,
+	// we choose the store which returned the highest shard count
+	var groups [][]*logproto.Shard
+	err := c.forStores(ctx, from, through, func(innerCtx context.Context, from, through model.Time, store Store) error {
+		shards, err := store.GetShards(innerCtx, userID, from, through, targetBytesPerShard, matchers...)
+		if err != nil {
+			return err
+		}
+		groups = append(groups, shards)
+		return nil
+	})
+
+	switch {
+	case err != nil:
+		return nil, err
+	case len(groups) == 1:
+		return groups[0], nil
+	case len(groups) == 0:
+		return nil, nil
+	default:
+		sort.Slice(groups, func(i, j int) bool {
+			return len(groups[i]) > len(groups[j])
+		})
+		return groups[0], nil
+	}
+
+}
+
 func (c CompositeStore) GetChunkFetcher(tm model.Time) *fetcher.Fetcher {
 	// find the schema with the lowest start _after_ tm
 	j := sort.Search(len(c.stores), func(j int) bool {

pkg/storage/stores/composite_store_entry.go

@@ -119,6 +119,21 @@ func (c *storeEntry) Stats(ctx context.Context, userID string, from, through mod
 	return c.indexReader.Stats(ctx, userID, from, through, matchers...)
 }
 
+func (c *storeEntry) GetShards(
+	ctx context.Context,
+	userID string,
+	from, through model.Time,
+	targetBytesPerShard uint64,
+	matchers ...*labels.Matcher,
+) ([]*logproto.Shard, error) {
+	_, err := c.validateQueryTimeRange(ctx, userID, &from, &through)
+	if err != nil {
+		return nil, err
+	}
+
+	return c.indexReader.GetShards(ctx, userID, from, through, targetBytesPerShard, matchers...)
+}
+
 func (c *storeEntry) Volume(ctx context.Context, userID string, from, through model.Time, limit int32, targetLabels []string, aggregateBy string, matchers ...*labels.Matcher) (*logproto.VolumeResponse, error) {
 	sp, ctx := opentracing.StartSpanFromContext(ctx, "SeriesStore.Volume")
 	defer sp.Finish()

pkg/storage/stores/composite_store_test.go

@@ -60,6 +60,10 @@ func (m mockStore) Volume(_ context.Context, _ string, _, _ model.Time, _ int32,
 	return nil, nil
 }
 
+func (m mockStore) GetShards(_ context.Context, _ string, _, _ model.Time, _ uint64, _ ...*labels.Matcher) ([]*logproto.Shard, error) {
+	return nil, nil
+}
+
 func (m mockStore) Stop() {}
 
 func TestCompositeStore(t *testing.T) {

pkg/storage/stores/index/index.go

@@ -28,6 +28,13 @@ type BaseReader interface {
 type StatsReader interface {
 	Stats(ctx context.Context, userID string, from, through model.Time, matchers ...*labels.Matcher) (*stats.Stats, error)
 	Volume(ctx context.Context, userID string, from, through model.Time, limit int32, targetLabels []string, aggregateBy string, matchers ...*labels.Matcher) (*logproto.VolumeResponse, error)
+	GetShards(
+		ctx context.Context,
+		userID string,
+		from, through model.Time,
+		targetBytesPerShard uint64,
+		matchers ...*labels.Matcher,
+	) ([]*logproto.Shard, error)
 }
 
 type Reader interface {
@@ -137,6 +144,24 @@ func (m MonitoredReaderWriter) Volume(ctx context.Context, userID string, from,
 	return vol, nil
 }
 
+func (m MonitoredReaderWriter) GetShards(
+	ctx context.Context,
+	userID string,
+	from, through model.Time,
+	targetBytesPerShard uint64,
+	matchers ...*labels.Matcher,
+) ([]*logproto.Shard, error) {
+	var shards []*logproto.Shard
+	if err := loki_instrument.TimeRequest(ctx, "shards", instrument.NewHistogramCollector(m.metrics.indexQueryLatency), instrument.ErrorCode, func(ctx context.Context) error {
+		var err error
+		shards, err = m.rw.GetShards(ctx, userID, from, through, targetBytesPerShard, matchers...)
+		return err
+	}); err != nil {
+		return nil, err
+	}
+	return shards, nil
+}
+
 func (m MonitoredReaderWriter) SetChunkFilterer(chunkFilter chunk.RequestChunkFilterer) {
 	m.rw.SetChunkFilterer(chunkFilter)
 }

pkg/storage/stores/series/series_index_gateway_store.go

@@ -15,13 +15,26 @@ import (
 	"github.com/grafana/loki/pkg/storage/stores/index/stats"
 )
 
+// NB(owen-d): mostly modeled off of the proto-generated `logproto.IndexGatewayClient`,
+// but decoupled from explicit GRPC dependencies to work well with streaming grpc methods
+type GatewayClient interface {
+	GetChunkRef(ctx context.Context, in *logproto.GetChunkRefRequest) (*logproto.GetChunkRefResponse, error)
+	GetSeries(ctx context.Context, in *logproto.GetSeriesRequest) (*logproto.GetSeriesResponse, error)
+	LabelNamesForMetricName(ctx context.Context, in *logproto.LabelNamesForMetricNameRequest) (*logproto.LabelResponse, error)
+	LabelValuesForMetricName(ctx context.Context, in *logproto.LabelValuesForMetricNameRequest) (*logproto.LabelResponse, error)
+	GetStats(ctx context.Context, in *logproto.IndexStatsRequest) (*logproto.IndexStatsResponse, error)
+	GetVolume(ctx context.Context, in *logproto.VolumeRequest) (*logproto.VolumeResponse, error)
+
+	GetShards(ctx context.Context, in *logproto.ShardsRequest) (*logproto.ShardsResponse, error)
+}
+
 // IndexGatewayClientStore implements pkg/storage/stores/index.ReaderWriter
 type IndexGatewayClientStore struct {
-	client logproto.IndexGatewayClient
+	client GatewayClient
 	logger log.Logger
 }
 
-func NewIndexGatewayClientStore(client logproto.IndexGatewayClient, logger log.Logger) *IndexGatewayClientStore {
+func NewIndexGatewayClientStore(client GatewayClient, logger log.Logger) *IndexGatewayClientStore {
 	return &IndexGatewayClientStore{
 		client: client,
 		logger: logger,
@@ -111,6 +124,25 @@ func (c *IndexGatewayClientStore) Volume(ctx context.Context, _ string, from, th
 	})
 }
 
+func (c *IndexGatewayClientStore) GetShards(
+	ctx context.Context,
+	_ string,
+	from, through model.Time,
+	targetBytesPerShard uint64,
+	matchers ...*labels.Matcher,
+) ([]*logproto.Shard, error) {
+	resp, err := c.client.GetShards(ctx, &logproto.ShardsRequest{
+		From:                from,
+		Through:             through,
+		Matchers:            (&syntax.MatchersExpr{Mts: matchers}).String(),
+		TargetBytesPerShard: targetBytesPerShard,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Shards, nil
+}
+
 func (c *IndexGatewayClientStore) SetChunkFilterer(_ chunk.RequestChunkFilterer) {
 	level.Warn(c.logger).Log("msg", "SetChunkFilterer called on index gateway client store, but it does not support it")
 }

pkg/storage/stores/series/series_index_store.go

@@ -2,6 +2,7 @@ package series
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sort"
 	"sync"
@@ -751,10 +752,20 @@ func (c *IndexReaderWriter) convertChunkIDsToChunkRefs(_ context.Context, userID
 
 // old index stores do not implement stats -- skip
 func (c *IndexReaderWriter) Stats(_ context.Context, _ string, _, _ model.Time, _ ...*labels.Matcher) (*stats.Stats, error) {
-	return nil, nil
+	return nil, errors.New("unimplemented Stats() on legacy index stores")
 }
 
 // old index stores do not implement label volume -- skip
 func (c *IndexReaderWriter) Volume(_ context.Context, _ string, _, _ model.Time, _ int32, _ []string, _ string, _ ...*labels.Matcher) (*logproto.VolumeResponse, error) {
-	return nil, nil
+	return nil, errors.New("unimplemented Volume() on legacy index stores")
+}
+
+func (c *IndexReaderWriter) GetShards(
+	ctx context.Context,
+	userID string,
+	from, through model.Time,
+	targetBytesPerShard uint64,
+	matchers ...*labels.Matcher,
+) ([]*logproto.Shard, error) {
+	return nil, errors.New("unimplemented GetShards() on legacy index stores")
 }