diff --git a/tsdb/tsm1/reader_index_iterator.go b/tsdb/tsm1/reader_index_iterator.go index 136e96961f..47864b4a13 100644 --- a/tsdb/tsm1/reader_index_iterator.go +++ b/tsdb/tsm1/reader_index_iterator.go @@ -59,6 +59,39 @@ func (t *TSMIndexIterator) Next() bool { return true } +// Seek points the iterator at the smallest key greater than or equal to the +// given key, returning true if it was an exact match. It returns false for +// ok if the key does not exist. +func (t *TSMIndexIterator) Seek(key []byte) (exact, ok bool) { + t.d.mu.RLock() + if n := len(t.d.ro.offsets); t.n != n { + t.err, t.ok = fmt.Errorf("Key count changed during iteration"), false + } + if t.err != nil { + t.d.mu.RUnlock() + return false, false + } + + t.peeked = false + t.first = false + + exact, t.ok = t.iter.Seek(key, t.b) + if !t.ok { + t.d.mu.RUnlock() + return false, false + } + + t.offset = t.iter.Offset() + t.eoffset = t.iter.EntryOffset(t.b) + t.d.mu.RUnlock() + + // reset lazy loaded state + t.key = nil + t.typ = 0 + t.entries = t.entries[:0] + return exact, true +} + // Peek reports the next key or nil if there is not one or an error happened. func (t *TSMIndexIterator) Peek() []byte { if !t.ok || t.err != nil { diff --git a/tsdb/tsm1/reader_index_iterator_test.go b/tsdb/tsm1/reader_index_iterator_test.go index 9d92aaede1..993a7c5636 100644 --- a/tsdb/tsm1/reader_index_iterator_test.go +++ b/tsdb/tsm1/reader_index_iterator_test.go @@ -51,6 +51,38 @@ func TestIndirectIndexIterator(t *testing.T) { checkEqual(t, iter.Next(), false) checkEqual(t, iter.Err(), error(nil)) + // check can seek and iterate index + iter = ind.Iterator(nil) + exact, ok := iter.Seek([]byte("cpu2")) + checkEqual(t, exact, true) + checkEqual(t, ok, true) + checkEqual(t, iter.Key(), []byte("cpu2")) + checkEqual(t, iter.Type(), BlockInteger) + checkEqual(t, iter.Entries(), []IndexEntry{ + {0, 10, 10, 20}, + {10, 20, 10, 20}, + }) + checkEqual(t, iter.Next(), true) + checkEqual(t, iter.Key(), []byte("mem")) + checkEqual(t, iter.Next(), false) + exact, ok = iter.Seek([]byte("cpu1")) + checkEqual(t, exact, true) + checkEqual(t, ok, true) + checkEqual(t, iter.Key(), []byte("cpu1")) + exact, ok = iter.Seek([]byte("cpu3")) + checkEqual(t, exact, false) + checkEqual(t, ok, true) + checkEqual(t, iter.Key(), []byte("mem")) + exact, ok = iter.Seek([]byte("cpu0")) + checkEqual(t, exact, false) + checkEqual(t, ok, true) + checkEqual(t, iter.Key(), []byte("cpu1")) + exact, ok = iter.Seek([]byte("zzz")) + checkEqual(t, exact, false) + checkEqual(t, ok, false) + checkEqual(t, iter.Next(), false) + checkEqual(t, iter.Err(), error(nil)) + // delete the cpu2 key and make sure it's skipped ind.Delete([][]byte{[]byte("cpu2")}) iter = ind.Iterator(nil)