store/bucket: wait until chunk loading ends in Close()

Chunk reader needs to wait until the chunk loading ends in Close()
because otherwise there will be a race between appending to r.chunkBytes
and reading from it.

Signed-off-by: Giedrius Statkevičius <[email protected]>

pkg/store/bucket.go

@@ -3153,6 +3153,8 @@ type bucketChunkReader struct {
 	mtx        sync.Mutex
 	stats      *queryStats
 	chunkBytes []*[]byte // Byte slice to return to the chunk pool on close.
+
+	chunkLoading sync.WaitGroup
 }
 
 func newBucketChunkReader(block *bucketBlock) *bucketChunkReader {
@@ -3170,6 +3172,9 @@ func (r *bucketChunkReader) reset() {
 }
 
 func (r *bucketChunkReader) Close() error {
+	// Wait until any chunk loading ends.
+	r.chunkLoading.Wait()
+
 	r.block.pendingReaders.Done()
 
 	for _, b := range r.chunkBytes {
@@ -3211,6 +3216,7 @@ func (r *bucketChunkReader) load(ctx context.Context, res []seriesEntry, aggrs [
 			r.stats.DataDownloadedSizeSum += units.Base2Bytes(p.End - p.Start)
 		}
 
+		r.chunkLoading.Add(len(parts))
 		for _, p := range parts {
 			seq := seq
 			p := p
@@ -3234,6 +3240,7 @@ func (r *bucketChunkReader) loadChunks(ctx context.Context, res []seriesEntry, a
 		}
 		r.stats.ChunksFetchDurationSum += time.Since(fetchBegin)
 		r.mtx.Unlock()
+		r.chunkLoading.Done()
 	}()
 
 	// Get a reader for the required range.