diff options
| -rw-r--r-- | cursor.go | 57 | ||||
| -rw-r--r-- | cursor_test.go | 52 |
2 files changed, 87 insertions, 22 deletions
@@ -55,26 +55,7 @@ func (c *Cursor) Last() (key []byte, value []byte) { // If the cursor is at the end of the bucket then a nil key and value are returned. func (c *Cursor) Next() (key []byte, value []byte) { _assert(c.bucket.tx.db != nil, "tx closed") - - // Attempt to move over one element until we're successful. - // Move up the stack as we hit the end of each page in our stack. - for i := len(c.stack) - 1; i >= 0; i-- { - elem := &c.stack[i] - if elem.index < elem.count()-1 { - elem.index++ - break - } - c.stack = c.stack[:i] - } - - // If we've hit the end then return nil. - if len(c.stack) == 0 { - return nil, nil - } - - // Move down the stack to find the first element of the first leaf under this branch. - c.first() - k, v, flags := c.keyValue() + k, v, flags := c.next() if (flags & uint32(bucketLeafFlag)) != 0 { return k, nil } @@ -116,6 +97,12 @@ func (c *Cursor) Prev() (key []byte, value []byte) { // follow, a nil key is returned. func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) { k, v, flags := c.seek(seek) + + // If we ended up after the last element of a page then move to the next one. + if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() { + k, v, flags = c.next() + } + if k == nil { return nil, nil } else if (flags & uint32(bucketLeafFlag)) != 0 { @@ -125,8 +112,7 @@ func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) { } // seek moves the cursor to a given key and returns it. -// If the key does not exist then the next key is used. If no keys -// follow, a nil value is returned. +// If the key does not exist then the next key is used. func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) { _assert(c.bucket.tx.db != nil, "tx closed") @@ -189,6 +175,33 @@ func (c *Cursor) last() { } } +// next moves to the next leaf element and returns the key and value. +// If the cursor is at the last leaf element then it stays there and returns nil. +func (c *Cursor) next() (key []byte, value []byte, flags uint32) { + // Attempt to move over one element until we're successful. + // Move up the stack as we hit the end of each page in our stack. + var i int + for i = len(c.stack) - 1; i >= 0; i-- { + elem := &c.stack[i] + if elem.index < elem.count()-1 { + elem.index++ + break + } + } + + // If we've hit the root page then stop and return. This will leave the + // cursor on the last element of the last page. + if i == -1 { + return nil, nil, 0 + } + + // Otherwise start from where we left off in the stack and find the + // first element of the first leaf page. + c.stack = c.stack[:i+1] + c.first() + return c.keyValue() +} + // search recursively performs a binary search against a given page/node until it finds a given key. func (c *Cursor) search(key []byte, pgid pgid) { p, n := c.bucket.pageNode(pgid) diff --git a/cursor_test.go b/cursor_test.go index 23b3a8e..b44bf53 100644 --- a/cursor_test.go +++ b/cursor_test.go @@ -1,6 +1,7 @@ package bolt import ( + "encoding/binary" "sort" "testing" "testing/quick" @@ -66,6 +67,57 @@ func TestCursor_Seek(t *testing.T) { }) } +// Ensure that a Tx cursor can seek to the appropriate keys when there are a +// large number of keys. This test also checks that seek will always move +// forward to the next key. +// +// Related: https://github.com/boltdb/bolt/pull/187 +func TestCursor_Seek_Large(t *testing.T) { + withOpenDB(func(db *DB, path string) { + var count = 10000 + + // Insert every other key between 0 and $count. + db.Update(func(tx *Tx) error { + b, _ := tx.CreateBucket([]byte("widgets")) + for i := 0; i < count; i += 100 { + for j := i; j < i+100; j += 2 { + k := make([]byte, 8) + binary.BigEndian.PutUint64(k, uint64(j)) + b.Put(k, make([]byte, 100)) + } + } + return nil + }) + + db.View(func(tx *Tx) error { + c := tx.Bucket([]byte("widgets")).Cursor() + for i := 0; i < count; i++ { + seek := make([]byte, 8) + binary.BigEndian.PutUint64(seek, uint64(i)) + + k, _ := c.Seek(seek) + + // The last seek is beyond the end of the the range so + // it should return nil. + if i == count-1 { + assert.Nil(t, k) + continue + } + + // Otherwise we should seek to the exact key or the next key. + num := binary.BigEndian.Uint64(k) + if i%2 == 0 { + assert.Equal(t, uint64(i), num) + } else { + assert.Equal(t, uint64(i+1), num) + } + } + + return nil + }) + }) +} + // Ensure that a cursor can iterate over an empty bucket without error. func TestCursor_EmptyBucket(t *testing.T) { withOpenDB(func(db *DB, path string) { |