diff --git a/tsdb/index/index.go b/tsdb/index/index.go index 09fb737bf..362105459 100644 --- a/tsdb/index/index.go +++ b/tsdb/index/index.go @@ -1559,8 +1559,10 @@ func (r *Reader) LabelNamesFor(ctx context.Context, postings Postings) ([]string 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 diff --git a/tsdb/index/index_test.go b/tsdb/index/index_test.go index 038caacf8..d81dd8696 100644 --- a/tsdb/index/index_test.go +++ b/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{}) {