Skip to content

Commit

Permalink
Refactor HeapIterator into Merge and Sort Iterator. (#5281)
Browse files Browse the repository at this point in the history 
* Refactor HeapIterator into merge and sort Iterator.

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

* lint

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
  • Loading branch information
@cyriltovena
cyriltovena committed Feb 1, 2022
1 parent 44ba390 commit 3af7b79
Show file tree
Hide file tree
Showing 18 changed files with 573 additions and 187 deletions.
10 changes: 5 additions & 5 deletions pkg/chunkenc/memchunk.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -818,7 +818,7 @@ func (c *MemChunk) Iterator(ctx context.Context, mintT, maxtT time.Time, directi
if ordered {
it = iter.NewNonOverlappingIterator(blockItrs, "")
} else {
it = iter.NewHeapIterator(ctx, blockItrs, direction)
it = iter.NewSortEntryIterator(blockItrs, direction)
}

return iter.NewTimeRangedIterator(
Expand Down Expand Up @@ -851,7 +851,7 @@ func (c *MemChunk) Iterator(ctx context.Context, mintT, maxtT time.Time, directi
if ordered {
return iter.NewNonOverlappingIterator(blockItrs, ""), nil
}
return iter.NewHeapIterator(ctx, blockItrs, direction), nil
return iter.NewSortEntryIterator(blockItrs, direction), nil
}

// Iterator implements Chunk.
Expand Down Expand Up @@ -886,7 +886,7 @@ func (c *MemChunk) SampleIterator(ctx context.Context, from, through time.Time,
if ordered {
it = iter.NewNonOverlappingSampleIterator(its, "")
} else {
it = iter.NewHeapSampleIterator(ctx, its)
it = iter.NewSortSampleIterator(its)
}

return iter.NewTimeRangedSampleIterator(
Expand Down Expand Up @@ -1041,7 +1041,7 @@ func (hb *headBlock) Iterator(ctx context.Context, direction logproto.Direction,
for _, stream := range streams {
streamsResult = append(streamsResult, *stream)
}
return iter.NewStreamsIterator(ctx, streamsResult, direction)
return iter.NewStreamsIterator(streamsResult, direction)
}

func (hb *headBlock) SampleIterator(ctx context.Context, mint, maxt int64, extractor log.StreamSampleExtractor) iter.SampleIterator {
Expand Down Expand Up @@ -1082,7 +1082,7 @@ func (hb *headBlock) SampleIterator(ctx context.Context, mint, maxt int64, extra
for _, s := range series {
seriesRes = append(seriesRes, *s)
}
return iter.SampleIteratorWithClose(iter.NewMultiSeriesIterator(ctx, seriesRes), func() error {
return iter.SampleIteratorWithClose(iter.NewMultiSeriesIterator(seriesRes), func() error {
for _, s := range series {
SamplesPool.Put(s.Samples)
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/chunkenc/unordered.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,7 +257,7 @@ func (hb *unorderedHeadBlock) Iterator(
for _, stream := range streams {
streamsResult = append(streamsResult, *stream)
}
return iter.NewStreamsIterator(ctx, streamsResult, direction)
return iter.NewStreamsIterator(streamsResult, direction)
}

// nolint:unused
Expand Down Expand Up @@ -308,7 +308,7 @@ func (hb *unorderedHeadBlock) SampleIterator(
for _, s := range series {
seriesRes = append(seriesRes, *s)
}
return iter.SampleIteratorWithClose(iter.NewMultiSeriesIterator(ctx, seriesRes), func() error {
return iter.SampleIteratorWithClose(iter.NewMultiSeriesIterator(seriesRes), func() error {
for _, s := range series {
SamplesPool.Put(s.Samples)
}
Expand Down
23 changes: 10 additions & 13 deletions pkg/ingester/ingester.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -561,7 +561,7 @@ func (i *Ingester) Query(req *logproto.QueryRequest, queryServer logproto.Querie
}

instance := i.GetOrCreateInstance(instanceID)
itrs, err := instance.Query(ctx, logql.SelectLogParams{QueryRequest: req})
it, err := instance.Query(ctx, logql.SelectLogParams{QueryRequest: req})
if err != nil {
return err
}
Expand All @@ -577,17 +577,15 @@ func (i *Ingester) Query(req *logproto.QueryRequest, queryServer logproto.Querie
}}
storeItr, err := i.store.SelectLogs(ctx, storeReq)
if err != nil {
errUtil.LogErrorWithContext(ctx, "closing iterator", it.Close)
return err
}

itrs = append(itrs, storeItr)
it = iter.NewMergeEntryIterator(ctx, []iter.EntryIterator{it, storeItr}, req.Direction)
}

heapItr := iter.NewHeapIterator(ctx, itrs, req.Direction)

defer errUtil.LogErrorWithContext(ctx, "closing iterator", heapItr.Close)
defer errUtil.LogErrorWithContext(ctx, "closing iterator", it.Close)

return sendBatches(ctx, heapItr, queryServer, req.Limit)
return sendBatches(ctx, it, queryServer, req.Limit)
}

// QuerySample the ingesters for series from logs matching a set of matchers.
Expand All @@ -601,7 +599,7 @@ func (i *Ingester) QuerySample(req *logproto.SampleQueryRequest, queryServer log
}

instance := i.GetOrCreateInstance(instanceID)
itrs, err := instance.QuerySample(ctx, logql.SelectSampleParams{SampleQueryRequest: req})
it, err := instance.QuerySample(ctx, logql.SelectSampleParams{SampleQueryRequest: req})
if err != nil {
return err
}
Expand All @@ -615,17 +613,16 @@ func (i *Ingester) QuerySample(req *logproto.SampleQueryRequest, queryServer log
}}
storeItr, err := i.store.SelectSamples(ctx, storeReq)
if err != nil {
errUtil.LogErrorWithContext(ctx, "closing iterator", it.Close)
return err
}

itrs = append(itrs, storeItr)
it = iter.NewMergeSampleIterator(ctx, []iter.SampleIterator{it, storeItr})
}

heapItr := iter.NewHeapSampleIterator(ctx, itrs)

defer errUtil.LogErrorWithContext(ctx, "closing iterator", heapItr.Close)
defer errUtil.LogErrorWithContext(ctx, "closing iterator", it.Close)

return sendSampleBatches(ctx, heapItr, queryServer)
return sendSampleBatches(ctx, it, queryServer)
}

// boltdbShipperMaxLookBack returns a max look back period only if active index type is boltdb-shipper.
Expand Down
8 changes: 4 additions & 4 deletions pkg/ingester/instance.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,7 @@ func (i *instance) getLabelsFromFingerprint(fp model.Fingerprint) labels.Labels
return s.labels
}

func (i *instance) Query(ctx context.Context, req logql.SelectLogParams) ([]iter.EntryIterator, error) {
func (i *instance) Query(ctx context.Context, req logql.SelectLogParams) (iter.EntryIterator, error) {
expr, err := req.LogSelector()
if err != nil {
return nil, err
Expand Down Expand Up @@ -341,10 +341,10 @@ func (i *instance) Query(ctx context.Context, req logql.SelectLogParams) ([]iter
return nil, err
}

return iters, nil
return iter.NewSortEntryIterator(iters, req.Direction), nil
}

func (i *instance) QuerySample(ctx context.Context, req logql.SelectSampleParams) ([]iter.SampleIterator, error) {
func (i *instance) QuerySample(ctx context.Context, req logql.SelectSampleParams) (iter.SampleIterator, error) {
expr, err := req.Expr()
if err != nil {
return nil, err
Expand Down Expand Up @@ -386,7 +386,7 @@ func (i *instance) QuerySample(ctx context.Context, req logql.SelectSampleParams
return nil, err
}

return iters, nil
return iter.NewSortSampleIterator(iters), nil
}

// Label returns the label names or values depending on the given request
Expand Down
9 changes: 3 additions & 6 deletions pkg/ingester/instance_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ import (
"github.com/prometheus/prometheus/model/labels"
"github.com/stretchr/testify/require"

"github.com/grafana/loki/pkg/iter"
"github.com/grafana/loki/pkg/logproto"
"github.com/grafana/loki/pkg/logql"
loki_runtime "github.com/grafana/loki/pkg/runtime"
Expand Down Expand Up @@ -497,7 +496,7 @@ func Test_Iterator(t *testing.T) {
}

// prepare iterators.
itrs, err := instance.Query(ctx,
it, err := instance.Query(ctx,
logql.SelectLogParams{
QueryRequest: &logproto.QueryRequest{
Selector: `{job="3"} | logfmt`,
Expand All @@ -509,12 +508,11 @@ func Test_Iterator(t *testing.T) {
},
)
require.NoError(t, err)
heapItr := iter.NewHeapIterator(ctx, itrs, direction)

// assert the order is preserved.
var res *logproto.QueryResponse
require.NoError(t,
sendBatches(ctx, heapItr,
sendBatches(ctx, it,
fakeQueryServer(
func(qr *logproto.QueryResponse) error {
res = qr
Expand Down Expand Up @@ -578,7 +576,7 @@ func Test_ChunkFilter(t *testing.T) {
}

// prepare iterators.
itrs, err := instance.Query(ctx,
it, err := instance.Query(ctx,
logql.SelectLogParams{
QueryRequest: &logproto.QueryRequest{
Selector: `{job="3"}`,
Expand All @@ -590,7 +588,6 @@ func Test_ChunkFilter(t *testing.T) {
},
)
require.NoError(t, err)
it := iter.NewHeapIterator(ctx, itrs, direction)
defer it.Close()

for it.Next() {
Expand Down
4 changes: 2 additions & 2 deletions pkg/ingester/stream.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -470,7 +470,7 @@ func (s *stream) Iterator(ctx context.Context, statsCtx *stats.Context, from, th
if ordered {
return iter.NewNonOverlappingIterator(iterators, ""), nil
}
return iter.NewHeapIterator(ctx, iterators, direction), nil
return iter.NewSortEntryIterator(iterators, direction), nil
}

// Returns an SampleIterator.
Expand Down Expand Up @@ -507,7 +507,7 @@ func (s *stream) SampleIterator(ctx context.Context, statsCtx *stats.Context, fr
if ordered {
return iter.NewNonOverlappingSampleIterator(iterators, ""), nil
}
return iter.NewHeapSampleIterator(ctx, iterators), nil
return iter.NewSortSampleIterator(iterators), nil
}

func (s *stream) addTailer(t *tailer) {
Expand Down
Loading

0 comments on commit 3af7b79

Please sign in to comment.