Handle context cancellation in some of the querier downstream requests

pkg/storage/batch.go

@@ -504,7 +504,7 @@ func (it *sampleBatchIterator) Sample() logproto.Sample {
 
 func (it *sampleBatchIterator) Next() bool {
 	// for loop to avoid recursion
-	for {
+	for it.ctx.Err() == nil {
 		if it.curr != nil && it.curr.Next() {
 			return true
 		}
@@ -527,6 +527,7 @@ func (it *sampleBatchIterator) Next() bool {
 			return false
 		}
 	}
+	return false
 }
 
 // newChunksIterator creates an iterator over a set of lazychunks.

pkg/storage/chunk/chunk_store_utils.go

@@ -145,6 +145,9 @@ func (c *Fetcher) worker() {
 // FetchChunks fetches a set of chunks from cache and store. Note that the keys passed in must be
 // lexicographically sorted, while the returned chunks are not in the same order as the passed in chunks.
 func (c *Fetcher) FetchChunks(ctx context.Context, chunks []Chunk, keys []string) ([]Chunk, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
 	log, ctx := spanlogger.New(ctx, "ChunkStore.FetchChunks")
 	defer log.Span.Finish()
 

pkg/storage/chunk/objectclient/client.go

@@ -102,6 +102,11 @@ func (o *Client) GetChunks(ctx context.Context, chunks []chunk.Chunk) ([]chunk.C
 }
 
 func (o *Client) getChunk(ctx context.Context, decodeContext *chunk.DecodeContext, c chunk.Chunk) (chunk.Chunk, error) {
+
+	if ctx.Err() != nil {
+		return chunk.Chunk{}, ctx.Err()
+	}
+
 	key := o.schema.ExternalKey(c)
 	if o.keyEncoder != nil {
 		key = o.keyEncoder(key)

pkg/storage/chunk/util/parallel_chunk_fetch.go

@@ -23,6 +23,10 @@ func GetParallelChunks(ctx context.Context, maxParallel int, chunks []chunk.Chun
 	defer log.Finish()
 	log.LogFields(otlog.Int("requested", len(chunks)))
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	queuedChunks := make(chan chunk.Chunk)
 
 	go func() {