// db helper functions
func testIterator(t *testing.T, itr *gorocksdb.Iterator, expectedValues map[string][]byte) {
itrResults := make(map[string][]byte)
itr.SeekToFirst()
for ; itr.Valid(); itr.Next() {
key := itr.Key()
value := itr.Value()
k := makeCopy(key.Data())
v := makeCopy(value.Data())
itrResults[string(k)] = v
}
if len(itrResults) != len(expectedValues) {
t.Fatalf("Expected [%d] results from iterator, found [%d]", len(expectedValues), len(itrResults))
}
for k, v := range expectedValues {
if !bytes.Equal(itrResults[k], v) {
t.Fatalf("Wrong value for key [%s]. Expected [%s], found [%s]", k, itrResults[k], v)
}
}
}
// 范围扫描
func (d *DB) Enumerate(iter *gorocksdb.Iterator, min, max []byte, direction IterDirection, fn func(i int, key, value []byte, quit *bool)) {
found := false
if direction == IterBackward {
if len(max) == 0 {
iter.SeekToLast()
} else {
iter.Seek(max)
}
} else {
if len(min) == 0 {
iter.SeekToFirst()
} else {
iter.Seek(min)
}
}
found = iter.Valid()
if !found {
return
}
i := -1
// 范围判断
if found && between(iter.Key().Data(), min, max) {
i++
quit := false
fn(i, iter.Key().Data(), iter.Value().Data(), &quit)
if quit {
return
}
}
for {
found = false
if direction == IterBackward {
iter.Prev()
} else {
iter.Next()
}
found = iter.Valid()
if found && between(iter.Key().Data(), min, max) {
i++
quit := false
fn(i, iter.Key().Data(), iter.Value().Data(), &quit)
if quit {
return
}
} else {
break
}
}
return
}
func (q *Queue) Dequeue(startId ...uint64) (uint64, []byte, error) {
store := q.store
var seekId uint64 = 1
if len(startId) > 0 {
seekId = startId[0]
}
if seekId < 1 {
seekId = 1
}
var it *rocks.Iterator
if q.useTailing {
it = q.tailIterator
if !it.Valid() {
// FIXME?(mijia): When Dequeue happens faster than enqueue, the tail iterator would be exhausted
// so we have seek it again.
it.Seek(q.key(seekId))
}
} else {
it = store.NewIteratorCF(store.ro, q.cfHandle)
defer it.Close()
it.Seek(q.key(seekId))
}
if !it.Valid() {
return 0, nil, EmptyQueue
}
wb := rocks.NewWriteBatch()
defer wb.Destroy()
key := makeSlice(it.Key())
value := makeSlice(it.Value())
wb.DeleteCF(q.cfHandle, key)
wb.MergeCF(q.cfHandle, q.metaKey("head"), oneBinary)
err := store.Write(store.wo, wb)
if err == nil {
atomic.AddUint64(&q.head, 1)
if q.useTailing {
it.Next()
}
}
id := q.id(key)
log.Debugf("[Queue] Dequeued data id=%d, err=%v", id, err)
return id, value, err
}
