[bloom-compactor] downloading chunks in batches

[vlad-diachenko]

What this PR does / why we need it:
Added chunks batches iterator to download chunks in batches instead of downloading all of them at once. Otherwise, when the stream contains a lot of chunks, it can lead to OOM.

Special notes for your reviewer:

Checklist

• Reviewed the CONTRIBUTING.md guide (required)
• Documentation added
• Tests updated
• CHANGELOG.md updated
  • If the change is worth mentioning in the release notes, add add-to-release-notes label
• Changes that require user attention or interaction to upgrade are documented in docs/sources/setup/upgrade/_index.md
• For Helm chart changes bump the Helm chart version in production/helm/loki/Chart.yaml and update production/helm/loki/CHANGELOG.md and production/helm/loki/README.md. Example PR
• If the change is deprecating or removing a configuration option, update the deprecated-config.yaml and deleted-config.yaml files respectively in the tools/deprecated-config-checker directory. Example PR

[chaudum]

Overall lgtm

[chaudum]

Same here

[vlad-diachenko]

I would prefer to pass the limits because it's extendable...

[chaudum]

Instead of passing down the limits, should we just pass the batchSize?
We also resolve the fpRate per tenant just before.

[vlad-diachenko]

I believe it's better to pass the limits because we will always add additional parameters and it will make us to add additional function's incoming parameters... I will remove fpRate and will resolve it inside the function.

[salvacorts]

nit: I think we can avoid modifying the argument of this method as well as its implementation at bloom_tokenizer.go.

We can call PopulateSeriesWithBloom for each call of the iterator Next(). E.g. on createPopulateFunc

batchesIterator, err := newChunkBatchesIterator(ctx, storeClient.chunk, chunkRefs, limits.BloomCompactorChunksBatchSize(job.tenantID))
if err != nil {
	return fmt.Errorf("error creating chunks batches iterator: %w", err)
}
for batchesIterator.Next() {
   if err := batchesIterator.Err() ...
   chunks := batchesIterator.At()
   err = bt.PopulateSeriesWithBloom(&bloomForChks, chunks)
}
err = bt.PopulateSeriesWithBloom(&bloomForChks, batchesIterator)

[salvacorts]

This would affect the metrics we expose inside PopulateSeriesWithBloom tho.

So what about PopulateSeriesWithBloom receiving an iterator of chunks (Iterator[chunk.Chunk]) and doing the buffering inside the iterator implementation. That way the batching logic is present only in the iterator.

func (c *chunksBatchesIterator) loadNextBatch() error {
	batchSize := c.batchSize
	chunksToDownloadCount := len(c.chunksToDownload)
	if chunksToDownloadCount < batchSize {
		batchSize = chunksToDownloadCount
	}
	chunksToDownload := c.chunksToDownload[:batchSize]
	c.chunksToDownload = c.chunksToDownload[batchSize:]

	newBatch, err := c.client.GetChunks(c.context, chunksToDownload)
	if err != nil {
		return err
	}

	c.currentBatch = newBatch
	return nil
}

func (c *chunksBatchesIterator) Next() bool {
	if len(c.currentBatch) == 0 {
		if len(c.chunksToDownload) == 0 {
			return false
		}
		if c.err = c.loadNextBatch(); c.err != nil {
			return false
		}
	}

	// Pop the first chunk from the current batch and set it as the current chunk.
	c.currentChunk = c.currentBatch[0]
	c.currentBatch = c.currentBatch[1:]

	return true
}

func (c *chunksBatchesIterator) At() chunk.Chunk {
	return c.currentChunk
}

[vlad-diachenko]

yes, it will do the same... but I do not believe that we win anything by changing it this way...
it just increases the complexity of the iterator.
let's say if we failed to download the batch, we will report an error, but this error is not connected to this particular chunk is returned by At() function. so, it might confuse somebody in the future...
not really against it, but it looks almost the same.

[vlad-diachenko]

wdyt?

[salvacorts]

Sure, not a strong opinion!