Remote read instrumentation tweaks

Signed-off-by: Justin Lei <lei.justin@gmail.com>

storage/remote/client.go

@@ -77,7 +77,7 @@ var (
 			Namespace: namespace,
 			Subsystem: subsystem,
 			Name:      "read_request_duration_seconds",
-			Help:      "Histogram of the latency for remote read requests.",
+			Help:      "Histogram of the latency for remote read requests. Note that for streamed responses this is only the duration of the initial call and does not include the processing of the stream.",
 			Buckets:   append(prometheus.DefBuckets, 25, 60),
 		},
 		[]string{remoteName, endpoint, "response_type"},
@@ -333,9 +333,19 @@ func (c *Client) Read(ctx context.Context, query *prompb.Query, sortSeries bool)
 		return c.handleSampledResponse(req, httpResp, sortSeries)
 	case strings.HasPrefix(contentType, "application/x-streamed-protobuf; proto=prometheus.ChunkedReadResponse"):
 		c.readQueriesDuration.WithLabelValues("chunked").Observe(time.Since(start).Seconds())
-		c.readQueriesTotal.WithLabelValues("chunked", strconv.Itoa(httpResp.StatusCode)).Inc()
-		ss := c.handleChunkedResponse(httpResp, query.StartTimestampMs, query.EndTimestampMs, cancel)
+		// We copy cancel here so we can nil out the original and prevent the context from being cancelled as soon as
+		// Read returns.
+		cancelCopy := cancel
 		cancel = nil
+
+		ss := c.handleChunkedResponse(httpResp, query.StartTimestampMs, query.EndTimestampMs, func(err error) {
+			code := strconv.Itoa(httpResp.StatusCode)
+			if err != io.EOF {
+				code = "aborted_stream"
+			}
+			c.readQueriesTotal.WithLabelValues("chunked", code).Inc()
+			cancelCopy()
+		})
 		return ss, nil
 	default:
 		c.readQueriesDuration.WithLabelValues("unsupported").Observe(time.Since(start).Seconds())
@@ -375,7 +385,7 @@ func (c *Client) handleSampledResponse(req *prompb.ReadRequest, httpResp *http.R
 	return FromQueryResult(sortSeries, res), nil
 }
 
-func (c *Client) handleChunkedResponse(httpResp *http.Response, mint, maxt int64, cancel context.CancelFunc) storage.SeriesSet {
+func (c *Client) handleChunkedResponse(httpResp *http.Response, mint, maxt int64, cancel func(error)) storage.SeriesSet {
 	s := NewChunkedReader(httpResp.Body, c.chunkedReadLimit, nil)
 	return NewChunkedSeriesSet(s, httpResp.Body, mint, maxt, cancel)
 }

storage/remote/codec.go

@@ -14,7 +14,6 @@
 package remote
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -547,13 +546,13 @@ type chunkedSeriesSet struct {
 	chunkedReader *ChunkedReader
 	respBody      io.ReadCloser
 	mint, maxt    int64
-	cancel        context.CancelFunc
+	cancel        func(error)
 
 	current storage.Series
 	err     error
 }
 
-func NewChunkedSeriesSet(chunkedReader *ChunkedReader, respBody io.ReadCloser, mint, maxt int64, cancel context.CancelFunc) storage.SeriesSet {
+func NewChunkedSeriesSet(chunkedReader *ChunkedReader, respBody io.ReadCloser, mint, maxt int64, cancel func(error)) storage.SeriesSet {
 	return &chunkedSeriesSet{
 		chunkedReader: chunkedReader,
 		respBody:      respBody,
@@ -576,7 +575,7 @@ func (s *chunkedSeriesSet) Next() bool {
 		}
 
 		_ = s.respBody.Close()
-		s.cancel()
+		s.cancel(err)
 
 		return false
 	}

storage/remote/codec_test.go

@@ -1005,7 +1005,7 @@ func TestChunkedSeriesSet(t *testing.T) {
 			require.NoError(t, err)
 		}
 
-		ss := NewChunkedSeriesSet(r, io.NopCloser(buf), 0, 14000, func() {})
+		ss := NewChunkedSeriesSet(r, io.NopCloser(buf), 0, 14000, func(error) {})
 		require.Nil(t, ss.Err())
 		require.Nil(t, ss.Warnings())
 
@@ -1060,7 +1060,7 @@ func TestChunkedSeriesSet(t *testing.T) {
 			require.NoError(t, err)
 		}
 
-		ss := NewChunkedSeriesSet(r, io.NopCloser(buf), 0, 14000, func() {})
+		ss := NewChunkedSeriesSet(r, io.NopCloser(buf), 0, 14000, func(error) {})
 		require.Nil(t, ss.Err())
 		require.Nil(t, ss.Warnings())