Merge pull request #14302 from yeya24/fix-check-ctx-cancel-count

fix check context cancellation not incrementing count

tsdb/index/index.go

@@ -1557,9 +1557,12 @@ func (r *Reader) LabelNamesFor(ctx context.Context, postings Postings) ([]string
 	i := 0
 	for postings.Next() {
 		id := postings.At()
+		i++
 
-		if i%checkContextEveryNIterations == 0 && ctx.Err() != nil {
-			return nil, ctx.Err()
+		if i%checkContextEveryNIterations == 0 {
+			if ctxErr := ctx.Err(); ctxErr != nil {
+				return nil, ctxErr
+			}
 		}
 
 		offset := id

tsdb/index/index_test.go

@@ -634,6 +634,31 @@ func TestReader_PostingsForLabelMatchingHonorsContextCancel(t *testing.T) {
 	require.Equal(t, failAfter, ctx.Count())
 }
 
+func TestReader_LabelNamesForHonorsContextCancel(t *testing.T) {
+	const seriesCount = 1000
+	var input indexWriterSeriesSlice
+	for i := 1; i <= seriesCount; i++ {
+		input = append(input, &indexWriterSeries{
+			labels: labels.FromStrings(labels.MetricName, fmt.Sprintf("%4d", i)),
+			chunks: []chunks.Meta{
+				{Ref: 1, MinTime: 0, MaxTime: 10},
+			},
+		})
+	}
+	ir, _, _ := createFileReader(context.Background(), t, input)
+
+	name, value := AllPostingsKey()
+	p, err := ir.Postings(context.Background(), name, value)
+	require.NoError(t, err)
+	// We check context cancellation every 128 iterations so 3 will fail after
+	// iterating 3 * 128 series.
+	failAfter := uint64(3)
+	ctx := &testutil.MockContextErrAfter{FailAfter: failAfter}
+	_, err = ir.LabelNamesFor(ctx, p)
+	require.Error(t, err)
+	require.Equal(t, failAfter, ctx.Count())
+}
+
 // createFileReader creates a temporary index file. It writes the provided input to this file.
 // It returns a Reader for this file, the file's name, and the symbol map.
 func createFileReader(ctx context.Context, tb testing.TB, input indexWriterSeriesSlice) (*Reader, string, map[string]struct{}) {