// SPDX-License-Identifier: AGPL-3.0-only
// Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/querier/blocks_store_queryable.go
// Provenance-includes-license: Apache-2.0
// Provenance-includes-copyright: The Cortex Authors.

package querier

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/gogo/protobuf/types"
	"github.com/grafana/dskit/cancellation"
	"github.com/grafana/dskit/grpcutil"
	"github.com/grafana/dskit/kv"
	"github.com/grafana/dskit/ring"
	"github.com/grafana/dskit/ring/client"
	"github.com/grafana/dskit/services"
	"github.com/grafana/dskit/tenant"
	"github.com/grafana/dskit/tracing"
	"github.com/oklog/ulid/v2"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/util/annotations"
	"golang.org/x/sync/errgroup"
	grpc_metadata "google.golang.org/grpc/metadata"

	"github.com/grafana/mimir/pkg/compartments"
	"github.com/grafana/mimir/pkg/mimirpb"
	"github.com/grafana/mimir/pkg/querier/stats"
	"github.com/grafana/mimir/pkg/storage/bucket"
	"github.com/grafana/mimir/pkg/storage/series"
	"github.com/grafana/mimir/pkg/storage/sharding"
	mimir_tsdb "github.com/grafana/mimir/pkg/storage/tsdb"
	"github.com/grafana/mimir/pkg/storage/tsdb/block"
	"github.com/grafana/mimir/pkg/storage/tsdb/bucketindex"
	"github.com/grafana/mimir/pkg/storegateway"
	"github.com/grafana/mimir/pkg/storegateway/hintspb"
	"github.com/grafana/mimir/pkg/storegateway/storegatewaypb"
	"github.com/grafana/mimir/pkg/storegateway/storepb"
	"github.com/grafana/mimir/pkg/util"
	"github.com/grafana/mimir/pkg/util/chunkinfologger"
	"github.com/grafana/mimir/pkg/util/globalerror"
	"github.com/grafana/mimir/pkg/util/limiter"
	util_log "github.com/grafana/mimir/pkg/util/log"
	"github.com/grafana/mimir/pkg/util/spanlogger"
)

var errAlreadyClosed = errors.New("querier already closed")

// BlocksStoreSet is the interface used to get the clients to query series on a set of blocks.
type BlocksStoreSet interface {
	services.Service

	// GetClientsFor returns the store gateway clients that should be used to
	// query the set of blocks in input. The exclude parameter is the map of
	// blocks -> store-gateway addresses that should be excluded.
	GetClientsFor(userID string, blocks bucketindex.Blocks, exclude map[ulid.ULID][]string) (map[BlocksStoreClient][]ulid.ULID, error)
}

// BlocksFinder is the interface used to find blocks for a given user and time range.
type BlocksFinder interface {
	services.Service

	// GetBlocks returns known blocks for userID containing samples within the range minT
	// and maxT (milliseconds, both included). Returned blocks are sorted by MaxTime descending.
	GetBlocks(ctx context.Context, userID string, minT, maxT int64) (bucketindex.Blocks, *bucketindex.Metadata, error)
}

// BlocksStoreClient is the interface that should be implemented by any client used
// to query a backend store-gateway.
type BlocksStoreClient interface {
	storegatewaypb.StoreGatewayClient

	// RemoteAddress returns the address of the remote store-gateway and is used to uniquely
	// identify a store-gateway backend instance.
	RemoteAddress() string

	// RemoteZone returns the zone of the remote store-gateway, or an empty string if the zone
	// is not set or is unknown.
	RemoteZone() string
}

// BlocksStoreLimits is the interface that should be implemented by the limits provider.
type BlocksStoreLimits interface {
	bucket.TenantConfigProvider

	MaxLabelsQueryLength(userID string) time.Duration
	MaxChunksPerQuery(userID string) int
	StoreGatewayTenantShardSize(userID string) int
	StoreGatewayTenantShardSizePerZone(userID string) int
}

type blocksStoreQueryableMetrics struct {
	storesHit prometheus.Histogram
	refetches prometheus.Histogram

	blocksFound                                       prometheus.Counter
	blocksQueried                                     prometheus.Counter
	blocksWithCompactorShardButIncompatibleQueryShard prometheus.Counter
	// The total number of chunks received from store-gateways that were used to evaluate queries
	chunksTotal prometheus.Counter

	// compartmentsHit counts the read compartments queried per query. It is nil (and unobserved) when
	// compartments are disabled.
	compartmentsHit prometheus.Histogram
}

func newBlocksStoreQueryableMetrics(compartmentsCfg compartments.Config, reg prometheus.Registerer) *blocksStoreQueryableMetrics {
	m := &blocksStoreQueryableMetrics{
		storesHit: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
			Name:    "cortex_querier_storegateway_instances_hit_per_query",
			Help:    "Number of store-gateway instances hit for a single query.",
			Buckets: []float64{0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048},
		}),
		refetches: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
			Name:    "cortex_querier_storegateway_refetches_per_query",
			Help:    "Number of re-fetches attempted while querying store-gateway instances due to missing blocks.",
			Buckets: []float64{0, 1, 2},
		}),

		blocksFound: promauto.With(reg).NewCounter(prometheus.CounterOpts{
			Name: "cortex_querier_blocks_found_total",
			Help: "Number of blocks found based on query time range.",
		}),
		blocksQueried: promauto.With(reg).NewCounter(prometheus.CounterOpts{
			Name: "cortex_querier_blocks_queried_total",
			Help: "Number of blocks queried to satisfy query. Compared to blocks found, some blocks may have been filtered out thanks to query and compactor sharding.",
		}),
		blocksWithCompactorShardButIncompatibleQueryShard: promauto.With(reg).NewCounter(prometheus.CounterOpts{
			Name: "cortex_querier_blocks_with_compactor_shard_but_incompatible_query_shard_total",
			Help: "Blocks that couldn't be checked for query and compactor sharding optimization due to incompatible shard counts.",
		}),
		chunksTotal: promauto.With(reg).NewCounter(prometheus.CounterOpts{
			Name: "cortex_querier_query_storegateway_chunks_total",
			Help: "Number of chunks received from store gateways at query time.",
		}),
	}

	// Only register the per-query compartments-hit histogram when compartments are enabled.
	if compartmentsCfg.Enabled {
		m.compartmentsHit = promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
			Name: "cortex_querier_compartments_hit_per_query",
			Help: "Number of read compartments queried for a single query.",
			// The "storage" label denotes which query backend was queried, matching the convention of
			// cortex_querier_queries_storage_type_total ("ingester" / "store-gateway").
			ConstLabels: prometheus.Labels{"storage": "store-gateway"},
			Buckets:     prometheus.LinearBuckets(1, 1, compartmentsCfg.Read.NumCompartments),
		})
	}

	return m
}

type blocksStoreCompartment struct {
	finder BlocksFinder
	stores BlocksStoreSet
}

// BlocksStoreQueryable is a queryable which queries blocks storage via
// the store-gateway.
type BlocksStoreQueryable struct {
	services.Service

	// compartments holds one block-query pipeline per read compartment, indexed by read compartment ID.
	// It always has at least one element (compartment 0 when compartments are disabled).
	compartments []blocksStoreCompartment

	// router narrows a query to the read compartments that can hold the selected metric names. It is nil
	// when compartments are disabled, in which case all queries go to the single compartment.
	router *compartments.Router

	consistency              *BlocksConsistency
	logger                   log.Logger
	queryStoreAfter          time.Duration
	metrics                  *blocksStoreQueryableMetrics
	limits                   BlocksStoreLimits
	streamingChunksBatchSize uint64

	// Subservices manager.
	subservicesManager *services.Manager
	subservicesWatcher *services.FailureWatcher
	dynamicReplication storegateway.DynamicReplication
}

func NewBlocksStoreQueryable(
	stores []blocksStoreCompartment,
	clientsPool *client.Pool,
	compartmentsCfg compartments.Config,
	dynamicReplication storegateway.DynamicReplication,
	consistency *BlocksConsistency,
	limits BlocksStoreLimits,
	queryStoreAfter time.Duration,
	streamingChunksBatchSize uint64,
	logger log.Logger,
	reg prometheus.Registerer,
) (*BlocksStoreQueryable, error) {
	if len(stores) == 0 {
		return nil, errors.New("blocks storage queryable requires at least one compartment")
	}
	if clientsPool == nil {
		return nil, errors.New("blocks storage queryable requires a store-gateway clients pool")
	}

	// The router narrows a query to the compartments owning the selected metric names; it is nil when
	// compartments are disabled (queries then go to the single compartment). Its compartment IDs index the
	// stores slice, so the counts must match to avoid an out-of-range access at query time.
	var router *compartments.Router
	if compartmentsCfg.Enabled {
		router = compartments.NewRouter(compartmentsCfg.Read.NumCompartments)
		if router.NumCompartments() != len(stores) {
			return nil, errors.Errorf("blocks storage queryable has %d compartments but the router expects %d", len(stores), router.NumCompartments())
		}
	}

	subservices := make([]services.Service, 0, len(stores)*2+1)
	for _, c := range stores {
		subservices = append(subservices, c.stores, c.finder)
	}
	// The store-gateway client pool is shared by all compartments, so it is started and stopped once here
	// rather than per compartment.
	subservices = append(subservices, clientsPool)

	subservicesManager, err := services.NewManager(subservices...)
	if err != nil {
		return nil, errors.Wrap(err, "register blocks storage queryable subservices")
	}

	q := &BlocksStoreQueryable{
		compartments:             stores,
		router:                   router,
		dynamicReplication:       dynamicReplication,
		consistency:              consistency,
		queryStoreAfter:          queryStoreAfter,
		logger:                   logger,
		subservicesManager:       subservicesManager,
		subservicesWatcher:       services.NewFailureWatcher(),
		metrics:                  newBlocksStoreQueryableMetrics(compartmentsCfg, reg),
		limits:                   limits,
		streamingChunksBatchSize: streamingChunksBatchSize,
	}

	q.Service = services.NewBasicService(q.starting, q.running, q.stopping)

	return q, nil
}

func NewBlocksStoreQueryableFromConfig(querierCfg Config, gatewayCfg storegateway.Config, storageCfg mimir_tsdb.BlocksStorageConfig, compartmentsCfg compartments.Config, limits BlocksStoreLimits, logger log.Logger, reg prometheus.Registerer) (*BlocksStoreQueryable, error) {
	// The configured bucket name must carry the read-compartment placeholder so each compartment resolves
	// to its own bucket; this is validated in the Mimir config for the components that query blocks.
	numCompartments := 1
	if compartmentsCfg.Enabled {
		numCompartments = compartmentsCfg.Read.NumCompartments
	}

	var dynamicReplication storegateway.DynamicReplication = storegateway.NewNopDynamicReplication(gatewayCfg.ShardingRing.ReplicationFactor)
	if gatewayCfg.DynamicReplication.Enabled {
		dynamicReplication = storegateway.NewMaxTimeDynamicReplication(
			gatewayCfg,
			// Keep syncing blocks to store-gateways for a grace period (3 times the sync interval) to
			// ensure they are not unloaded while they are still being queried.
			mimir_tsdb.NewBlockDiscoveryDelayMultiplier*storageCfg.BucketStore.SyncInterval,
		)
	}

	consistency := NewBlocksConsistency(
		// Exclude blocks which have been recently uploaded, in order to give enough time to store-gateways
		// to discover and load them (3 times the sync interval).
		mimir_tsdb.NewBlockDiscoveryDelayMultiplier*storageCfg.BucketStore.SyncInterval,
		reg,
	)

	// A single store-gateway ring KV backend is shared by all compartments; each compartment's ring reads
	// its own key from it.
	storesRingCfg := gatewayCfg.ShardingRing.ToRingConfig()
	storesRingBackend, err := kv.NewClient(
		storesRingCfg.KVStore,
		ring.GetCodec(),
		kv.RegistererWithKVName(prometheus.WrapRegistererWithPrefix("cortex_", reg), "querier-store-gateway"),
		logger,
	)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create store-gateway ring backend")
	}

	finders := make([]BlocksFinder, numCompartments)
	storesRings := make([]*ring.Ring, numCompartments)
	for idx := 0; idx < numCompartments; idx++ {
		var (
			component = "querier"
			bucketCfg = storageCfg.Bucket
			ringName  = storegateway.RingNameForClient
			ringKey   = storegateway.RingKey
		)

		// cacheBucketID prefixes every metadata cache key. It should be "blocks", but without compartments
		// we keep it empty to not cause a massive cache invalidation when rolling out the Mimir version
		// introducing the bucket ID; this is fine as far as every other caching bucket uses its own unique
		// ID.
		//
		// With compartments each compartment reads a different bucket but the cache could be shared, so its
		// keys must be scoped by a compartment-specific ID to avoid collisions and subtle bugs.
		cacheBucketID := ""

		if compartmentsCfg.Enabled {
			// Each read compartment has its own bucket, store-gateway ring and metadata cache within this
			// single process.
			//
			// In the metrics, we identify them by a "-rc-<id>" suffix on each subsystem's own identity — the
			// metrics "component" label and the ring name/key — rather than an extra label which would change
			// a metric's label set and panic on collision with the same metric registered by another component
			//(e.g. the usage-stats reporter's bucket client, or the ingester ring).
			component = compartments.WithReadCompartmentSuffix(component, idx)
			bucketCfg = storageCfg.Bucket.ReadCompartmentConfig(idx)
			cacheBucketID = compartments.WithReadCompartmentSuffix("blocks", idx)
			ringName = compartments.WithReadCompartmentSuffix(ringName, idx)
			ringKey = compartments.WithReadCompartmentSuffix(ringKey, idx)
		}

		finder, err := newBlocksStoreQueryableFinder(component, cacheBucketID, bucketCfg, storageCfg, limits, logger, reg)
		if err != nil {
			return nil, err
		}

		storesRing, err := ring.NewWithStoreClientAndStrategy(storesRingCfg, ringName, ringKey, storesRingBackend, ring.NewIgnoreUnhealthyInstancesReplicationStrategy(), prometheus.WrapRegistererWithPrefix("cortex_", reg), logger)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create store-gateway ring client")
		}

		finders[idx] = finder
		storesRings[idx] = storesRing
	}

	// One store-gateway client pool is shared by all compartments. It is a connection cache keyed by
	// instance address, while per-compartment block ownership is resolved by each compartment's ring, so
	// it doesn't need to be per-compartment. Registering it once keeps its metrics free of a
	// per-compartment label that would otherwise collide with the same gRPC-client metrics registered by
	// other components running in the same process. Its discovery is the union of all compartment rings,
	// so its stale-client cleanup keeps clients for every compartment's store-gateways.
	readRings := make([]ring.ReadRing, len(storesRings))
	for i, r := range storesRings {
		readRings[i] = r
	}
	clientsPool := newStoreGatewayClientPool(client.NewRingsServiceDiscovery(readRings), querierCfg.StoreGatewayClient, logger, reg)

	stores := make([]blocksStoreCompartment, numCompartments)
	for idx := range stores {
		storeSet, err := newBlocksStoreReplicationSet(storesRings[idx], clientsPool, randomLoadBalancing, dynamicReplication, querierCfg.PreferAvailabilityZones, limits)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create store set")
		}

		stores[idx] = blocksStoreCompartment{finder: finders[idx], stores: storeSet}
	}

	streamingBufferSize := querierCfg.StreamingChunksPerStoreGatewaySeriesBufferSize

	return NewBlocksStoreQueryable(stores, clientsPool, compartmentsCfg, dynamicReplication, consistency, limits, querierCfg.QueryStoreAfter, streamingBufferSize, logger, reg)
}

// newBlocksStoreQueryableFinder creates a BucketIndexBlocksFinder over the given bucket config.
func newBlocksStoreQueryableFinder(component, cacheBucketID string, bucketCfg bucket.Config, storageCfg mimir_tsdb.BlocksStorageConfig, limits BlocksStoreLimits, logger log.Logger, reg prometheus.Registerer) (BlocksFinder, error) {
	bucketClient, err := bucket.NewClient(context.Background(), bucketCfg, component, logger, reg)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create bucket client")
	}

	componentReg := prometheus.WrapRegistererWith(prometheus.Labels{"component": component}, reg)

	cachingBucket, err := mimir_tsdb.NewMetadataCachingBucket(
		cacheBucketID,
		storageCfg.BucketStore.MetadataCache,
		bucketClient,
		logger,
		componentReg,
	)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create caching bucket")
	}

	return NewBucketIndexBlocksFinder(BucketIndexBlocksFinderConfig{
		IndexLoader: bucketindex.LoaderConfig{
			CheckInterval:         time.Minute,
			UpdateOnStaleInterval: storageCfg.BucketStore.SyncInterval,
			UpdateOnErrorInterval: storageCfg.BucketStore.BucketIndex.UpdateOnErrorInterval,
			IdleTimeout:           storageCfg.BucketStore.BucketIndex.IdleTimeout,
		},
		MaxStalePeriod:           storageCfg.BucketStore.BucketIndex.MaxStalePeriod,
		IgnoreDeletionMarksDelay: storageCfg.BucketStore.IgnoreDeletionMarksWhileQueryingDelay,
	}, cachingBucket, limits, logger, componentReg), nil
}

func (q *BlocksStoreQueryable) starting(ctx context.Context) error {
	q.subservicesWatcher.WatchManager(q.subservicesManager)

	if err := services.StartManagerAndAwaitHealthy(ctx, q.subservicesManager); err != nil {
		return errors.Wrap(err, "unable to start blocks storage queryable subservices")
	}

	return nil
}

func (q *BlocksStoreQueryable) running(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		case err := <-q.subservicesWatcher.Chan():
			return errors.Wrap(err, "block storage queryable subservice failed")
		}
	}
}

func (q *BlocksStoreQueryable) stopping(_ error) error {
	return services.StopManagerAndAwaitStopped(context.Background(), q.subservicesManager)
}

// Querier returns a new Querier on the storage.
func (q *BlocksStoreQueryable) Querier(mint, maxt int64) (storage.Querier, error) {
	if s := q.State(); s != services.Running {
		return nil, errors.Errorf("BlocksStoreQueryable is not running: %v", s)
	}

	return &blocksStoreQuerier{
		minT:                     mint,
		maxT:                     maxt,
		compartments:             q.compartments,
		router:                   q.router,
		dynamicReplication:       q.dynamicReplication,
		metrics:                  q.metrics,
		limits:                   q.limits,
		streamingChunksBatchSize: q.streamingChunksBatchSize,
		consistency:              q.consistency,
		logger:                   q.logger,
		queryStoreAfter:          q.queryStoreAfter,
	}, nil
}

type blocksStoreQuerier struct {
	minT, maxT               int64
	compartments             []blocksStoreCompartment
	router                   *compartments.Router
	dynamicReplication       storegateway.DynamicReplication
	metrics                  *blocksStoreQueryableMetrics
	consistency              *BlocksConsistency
	limits                   BlocksStoreLimits
	streamingChunksBatchSize uint64
	logger                   log.Logger

	// If set, the querier manipulates the max time to not be greater than
	// "now - queryStoreAfter" so that most recent blocks are not queried.
	queryStoreAfter time.Duration

	streamReadersMtx sync.Mutex
	closed           bool
	streamReaders    []*storeGatewayStreamReader
}

// Select implements storage.Querier interface.
// The bool passed is ignored because the series is always sorted.
func (q *blocksStoreQuerier) Select(ctx context.Context, _ bool, sp *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
	tenantID, err := tenant.TenantID(ctx)
	if err != nil {
		return storage.ErrSeriesSet(err)
	}

	return q.selectSorted(ctx, sp, tenantID, matchers...)
}

func (q *blocksStoreQuerier) LabelNames(ctx context.Context, hints *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
	spanLog, ctx := spanlogger.New(ctx, q.logger, tracer, "blocksStoreQuerier.LabelNames")
	defer spanLog.Finish()

	tenantID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, nil, err
	}

	minT, maxT := q.minT, q.maxT

	spanLog.DebugLog("start", util.TimeFromMillis(minT).UTC().String(), "end",
		util.TimeFromMillis(maxT).UTC().String(), "matchers", util.MatchersStringer(matchers))

	// Clamp minT; we cannot push this down into queryWithConsistencyCheck as not all its callers need to clamp minT
	maxQueryLength := q.limits.MaxLabelsQueryLength(tenantID)
	if maxQueryLength != 0 {
		minT = clampToMaxLabelQueryLength(spanLog, minT, maxT, time.Now().UnixMilli(), maxQueryLength.Milliseconds())
	}

	var (
		convertedMatchers = convertMatchersToLabelMatcher(matchers)

		resMtx      sync.Mutex
		resNameSets = [][]string{}
		resWarnings annotations.Annotations
	)

	queryF := func(ctx context.Context, clients map[BlocksStoreClient][]ulid.ULID, minT, maxT int64, indexMeta *bucketindex.Metadata) ([]ulid.ULID, error) {
		nameSets, warnings, queriedBlocks, err := q.fetchLabelNamesFromStore(ctx, clients, minT, maxT, tenantID, hints, convertedMatchers, indexMeta)
		if err != nil {
			return nil, err
		}

		resMtx.Lock()
		resNameSets = append(resNameSets, nameSets...)
		resWarnings.Merge(warnings)
		resMtx.Unlock()

		return queriedBlocks, nil
	}

	if err := q.queryCompartmentsWithConsistencyCheck(ctx, q.targetCompartments(tenantID, matchers), spanLog, minT, maxT, tenantID, nil, queryF); err != nil {
		return nil, nil, err
	}

	return util.MergeSlices(resNameSets...), resWarnings, nil
}

func (q *blocksStoreQuerier) LabelValues(ctx context.Context, name string, hints *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
	spanLog, ctx := spanlogger.New(ctx, q.logger, tracer, "blocksStoreQuerier.LabelValues")
	defer spanLog.Finish()

	tenantID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, nil, err
	}

	minT, maxT := q.minT, q.maxT

	spanLog.DebugLog("start", util.TimeFromMillis(minT).UTC().String(), "end",
		util.TimeFromMillis(maxT).UTC().String(), "matchers", util.MatchersStringer(matchers))

	// Clamp minT; we cannot push this down into queryWithConsistencyCheck as not all its callers need to clamp minT
	maxQueryLength := q.limits.MaxLabelsQueryLength(tenantID)
	if maxQueryLength != 0 {
		minT = clampToMaxLabelQueryLength(spanLog, minT, maxT, time.Now().UnixMilli(), maxQueryLength.Milliseconds())
	}

	var (
		resMtx       sync.Mutex
		resValueSets = [][]string{}
		resWarnings  annotations.Annotations
	)

	queryF := func(ctx context.Context, clients map[BlocksStoreClient][]ulid.ULID, minT, maxT int64, indexMeta *bucketindex.Metadata) ([]ulid.ULID, error) {
		valueSets, warnings, queriedBlocks, err := q.fetchLabelValuesFromStore(ctx, name, clients, minT, maxT, tenantID, hints, matchers, indexMeta)
		if err != nil {
			return nil, err
		}

		resMtx.Lock()
		resValueSets = append(resValueSets, valueSets...)
		resWarnings.Merge(warnings)
		resMtx.Unlock()

		return queriedBlocks, nil
	}

	if err := q.queryCompartmentsWithConsistencyCheck(ctx, q.targetCompartments(tenantID, matchers), spanLog, minT, maxT, tenantID, nil, queryF); err != nil {
		return nil, nil, err
	}

	return util.MergeSlices(resValueSets...), resWarnings, nil
}

func (q *blocksStoreQuerier) Close() error {
	q.streamReadersMtx.Lock()
	defer q.streamReadersMtx.Unlock()

	q.closed = true

	for _, r := range q.streamReaders {
		r.FreeBuffer()
	}

	return nil
}

func (q *blocksStoreQuerier) selectSorted(ctx context.Context, sp *storage.SelectHints, tenantID string, matchers ...*labels.Matcher) storage.SeriesSet {
	spanLog, ctx := spanlogger.New(ctx, q.logger, tracer, "blocksStoreQuerier.selectSorted")
	defer spanLog.Finish()

	minT, maxT := sp.Start, sp.End

	var (
		convertedMatchers = convertMatchersToLabelMatcher(matchers)
		queryLimiter      = limiter.QueryLimiterFromContextWithFallback(ctx)

		resMtx             sync.Mutex
		resSeriesSets      = []storage.SeriesSet(nil)
		resWarnings        annotations.Annotations
		resStreamReaders   []*storeGatewayStreamReader
		resChunkEstimators []func() int
	)

	shard, _, err := sharding.ShardFromMatchers(matchers)
	if err != nil {
		return storage.ErrSeriesSet(err)
	}

	memoryTracker, err := limiter.MemoryConsumptionTrackerFromContext(ctx)
	if err != nil {
		return storage.ErrSeriesSet(err)
	}

	queryF := func(ctx context.Context, clients map[BlocksStoreClient][]ulid.ULID, minT, maxT int64, indexMeta *bucketindex.Metadata) ([]ulid.ULID, error) {
		seriesSets, queriedBlocks, warnings, streamReaders, chunkEstimator, err := q.fetchSeriesFromStores(ctx, sp, clients, minT, maxT, tenantID, convertedMatchers, memoryTracker, indexMeta)
		if err != nil {
			return nil, err
		}

		resMtx.Lock()
		resSeriesSets = append(resSeriesSets, seriesSets...)
		resWarnings.Merge(warnings)
		resStreamReaders = append(resStreamReaders, streamReaders...)
		resChunkEstimators = append(resChunkEstimators, chunkEstimator)
		resMtx.Unlock()

		return queriedBlocks, nil
	}

	err = q.queryCompartmentsWithConsistencyCheck(ctx, q.targetCompartments(tenantID, matchers), spanLog, minT, maxT, tenantID, shard, queryF)
	if err != nil {
		// Streams opened for already-queried compartments have not been handed to startBuffering yet, so
		// close them here to avoid leaking store-gateway streams when a compartment fails.
		for _, r := range resStreamReaders {
			r.Close()
		}
		return storage.ErrSeriesSet(err)
	}

	if len(resStreamReaders) > 0 {
		spanLog.DebugLog("msg", "starting streaming")

		err := q.startBuffering(resStreamReaders)
		if err != nil {
			return storage.ErrSeriesSet(err)
		}

		spanLog.DebugLog("msg", "streaming started, waiting for chunks estimates")

		chunksEstimate := 0
		for _, chunkEstimator := range resChunkEstimators {
			chunksEstimate += chunkEstimator()
		}

		spanLog.DebugLog("msg", "received chunks estimate from all store-gateways", "chunks_estimate", chunksEstimate)

		if err := queryLimiter.AddEstimatedChunks(chunksEstimate); err != nil {
			return storage.ErrSeriesSet(err)
		}
	}

	return series.NewSeriesSetWithWarnings(
		storage.NewMergeSeriesSet(resSeriesSets, 0, storage.ChainedSeriesMerge),
		resWarnings)
}

func (q *blocksStoreQuerier) startBuffering(streamReaders []*storeGatewayStreamReader) error {
	q.streamReadersMtx.Lock()
	defer q.streamReadersMtx.Unlock()

	if q.closed {
		// We were closed while loading, give up now.
		for _, r := range streamReaders {
			r.Close()
		}

		return errAlreadyClosed
	}

	q.streamReaders = append(q.streamReaders, streamReaders...)

	// If this was a streaming call, start fetching streaming chunks here.
	for _, r := range streamReaders {
		r.StartBuffering()
	}

	return nil
}

type queryFunc func(ctx context.Context, clients map[BlocksStoreClient][]ulid.ULID, minT, maxT int64, indexMeta *bucketindex.Metadata) ([]ulid.ULID, error)

// storeGatewayQueryStats holds the per-query store-gateway histogram values. queryWithConsistencyCheck
// returns it (rather than observing the histograms itself) so the caller observes them once per query.
type storeGatewayQueryStats struct {
	storesHit int  // number of distinct store-gateway instances queried
	refetches int  // number of retries due to missing blocks
	queried   bool // whether the block store was actually queried (true only on the success path)
}

func (q *blocksStoreQuerier) queryWithConsistencyCheck(
	ctx context.Context, spanLog *spanlogger.SpanLogger, finder BlocksFinder, stores BlocksStoreSet, minT, maxT int64, tenantID string, shard *sharding.ShardSelector, queryF queryFunc,
) (stats storeGatewayQueryStats, returnErr error) {
	now := time.Now()

	if !ShouldQueryBlockStore(q.queryStoreAfter, now, minT) {
		spanLog.DebugLog("msg", "not querying block store; query time range begins after the query-store-after limit")
		return storeGatewayQueryStats{}, nil
	}

	maxT = clampMaxTime(spanLog, maxT, now.UnixMilli(), -q.queryStoreAfter, "query store after")

	// Find the list of blocks we need to query given the time range.
	knownBlocks, indexMeta, err := finder.GetBlocks(ctx, tenantID, minT, maxT)
	if err != nil {
		return storeGatewayQueryStats{}, err
	}

	if len(knownBlocks) == 0 {
		spanLog.DebugLog("msg", "no blocks found")
		return storeGatewayQueryStats{}, nil
	}

	q.metrics.blocksFound.Add(float64(len(knownBlocks)))

	if shard != nil && shard.ShardCount > 0 {
		spanLog.DebugLog("msg", "filtering blocks due to sharding", "blocksBeforeFiltering", knownBlocks.String(), "shardID", shard.LabelValue())

		result, incompatibleBlocks := filterBlocksByShard(knownBlocks, shard.ShardIndex, shard.ShardCount)

		spanLog.DebugLog("msg", "result of filtering blocks", "before", len(knownBlocks), "after", len(result), "filtered", len(knownBlocks)-len(result), "incompatible", incompatibleBlocks)
		q.metrics.blocksWithCompactorShardButIncompatibleQueryShard.Add(float64(incompatibleBlocks))

		knownBlocks = result
	}

	q.metrics.blocksQueried.Add(float64(len(knownBlocks)))

	spanLog.DebugLog("msg", "found blocks to query", "expected", knownBlocks.String())

	var (
		// At the beginning the list of blocks to query are all known blocks.
		remainingBlocks = knownBlocks
		attemptedBlocks = map[ulid.ULID][]string{}
		touchedStores   = map[string]struct{}{}
	)

	consistencyTracker := q.consistency.NewTracker(knownBlocks, spanLog)
	defer func() {
		// Do not track consistency check metrics if query failed with an error unrelated to consistency check (e.g. context canceled),
		// because it means we interrupted the requests, and we don't know whether consistency check would have succeeded
		// or failed.
		if returnErr != nil && !errors.Is(returnErr, &storeConsistencyCheckFailedErr{}) {
			return
		}

		consistencyTracker.Complete()
	}()

	for attempt := 1; attempt <= q.dynamicReplication.MaxReplicationFactor(); attempt++ {
		// Find the set of store-gateway instances having the blocks. The exclude parameter is the
		// map of blocks queried so far, with the list of store-gateway addresses for each block.
		clients, err := stores.GetClientsFor(tenantID, remainingBlocks, attemptedBlocks)
		if err != nil {
			// If it's a retry and we get an error, it means there are no more store-gateways left
			// from which running another attempt, so we're just stopping retrying.
			if attempt > 1 {
				level.Warn(spanLog).Log("msg", "unable to get store-gateway clients while retrying to fetch missing blocks", "err", err)
				break
			}

			return storeGatewayQueryStats{}, err
		}
		spanLog.DebugLog("msg", "found store-gateway instances to query", "num instances", len(clients), "attempt", attempt)

		// Fetch series from stores. If an error occur we do not retry because retries
		// are only meant to cover missing blocks.
		queriedBlocks, err := queryF(ctx, clients, minT, maxT, indexMeta)
		if err != nil {
			return storeGatewayQueryStats{}, err
		}
		spanLog.DebugLog("msg", "received series from all store-gateways", "queried blocks", strings.Join(convertULIDsToString(queriedBlocks), " "))

		// Update the map of blocks we attempted to query.
		for client, blockIDs := range clients {
			touchedStores[client.RemoteAddress()] = struct{}{}

			for _, blockID := range blockIDs {
				attemptedBlocks[blockID] = append(attemptedBlocks[blockID], client.RemoteAddress())
			}
		}

		// Ensure all expected blocks have been queried (during all tries done so far).
		// The next attempt should just query the missing blocks.
		remainingBlocks = consistencyTracker.Check(queriedBlocks)
		if len(remainingBlocks) == 0 {
			return storeGatewayQueryStats{storesHit: len(touchedStores), refetches: attempt - 1, queried: true}, nil
		}

		spanLog.DebugLog("msg", "couldn't query all blocks", "attempt", attempt, "missing blocks", strings.Join(convertULIDsToString(remainingBlocks.GetULIDs()), " "))
	}

	// We've not been able to query all expected blocks after all retries.
	err = newStoreConsistencyCheckFailedError(remainingBlocks.GetULIDs())
	level.Warn(util_log.WithContext(ctx, spanLog)).Log("msg", "failed consistency check after all attempts", "err", err)
	return storeGatewayQueryStats{}, err
}

// targetCompartments returns the read compartments to query for the given matchers. With compartments
// disabled it returns the single compartment 0.
func (q *blocksStoreQuerier) targetCompartments(tenantID string, matchers []*labels.Matcher) []int {
	if q.router == nil {
		return []int{0}
	}

	return q.router.CompartmentsForMatchers(tenantID, matchers)
}

// forEachCompartment runs fn for each targeted compartment. With a single target it runs inline (no
// goroutine), so the compartments-disabled path is unchanged. With several, it runs them concurrently
// and returns the first error, cancelling the in-flight siblings on error.
func (q *blocksStoreQuerier) forEachCompartment(ctx context.Context, targets []int, fn func(ctx context.Context, c blocksStoreCompartment) error) error {
	if len(targets) == 1 {
		return fn(ctx, q.compartments[targets[0]])
	}

	// fn may open store-gateway streams that are read lazily after this returns, so the context passed to
	// fn must outlive this call and only be cancelled on error, never on normal completion (mirrors the
	// contract in fetchSeriesFromStores; errgroup.WithContext would cancel it when Wait returns).
	ctx, cancel := context.WithCancelCause(ctx) //nolint:govet
	g := &errgroup.Group{}
	for _, c := range targets {
		g.Go(func() error {
			if err := fn(ctx, q.compartments[c]); err != nil {
				cancel(err)
				return err
			}
			return nil
		})
	}
	return g.Wait() //nolint:govet // OK to return without cancelling ctx on success, see comment above.
}

// queryCompartmentsWithConsistencyCheck runs queryWithConsistencyCheck against each targeted compartment
// (concurrently when there is more than one) and, on success, observes the per-query metrics once for the
// whole query, summed across the compartments that were queried.
func (q *blocksStoreQuerier) queryCompartmentsWithConsistencyCheck(ctx context.Context, targets []int, spanLog *spanlogger.SpanLogger, minT, maxT int64, tenantID string, shard *sharding.ShardSelector, queryF queryFunc) error {
	var (
		statsMtx   sync.Mutex
		storesHit  int
		refetches  int
		anyQueried bool
	)

	err := q.forEachCompartment(ctx, targets, func(ctx context.Context, c blocksStoreCompartment) error {
		stats, err := q.queryWithConsistencyCheck(ctx, spanLog, c.finder, c.stores, minT, maxT, tenantID, shard, queryF)
		if err != nil {
			return err
		}

		statsMtx.Lock()
		storesHit += stats.storesHit
		refetches += stats.refetches
		anyQueried = anyQueried || stats.queried
		statsMtx.Unlock()
		return nil
	})
	if err != nil {
		return err
	}

	q.metrics.storesHit.Observe(float64(storesHit))
	if q.metrics.compartmentsHit != nil {
		q.metrics.compartmentsHit.Observe(float64(len(targets)))
	}
	if anyQueried {
		// refetches is only meaningful (and was historically only observed) when the block store was queried.
		q.metrics.refetches.Observe(float64(refetches))
	}
	return nil
}

type storeConsistencyCheckFailedErr struct {
	remainingBlocks []ulid.ULID
}

func newStoreConsistencyCheckFailedError(remainingBlocks []ulid.ULID) *storeConsistencyCheckFailedErr {
	// Sort the blocks, so it's easier to test the error strings.
	slices.SortFunc(remainingBlocks, func(i, j ulid.ULID) int {
		return i.Compare(j)
	})

	return &storeConsistencyCheckFailedErr{
		remainingBlocks: remainingBlocks,
	}
}

func (e *storeConsistencyCheckFailedErr) Error() string {
	return fmt.Sprintf("%s. The failed blocks are: %s", globalerror.StoreConsistencyCheckFailed.Message("failed to fetch some blocks"), strings.Join(convertULIDsToString(e.remainingBlocks), " "))
}

// Is implements support for errors.Is.
func (e *storeConsistencyCheckFailedErr) Is(err error) bool {
	var target *storeConsistencyCheckFailedErr
	return errors.As(err, &target)
}

// filterBlocksByShard removes blocks that can be safely ignored when using query sharding.
// We know that block can be safely ignored, if it was compacted using split-and-merge
// compactor, and it has a valid compactor shard ID. We exploit the fact that split-and-merge
// compactor and query-sharding use the same series-sharding algorithm.
//
// This function modifies input slice.
//
// This function also returns number of "incompatible" blocks -- blocks with compactor shard ID,
// but with compactor shard and query shard being incompatible for optimization.
func filterBlocksByShard(blocks bucketindex.Blocks, queryShardIndex, queryShardCount uint64) (_ bucketindex.Blocks, incompatibleBlocks int) {
	for ix := 0; ix < len(blocks); {
		b := blocks[ix]
		if b.CompactorShardID == "" {
			ix++
			continue
		}

		compactorShardIndex, compactorShardCount, err := sharding.ParseShardIDLabelValue(b.CompactorShardID)
		if err != nil {
			// Cannot parse compactor shardID, we must query this block.
			ix++
			continue
		}

		res, divisible := canBlockWithCompactorShardIndexContainQueryShard(queryShardIndex, queryShardCount, compactorShardIndex, compactorShardCount)
		if !divisible {
			incompatibleBlocks++
		}

		if res {
			ix++
			continue
		}

		// Series shard is NOT included in this block, we can remove this block.
		blocks = append(blocks[:ix], blocks[ix+1:]...)
	}

	return blocks, incompatibleBlocks
}

// canBlockWithCompactorShardIndexContainQueryShard returns false if block with
// given compactor shard ID can *definitely NOT* contain series for given query shard.
// Returns true otherwise (we don't know if block *does* contain such series,
// but we cannot rule it out).
//
// In other words, if this function returns false, block with given compactorShardID
// doesn't need to be searched for series from given query shard.
//
// In addition this function also returns whether query and compactor shard counts were
// divisible by each other (one way or the other).
func canBlockWithCompactorShardIndexContainQueryShard(queryShardIndex, queryShardCount, compactorShardIndex,
	compactorShardCount uint64) (result bool, divisibleShardCounts bool) {
	// If queryShardCount = compactorShardCount * K for integer K, then we know that series in queryShardIndex
	// can only be in the block for which (queryShardIndex % compactorShardCount == compactorShardIndex).
	//
	// For example if queryShardCount = 8 and compactorShardCount = 4, then series that should be returned
	// for queryShardIndex 5 can only be in block with compactorShardIndex = 1.
	if queryShardCount >= compactorShardCount && queryShardCount%compactorShardCount == 0 {
		wantedCompactorShardIndex := queryShardIndex % compactorShardCount

		return compactorShardIndex == wantedCompactorShardIndex, true
	}

	// If compactorShardCount = queryShardCount * K for some integer K, then series in queryShardIndex
	// can only be in K blocks for which queryShardIndex % compactorShardCount == compactorShardIndex.
	//
	// For example if queryShardCount = 4, and compactorShardCount = 8, then series that should be returned for
	// queryShardIndex 3 can only be in blocks with compactorShardIndex 3 and 7.
	if compactorShardCount >= queryShardCount && compactorShardCount%queryShardCount == 0 {
		wantedQueryShardIndex := compactorShardIndex % queryShardCount

		return queryShardIndex == wantedQueryShardIndex, true
	}

	return true, false
}

// fetchSeriesFromStores fetches series satisfying convertedMatchers and in the
// time range [minT, maxT) from all store-gateways in clients. Series are fetched
// from the given set of store-gateways concurrently. In successful case, i.e.,
// when all the concurrent fetches terminate with no exception, fetchSeriesFromStores returns:
//  1. a slice of fetched storage.SeriesSet
//  2. a slice of ulid.ULID corresponding to the queried blocks
//  3. annotations.Annotations encountered during the operation
//
// In case of a serious error during any of the concurrent executions, the error is returned.
// Errors while creating storepb.SeriesRequest, context cancellation, and unprocessable
// requests to the store-gateways (e.g., if a chunk or series limit is hit) are
// considered serious errors. All other errors are not returned, but they give rise to fetch retrials.
//
// In case of a successful run, fetchSeriesFromStores returns a startStreamingChunks function to start streaming
// chunks for the fetched series if it was a streaming call for series+chunks. startStreamingChunks must be called
// before iterating on the series.
func (q *blocksStoreQuerier) fetchSeriesFromStores(
	ctx context.Context,
	sp *storage.SelectHints,
	clients map[BlocksStoreClient][]ulid.ULID,
	minT int64,
	maxT int64,
	tenantID string,
	matchers []storepb.LabelMatcher,
	memoryTracker *limiter.MemoryConsumptionTracker,
	indexMeta *bucketindex.Metadata,
) (_ []storage.SeriesSet, _ []ulid.ULID, _ annotations.Annotations, streamReaders []*storeGatewayStreamReader, estimateChunks func() int, _ error) {
	reqCtx := grpcContextWithBucketStoreRequestMeta(ctx, tenantID, indexMeta)

	// We deliberately only cancel this context if any store-gateway call fails, to ensure that all streams are aborted promptly.
	// When all calls succeed, we rely on the parent context being cancelled, otherwise we'd abort all the store-gateway streams returned by this method, which makes them unusable.
	reqCtx, cancelReqCtx := context.WithCancelCause(reqCtx) //nolint:govet

	var (
		g, gCtx       = errgroup.WithContext(reqCtx)
		mtx           = sync.Mutex{}
		seriesSets    = []storage.SeriesSet(nil)
		warnings      annotations.Annotations
		queriedBlocks = []ulid.ULID(nil)
		spanLog       = spanlogger.FromContext(ctx, q.logger)
		queryLimiter  = limiter.QueryLimiterFromContextWithFallback(ctx)
		reqStats      = stats.FromContext(ctx)
		streams       []storegatewaypb.StoreGateway_SeriesClient
	)

	debugQuery := chunkinfologger.IsChunkInfoLoggingEnabled(ctx)

	// Concurrently fetch series from all clients.
	for c, blockIDs := range clients {
		g.Go(func() error {
			clientSpanLog, clientCtx := spanlogger.New(reqCtx, spanLog, tracer, "blocksStoreQuerier.fetchSeriesFromStores")
			defer clientSpanLog.Finish()
			clientSpanLog.SetTag("store_gateway_address", c.RemoteAddress())

			var skipChunks bool
			if sp != nil {
				// See: https://github.com/prometheus/prometheus/pull/8050
				// TODO(goutham): we should ideally be passing the hints down to the storage layer
				// and let the TSDB return us data with no chunks as in prometheus#8050.
				// But this is an acceptable workaround for now.
				skipChunks = sp.Func == "series"
			}

			req, err := createSeriesRequest(minT, maxT, matchers, skipChunks, sp, blockIDs, q.streamingChunksBatchSize)
			if err != nil {
				return errors.Wrapf(err, "failed to create series request")
			}

			stream, err := c.Series(clientCtx, req)
			if err == nil {
				mtx.Lock()
				streams = append(streams, stream)
				mtx.Unlock()
				err = gCtx.Err()
			}
			if err != nil {
				if shouldRetry(err) {
					level.Warn(clientSpanLog).Log("msg", "failed to fetch series", "remote", c.RemoteAddress(), "err", err)
					return nil
				}

				return err
			}

			myStreamingSeriesLabels := []labels.Labels(nil)
			var myWarnings annotations.Annotations
			myQueriedBlocks := []ulid.ULID(nil)
			indexBytesFetched := uint64(0)

			deduplicator, err := limiter.SeriesLabelsDeduplicatorFromContext(ctx)
			if err != nil {
				return err
			}

			for {
				// Ensure the context hasn't been canceled in the meanwhile (eg. an error occurred
				// in another goroutine).
				if gCtx.Err() != nil {
					return gCtx.Err()
				}

				var err error
				var isEOS bool
				var shouldRetry bool

				myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, isEOS, shouldRetry, err = q.receiveMessage(
					c, stream, queryLimiter, memoryTracker, deduplicator, myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched,
				)
				if errors.Is(err, io.EOF) {
					util.CloseAndExhaust[*storepb.SeriesResponse](stream) //nolint:errcheck
					break
				}
				if shouldRetry {
					level.Warn(clientSpanLog).Log("msg", "failed to receive series", "remote", c.RemoteAddress(), "err", err)
					return nil
				}
				if err != nil {
					return err
				}

				if isEOS {
					// If we aren't expecting any series from this stream, close it now.
					if len(myStreamingSeriesLabels) == 0 {
						util.CloseAndExhaust[*storepb.SeriesResponse](stream) //nolint:errcheck
					}

					// We expect "end of stream" to be sent after the hints and the stats have been sent, so we can break out of the loop now.
					break
				}
			}

			// debug
			var chunkInfo *chunkinfologger.ChunkInfoLogger
			if debugQuery {
				traceID, spanID, _ := tracing.ExtractTraceSpanID(ctx)
				chunkInfo = chunkinfologger.NewChunkInfoLogger("store-gateway message", traceID, spanID, q.logger, chunkinfologger.ChunkInfoLoggingFromContext(ctx))
				chunkInfo.LogSelect("store-gateway", minT, maxT)
			}

			reqStats.AddFetchedIndexBytes(indexBytesFetched)
			var streamReader *storeGatewayStreamReader
			if len(myStreamingSeriesLabels) > 0 {
				// FetchedChunks and FetchedChunkBytes are added by the SeriesChunksStreamReader.
				reqStats.AddFetchedSeries(uint64(len(myStreamingSeriesLabels)))

				if req.SkipChunks {
					// If we aren't creating a stream reader for reading chunks, we need to close the stream
					// ourselves. It's safe to close the stream multiple times so we don't worry about closing
					// all streams below when there's an error.
					if err := util.CloseAndExhaust[*storepb.SeriesResponse](stream); err != nil {
						level.Warn(clientSpanLog).Log("msg", "closing store-gateway client stream failed", "err", err)
					}
				} else {
					streamReader = newStoreGatewayStreamReader(clientCtx, stream, len(myStreamingSeriesLabels), queryLimiter, memoryTracker, reqStats, q.metrics, q.logger)
				}

				clientSpanLog.DebugLog(
					"msg", "received streaming series from store-gateway",
					"instance", c.RemoteAddress(),
					"fetched series", len(myStreamingSeriesLabels),
					"fetched index bytes", indexBytesFetched,
					"requested blocks", strings.Join(convertULIDsToString(blockIDs), " "),
					"queried blocks", strings.Join(convertULIDsToString(myQueriedBlocks), " "),
				)
			} else {
				clientSpanLog.DebugLog(
					"msg", "received no series from store-gateway",
					"instance", c.RemoteAddress(),
					"requested blocks", strings.Join(convertULIDsToString(blockIDs), " "),
					"queried blocks", strings.Join(convertULIDsToString(myQueriedBlocks), " "),
				)
			}

			// Store the result.
			mtx.Lock()
			if len(myStreamingSeriesLabels) > 0 {
				if chunkInfo != nil {
					chunkInfo.SetMsg("store-gateway streaming")
				}

				if req.SkipChunks {
					noChunkSeries := make([]storage.Series, 0, len(myStreamingSeriesLabels))
					for _, lbls := range myStreamingSeriesLabels {
						noChunkSeries = append(noChunkSeries, series.NewConcreteSeries(lbls, nil, nil))
					}

					seriesSets = append(seriesSets, series.NewConcreteSeriesSetFromSortedSeries(noChunkSeries))
				} else {
					seriesSets = append(seriesSets, &blockStreamingQuerierSeriesSet{
						series:        myStreamingSeriesLabels,
						streamReader:  streamReader,
						chunkInfo:     chunkInfo,
						remoteAddress: c.RemoteAddress(),
					})
					streamReaders = append(streamReaders, streamReader)
				}
			}
			warnings.Merge(myWarnings)
			queriedBlocks = append(queriedBlocks, myQueriedBlocks...)
			mtx.Unlock()

			return nil
		})
	}

	// Wait until all client requests complete.
	if err := g.Wait(); err != nil {
		cancelReqCtx(cancellation.NewErrorf("cancelling queries because query to at least one store-gateway failed: %w", err))

		for _, stream := range streams {
			if err := util.CloseAndExhaust[*storepb.SeriesResponse](stream); err != nil {
				level.Warn(q.logger).Log("msg", "closing store-gateway client stream failed", "err", err)
			}
		}
		return nil, nil, nil, nil, nil, err
	}

	estimateChunks = func() int {
		totalChunks := 0

		for _, sr := range streamReaders {
			totalChunks += sr.EstimateChunkCount()
		}

		return totalChunks
	}

	return seriesSets, queriedBlocks, warnings, streamReaders, estimateChunks, nil //nolint:govet // It's OK to return without cancelling reqCtx, see comment above.
}

func (q *blocksStoreQuerier) receiveMessage(c BlocksStoreClient, stream storegatewaypb.StoreGateway_SeriesClient, queryLimiter *limiter.QueryLimiter, memoryTracker *limiter.MemoryConsumptionTracker, deduplicator limiter.SeriesLabelsDeduplicator, myWarnings annotations.Annotations, myQueriedBlocks []ulid.ULID, myStreamingSeriesLabels []labels.Labels, indexBytesFetched uint64) (annotations.Annotations, []ulid.ULID, []labels.Labels, uint64, bool, bool, error) {
	resp, err := stream.Recv()
	if err != nil {
		if errors.Is(err, io.EOF) {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, err
		}

		if shouldRetry(err) {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, true, err
		}

		return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, err
	}
	defer resp.FreeBuffer()

	// Response may either contain streaming series, warning, or hints.
	if w := resp.GetWarning(); w != "" {
		myWarnings.Add(errors.New(w))
	}

	// Check both the opaque hints type and the non-opaque hints type. Store-gateways will only send
	// a single one of these as part of a series request. We need to maintain support for either one
	// in queriers until after store-gateways are switched to only use the non-opaque type.
	if h := resp.GetResponseHints(); h != nil {
		ids, err := convertBlockHintsToULIDs(h.QueriedBlocks)
		if err != nil {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, errors.Wrapf(err, "failed to parse queried block IDs from received hints")
		}

		myQueriedBlocks = append(myQueriedBlocks, ids...)
	}

	if h := resp.GetHints(); h != nil {
		// Note that we use a different but equivalent hints type for the opaque field.
		opaqueResHints := hintspb.SeriesResponseHints{}
		if err := types.UnmarshalAny(h, &opaqueResHints); err != nil {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, errors.Wrapf(err, "failed to unmarshal series hints from %s", c.RemoteAddress())
		}

		ids, err := convertBlockHintsToULIDsOpaque(opaqueResHints.QueriedBlocks)
		if err != nil {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, errors.Wrapf(err, "failed to parse queried block IDs from received hints")
		}

		myQueriedBlocks = append(myQueriedBlocks, ids...)
	}

	if s := resp.GetStats(); s != nil {
		indexBytesFetched += s.FetchedIndexBytes
	}

	if ss := resp.GetStreamingSeries(); ss != nil {
		myStreamingSeriesLabels = slices.Grow(myStreamingSeriesLabels, len(ss.Series))

		for _, s := range ss.Series {
			ls := mimirpb.FromLabelAdaptersToLabelsWithCopy(s.Labels)

			uniqueSeriesLabels, err := deduplicator.Deduplicate(ls, memoryTracker)
			if err != nil {
				return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, err
			}

			// Add series fingerprint to query limiter; will return error if we are over the limit
			if limitErr := queryLimiter.AddSeries(uniqueSeriesLabels); limitErr != nil {
				return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, limitErr
			}

			myStreamingSeriesLabels = append(myStreamingSeriesLabels, uniqueSeriesLabels)
		}

		if ss.IsEndOfSeriesStream {
			return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, true, false, nil
		}
	}

	return myWarnings, myQueriedBlocks, myStreamingSeriesLabels, indexBytesFetched, false, false, nil
}

func shouldRetry(err error) bool {
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false
	}

	if st, ok := grpcutil.ErrorToStatus(err); ok {
		return int(st.Code()) != http.StatusUnprocessableEntity
	}

	return true
}

func (q *blocksStoreQuerier) fetchLabelNamesFromStore(
	ctx context.Context,
	clients map[BlocksStoreClient][]ulid.ULID,
	minT int64,
	maxT int64,
	tenantID string,
	hints *storage.LabelHints,
	matchers []storepb.LabelMatcher,
	indexMeta *bucketindex.Metadata,
) ([][]string, annotations.Annotations, []ulid.ULID, error) {
	reqCtx := grpcContextWithBucketStoreRequestMeta(ctx, tenantID, indexMeta)

	var (
		g, gCtx       = errgroup.WithContext(reqCtx)
		mtx           = sync.Mutex{}
		nameSets      = [][]string{}
		warnings      annotations.Annotations
		queriedBlocks = []ulid.ULID(nil)
		spanLog       = spanlogger.FromContext(ctx, q.logger)
	)

	// Concurrently fetch series from all clients.
	for c, blockIDs := range clients {
		g.Go(func() error {
			req, err := createLabelNamesRequest(minT, maxT, blockIDs, hints, matchers)
			if err != nil {
				return errors.Wrapf(err, "failed to create label names request")
			}

			namesResp, err := c.LabelNames(gCtx, req)
			if err != nil {
				if shouldRetry(err) {
					level.Warn(spanLog).Log("msg", "failed to fetch label names; error is retriable", "remote", c.RemoteAddress(), "err", err)
					return nil
				}
				return fmt.Errorf("non-retriable error while fetching label names from store: %w", err)
			}

			myQueriedBlocks := []ulid.ULID(nil)
			if namesResp.ResponseHints != nil {
				ids, err := convertBlockHintsToULIDs(namesResp.ResponseHints.QueriedBlocks)
				if err != nil {
					return errors.Wrapf(err, "failed to parse queried block IDs from received hints")
				}

				myQueriedBlocks = ids
			} else if namesResp.Hints != nil { //nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
				// Note that we use a different but equivalent hints type for the opaque field.
				resHints := hintspb.LabelNamesResponseHints{}
				//nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
				if err := types.UnmarshalAny(namesResp.Hints, &resHints); err != nil {
					return errors.Wrapf(err, "failed to unmarshal label names hints from %s", c.RemoteAddress())
				}

				ids, err := convertBlockHintsToULIDsOpaque(resHints.QueriedBlocks)
				if err != nil {
					return errors.Wrapf(err, "failed to parse queried block IDs from received hints")
				}

				myQueriedBlocks = ids
			}

			spanLog.DebugLog("msg", "received label names from store-gateway",
				"instance", c,
				"num labels", len(namesResp.Names),
				"requested blocks", strings.Join(convertULIDsToString(blockIDs), " "),
				"queried blocks", strings.Join(convertULIDsToString(myQueriedBlocks), " "))

			// Store the result.
			mtx.Lock()
			nameSets = append(nameSets, namesResp.Names)
			for _, w := range namesResp.Warnings {
				warnings.Add(errors.New(w))
			}
			queriedBlocks = append(queriedBlocks, myQueriedBlocks...)
			mtx.Unlock()

			return nil
		})
	}

	// Wait until all client requests complete.
	if err := g.Wait(); err != nil {
		return nil, nil, nil, err
	}

	return nameSets, warnings, queriedBlocks, nil
}

func (q *blocksStoreQuerier) fetchLabelValuesFromStore(
	ctx context.Context,
	name string,
	clients map[BlocksStoreClient][]ulid.ULID,
	minT int64,
	maxT int64,
	tenantID string,
	hints *storage.LabelHints,
	matchers []*labels.Matcher,
	indexMeta *bucketindex.Metadata,
) ([][]string, annotations.Annotations, []ulid.ULID, error) {
	reqCtx := grpcContextWithBucketStoreRequestMeta(ctx, tenantID, indexMeta)

	var (
		g, gCtx       = errgroup.WithContext(reqCtx)
		mtx           = sync.Mutex{}
		valueSets     = [][]string{}
		warnings      annotations.Annotations
		queriedBlocks = []ulid.ULID(nil)
		spanLog       = spanlogger.FromContext(ctx, q.logger)
	)

	// Concurrently fetch series from all clients.
	for c, blockIDs := range clients {
		g.Go(func() error {
			req, err := createLabelValuesRequest(minT, maxT, name, blockIDs, hints, matchers...)
			if err != nil {
				return errors.Wrapf(err, "failed to create label values request")
			}

			valuesResp, err := c.LabelValues(gCtx, req)
			if err != nil {
				if shouldRetry(err) {
					level.Warn(spanLog).Log("msg", "failed to fetch label values; error is retriable", "remote", c.RemoteAddress(), "err", err)
					return nil
				}
				return fmt.Errorf("non-retriable error while fetching label values from store: %w", err)
			}

			myQueriedBlocks := []ulid.ULID(nil)
			if valuesResp.ResponseHints != nil {
				ids, err := convertBlockHintsToULIDs(valuesResp.ResponseHints.QueriedBlocks)
				if err != nil {
					return errors.Wrapf(err, "failed to parse queried block IDs from received hints")
				}

				myQueriedBlocks = ids
			} else if valuesResp.Hints != nil { //nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
				// Note that we use a different but equivalent hints type for the opaque field.
				resHints := hintspb.LabelValuesResponseHints{}
				//nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
				if err := types.UnmarshalAny(valuesResp.Hints, &resHints); err != nil {
					return errors.Wrapf(err, "failed to unmarshal label values hints from %s", c.RemoteAddress())
				}

				ids, err := convertBlockHintsToULIDsOpaque(resHints.QueriedBlocks)
				if err != nil {
					return errors.Wrapf(err, "failed to parse queried block IDs from received hints")
				}

				myQueriedBlocks = ids
			}

			spanLog.DebugLog("msg", "received label values from store-gateway",
				"instance", c.RemoteAddress(),
				"num values", len(valuesResp.Values),
				"requested blocks", strings.Join(convertULIDsToString(blockIDs), " "),
				"queried blocks", strings.Join(convertULIDsToString(myQueriedBlocks), " "))

			// Values returned need not be sorted, but we need them to be sorted so we can merge.
			slices.Sort(valuesResp.Values)

			// Store the result.
			mtx.Lock()
			valueSets = append(valueSets, valuesResp.Values)
			for _, w := range valuesResp.Warnings {
				warnings.Add(errors.New(w))
			}
			queriedBlocks = append(queriedBlocks, myQueriedBlocks...)
			mtx.Unlock()

			return nil
		})
	}

	// Wait until all client requests complete.
	if err := g.Wait(); err != nil {
		return nil, nil, nil, err
	}

	return valueSets, warnings, queriedBlocks, nil
}

func grpcContextWithBucketStoreRequestMeta(ctx context.Context, tenantID string, indexMeta *bucketindex.Metadata) context.Context {
	return grpc_metadata.AppendToOutgoingContext(
		ctx,
		storegateway.GrpcContextMetadataTenantID,
		tenantID,
		storegateway.GrpcContextMetadataBucketIndexUpdatedAt,
		strconv.FormatInt(indexMeta.UpdatedAt, 10),
	)
}

func createSeriesRequest(minT, maxT int64, matchers []storepb.LabelMatcher, skipChunks bool, hints *storage.SelectHints, blockIDs []ulid.ULID, streamingBatchSize uint64) (*storepb.SeriesRequest, error) {
	var limit int64
	if hints != nil && hints.Limit > 0 {
		limit = int64(hints.Limit)
	}

	// Selectively query only specific blocks.
	requestHints := &storepb.SeriesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	// We send both opaque and non-opaque request hints to store-gateways. New store-gateways
	// will prefer the non-opaque type and ignore the opaque type. Older store-gateways will
	// ignore the unknown non-opaque type and instead use the opaque type.
	opaqueHints := &hintspb.SeriesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	anyHints, err := types.MarshalAny(opaqueHints)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to marshal series request hints")
	}

	if skipChunks {
		// We don't do the streaming call if we are not requesting the chunks. Note that setting
		// a batch size of 0 is ignored in newer store-gateways as the streaming code path is always
		// used. We set this to 0 anyway here so that newer queriers will continue to work with
		// older store-gateways that have not been updated to ignore the value 0.
		streamingBatchSize = 0
	}

	return &storepb.SeriesRequest{
		MinTime:  minT,
		MaxTime:  maxT,
		Matchers: matchers,
		//nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
		Hints:                    anyHints,
		RequestHints:             requestHints,
		SkipChunks:               skipChunks,
		StreamingChunksBatchSize: streamingBatchSize,
		Limit:                    limit,
	}, nil
}

func createLabelNamesRequest(minT, maxT int64, blockIDs []ulid.ULID, hints *storage.LabelHints, matchers []storepb.LabelMatcher) (*storepb.LabelNamesRequest, error) {
	var limit int64
	if hints != nil && hints.Limit > 0 {
		limit = int64(hints.Limit)
	}
	req := &storepb.LabelNamesRequest{
		Start:    minT,
		End:      maxT,
		Matchers: matchers,
		Limit:    limit,
	}

	// Selectively query only specific blocks.
	requestHints := &storepb.LabelNamesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	// We send both opaque and non-opaque request hints to store-gateways. New store-gateways
	// will prefer the non-opaque type and ignore the opaque type. Older store-gateways will
	// ignore the unknown non-opaque type and instead use the opaque type.
	opaqueHints := &hintspb.LabelNamesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	anyHints, err := types.MarshalAny(opaqueHints)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to marshal label names request hints")
	}

	//nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
	req.Hints = anyHints
	req.RequestHints = requestHints

	return req, nil
}

func createLabelValuesRequest(minT, maxT int64, label string, blockIDs []ulid.ULID, hints *storage.LabelHints, matchers ...*labels.Matcher) (*storepb.LabelValuesRequest, error) {
	var limit int64
	if hints != nil && hints.Limit > 0 {
		limit = int64(hints.Limit)
	}

	req := &storepb.LabelValuesRequest{
		Start:    minT,
		End:      maxT,
		Label:    label,
		Matchers: convertMatchersToLabelMatcher(matchers),
		Limit:    limit,
	}

	// Selectively query only specific blocks.
	requestHints := &storepb.LabelValuesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	// We send both opaque and non-opaque request hints to store-gateways. New store-gateways
	// will prefer the non-opaque type and ignore the opaque type. Older store-gateways will
	// ignore the unknown non-opaque type and instead use the opaque type.
	opaqueHints := &hintspb.LabelValuesRequestHints{
		BlockMatchers: []storepb.LabelMatcher{
			{
				Type:  storepb.LabelMatcher_RE,
				Name:  block.BlockIDLabel,
				Value: strings.Join(convertULIDsToString(blockIDs), "|"),
			},
		},
	}

	anyHints, err := types.MarshalAny(opaqueHints)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to marshal label values request hints")
	}

	//nolint:staticcheck // Ignore SA1019. This use will be removed in Mimir 3.2
	req.Hints = anyHints
	req.RequestHints = requestHints

	return req, nil
}

func convertULIDsToString(ids []ulid.ULID) []string {
	res := make([]string, len(ids))
	for idx, id := range ids {
		res[idx] = id.String()
	}
	return res
}

func convertBlockHintsToULIDs(hints []storepb.Block) ([]ulid.ULID, error) {
	res := make([]ulid.ULID, len(hints))

	for idx, hint := range hints {
		blockID, err := ulid.Parse(hint.Id)
		if err != nil {
			return nil, err
		}

		res[idx] = blockID
	}

	return res, nil
}

func convertBlockHintsToULIDsOpaque(hints []hintspb.Block) ([]ulid.ULID, error) {
	res := make([]ulid.ULID, len(hints))

	for idx, hint := range hints {
		blockID, err := ulid.Parse(hint.Id)
		if err != nil {
			return nil, err
		}

		res[idx] = blockID
	}

	return res, nil
}
