store-gateway: use streaming for skipchunks requests

This changes store-gateways to use the streaming code path for requests
where chunks are not sent. This provides the same results to queriers
since they already know how to parse the streaming response and allows
us to remove the non-streaming code path in the store-gateway.

Notably however, this change does not make remove any of the non-
streaming code from queriers. This is required to avoid failing requests
during a rollout. When queriers make requests without chunks to new
store-gateways, the streaming code path will be used and queriers will
be able to parse the result. When requests without chunks are made to
older store-gateways the non-streaming code path will be used which
queriers are still able to parse the result of.

Part of #12673

Signed-off-by: Nick Pillitteri <[email protected]>

Author: 56quarters

CHANGELOG.md

@@ -16,7 +16,7 @@
 * [CHANGE] Query-frontend: Use the Mimir Query Engine (MQE) by default. #12361
 * [CHANGE] Query-frontend: Remove the CLI flags `-querier.frontend-address`, `-querier.max-outstanding-requests-per-tenant`, and `-query-frontend.querier-forget-delay` and corresponding YAML configurations. This is part of a change that makes the query-scheduler a required component. This removes the ability to run the query-frontend with an embedded query-scheduler. Instead, you must run a dedicated query-scheduler component. #12200
 * [CHANGE] Ingester: Remove deprecated `-ingester.stream-chunks-when-using-blocks` CLI flag and `ingester_stream_chunks_when_using_blocks` runtime configuration option. #12615
-* [CHANGE] Querier: Require non-zero values for `-querier.streaming-chunks-per-ingester-buffer-size` and `-querier.streaming-chunks-per-store-gateway-buffer-size` CLI flags and corresponding YAML configurations. This is part of a change that makes streaming required between queriers and ingesters and store-gateways. Streaming has been the default since Mimir 2.14. #12790
+* [CHANGE] Querier: Require non-zero values for `-querier.streaming-chunks-per-ingester-buffer-size` and `-querier.streaming-chunks-per-store-gateway-buffer-size` CLI flags and corresponding YAML configurations. This is part of a change that makes streaming required between queriers and ingesters and store-gateways. Streaming has been the default since Mimir 2.14. #12790 #12818
 * [CHANGE] Remove support for the experimental read-write deployment mode. #12584
 * [CHANGE] Store-gateway: Update default value of `-store-gateway.dynamic-replication.multiple` to `5` to increase replication of recent blocks. #12433
 * [CHANGE] Cost attribution: Reduce the default maximum per-user cardinality of cost attribution labels to 2000. #12625

pkg/querier/blocks_store_queryable.go

@@ -1236,9 +1236,13 @@ func createSeriesRequest(minT, maxT int64, matchers []storepb.LabelMatcher, skip
 	}
 
 	if skipChunks {
-		// We don't do the streaming call if we are not requesting the chunks.
+		// 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,

pkg/storegateway/bucket.go

@@ -67,6 +67,7 @@ const (
 	labelDecode = "decode"
 
 	targetQueryStreamBatchMessageSize = 1 * 1024 * 1024
+	fallbackStreamingBatchSize        = 256
 )
 
 type BucketStoreStats struct {
@@ -555,10 +556,17 @@ type seriesChunks struct {
 
 // Series implements the storegatewaypb.StoreGatewayServer interface.
 func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storegatewaypb.StoreGateway_SeriesServer) (err error) {
-	if req.SkipChunks {
-		// We don't do the streaming call if we are not requesting the chunks.
-		req.StreamingChunksBatchSize = 0
+	ctx := srv.Context()
+	spanLogger := spanlogger.FromContext(ctx, s.logger)
+
+	// Previously, a batch size of 0 was used to indicate this was a Prometheus series request
+	// that didn't need to load any chunk data (in addition to the SkipChunks field) and used
+	// the non-streaming path. Now that Prometheus series requests use the streaming path we
+	// need to make sure that the batch size is always non-zero.
+	if req.StreamingChunksBatchSize == 0 {
+		req.StreamingChunksBatchSize = fallbackStreamingBatchSize
 	}
+
 	defer func() { err = mapSeriesError(err) }()
 
 	matchers, err := storepb.MatchersToPromMatchers(req.Matchers...)
@@ -572,15 +580,11 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storegatewaypb.Stor
 		return status.Error(codes.InvalidArgument, errors.Wrap(err, "parse query sharding label").Error())
 	}
 
-	var (
-		spanLogger       = spanlogger.FromContext(srv.Context(), s.logger)
-		ctx              = srv.Context()
-		stats            = newSafeQueryStats()
-		reqBlockMatchers []*labels.Matcher
-	)
+	stats := newSafeQueryStats()
 	defer s.recordSeriesCallResult(stats)
 	defer s.recordRequestAmbientTime(stats, time.Now())
 
+	var reqBlockMatchers []*labels.Matcher
 	if req.Hints != nil {
 		reqHints := &hintspb.SeriesRequestHints{}
 		if err := types.UnmarshalAny(req.Hints, reqHints); err != nil {
@@ -597,17 +601,16 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storegatewaypb.Stor
 
 	blocks, indexReaders, chunkReaders := s.openBlocksForReading(ctx, req.SkipChunks, req.MinTime, req.MaxTime, reqBlockMatchers, stats)
 	// We must keep the readers open until all their data has been sent.
-	for _, r := range indexReaders {
-		defer runutil.CloseWithLogOnErr(s.logger, r, "close block index reader")
-	}
-	for _, r := range chunkReaders {
-		defer runutil.CloseWithLogOnErr(s.logger, r, "close block chunk reader")
-	}
-
-	var readers *bucketChunkReaders
-	if !req.SkipChunks {
-		readers = newChunkReaders(chunkReaders)
-	}
+	defer func() {
+		for _, r := range indexReaders {
+			runutil.CloseWithLogOnErr(s.logger, r, "close block index reader")
+		}
+	}()
+	defer func() {
+		for _, r := range chunkReaders {
+			runutil.CloseWithLogOnErr(s.logger, r, "close block chunk reader")
+		}
+	}()
 
 	// Wait for the query gate only after opening blocks. Opening blocks is usually fast (~1ms),
 	// but sometimes it can take minutes if the block isn't loaded and there is a surge in queries for unloaded blocks.
@@ -617,92 +620,80 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, srv storegatewaypb.Stor
 	}
 	defer done()
 
+	// Send hints about the blocks loaded for this query before sending series or stats.
+	resHints := buildSeriesResponseHints(blocks)
+	if err = s.sendHints(srv, resHints); err != nil {
+		return err
+	}
+
+	streamingSeriesCount := 0
 	var (
+		seriesSet          storepb.SeriesSet
 		streamingIterators *streamingSeriesIterators
-		resHints           = &hintspb.SeriesResponseHints{}
+		seriesLoadStart    = time.Now()
+		chunksLimiter      = s.chunksLimiterFactory(s.metrics.queriesDropped.WithLabelValues("chunks"))
+		seriesLimiter      = s.seriesLimiterFactory(s.metrics.queriesDropped.WithLabelValues("series"))
 	)
-	for _, b := range blocks {
-		resHints.AddQueriedBlock(b.meta.ULID)
-
-		if b.meta.Compaction.Level == 1 && b.meta.Thanos.Source == block.ReceiveSource && !b.queried.Load() {
-			level.Debug(s.logger).Log("msg", "queried non-compacted block", "blockId", b.meta.ULID, "ooo", b.meta.Compaction.FromOutOfOrder())
-		}
 
-		b.queried.Store(true)
-	}
-	if err := s.sendHints(srv, resHints); err != nil {
+	seriesSet, streamingIterators, err = s.createIteratorForChunksStreamingLabelsPhase(ctx, req, blocks, indexReaders, shardSelector, matchers, chunksLimiter, seriesLimiter, stats)
+	if err != nil {
 		return err
 	}
 
-	streamingSeriesCount := 0
-	if req.StreamingChunksBatchSize > 0 {
-		var (
-			seriesSet       storepb.SeriesSet
-			seriesLoadStart = time.Now()
-			chunksLimiter   = s.chunksLimiterFactory(s.metrics.queriesDropped.WithLabelValues("chunks"))
-			seriesLimiter   = s.seriesLimiterFactory(s.metrics.queriesDropped.WithLabelValues("series"))
-		)
-
-		seriesSet, streamingIterators, err = s.createIteratorForChunksStreamingLabelsPhase(ctx, req, blocks, indexReaders, shardSelector, matchers, chunksLimiter, seriesLimiter, stats)
-		if err != nil {
-			return err
-		}
-
-		streamingSeriesCount, err = s.sendStreamingSeriesLabelsAndStats(req, srv, stats, seriesSet)
-		if err != nil {
-			return err
-		}
-		spanLogger.DebugLog(
-			"msg", "sent streaming series",
-			"num_series", streamingSeriesCount,
-			"duration", time.Since(seriesLoadStart),
-		)
+	streamingSeriesCount, err = s.sendStreamingSeriesLabelsAndStats(req, srv, stats, seriesSet)
+	if err != nil {
+		return err
+	}
+	spanLogger.DebugLog(
+		"msg", "sent streaming series",
+		"num_series", streamingSeriesCount,
+		"duration", time.Since(seriesLoadStart),
+	)
 
-		if streamingSeriesCount == 0 {
-			// There is no series to send chunks for.
-			return nil
-		}
+	if streamingSeriesCount == 0 {
+		// There is no series to send chunks for.
+		return nil
 	}
 
 	// We create the limiter twice in the case of streaming so that we don't double count the series
 	// and hit the limit prematurely.
-	chunksLimiter := s.chunksLimiterFactory(s.metrics.queriesDropped.WithLabelValues("chunks"))
-	seriesLimiter := s.seriesLimiterFactory(s.metrics.queriesDropped.WithLabelValues("series"))
+	chunksLimiter = s.chunksLimiterFactory(s.metrics.queriesDropped.WithLabelValues("chunks"))
+	seriesLimiter = s.seriesLimiterFactory(s.metrics.queriesDropped.WithLabelValues("series"))
 
 	start := time.Now()
-	if req.StreamingChunksBatchSize > 0 {
+	if !req.SkipChunks {
+		readers := newChunkReaders(chunkReaders)
 		seriesChunkIt := s.createIteratorForChunksStreamingChunksPhase(ctx, readers, stats, chunksLimiter, seriesLimiter, streamingIterators)
 		err = s.sendStreamingChunks(req, srv, seriesChunkIt, stats, streamingSeriesCount)
-	} else {
-		var seriesSet storepb.SeriesSet
-		seriesSet, err = s.createIteratorForNonChunksStreamingRequest(ctx, req, blocks, indexReaders, readers, shardSelector, matchers, chunksLimiter, seriesLimiter, stats)
 		if err != nil {
-			return err
+			return
 		}
-		err = s.sendSeriesChunks(req, srv, seriesSet, stats)
-	}
-	if err != nil {
-		return
 	}
 
 	numSeries, numChunks := stats.seriesAndChunksCount()
 	debugMessage := "sent series"
-	if req.StreamingChunksBatchSize > 0 {
+	if !req.SkipChunks {
 		debugMessage = "sent streaming chunks"
 	}
+
 	spanLogger.DebugLog(
 		"msg", debugMessage,
 		"num_series", numSeries,
 		"num_chunks", numChunks,
 		"duration", time.Since(start),
 	)
 
-	if req.StreamingChunksBatchSize == 0 {
-		// Stats were not sent before, so send it now.
-		return s.sendStats(srv, stats)
+	return nil
+}
+
+func buildSeriesResponseHints(blocks []*bucketBlock) *hintspb.SeriesResponseHints {
+	resHints := &hintspb.SeriesResponseHints{}
+	for _, b := range blocks {
+		resHints.AddQueriedBlock(b.meta.ULID)
+		b.queried.Store(true)
 	}
 
-	return nil
+	return resHints
 }
 
 func mapSeriesError(err error) error {
@@ -918,53 +909,6 @@ func (s *BucketStore) sendStreamingChunks(
 	return it.Err()
 }
 
-func (s *BucketStore) sendSeriesChunks(
-	req *storepb.SeriesRequest,
-	srv storegatewaypb.StoreGateway_SeriesServer,
-	seriesSet storepb.SeriesSet,
-	stats *safeQueryStats,
-) error {
-	var (
-		encodeDuration           time.Duration
-		sendDuration             time.Duration
-		seriesCount, chunksCount int
-	)
-
-	defer stats.update(func(stats *queryStats) {
-		stats.mergedSeriesCount += seriesCount
-		stats.mergedChunksCount += chunksCount
-
-		stats.streamingSeriesEncodeResponseDuration += encodeDuration
-		stats.streamingSeriesSendResponseDuration += sendDuration
-	})
-
-	for seriesSet.Next() {
-		seriesCount++
-		// IMPORTANT: do not retain the memory returned by seriesSet.At() beyond this loop cycle
-		// because the subsequent call to seriesSet.Next() may release it. But it is safe to hold
-		// onto lset because the labels are not released.
-		lset, chks := seriesSet.At()
-		series := storepb.Series{
-			Labels: mimirpb.FromLabelsToLabelAdapters(lset),
-		}
-		if !req.SkipChunks {
-			series.Chunks = chks
-			chunksCount += len(chks)
-			s.metrics.chunkSizeBytes.Observe(float64(chunksSize(chks)))
-		}
-
-		err := s.sendMessage("series", srv, storepb.NewSeriesResponse(&series), &encodeDuration, &sendDuration)
-		if err != nil {
-			return err
-		}
-	}
-	if seriesSet.Err() != nil {
-		return errors.Wrap(seriesSet.Err(), "expand series set")
-	}
-
-	return nil
-}
-
 func (s *BucketStore) sendMessage(typ string, srv storegatewaypb.StoreGateway_SeriesServer, msg interface{}, encodeDuration, sendDuration *time.Duration) error {
 	// We encode it ourselves into a PreparedMsg in order to measure the time it takes.
 	encodeBegin := time.Now()
@@ -1032,38 +976,6 @@ func chunksSize(chks []storepb.AggrChunk) (size int) {
 	return size
 }
 
-// createIteratorForNonChunksStreamingRequest is used when the streaming feature is not enabled.
-func (s *BucketStore) createIteratorForNonChunksStreamingRequest(
-	ctx context.Context,
-	req *storepb.SeriesRequest,
-	blocks []*bucketBlock,
-	indexReaders map[ulid.ULID]*bucketIndexReader,
-	chunkReaders *bucketChunkReaders,
-	shardSelector *sharding.ShardSelector,
-	matchers []*labels.Matcher,
-	chunksLimiter ChunksLimiter,
-	seriesLimiter SeriesLimiter,
-	stats *safeQueryStats,
-) (storepb.SeriesSet, error) {
-	strategy := defaultStrategy
-	if req.SkipChunks {
-		strategy = noChunkRefs
-	}
-	it, err := s.getSeriesIteratorFromBlocks(ctx, req, blocks, indexReaders, shardSelector, matchers, chunksLimiter, seriesLimiter, stats, strategy, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	var set storepb.SeriesSet
-	if !req.SkipChunks {
-		ss := newChunksPreloadingIterator(ctx, s.logger, s.userID, *chunkReaders, it, s.maxSeriesPerBatch, stats)
-		set = newSeriesChunksSeriesSet(ss)
-	} else {
-		set = newSeriesSetWithoutChunks(ctx, it, stats)
-	}
-	return set, nil
-}
-
 // createIteratorForChunksStreamingLabelsPhase is used when streaming feature is enabled.
 // It returns a series set that only contains the series labels without any chunks information.
 // The streamingSeriesIterators should be re-used when getting chunks to save on computation.

pkg/storegateway/bucket_e2e_test.go

@@ -438,8 +438,8 @@ func testBucketStore_e2e(t *testing.T, ctx context.Context, s *storeSuite, addit
 		},
 	}
 	for i, tcase := range append(testCases, additionalCases...) {
-		for _, streamingBatchSize := range []int{0, 1, 5, 256} {
-			if ok := t.Run(fmt.Sprintf("%d,streamingBatchSize=%d", i, streamingBatchSize), func(t *testing.T) {
+		for _, streamingBatchSize := range []int{1, 5, 256} {
+			t.Run(fmt.Sprintf("%d,streamingBatchSize=%d", i, streamingBatchSize), func(t *testing.T) {
 				tcase.req.StreamingChunksBatchSize = uint64(streamingBatchSize)
 				seriesSet, _, _, _, err := srv.Series(context.Background(), tcase.req)
 				require.NoError(t, err)
@@ -451,9 +451,7 @@ func testBucketStore_e2e(t *testing.T, ctx context.Context, s *storeSuite, addit
 					assert.Equal(t, tcase.expectedChunkLen, len(s.Chunks))
 				}
 				assertQueryStatsMetricsRecorded(t, len(tcase.expected), tcase.expectedChunkLen, s.metricsRegistry)
-			}); !ok {
-				return
-			}
+			})
 		}
 	}
 }

pkg/storegateway/bucket_store_server_test.go

@@ -143,11 +143,6 @@ func (s *storeTestServer) Series(ctx context.Context, req *storepb.SeriesRequest
 		}
 
 		if recvSeries := res.GetStreamingSeries(); recvSeries != nil {
-			if req.StreamingChunksBatchSize == 0 || req.SkipChunks {
-				err = errors.New("got a streaming series when streaming was disabled")
-				return
-			}
-
 			var recvSeriesData []byte
 
 			// We prefer to stay on the safest side at this stage
@@ -172,7 +167,7 @@ func (s *storeTestServer) Series(ctx context.Context, req *storepb.SeriesRequest
 		}
 	}
 
-	if req.StreamingChunksBatchSize > 0 && !req.SkipChunks {
+	if !req.SkipChunks {
 		res, err = stream.Recv()
 		if err != nil {
 			if errors.Is(err, io.EOF) {
@@ -250,6 +245,14 @@ func (s *storeTestServer) Series(ctx context.Context, req *storepb.SeriesRequest
 		if errors.Is(err, io.EOF) {
 			err = nil
 		}
+	} else {
+		// The store-gateway returns a streaming response even when SkipChunks = true. In that
+		// case, convert the streaming series responses into the expected series set.
+		for _, streamingSeries := range streamingSeriesSet {
+			seriesSet = append(seriesSet, &storepb.Series{
+				Labels: streamingSeries.Labels,
+			})
+		}
 	}
 
 	return