fix read reader multiple times in QFE

Signed-off-by: Ben Ye <[email protected]>

Author: yeya24

integration/e2ecortex/client.go

@@ -122,7 +122,7 @@ func (c *Client) Query(query string, ts time.Time) (model.Value, error) {
 	return value, err
 }
 
-// Query runs a query range.
+// QueryRange runs a query range.
 func (c *Client) QueryRange(query string, start, end time.Time, step time.Duration) (model.Value, error) {
 	value, _, err := c.querierClient.QueryRange(context.Background(), query, promv1.Range{
 		Start: start,

integration/query_frontend_test.go

@@ -154,6 +154,46 @@ func TestQueryFrontendTLSWithBlocksStorageViaFlags(t *testing.T) {
 	})
 }
 
+func TestQueryFrontendWithVerticalSharding(t *testing.T) {
+	runQueryFrontendTest(t, queryFrontendTestConfig{
+		testMissingMetricName: false,
+		querySchedulerEnabled: false,
+		queryStatsEnabled:     true,
+		setup: func(t *testing.T, s *e2e.Scenario) (configFile string, flags map[string]string) {
+			require.NoError(t, writeFileToSharedDir(s, cortexConfigFile, []byte(BlocksStorageConfig)))
+
+			minio := e2edb.NewMinio(9000, BlocksStorageFlags()["-blocks-storage.s3.bucket-name"])
+			require.NoError(t, s.StartAndWaitReady(minio))
+
+			// Enable vertical sharding.
+			flags := mergeFlags(e2e.EmptyFlags(), map[string]string{
+				"-frontend.query-vertical-shard-size": "2",
+			})
+			return cortexConfigFile, e2e.EmptyFlags()
+		},
+	})
+}
+
+func TestQueryFrontendWithVerticalShardingQueryScheduler(t *testing.T) {
+	runQueryFrontendTest(t, queryFrontendTestConfig{
+		testMissingMetricName: false,
+		querySchedulerEnabled: true,
+		queryStatsEnabled:     true,
+		setup: func(t *testing.T, s *e2e.Scenario) (configFile string, flags map[string]string) {
+			require.NoError(t, writeFileToSharedDir(s, cortexConfigFile, []byte(BlocksStorageConfig)))
+
+			minio := e2edb.NewMinio(9000, BlocksStorageFlags()["-blocks-storage.s3.bucket-name"])
+			require.NoError(t, s.StartAndWaitReady(minio))
+
+			// Enable vertical sharding.
+			flags := mergeFlags(e2e.EmptyFlags(), map[string]string{
+				"-frontend.query-vertical-shard-size": "2",
+			})
+			return cortexConfigFile, e2e.EmptyFlags()
+		},
+	})
+}
+
 func runQueryFrontendTest(t *testing.T, cfg queryFrontendTestConfig) {
 	const numUsers = 10
 	const numQueriesPerUser = 10
@@ -304,7 +344,7 @@ func runQueryFrontendTest(t *testing.T, cfg queryFrontendTestConfig) {
 
 	require.NoError(t, queryFrontend.WaitSumMetrics(e2e.Equals(numUsers*numQueriesPerUser+extra), "cortex_query_frontend_queries_total"))
 
-	// The number of received request is greater then the query requests because include
+	// The number of received request is greater than the query requests because include
 	// requests to /metrics and /ready.
 	require.NoError(t, queryFrontend.WaitSumMetricsWithOptions(e2e.Greater(numUsers*numQueriesPerUser), []string{"cortex_request_duration_seconds"}, e2e.WithMetricCount))
 	require.NoError(t, querier.WaitSumMetricsWithOptions(e2e.Greater(numUsers*numQueriesPerUser), []string{"cortex_request_duration_seconds"}, e2e.WithMetricCount))

pkg/frontend/transport/handler.go

@@ -137,6 +137,12 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var buf bytes.Buffer
 	r.Body = http.MaxBytesReader(w, r.Body, f.cfg.MaxBodySize)
 	r.Body = io.NopCloser(io.TeeReader(r.Body, &buf))
+	// We parse form here so that we can use buf as body, in order to
+	// prevent https://github.com/cortexproject/cortex/issues/5201.
+	if err := r.ParseForm(); err != nil {
+		level.Warn(util_log.WithContext(r.Context(), f.log)).Log("msg", "unable to parse request form", "err", err)
+	}
+	r.Body = io.NopCloser(&buf)
 
 	startTime := time.Now()
 	resp, err := f.roundTripper.RoundTrip(r)