continuoustest: query with strong read consistency (grafana#7540)

narqo · web-flow · commit 48ece79fb706 · 2024-03-11T09:54:38.000+01:00
Signed-off-by: Vladimir Varankin &lt;vladimir.varankin@grafana.com&gt;
diff --git a/pkg/continuoustest/client.go b/pkg/continuoustest/client.go
@@ -17,6 +17,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/prompb"
 
+	querierapi "github.com/grafana/mimir/pkg/querier/api"
 	"github.com/grafana/mimir/pkg/util/instrumentation"
 	util_math "github.com/grafana/mimir/pkg/util/math"
 )
@@ -153,6 +154,8 @@ func (c *Client) QueryRange(ctx context.Context, query string, start, end time.T
 	ctx, cancel := context.WithTimeout(ctx, c.cfg.ReadTimeout)
 	defer cancel()
 
+	ctx = querierapi.ContextWithReadConsistency(ctx, querierapi.ReadConsistencyStrong)
+
 	value, _, err := c.readClient.QueryRange(ctx, query, v1.Range{
 		Start: start,
 		End:   end,
@@ -180,6 +183,8 @@ func (c *Client) Query(ctx context.Context, query string, ts time.Time, options
 	ctx, cancel := context.WithTimeout(ctx, c.cfg.ReadTimeout)
 	defer cancel()
 
+	ctx = querierapi.ContextWithReadConsistency(ctx, querierapi.ReadConsistencyStrong)
+
 	value, _, err := c.readClient.Query(ctx, query, ts)
 	if err != nil {
 		return nil, err
@@ -269,5 +274,10 @@ func (rt *clientRoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
 	} else {
 		req.Header.Set("X-Scope-OrgID", rt.tenantID)
 	}
+
+	if lvl, ok := querierapi.ReadConsistencyFromContext(req.Context()); ok {
+		req.Header.Add(querierapi.ReadConsistencyHeader, lvl)
+	}
+
 	return rt.rt.RoundTrip(req)
 }
diff --git a/pkg/continuoustest/client_test.go b/pkg/continuoustest/client_test.go
@@ -21,6 +21,8 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/pdata/pmetric/pmetricotlp"
+
+	"github.com/grafana/mimir/pkg/querier/api"
 )
 
 func TestOTLPHttpClient_WriteSeries(t *testing.T) {
@@ -203,6 +205,9 @@ func TestClient_QueryRange(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		receivedRequests = append(receivedRequests, request)
 
+		// Read requests must go through strong read consistency
+		require.Equal(t, api.ReadConsistencyStrong, request.Header.Get(api.ReadConsistencyHeader))
+
 		writer.WriteHeader(http.StatusOK)
 		_, err := writer.Write([]byte(`{"status":"success","data":{"resultType":"matrix","result":[]}}`))
 		require.NoError(t, err)
@@ -248,6 +253,9 @@ func TestClient_Query(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		receivedRequests = append(receivedRequests, request)
 
+		// Read requests must go through strong read consistency
+		require.Equal(t, api.ReadConsistencyStrong, request.Header.Get(api.ReadConsistencyHeader))
+
 		writer.WriteHeader(http.StatusOK)
 		_, err := writer.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`))
 		require.NoError(t, err)
