querymiddleware: Fix race condition in shardActiveSeriesMiddleware

CHANGELOG.md

@@ -78,6 +78,7 @@
   * `active_series`: active series query
   * `other`: any other request
 * [BUGFIX] Fix performance regression introduced in Mimir 2.11.0 when uploading blocks to AWS S3. #7240
+* [BUGFIX] Query-frontend: fix race condition when sharding active series is enabled (see above) and response is compressed with snappy. #7290
 
 ### Mixin
 

pkg/frontend/querymiddleware/shard_active_series.go

@@ -56,15 +56,13 @@ type shardActiveSeriesMiddleware struct {
 	upstream http.RoundTripper
 	limits   Limits
 	logger   log.Logger
-	encoder  *s2.Writer
 }
 
 func newShardActiveSeriesMiddleware(upstream http.RoundTripper, limits Limits, logger log.Logger) http.RoundTripper {
 	return &shardActiveSeriesMiddleware{
 		upstream: upstream,
 		limits:   limits,
 		logger:   logger,
-		encoder:  s2.NewWriter(nil),
 	}
 }
 
@@ -340,19 +338,17 @@ func (s *shardActiveSeriesMiddleware) mergeResponses(ctx context.Context, respon
 }
 
 func (s *shardActiveSeriesMiddleware) writeMergedResponse(ctx context.Context, check func() error, w io.WriteCloser, items <-chan *labels.Builder, encodingType string) {
-	defer func(encoder, w io.Closer) {
-		_ = encoder.Close()
-		_ = w.Close()
-	}(s.encoder, w)
+	defer w.Close()
 
 	span, _ := opentracing.StartSpanFromContext(ctx, "shardActiveSeries.writeMergedResponse")
 	defer span.Finish()
 
 	var out io.Writer = w
 	if encodingType == encodingTypeSnappyFramed {
 		span.LogFields(otlog.String("encoding", encodingTypeSnappyFramed))
-		out = s.encoder
-		s.encoder.Reset(w)
+		enc := s2.NewWriter(w)
+		defer enc.Close()
+		out = enc
 	} else {
 		span.LogFields(otlog.String("encoding", "none"))
 	}

pkg/frontend/querymiddleware/shard_active_series_test.go

@@ -267,7 +267,6 @@ func Test_shardActiveSeriesMiddleware_RoundTrip(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-
 			// Stub upstream with valid or invalid responses.
 			var requestCount atomic.Int32
 			upstream := RoundTripFunc(func(r *http.Request) (*http.Response, error) {
@@ -359,69 +358,6 @@ func Test_shardActiveSeriesMiddleware_RoundTrip(t *testing.T) {
 	}
 }
 
-func Test_shardActiveSeriesMiddleware_RoundTrip_ResponseBodyStreamed(t *testing.T) {
-	// This value needs to be set at least as large as the buffer size used by the
-	// implementation for this test to make sense.
-	const bufferSize = 512
-	const shardCount = 2
-
-	// Stub upstream with two responses that are larger than the buffer size and
-	// retain a reference to the response bodies. The responses use a custom body
-	// type that counts the number of bytes read, so we can assert on that later in
-	// the test.
-	var upstreamResponseBodies [shardCount]*bodyReadBytesCounter
-	var responseSize [shardCount]int
-	upstream := RoundTripFunc(func(r *http.Request) (*http.Response, error) {
-		// Extract requested shard index
-		require.NoError(t, r.ParseForm())
-		req, err := cardinality.DecodeActiveSeriesRequestFromValues(r.Form)
-		require.NoError(t, err)
-		shard, _, err := sharding.ShardFromMatchers(req.Matchers)
-		require.NoError(t, err)
-		require.NotNil(t, shard, "this test requires a shard to be requested")
-
-		// Make sure the response body is big enough to not be buffered entirely.
-		response := fmt.Sprintf(fmt.Sprintf(`{"data": [{"__name__": "metric-%%0%dd"}]}`, bufferSize), shard.ShardIndex)
-		body := &bodyReadBytesCounter{body: io.NopCloser(strings.NewReader(response))}
-		upstreamResponseBodies[shard.ShardIndex] = body
-		responseSize[shard.ShardIndex] = len(response)
-
-		return &http.Response{StatusCode: http.StatusOK, Body: body}, nil
-	})
-
-	s := newShardActiveSeriesMiddleware(
-		upstream,
-		mockLimits{maxShardedQueries: shardCount, totalShards: shardCount},
-		log.NewNopLogger(),
-	)
-
-	r := httptest.NewRequest("POST", "/active_series", strings.NewReader(`selector={__name__=~"metric-.*"}`))
-	r.Header.Add("Content-Type", "application/x-www-form-urlencoded")
-
-	resp, err := s.RoundTrip(r.WithContext(user.InjectOrgID(r.Context(), "test")))
-	require.NoError(t, err)
-	defer func(body io.ReadCloser) {
-		_, _ = io.ReadAll(body)
-		_ = body.Close()
-	}(resp.Body)
-
-	// Check that upstream responses have been read only up to a max of the buffer size.
-	for _, body := range upstreamResponseBodies {
-		bytesRead := int(body.BytesRead())
-		require.GreaterOrEqual(t, bufferSize, bytesRead)
-	}
-
-	// Read and close the response body.
-	_, _ = io.ReadAll(resp.Body)
-	_ = resp.Body.Close()
-
-	// Check that upstream responses have been fully read now.
-	for i, body := range upstreamResponseBodies {
-		bytesRead := int(body.BytesRead())
-		assert.Equal(t, responseSize[i], bytesRead)
-	}
-}
-
 func BenchmarkActiveSeriesMiddlewareMergeResponses(b *testing.B) {
 	type activeSeriesResponse struct {
 		Data []labels.Labels `json:"data"`
@@ -465,26 +401,6 @@ func BenchmarkActiveSeriesMiddlewareMergeResponses(b *testing.B) {
 	}
 }
 
-// bodyReadBytesCounter is a wrapper around a response body that counts the number of bytes read from it.
-type bodyReadBytesCounter struct {
-	body      io.ReadCloser
-	bytesRead atomic.Uint64
-}
-
-func (b *bodyReadBytesCounter) Read(p []byte) (n int, err error) {
-	read, err := b.body.Read(p)
-	b.bytesRead.Add(uint64(read))
-	return read, err
-}
-
-func (b *bodyReadBytesCounter) Close() error {
-	return b.body.Close()
-}
-
-func (b *bodyReadBytesCounter) BytesRead() uint64 {
-	return b.bytesRead.Load()
-}
-
 type result struct {
 	Data   []labels.Labels `json:"data"`
 	Status string          `json:"status,omitempty"`