120 lines
2.9 KiB
Go
120 lines
2.9 KiB
Go
package tsdb
|
|
|
|
import (
|
|
"bytes"
|
|
"container/heap"
|
|
)
|
|
|
|
// MultiCursor returns a single cursor that combines the results of all cursors in order.
|
|
//
|
|
// If the same key is returned from multiple cursors then the first cursor
|
|
// specified will take precendence. A key will only be returned once from the
|
|
// returned cursor.
|
|
func MultiCursor(cursors ...Cursor) Cursor {
|
|
return &multiCursor{cursors: cursors}
|
|
}
|
|
|
|
// multiCursor represents a cursor that combines multiple cursors into one.
|
|
type multiCursor struct {
|
|
cursors []Cursor
|
|
heap cursorHeap
|
|
prev []byte
|
|
}
|
|
|
|
// Seek moves the cursor to a given key.
|
|
func (mc *multiCursor) Seek(seek []byte) (key, value []byte) {
|
|
// Initialize heap.
|
|
h := make(cursorHeap, 0, len(mc.cursors))
|
|
for i, c := range mc.cursors {
|
|
// Move cursor to position. Skip if it's empty.
|
|
k, v := c.Seek(seek)
|
|
if k == nil {
|
|
continue
|
|
}
|
|
|
|
// Append cursor to heap.
|
|
h = append(h, &cursorHeapItem{
|
|
key: k,
|
|
value: v,
|
|
cursor: c,
|
|
priority: len(mc.cursors) - i,
|
|
})
|
|
}
|
|
|
|
heap.Init(&h)
|
|
mc.heap = h
|
|
mc.prev = nil
|
|
|
|
return mc.pop()
|
|
}
|
|
|
|
// Next returns the next key/value from the cursor.
|
|
func (mc *multiCursor) Next() (key, value []byte) { return mc.pop() }
|
|
|
|
// pop returns the next item from the heap.
|
|
// Reads the next key/value from item's cursor and puts it back on the heap.
|
|
func (mc *multiCursor) pop() (key, value []byte) {
|
|
// Read items until we have a key that doesn't match the previously read one.
|
|
// This is to perform deduplication when there's multiple items with the same key.
|
|
// The highest priority cursor will be read first and then remaining keys will be dropped.
|
|
for {
|
|
// Return nil if there are no more items left.
|
|
if len(mc.heap) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
// Read the next item from the heap.
|
|
item := heap.Pop(&mc.heap).(*cursorHeapItem)
|
|
|
|
// Save the key/value for return.
|
|
key, value = item.key, item.value
|
|
|
|
// Read the next item from the cursor. Push back to heap if one exists.
|
|
if item.key, item.value = item.cursor.Next(); item.key != nil {
|
|
heap.Push(&mc.heap, item)
|
|
}
|
|
|
|
// Skip if this key matches the previously returned one.
|
|
if bytes.Equal(mc.prev, key) {
|
|
continue
|
|
}
|
|
|
|
mc.prev = key
|
|
return
|
|
}
|
|
}
|
|
|
|
// cursorHeap represents a heap of cursorHeapItems.
|
|
type cursorHeap []*cursorHeapItem
|
|
|
|
func (h cursorHeap) Len() int { return len(h) }
|
|
func (h cursorHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
|
|
func (h cursorHeap) Less(i, j int) bool {
|
|
if cmp := bytes.Compare(h[i].key, h[j].key); cmp == -1 {
|
|
return true
|
|
} else if cmp == 0 {
|
|
return h[i].priority > h[j].priority
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (h *cursorHeap) Push(x interface{}) {
|
|
*h = append(*h, x.(*cursorHeapItem))
|
|
}
|
|
|
|
func (h *cursorHeap) Pop() interface{} {
|
|
old := *h
|
|
n := len(old)
|
|
item := old[n-1]
|
|
*h = old[0 : n-1]
|
|
return item
|
|
}
|
|
|
|
// cursorHeapItem is something we manage in a priority queue.
|
|
type cursorHeapItem struct {
|
|
key []byte
|
|
value []byte
|
|
cursor Cursor
|
|
priority int
|
|
}
|