Speed up EntrySortIterator by 20%.

pkg/iter/entry_iterator.go

@@ -4,6 +4,7 @@ import (
 	"container/heap"
 	"context"
 	"io"
+	"sort"
 	"sync"
 	"time"
 
@@ -323,12 +324,9 @@ func (i *mergeEntryIterator) Len() int {
 }
 
 type entrySortIterator struct {
-	heap interface {
-		heap.Interface
-		Peek() EntryIterator
-	}
 	is         []EntryIterator
 	prefetched bool
+	dir        logproto.Direction
 
 	currEntry entryWithLabels
 	errs      []error
@@ -344,28 +342,61 @@ func NewSortEntryIterator(is []EntryIterator, direction logproto.Direction) Entr
 	if len(is) == 1 {
 		return is[0]
 	}
-	result := &entrySortIterator{is: is}
-	switch direction {
-	case logproto.BACKWARD:
-		result.heap = &iteratorMaxHeap{iteratorHeap: make([]EntryIterator, 0, len(is))}
-	case logproto.FORWARD:
-		result.heap = &iteratorMinHeap{iteratorHeap: make([]EntryIterator, 0, len(is))}
-	default:
-		panic("bad direction")
-	}
+
+	result := &entrySortIterator{is: is, dir: direction}
 	return result
 }
 
-// init initialize the underlaying heap
+func (it *entrySortIterator) less(i, j int) bool {
+	t1, t2 := it.is[i].Entry().Timestamp, it.is[j].Entry().Timestamp
+	if it.dir == logproto.BACKWARD {
+		t2, t1 = t1, t2
+
+	}
+
+	un1 := t1.UnixNano()
+	un2 := t2.UnixNano()
+
+	switch {
+	case un1 < un2:
+		return true
+	case un1 > un2:
+		return false
+	default: // un1 == un2:
+		return it.is[i].Labels() < it.is[j].Labels()
+	}
+}
+
+func (it *entrySortIterator) lessThan(t1 time.Time, l1 string, j int) bool {
+	t2 := it.is[j].Entry().Timestamp
+	if it.dir == logproto.BACKWARD {
+		t2, t1 = t1, t2
+	}
+
+	un1 := t1.UnixNano()
+	un2 := t2.UnixNano()
+
+	switch {
+	case un1 < un2:
+		return true
+	case un1 > un2:
+		return false
+	default: // un1 == un2:
+		return l1 < it.is[j].Labels()
+	}
+}
+
+// init throws out empty iterators and sorts them.
 func (i *entrySortIterator) init() {
 	if i.prefetched {
 		return
 	}
 
 	i.prefetched = true
+	tmp := make([]EntryIterator, 0, len(i.is))
 	for _, it := range i.is {
 		if it.Next() {
-			i.heap.Push(it)
+			tmp = append(tmp, it)
 			continue
 		}
 
@@ -374,35 +405,55 @@ func (i *entrySortIterator) init() {
 		}
 		util.LogError("closing iterator", it.Close)
 	}
-	heap.Init(i.heap)
+	i.is = tmp
+	sort.Slice(i.is, i.less)
+}
 
-	// We can now clear the list of input iterators to merge, given they have all
-	// been processed and the non empty ones have been pushed to the heap
-	i.is = nil
+func (i *entrySortIterator) fix() {
+	t1 := i.is[0].Entry().Timestamp
+	l1 := i.is[0].Labels()
+
+	// shortcut
+	if len(i.is) <= 1 || i.lessThan(t1, l1, 1) {
+		return
+	}
+
+	// First element is out of place. So we reposition it.
+	index := sort.Search(len(i.is), func(in int) bool { return i.lessThan(t1, l1, in) })
+
+	head := i.is[0]
+	if index == len(i.is) {
+		i.is = append(i.is[1:], head)
+	} else {
+		i.is = append(i.is[1:index+1], i.is[index:]...)
+		i.is[index] = head
+	}
 }
 
 func (i *entrySortIterator) Next() bool {
 	i.init()
 
-	if i.heap.Len() == 0 {
+	if len(i.is) == 0 {
 		return false
 	}
 
-	next := i.heap.Peek()
+	next := i.is[0]
 	i.currEntry.entry = next.Entry()
 	i.currEntry.labels = next.Labels()
 	// if the top iterator is empty, we remove it.
 	if !next.Next() {
-		heap.Pop(i.heap)
+		i.is = i.is[1:]
 		if err := next.Error(); err != nil {
 			i.errs = append(i.errs, err)
 		}
 		util.LogError("closing iterator", next.Close)
 		return true
 	}
-	if i.heap.Len() > 1 {
-		heap.Fix(i.heap, 0)
+
+	if len(i.is) > 1 {
+		i.fix()
 	}
+
 	return true
 }
 
@@ -426,11 +477,6 @@ func (i *entrySortIterator) Error() error {
 }
 
 func (i *entrySortIterator) Close() error {
-	for i.heap.Len() > 0 {
-		if err := i.heap.Pop().(EntryIterator).Close(); err != nil {
-			return err
-		}
-	}
 	return nil
 }
 

pkg/iter/entry_iterator_test.go

@@ -676,6 +676,7 @@ func BenchmarkSortIterator(b *testing.B) {
 	})
 
 	b.Run("sort", func(b *testing.B) {
+		b.ReportAllocs()
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			b.StopTimer()

vendor/modules.txt

@@ -1028,6 +1028,9 @@ github.com/prometheus/prometheus/util/teststorage
 github.com/prometheus/prometheus/util/testutil
 github.com/prometheus/prometheus/util/treecache
 github.com/prometheus/prometheus/web/api/v1
+# github.com/psilva261/timsort/v2 v2.0.0
+## explicit; go 1.13
+github.com/psilva261/timsort/v2
 # github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 ## explicit
 github.com/rcrowley/go-metrics