Method queryIngesterStream

pkg/distributor/query.go:225–340

queryIngesterStream queries the ingesters using the new streaming API.

(ctx context.Context, replicationSet ring.ReplicationSet, req *ingester_client.QueryRequest, partialDataEnabled bool)

// queryIngesterStream queries the ingesters using the new streaming API.
func (d *Distributor) queryIngesterStream(ctx context.Context, replicationSet ring.ReplicationSet, req *ingester_client.QueryRequest, partialDataEnabled bool) (*ingester_client.QueryStreamResponse, error) {
	var (
		queryLimiter = limiter.QueryLimiterFromContextWithFallback(ctx)
		reqStats     = stats.FromContext(ctx)
	)

	// Fetch samples from multiple ingesters
	results, err := replicationSet.Do(ctx, d.cfg.ExtraQueryDelay, false, partialDataEnabled, func(ctx context.Context, ing *ring.InstanceDesc) (any, error) {
		client, err := d.ingesterPool.GetClientFor(ing.Addr)
		if err != nil {
			return nil, err
		}

		ingesterId, err := d.ingestersRing.GetInstanceIdByAddr(ing.Addr)
		if err != nil {
			level.Warn(d.log).Log("msg", "instance not found in the ring", "addr", ing.Addr, "err", err)
		}

		d.ingesterQueries.WithLabelValues(ingesterId).Inc()

		stream, err := client.(ingester_client.IngesterClient).QueryStream(ctx, req)
		if err != nil {
			d.ingesterQueryFailures.WithLabelValues(ingesterId).Inc()
			return nil, err
		}
		defer stream.CloseSend() //nolint:errcheck

		result := &ingester_client.QueryStreamResponse{}
		for {
			resp, err := stream.Recv()
			if err == io.EOF {
				break
			} else if err != nil {
				// Do not track a failure if the context was canceled.
				if !grpcutil.IsGRPCContextCanceled(err) {
					d.ingesterQueryFailures.WithLabelValues(ingesterId).Inc()
				}

				return nil, err
			}

			// Enforce the max chunks limits.
			if chunkLimitErr := queryLimiter.AddChunks(resp.ChunksCount()); chunkLimitErr != nil {
				return nil, validation.LimitError(chunkLimitErr.Error())
			}

			s := make([][]cortexpb.LabelAdapter, 0, len(resp.Chunkseries))
			for _, series := range resp.Chunkseries {
				s = append(s, series.Labels)
			}

			if limitErr := queryLimiter.AddSeries(s...); limitErr != nil {
				return nil, validation.LimitError(limitErr.Error())
			}

			if chunkBytesLimitErr := queryLimiter.AddChunkBytes(resp.ChunksSize()); chunkBytesLimitErr != nil {
				return nil, validation.LimitError(chunkBytesLimitErr.Error())
			}