Skip to content

Commit

Permalink
e3: simplify findShortenedKey (#11915)
Browse files Browse the repository at this point in the history
  • Loading branch information
AskAlexSharov authored Sep 9, 2024
1 parent ebfa41c commit 5f42f0d
Show file tree
Hide file tree
Showing 5 changed files with 44 additions and 56 deletions.
10 changes: 5 additions & 5 deletions erigon-lib/state/bps_tree.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ type indexSeekerIterator interface {
KVFromGetter(g *seg.Reader) ([]byte, []byte, error)
}

type dataLookupFunc func(di uint64, g *seg.Reader) ([]byte, []byte, error)
type dataLookupFunc func(di uint64, g *seg.Reader) ([]byte, []byte, uint64, error)
type keyCmpFunc func(k []byte, di uint64, g *seg.Reader, copyBuf []byte) (int, []byte, error)

// M limits amount of child for tree node.
Expand Down Expand Up @@ -111,7 +111,7 @@ func (it *BpsTreeIterator) KVFromGetter(g *seg.Reader) ([]byte, []byte, error) {
return nil, nil, errors.New("iterator is nil")
}
//fmt.Printf("kv from %p getter %p tree %p offt %d\n", it, g, it.t, it.i)
k, v, err := it.t.dataLookupFunc(it.i, g)
k, v, _, err := it.t.dataLookupFunc(it.i, g)
if err != nil {
if errors.Is(err, ErrBtIndexLookupBounds) {
return nil, nil, nil
Expand Down Expand Up @@ -287,7 +287,7 @@ func (b *BpsTree) Seek(g *seg.Reader, seekKey []byte) (key, value []byte, di uin
fmt.Printf("seek %x\n", seekKey)
}
if len(seekKey) == 0 && b.offt.Count() > 0 {
key, value, err = b.dataLookupFunc(0, g)
key, value, _, err = b.dataLookupFunc(0, g)
if err != nil {
return nil, nil, 0, false, err
}
Expand Down Expand Up @@ -348,7 +348,7 @@ func (b *BpsTree) Seek(g *seg.Reader, seekKey []byte) (key, value []byte, di uin
if l == r {
m = l
}
key, value, err = b.dataLookupFunc(m, g)
key, value, _, err = b.dataLookupFunc(m, g)
if err != nil {
return nil, nil, 0, false, err
}
Expand All @@ -363,7 +363,7 @@ func (b *BpsTree) Get(g *seg.Reader, key []byte) (k []byte, ok bool, i uint64, e
fmt.Printf("get %x\n", key)
}
if len(key) == 0 && b.offt.Count() > 0 {
k0, v0, err := b.dataLookupFunc(0, g)
k0, v0, _, err := b.dataLookupFunc(0, g)
if err != nil || k0 != nil {
return nil, false, 0, err
}
Expand Down
46 changes: 21 additions & 25 deletions erigon-lib/state/btree_index.go
Original file line number Diff line number Diff line change
Expand Up @@ -105,10 +105,6 @@ func (c *Cursor) Di() uint64 {
return c.d
}

func (c *Cursor) offsetInFile() uint64 {
return c.btt.ef.Get(c.d)
}

func (c *Cursor) Value() []byte {
return c.value
}
Expand All @@ -118,7 +114,7 @@ func (c *Cursor) Next() bool {
return false
}

key, value, err := c.btt.dataLookup(c.d, c.getter)
key, value, _, err := c.btt.dataLookup(c.d, c.getter)
if err != nil {
return false
}
Expand Down Expand Up @@ -543,7 +539,7 @@ func (a *btAlloc) WarmUp(gr *seg.Reader) error {
break
}

kb, v, err := a.dataLookup(s.d, gr)
kb, v, _, err := a.dataLookup(s.d, gr)
if err != nil {
fmt.Printf("d %d not found %v\n", s.d, err)
}
Expand Down Expand Up @@ -898,23 +894,23 @@ func OpenBtreeIndexWithDecompressor(indexPath string, M uint64, kv *seg.Decompre

// dataLookup fetches key and value from data file by di (data index)
// di starts from 0 so di is never >= keyCount
func (b *BtIndex) dataLookup(di uint64, g *seg.Reader) ([]byte, []byte, error) {
func (b *BtIndex) dataLookup(di uint64, g *seg.Reader) (k, v []byte, offset uint64, err error) {
if di >= b.ef.Count() {
return nil, nil, fmt.Errorf("%w: keyCount=%d, but key %d requested. file: %s", ErrBtIndexLookupBounds, b.ef.Count(), di, b.FileName())
return nil, nil, 0, fmt.Errorf("%w: keyCount=%d, but key %d requested. file: %s", ErrBtIndexLookupBounds, b.ef.Count(), di, b.FileName())
}

offset := b.ef.Get(di)
offset = b.ef.Get(di)
g.Reset(offset)
if !g.HasNext() {
return nil, nil, fmt.Errorf("pair %d/%d key not found, file: %s/%s", di, b.ef.Count(), b.FileName(), g.FileName())
return nil, nil, 0, fmt.Errorf("pair %d/%d key not found, file: %s/%s", di, b.ef.Count(), b.FileName(), g.FileName())
}

k, _ := g.Next(nil)
k, _ = g.Next(nil)
if !g.HasNext() {
return nil, nil, fmt.Errorf("pair %d/%d value not found, file: %s/%s", di, b.ef.Count(), b.FileName(), g.FileName())
return nil, nil, 0, fmt.Errorf("pair %d/%d value not found, file: %s/%s", di, b.ef.Count(), b.FileName(), g.FileName())
}
v, _ := g.Next(nil)
return k, v, nil
v, _ = g.Next(nil)
return k, v, offset, nil
}

// comparing `k` with item of index `di`. using buffer `kBuf` to avoid allocations
Expand Down Expand Up @@ -988,14 +984,14 @@ func (b *BtIndex) Close() {
}

// Get - exact match of key. `k == nil` - means not found
func (b *BtIndex) Get(lookup []byte, gr *seg.Reader) (k, v []byte, found bool, err error) {
func (b *BtIndex) Get(lookup []byte, gr *seg.Reader) (k, v []byte, offsetInFile uint64, found bool, err error) {
// TODO: optimize by "push-down" - instead of using seek+compare, alloc can have method Get which will return nil if key doesn't exists
// alternativaly: can allocate cursor on-stack
// it := Iter{} // allocation on stack
// it.Initialize(file)

if b.Empty() {
return k, v, false, nil
return k, v, 0, false, nil
}

var index uint64
Expand All @@ -1015,29 +1011,29 @@ func (b *BtIndex) Get(lookup []byte, gr *seg.Reader) (k, v []byte, found bool, e
k, found, index, err = b.bplus.Get(gr, lookup)
} else {
if b.alloc == nil {
return k, v, false, err
return k, v, 0, false, err
}
k, found, index, err = b.alloc.Get(gr, lookup)
}
if err != nil || !found {
if errors.Is(err, ErrBtIndexLookupBounds) {
return k, v, false, nil
return k, v, offsetInFile, false, nil
}
return nil, nil, false, err
return nil, nil, 0, false, err
}

// this comparation should be done by index get method, and in case of mismatch, key is not found
//if !bytes.Equal(k, lookup) {
// return k, v, false, nil
//}
k, v, err = b.dataLookup(index, gr)
k, v, offsetInFile, err = b.dataLookup(index, gr)
if err != nil {
if errors.Is(err, ErrBtIndexLookupBounds) {
return k, v, false, nil
return k, v, offsetInFile, false, nil
}
return k, v, false, err
return k, v, offsetInFile, false, err
}
return k, v, true, nil
return k, v, offsetInFile, true, nil
}

// Seek moves cursor to position where key >= x.
Expand Down Expand Up @@ -1070,7 +1066,7 @@ func (b *BtIndex) Seek(g *seg.Reader, x []byte) (*Cursor, error) {
return nil, err
}

k, v, err := b.dataLookup(dt, g)
k, v, _, err := b.dataLookup(dt, g)
if err != nil {
if errors.Is(err, ErrBtIndexLookupBounds) {
return nil, nil
Expand All @@ -1081,7 +1077,7 @@ func (b *BtIndex) Seek(g *seg.Reader, x []byte) (*Cursor, error) {
}

func (b *BtIndex) OrdinalLookup(getter *seg.Reader, i uint64) *Cursor {
k, v, err := b.dataLookup(i, getter)
k, v, _, err := b.dataLookup(i, getter)
if err != nil {
return nil
}
Expand Down
28 changes: 14 additions & 14 deletions erigon-lib/state/btree_index_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,14 +104,14 @@ func Test_BtreeIndex_Seek(t *testing.T) {
getter := seg.NewReader(kv.MakeGetter(), compressFlags)

t.Run("seek beyond the last key", func(t *testing.T) {
_, _, err := bt.dataLookup(bt.ef.Count()+1, getter)
_, _, _, err := bt.dataLookup(bt.ef.Count()+1, getter)
require.ErrorIs(t, err, ErrBtIndexLookupBounds)

_, _, err = bt.dataLookup(bt.ef.Count(), getter)
_, _, _, err = bt.dataLookup(bt.ef.Count(), getter)
require.ErrorIs(t, err, ErrBtIndexLookupBounds)
require.Error(t, err)

_, _, err = bt.dataLookup(bt.ef.Count()-1, getter)
_, _, _, err = bt.dataLookup(bt.ef.Count()-1, getter)
require.NoError(t, err)

cur, err := bt.Seek(getter, common.FromHex("0xffffffffffffff")) //seek beyeon the last key
Expand Down Expand Up @@ -224,14 +224,14 @@ func Test_BtreeIndex_Seek2(t *testing.T) {
getter := seg.NewReader(kv.MakeGetter(), compressFlags)

t.Run("seek beyond the last key", func(t *testing.T) {
_, _, err := bt.dataLookup(bt.ef.Count()+1, getter)
_, _, _, err := bt.dataLookup(bt.ef.Count()+1, getter)
require.ErrorIs(t, err, ErrBtIndexLookupBounds)

_, _, err = bt.dataLookup(bt.ef.Count(), getter)
_, _, _, err = bt.dataLookup(bt.ef.Count(), getter)
require.ErrorIs(t, err, ErrBtIndexLookupBounds)
require.Error(t, err)

_, _, err = bt.dataLookup(bt.ef.Count()-1, getter)
_, _, _, err = bt.dataLookup(bt.ef.Count()-1, getter)
require.NoError(t, err)

cur, err := bt.Seek(getter, common.FromHex("0xffffffffffffff")) //seek beyeon the last key
Expand Down Expand Up @@ -337,23 +337,23 @@ type mockIndexReader struct {
ef *eliasfano32.EliasFano
}

func (b *mockIndexReader) dataLookup(di uint64, g *seg.Reader) ([]byte, []byte, error) {
func (b *mockIndexReader) dataLookup(di uint64, g *seg.Reader) (k, v []byte, offset uint64, err error) {
if di >= b.ef.Count() {
return nil, nil, fmt.Errorf("%w: keyCount=%d, but key %d requested. file: %s", ErrBtIndexLookupBounds, b.ef.Count(), di, g.FileName())
return nil, nil, 0, fmt.Errorf("%w: keyCount=%d, but key %d requested. file: %s", ErrBtIndexLookupBounds, b.ef.Count(), di, g.FileName())
}

offset := b.ef.Get(di)
offset = b.ef.Get(di)
g.Reset(offset)
if !g.HasNext() {
return nil, nil, fmt.Errorf("pair %d/%d key not found, file: %s", di, b.ef.Count(), g.FileName())
return nil, nil, 0, fmt.Errorf("pair %d/%d key not found, file: %s", di, b.ef.Count(), g.FileName())
}

k, _ := g.Next(nil)
k, _ = g.Next(nil)
if !g.HasNext() {
return nil, nil, fmt.Errorf("pair %d/%d value not found, file: %s", di, b.ef.Count(), g.FileName())
return nil, nil, 0, fmt.Errorf("pair %d/%d value not found, file: %s", di, b.ef.Count(), g.FileName())
}
v, _ := g.Next(nil)
return k, v, nil
v, _ = g.Next(nil)
return k, v, offset, nil
}

// comparing `k` with item of index `di`. using buffer `kBuf` to avoid allocations
Expand Down
2 changes: 1 addition & 1 deletion erigon-lib/state/domain.go
Original file line number Diff line number Diff line change
Expand Up @@ -770,7 +770,7 @@ func (dt *DomainRoTx) getFromFile(i int, filekey []byte) ([]byte, bool, error) {
return v, true, nil
}

_, v, ok, err := dt.statelessBtree(i).Get(filekey, g)
_, v, _, ok, err := dt.statelessBtree(i).Get(filekey, g)
if err != nil || !ok {
return nil, false, err
}
Expand Down
14 changes: 3 additions & 11 deletions erigon-lib/state/domain_committed.go
Original file line number Diff line number Diff line change
Expand Up @@ -134,23 +134,15 @@ func (dt *DomainRoTx) findShortenedKey(fullKey []byte, itemGetter *seg.Reader, i
if item.bindex == nil {
dt.d.logger.Warn("[agg] commitment branch key replacement: file doesn't have index", "name", item.decompressor.FileName())
}
cur, err := item.bindex.Seek(itemGetter, fullKey)
_, _, offsetInFile, ok, err := item.bindex.Get(fullKey, itemGetter)
if err != nil {
dt.d.logger.Warn("[agg] commitment branch key replacement seek failed",
"key", fmt.Sprintf("%x", fullKey), "idx", "bt", "err", err, "file", item.decompressor.FileName())
}

if cur == nil || !bytes.Equal(cur.Key(), fullKey) {
return nil, false
}

offset := cur.offsetInFile()
if uint64(itemGetter.Size()) <= offset {
dt.d.logger.Warn("commitment branch key replacement seek gone too far",
"key", fmt.Sprintf("%x", fullKey), "offset", offset, "size", itemGetter.Size(), "file", item.decompressor.FileName())
if !ok {
return nil, false
}
return encodeShorterKey(nil, offset), true
return encodeShorterKey(nil, offsetInFile), true
}
return nil, false
}
Expand Down

0 comments on commit 5f42f0d

Please sign in to comment.