github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/queryrange/queryrangebase/util.go (about)

     1  package queryrangebase
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  
     7  	"github.com/weaveworks/common/httpgrpc"
     8  
     9  	"github.com/grafana/dskit/tenant"
    10  
    11  	"github.com/grafana/loki/pkg/util/validation"
    12  )
    13  
    14  // RequestResponse contains a request response and the respective request that was used.
    15  type RequestResponse struct {
    16  	Request  Request
    17  	Response Response
    18  }
    19  
    20  // DoRequests executes a list of requests in parallel. The limits parameters is used to limit parallelism per single request.
    21  func DoRequests(ctx context.Context, downstream Handler, reqs []Request, limits Limits) ([]RequestResponse, error) {
    22  	tenantIDs, err := tenant.TenantIDs(ctx)
    23  	if err != nil {
    24  		return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
    25  	}
    26  
    27  	// If one of the requests fail, we want to be able to cancel the rest of them.
    28  	ctx, cancel := context.WithCancel(ctx)
    29  	defer cancel()
    30  
    31  	// Feed all requests to a bounded intermediate channel to limit parallelism.
    32  	intermediate := make(chan Request)
    33  	go func() {
    34  		for _, req := range reqs {
    35  			intermediate <- req
    36  		}
    37  		close(intermediate)
    38  	}()
    39  
    40  	respChan, errChan := make(chan RequestResponse), make(chan error)
    41  	parallelism := validation.SmallestPositiveIntPerTenant(tenantIDs, limits.MaxQueryParallelism)
    42  	if parallelism > len(reqs) {
    43  		parallelism = len(reqs)
    44  	}
    45  	for i := 0; i < parallelism; i++ {
    46  		go func() {
    47  			for req := range intermediate {
    48  				resp, err := downstream.Do(ctx, req)
    49  				if err != nil {
    50  					errChan <- err
    51  				} else {
    52  					respChan <- RequestResponse{req, resp}
    53  				}
    54  			}
    55  		}()
    56  	}
    57  
    58  	resps := make([]RequestResponse, 0, len(reqs))
    59  	var firstErr error
    60  	for range reqs {
    61  		select {
    62  		case resp := <-respChan:
    63  			resps = append(resps, resp)
    64  		case err := <-errChan:
    65  			if firstErr == nil {
    66  				cancel()
    67  				firstErr = err
    68  			}
    69  		}
    70  	}
    71  
    72  	return resps, firstErr
    73  }