From 53091126c28f518c154c365a47d617c9ee7634de Mon Sep 17 00:00:00 2001
From: Julien Pivotto <roidelapluie@o11y.eu>
Date: Wed, 13 Mar 2024 15:36:53 +0100
Subject: [PATCH] Chunked remote read: close the querier earlier

I have seen prometheis instances misebehaving because of broken chinked remote
read requests.

In order to avoid OOM's when this happens, I propose to close the
queries used by the streamed remote read requests earlier.

Signed-off-by: Julien Pivotto <roidelapluie@o11y.eu>
---
 storage/remote/read_handler.go | 53 ++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/storage/remote/read_handler.go b/storage/remote/read_handler.go
index ffc64c9c3..2a00ce897 100644
--- a/storage/remote/read_handler.go
+++ b/storage/remote/read_handler.go
@@ -202,34 +202,16 @@ func (h *readHandler) remoteReadStreamedXORChunks(ctx context.Context, w http.Re
 				return err
 			}
 
-			querier, err := h.queryable.ChunkQuerier(query.StartTimestampMs, query.EndTimestampMs)
-			if err != nil {
+			chunks := h.getChunkSeriesSet(ctx, query, filteredMatchers)
+			if err := chunks.Err(); err != nil {
 				return err
 			}
-			defer func() {
-				if err := querier.Close(); err != nil {
-					level.Warn(h.logger).Log("msg", "Error on chunk querier close", "err", err.Error())
-				}
-			}()
-
-			var hints *storage.SelectHints
-			if query.Hints != nil {
-				hints = &storage.SelectHints{
-					Start:    query.Hints.StartMs,
-					End:      query.Hints.EndMs,
-					Step:     query.Hints.StepMs,
-					Func:     query.Hints.Func,
-					Grouping: query.Hints.Grouping,
-					Range:    query.Hints.RangeMs,
-					By:       query.Hints.By,
-				}
-			}
 
 			ws, err := StreamChunkedReadResponses(
 				NewChunkedWriter(w, f),
 				int64(i),
 				// The streaming API has to provide the series sorted.
-				querier.Select(ctx, true, hints, filteredMatchers...),
+				chunks,
 				sortedExternalLabels,
 				h.remoteReadMaxBytesInFrame,
 				h.marshalPool,
@@ -254,6 +236,35 @@ func (h *readHandler) remoteReadStreamedXORChunks(ctx context.Context, w http.Re
 	}
 }
 
+// getChunkSeriesSet executes a query to retrieve a ChunkSeriesSet,
+// encapsulating the operation in its own function to ensure timely release of
+// the querier resources.
+func (h *readHandler) getChunkSeriesSet(ctx context.Context, query *prompb.Query, filteredMatchers []*labels.Matcher) storage.ChunkSeriesSet {
+	querier, err := h.queryable.ChunkQuerier(query.StartTimestampMs, query.EndTimestampMs)
+	if err != nil {
+		return storage.ErrChunkSeriesSet(err)
+	}
+	defer func() {
+		if err := querier.Close(); err != nil {
+			level.Warn(h.logger).Log("msg", "Error on chunk querier close", "err", err.Error())
+		}
+	}()
+
+	var hints *storage.SelectHints
+	if query.Hints != nil {
+		hints = &storage.SelectHints{
+			Start:    query.Hints.StartMs,
+			End:      query.Hints.EndMs,
+			Step:     query.Hints.StepMs,
+			Func:     query.Hints.Func,
+			Grouping: query.Hints.Grouping,
+			Range:    query.Hints.RangeMs,
+			By:       query.Hints.By,
+		}
+	}
+	return querier.Select(ctx, true, hints, filteredMatchers...)
+}
+
 // filterExtLabelsFromMatchers change equality matchers which match external labels
 // to a matcher that looks for an empty label,
 // as that label should not be present in the storage.