Fixes a swallowed context err in the batch storage. (#3025)

Fixes #2999

There's still a potential cases where, the context is done (and so the Err is not nil) and we're finished
with iterating over the list of batches. I didn't introduce the error here because I think I would prefer a success in this case.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

pkg/logql/engine.go

@@ -150,7 +150,10 @@ func (q *query) Exec(ctx context.Context) (Result, error) {
 	status := "200"
 	if err != nil {
 		status = "500"
-		if errors.Is(err, ErrParse) || errors.Is(err, ErrPipeline) || errors.Is(err, ErrLimit) {
+		if errors.Is(err, ErrParse) ||
+			errors.Is(err, ErrPipeline) ||
+			errors.Is(err, ErrLimit) ||
+			errors.Is(err, context.Canceled) {
 			status = "400"
 		}
 	}

pkg/storage/batch.go

@@ -143,6 +143,7 @@ func (it *batchChunkIterator) loop(ctx context.Context) {
 		}
 		select {
 		case <-ctx.Done():
+			it.next <- &chunkBatch{err: ctx.Err()}
 			close(it.next)
 			return
 		case it.next <- it.nextBatch():

pkg/storage/batch_test.go

@@ -1493,6 +1493,35 @@ func Test_IsInvalidChunkError(t *testing.T) {
 	}
 }
 
+func TestBatchCancel(t *testing.T) {
+	chunk := func(from time.Time) *LazyChunk {
+		return newLazyChunk(logproto.Stream{
+			Labels: fooLabelsWithName,
+			Entries: []logproto.Entry{
+				{
+					Timestamp: from,
+					Line:      "1",
+				},
+				{
+					Timestamp: from.Add(time.Millisecond),
+					Line:      "2",
+				},
+			},
+		})
+	}
+	chunks := []*LazyChunk{
+		chunk(from), chunk(from.Add(10 * time.Millisecond)), chunk(from.Add(30 * time.Millisecond)),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	it, err := newLogBatchIterator(ctx, NilMetrics, chunks, 1, newMatchers(fooLabels), log.NewNoopPipeline(), logproto.FORWARD, from, time.Now())
+	require.NoError(t, err)
+	defer require.NoError(t, it.Close())
+	for it.Next() {
+	}
+	require.Equal(t, context.Canceled, it.Error())
+}
+
 var entry logproto.Entry
 
 func Benchmark_store_OverlappingChunks(b *testing.B) {