From e7db2e30a41b71e21f2477c0a756884bfba4579f Mon Sep 17 00:00:00 2001
From: Ben Ye <benye@amazon.com>
Date: Sat, 15 Jun 2024 11:43:26 -0700
Subject: [PATCH 1/2] fix check context cancellation not incrementing count

Signed-off-by: Ben Ye <benye@amazon.com>
---
 tsdb/index/index.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tsdb/index/index.go b/tsdb/index/index.go
index 8172b81ce..09fb737bf 100644
--- a/tsdb/index/index.go
+++ b/tsdb/index/index.go
@@ -1557,6 +1557,7 @@ 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()

From 0e6fca8e76baf59e4a5f67f398dbf08f654b8013 Mon Sep 17 00:00:00 2001
From: Ben Ye <benye@amazon.com>
Date: Sun, 16 Jun 2024 12:09:42 -0700
Subject: [PATCH 2/2] add unit test

Signed-off-by: Ben Ye <benye@amazon.com>
---
 tsdb/index/index.go      |  6 ++++--
 tsdb/index/index_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

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{}) {