Correctly apply parallelism limits when doing querysharding. (#253)

* Correctly apply parallelism limits when doing querysharding.

Previously we were only splitting by day/time in the frontend, so applying the max parallelism was easy.
Now we can also apply parallelism at the querysharding level, and this means we can easily by-pass the `MaxQueryParallelism` limits.

This PR apply the limits at a lower level and so fixes the problem of overscheduling per query which we currently have since querysharding is activable.

This is inspired by work we've done in Loki.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

* Update pkg/querier/queryrange/limits.go

Co-authored-by: Christian Simon <simon@swine.de>

* Update pkg/querier/queryrange/limits_test.go

Co-authored-by: Christian Simon <simon@swine.de>

* Checks for error in tests.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

* Improve concurrency handling.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

* Added unit test on context cancellation

Signed-off-by: Marco Pracucci <marco@pracucci.com>

* Simplify the code.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

Co-authored-by: Christian Simon <simon@swine.de>
Co-authored-by: Marco Pracucci <marco@pracucci.com>

Author: 3 people

pkg/querier/queryrange/limits.go

@@ -8,9 +8,11 @@ package queryrange
 import (
 	"context"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/go-kit/kit/log/level"
+	"github.com/opentracing/opentracing-go"
 	"github.com/prometheus/prometheus/pkg/timestamp"
 	"github.com/weaveworks/common/httpgrpc"
 
@@ -103,3 +105,116 @@ func (l limitsMiddleware) Do(ctx context.Context, r Request) (Response, error) {
 
 	return l.next.Do(ctx, r)
 }
+
+type limitedRoundTripper struct {
+	downstream Handler
+	limits     Limits
+
+	codec      Codec
+	middleware Middleware
+}
+
+// NewLimitedRoundTripper creates a new roundtripper that enforces MaxQueryParallelism to the `next` roundtripper across `middlewares`.
+func NewLimitedRoundTripper(next http.RoundTripper, codec Codec, limits Limits, middlewares ...Middleware) http.RoundTripper {
+	transport := limitedRoundTripper{
+		downstream: roundTripperHandler{
+			next:  next,
+			codec: codec,
+		},
+		codec:      codec,
+		limits:     limits,
+		middleware: MergeMiddlewares(middlewares...),
+	}
+	return transport
+}
+
+type subRequest struct {
+	req    Request
+	ctx    context.Context
+	result chan result
+}
+
+type result struct {
+	response Response
+	err      error
+}
+
+func newSubRequest(ctx context.Context, req Request) subRequest {
+	return subRequest{
+		req:    req,
+		ctx:    ctx,
+		result: make(chan result, 1),
+	}
+}
+
+func (rt limitedRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	var (
+		wg           sync.WaitGroup
+		intermediate = make(chan subRequest)
+		ctx, cancel  = context.WithCancel(r.Context())
+	)
+	defer func() {
+		cancel()
+		wg.Wait()
+	}()
+
+	request, err := rt.codec.DecodeRequest(ctx, r)
+	if err != nil {
+		return nil, err
+	}
+
+	if span := opentracing.SpanFromContext(ctx); span != nil {
+		request.LogToSpan(span)
+	}
+
+	tenantIDs, err := tenant.TenantIDs(ctx)
+	if err != nil {
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
+	}
+
+	// Creates workers that will process the sub-requests in parallel for this query.
+	// The amount of workers is limited by the MaxQueryParallelism tenant setting.
+	parallelism := validation.SmallestPositiveIntPerTenant(tenantIDs, rt.limits.MaxQueryParallelism)
+	for i := 0; i < parallelism; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case w := <-intermediate:
+					resp, err := rt.downstream.Do(w.ctx, w.req)
+					w.result <- result{response: resp, err: err}
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
+	}
+
+	// Wraps middlewares with a final handler, which will receive requests in
+	// parallel from upstream handlers. Then each requests gets scheduled to a
+	// different worker via the `intermediate` channel, so the maximum
+	// parallelism is limited. This worker will then call `Do` on the resulting
+	// handler.
+	response, err := rt.middleware.Wrap(
+		HandlerFunc(func(ctx context.Context, r Request) (Response, error) {
+			s := newSubRequest(ctx, r)
+			select {
+			case intermediate <- s:
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+
+			select {
+			case response := <-s.result:
+				return response.response, response.err
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+		})).Do(ctx, request)
+	if err != nil {
+		return nil, err
+	}
+
+	return rt.codec.EncodeResponse(ctx, response)
+}

pkg/querier/queryrange/limits_test.go

@@ -7,13 +7,16 @@ package queryrange
 
 import (
 	"context"
+	"net/http"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/weaveworks/common/user"
+	"go.uber.org/atomic"
 
 	"github.com/grafana/mimir/pkg/util"
 )
@@ -192,10 +195,11 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 }
 
 type mockLimits struct {
-	maxQueryLookback  time.Duration
-	maxQueryLength    time.Duration
-	maxCacheFreshness time.Duration
-	totalShards       int
+	maxQueryLookback    time.Duration
+	maxQueryLength      time.Duration
+	maxCacheFreshness   time.Duration
+	maxQueryParallelism int
+	totalShards         int
 }
 
 func (m mockLimits) MaxQueryLookback(string) time.Duration {
@@ -206,8 +210,11 @@ func (m mockLimits) MaxQueryLength(string) time.Duration {
 	return m.maxQueryLength
 }
 
-func (mockLimits) MaxQueryParallelism(string) int {
-	return 14 // Flag default.
+func (m mockLimits) MaxQueryParallelism(string) int {
+	if m.maxQueryParallelism == 0 {
+		return 14 // Flag default.
+	}
+	return m.maxQueryParallelism
 }
 
 func (m mockLimits) MaxCacheFreshness(string) time.Duration {
@@ -226,3 +233,149 @@ func (m *mockHandler) Do(ctx context.Context, req Request) (Response, error) {
 	args := m.Called(ctx, req)
 	return args.Get(0).(Response), args.Error(1)
 }
+
+func TestLimitedRoundTripper_MaxQueryParallelism(t *testing.T) {
+	var (
+		maxQueryParallelism = 2
+		count               atomic.Int32
+		max                 atomic.Int32
+		downstream          = RoundTripFunc(func(_ *http.Request) (*http.Response, error) {
+			cur := count.Inc()
+			if cur > max.Load() {
+				max.Store(cur)
+			}
+			defer count.Dec()
+			// simulate some work
+			time.Sleep(20 * time.Millisecond)
+			return &http.Response{
+				Body: http.NoBody,
+			}, nil
+		})
+		ctx = user.InjectOrgID(context.Background(), "foo")
+	)
+
+	r, err := PrometheusCodec.EncodeRequest(ctx, &PrometheusRequest{
+		Path:  "/query_range",
+		Start: time.Now().Add(time.Hour).Unix(),
+		End:   util.TimeToMillis(time.Now()),
+		Step:  int64(1 * time.Second * time.Millisecond),
+		Query: `foo`,
+	})
+	require.Nil(t, err)
+
+	_, err = NewLimitedRoundTripper(downstream, PrometheusCodec, mockLimits{maxQueryParallelism: maxQueryParallelism},
+		MiddlewareFunc(func(next Handler) Handler {
+			return HandlerFunc(func(c context.Context, _ Request) (Response, error) {
+				var wg sync.WaitGroup
+				for i := 0; i < maxQueryParallelism+20; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						_, _ = next.Do(c, &PrometheusRequest{})
+					}()
+				}
+				wg.Wait()
+				return NewEmptyPrometheusResponse(), nil
+			})
+		}),
+	).RoundTrip(r)
+	require.NoError(t, err)
+	maxFound := int(max.Load())
+	require.LessOrEqual(t, maxFound, maxQueryParallelism, "max query parallelism: ", maxFound, " went over the configured one:", maxQueryParallelism)
+}
+
+func TestLimitedRoundTripper_MaxQueryParallelismLateScheduling(t *testing.T) {
+	var (
+		maxQueryParallelism = 2
+		downstream          = RoundTripFunc(func(_ *http.Request) (*http.Response, error) {
+			// simulate some work
+			time.Sleep(20 * time.Millisecond)
+			return &http.Response{
+				Body: http.NoBody,
+			}, nil
+		})
+		ctx = user.InjectOrgID(context.Background(), "foo")
+	)
+
+	r, err := PrometheusCodec.EncodeRequest(ctx, &PrometheusRequest{
+		Path:  "/query_range",
+		Start: time.Now().Add(time.Hour).Unix(),
+		End:   util.TimeToMillis(time.Now()),
+		Step:  int64(1 * time.Second * time.Millisecond),
+		Query: `foo`,
+	})
+	require.Nil(t, err)
+
+	_, err = NewLimitedRoundTripper(downstream, PrometheusCodec, mockLimits{maxQueryParallelism: maxQueryParallelism},
+		MiddlewareFunc(func(next Handler) Handler {
+			return HandlerFunc(func(c context.Context, _ Request) (Response, error) {
+				// fire up work and we don't wait.
+				for i := 0; i < 10; i++ {
+					go func() {
+						_, _ = next.Do(c, &PrometheusRequest{})
+					}()
+				}
+				return NewEmptyPrometheusResponse(), nil
+			})
+		}),
+	).RoundTrip(r)
+	require.NoError(t, err)
+}
+
+func TestLimitedRoundTripper_OriginalRequestContextCancellation(t *testing.T) {
+	var (
+		maxQueryParallelism = 2
+		downstream          = RoundTripFunc(func(req *http.Request) (*http.Response, error) {
+			// Sleep for a long time or until the request context is canceled.
+			select {
+			case <-time.After(time.Minute):
+				return &http.Response{Body: http.NoBody}, nil
+			case <-req.Context().Done():
+				return nil, req.Context().Err()
+			}
+		})
+		reqCtx, reqCancel = context.WithCancel(user.InjectOrgID(context.Background(), "foo"))
+	)
+
+	r, err := PrometheusCodec.EncodeRequest(reqCtx, &PrometheusRequest{
+		Path:  "/query_range",
+		Start: time.Now().Add(time.Hour).Unix(),
+		End:   util.TimeToMillis(time.Now()),
+		Step:  int64(1 * time.Second * time.Millisecond),
+		Query: `foo`,
+	})
+	require.Nil(t, err)
+
+	_, err = NewLimitedRoundTripper(downstream, PrometheusCodec, mockLimits{maxQueryParallelism: maxQueryParallelism},
+		MiddlewareFunc(func(next Handler) Handler {
+			return HandlerFunc(func(c context.Context, _ Request) (Response, error) {
+				var wg sync.WaitGroup
+
+				// Fire up some work. Each sub-request will either be blocked in the sleep or in the queue
+				// waiting to be scheduled.
+				for i := 0; i < maxQueryParallelism+20; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						_, _ = next.Do(c, &PrometheusRequest{})
+					}()
+				}
+
+				// Give it a bit a time to get the first sub-requests running.
+				time.Sleep(100 * time.Millisecond)
+
+				// Cancel the original request context.
+				reqCancel()
+
+				// Wait until all sub-requests have done. We expect all of them to cancel asap,
+				// so it should take a very short time.
+				waitStart := time.Now()
+				wg.Wait()
+				assert.Less(t, time.Since(waitStart).Milliseconds(), int64(100))
+
+				return NewEmptyPrometheusResponse(), nil
+			})
+		}),
+	).RoundTrip(r)
+	require.NoError(t, err)
+}

pkg/querier/queryrange/roundtrip.go

@@ -202,7 +202,7 @@ func NewTripperware(
 	return func(next http.RoundTripper) http.RoundTripper {
 		// Finally, if the user selected any query range middleware, stitch it in.
 		if len(queryRangeMiddleware) > 0 {
-			queryrange := NewRoundTripper(next, codec, queryRangeMiddleware...)
+			queryrange := NewLimitedRoundTripper(next, codec, limits, queryRangeMiddleware...)
 			return RoundTripFunc(func(r *http.Request) (*http.Response, error) {
 				isQueryRange := strings.HasSuffix(r.URL.Path, "/query_range")
 				op := "query"
@@ -230,20 +230,20 @@ func NewTripperware(
 }
 
 type roundTripper struct {
-	next    http.RoundTripper
 	handler Handler
 	codec   Codec
 }
 
 // NewRoundTripper merges a set of middlewares into an handler, then inject it into the `next` roundtripper
 // using the codec to translate requests and responses.
 func NewRoundTripper(next http.RoundTripper, codec Codec, middlewares ...Middleware) http.RoundTripper {
-	transport := roundTripper{
-		next:  next,
+	return roundTripper{
+		handler: MergeMiddlewares(middlewares...).Wrap(roundTripperHandler{
+			next:  next,
+			codec: codec,
+		}),
 		codec: codec,
 	}
-	transport.handler = MergeMiddlewares(middlewares...).Wrap(&transport)
-	return transport
 }
 
 func (q roundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
@@ -264,8 +264,15 @@ func (q roundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
 	return q.codec.EncodeResponse(r.Context(), response)
 }
 
+// roundTripperHandler is a handler that roundtrips requests to next roundtripper.
+// It basically encodes a Request from Handler.Do and decode response from next roundtripper.
+type roundTripperHandler struct {
+	next  http.RoundTripper
+	codec Codec
+}
+
 // Do implements Handler.
-func (q roundTripper) Do(ctx context.Context, r Request) (Response, error) {
+func (q roundTripperHandler) Do(ctx context.Context, r Request) (Response, error) {
 	request, err := q.codec.EncodeRequest(ctx, r)
 	if err != nil {
 		return nil, err