github.com/Tri-stone/burrow@v0.25.0/storage/multi_iterator.go (about)

     1  package storage
     2  
     3  import (
     4  	"bytes"
     5  	"container/heap"
     6  )
     7  
     8  type MultiIterator struct {
     9  	start []byte
    10  	end   []byte
    11  	// Acts as priority queue based on sort order of current key in each iterator
    12  	iterators     []KVIterator
    13  	iteratorOrder map[KVIterator]int
    14  	lessComp      int
    15  }
    16  
    17  // MultiIterator iterates in order over a series o
    18  func NewMultiIterator(reverse bool, iterators ...KVIterator) *MultiIterator {
    19  	// reuse backing array
    20  	lessComp := -1
    21  	if reverse {
    22  		lessComp = 1
    23  	}
    24  	mi := &MultiIterator{
    25  		iterators:     iterators,
    26  		iteratorOrder: make(map[KVIterator]int),
    27  		lessComp:      lessComp,
    28  	}
    29  	mi.init()
    30  	return mi
    31  }
    32  
    33  func (mi *MultiIterator) init() {
    34  	validIterators := mi.iterators[:0]
    35  	for i, it := range mi.iterators {
    36  		mi.iteratorOrder[it] = i
    37  		if it.Valid() {
    38  			validIterators = append(validIterators, it)
    39  			start, end := it.Domain()
    40  			if i == 0 || CompareKeys(start, mi.start) == mi.lessComp {
    41  				mi.start = start
    42  			}
    43  			if i == 0 || CompareKeys(mi.end, end) == mi.lessComp {
    44  				mi.end = end
    45  			}
    46  		} else {
    47  			// Not clear if this is necessary - fairly sure it is permitted so can't hurt
    48  			it.Close()
    49  		}
    50  	}
    51  	mi.iterators = validIterators
    52  	heap.Init(mi)
    53  }
    54  
    55  // sort.Interface implementation
    56  func (mi *MultiIterator) Len() int {
    57  	return len(mi.iterators)
    58  }
    59  
    60  func (mi *MultiIterator) Less(i, j int) bool {
    61  	comp := bytes.Compare(mi.iterators[i].Key(), mi.iterators[j].Key())
    62  	// Use order iterators passed to NewMultiIterator if keys are equal1
    63  	return comp == mi.lessComp || (comp == 0 && mi.iteratorOrder[mi.iterators[i]] < mi.iteratorOrder[mi.iterators[j]])
    64  }
    65  
    66  func (mi *MultiIterator) Swap(i, j int) {
    67  	mi.iterators[i], mi.iterators[j] = mi.iterators[j], mi.iterators[i]
    68  }
    69  
    70  func (mi *MultiIterator) Push(x interface{}) {
    71  	mi.iterators = append(mi.iterators, x.(KVIterator))
    72  }
    73  
    74  func (mi *MultiIterator) Pop() interface{} {
    75  	n := len(mi.iterators) - 1
    76  	it := mi.iterators[n]
    77  	mi.iterators = mi.iterators[:n]
    78  	return it
    79  }
    80  
    81  func (mi *MultiIterator) Domain() ([]byte, []byte) {
    82  	return mi.start, mi.end
    83  }
    84  
    85  func (mi *MultiIterator) Valid() bool {
    86  	return len(mi.iterators) > 0
    87  }
    88  
    89  func (mi *MultiIterator) Next() {
    90  	// Always advance the lowest iterator - the same one we serve the KV pair from
    91  	it := heap.Pop(mi).(KVIterator)
    92  	it.Next()
    93  	if it.Valid() {
    94  		heap.Push(mi, it)
    95  	}
    96  }
    97  
    98  func (mi *MultiIterator) Key() []byte {
    99  	return mi.Peek().Key()
   100  }
   101  
   102  func (mi *MultiIterator) Value() []byte {
   103  	return mi.Peek().Value()
   104  }
   105  
   106  func (mi *MultiIterator) Peek() KVIterator {
   107  	return mi.iterators[0]
   108  }
   109  
   110  func (mi *MultiIterator) Close() {
   111  	// Close any remaining valid iterators
   112  	for _, it := range mi.iterators {
   113  		it.Close()
   114  	}
   115  }