-
Notifications
You must be signed in to change notification settings - Fork 3.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
added chunks batches iterator to download chunks in batches instead o…
…f downloading all of them at once Signed-off-by: Vladyslav Diachenko <vlad.diachenko@grafana.com>
- Loading branch information
1 parent
5559b26
commit 4dcf3c8
Showing
13 changed files
with
274 additions
and
99 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
package bloomcompactor | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
|
||
"github.com/grafana/loki/pkg/storage/chunk" | ||
) | ||
|
||
type chunksBatchesIterator struct { | ||
context context.Context | ||
client chunkClient | ||
chunksToDownload []chunk.Chunk | ||
batchSize uint | ||
|
||
currentBatch []chunk.Chunk | ||
err error | ||
} | ||
|
||
func newChunkBatchesIterator(context context.Context, client chunkClient, chunksToDownload []chunk.Chunk, batchSize uint) (*chunksBatchesIterator, error) { | ||
if batchSize == 0 { | ||
return nil, errors.New("batchSize must be greater than 0") | ||
} | ||
return &chunksBatchesIterator{context: context, client: client, chunksToDownload: chunksToDownload, batchSize: batchSize}, nil | ||
} | ||
|
||
func (c *chunksBatchesIterator) Next() bool { | ||
if len(c.chunksToDownload) == 0 { | ||
return false | ||
} | ||
batchSize := c.batchSize | ||
chunksToDownloadCount := uint(len(c.chunksToDownload)) | ||
if chunksToDownloadCount < batchSize { | ||
batchSize = chunksToDownloadCount | ||
} | ||
chunksToDownload := c.chunksToDownload[:batchSize] | ||
c.chunksToDownload = c.chunksToDownload[batchSize:] | ||
c.currentBatch, c.err = c.client.GetChunks(c.context, chunksToDownload) | ||
return true | ||
} | ||
|
||
func (c *chunksBatchesIterator) Err() error { | ||
return c.err | ||
} | ||
|
||
func (c *chunksBatchesIterator) At() []chunk.Chunk { | ||
return c.currentBatch | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,96 @@ | ||
package bloomcompactor | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"testing" | ||
|
||
"github.com/stretchr/testify/require" | ||
|
||
"github.com/grafana/loki/pkg/storage/chunk" | ||
tsdbindex "github.com/grafana/loki/pkg/storage/stores/shipper/indexshipper/tsdb/index" | ||
) | ||
|
||
func Test_chunksBatchesIterator(t *testing.T) { | ||
tests := map[string]struct { | ||
batchSize uint | ||
chunksToDownload []chunk.Chunk | ||
constructorError error | ||
|
||
hadNextCount int | ||
}{ | ||
"expected error if batch size is set to 0": { | ||
batchSize: 0, | ||
constructorError: errors.New("batchSize must be greater than 0"), | ||
}, | ||
"expected no error if there are no chunks": { | ||
hadNextCount: 0, | ||
batchSize: 10, | ||
}, | ||
"expected 1 call to the client": { | ||
chunksToDownload: createFakeChunks(10), | ||
hadNextCount: 1, | ||
batchSize: 20, | ||
}, | ||
"expected 1 call to the client(2)": { | ||
chunksToDownload: createFakeChunks(10), | ||
hadNextCount: 1, | ||
batchSize: 10, | ||
}, | ||
"expected 2 calls to the client": { | ||
chunksToDownload: createFakeChunks(10), | ||
hadNextCount: 2, | ||
batchSize: 6, | ||
}, | ||
"expected 10 calls to the client": { | ||
chunksToDownload: createFakeChunks(10), | ||
hadNextCount: 10, | ||
batchSize: 1, | ||
}, | ||
} | ||
for name, data := range tests { | ||
t.Run(name, func(t *testing.T) { | ||
client := &fakeClient{} | ||
iterator, err := newChunkBatchesIterator(context.Background(), client, data.chunksToDownload, data.batchSize) | ||
if data.constructorError != nil { | ||
require.Equal(t, err, data.constructorError) | ||
return | ||
} | ||
hadNextCount := 0 | ||
var downloadedChunks []chunk.Chunk | ||
for iterator.Next() { | ||
require.NoError(t, iterator.Err()) | ||
hadNextCount++ | ||
downloaded := iterator.At() | ||
downloadedChunks = append(downloadedChunks, downloaded...) | ||
require.LessOrEqual(t, uint(len(downloaded)), data.batchSize) | ||
} | ||
require.Equal(t, data.chunksToDownload, downloadedChunks) | ||
require.Equal(t, data.hadNextCount, client.callsCount) | ||
require.Equal(t, data.hadNextCount, hadNextCount) | ||
}) | ||
} | ||
} | ||
|
||
func createFakeChunks(count int) []chunk.Chunk { | ||
metas := make([]tsdbindex.ChunkMeta, 0, count) | ||
for i := 0; i < count; i++ { | ||
metas = append(metas, tsdbindex.ChunkMeta{ | ||
Checksum: uint32(i), | ||
MinTime: int64(i), | ||
MaxTime: int64(i + 100), | ||
KB: uint32(i * 100), | ||
Entries: uint32(i * 10), | ||
}) | ||
} | ||
return makeChunkRefs(metas, "fake", 0xFFFF) | ||
} | ||
|
||
type fakeClient struct { | ||
callsCount int | ||
} | ||
|
||
func (f *fakeClient) GetChunks(ctx context.Context, chunks []chunk.Chunk) ([]chunk.Chunk, error) { | ||
f.callsCount++ | ||
return chunks, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.