// 范围扫描 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 }
// 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) } } }