TxPool: btree with generics (#514)

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/flanglet/kanzi-go v1.9.1-0.20211212184056-72dda96261ee
 	github.com/go-stack/stack v1.8.1
 	github.com/gofrs/flock v0.8.1
-	github.com/google/btree v1.0.1
+	github.com/google/btree v1.1.2
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/hashicorp/golang-lru v0.5.5-0.20210104140557-80c98217689d
 	github.com/holiman/uint256 v1.2.0

go.sum

@@ -55,8 +55,8 @@ github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
-github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
+github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU=
-github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
+github.com/google/btree v1.1.2/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=

kv/kvcache/cache.go

@@ -99,8 +99,8 @@ type Coherent struct {
 }
 
 type CoherentRoot struct {
-	cache           *btree.BTree
+	cache           *btree.BTreeG[*Element]
-	codeCache       *btree.BTree
+	codeCache       *btree.BTreeG[*Element]
 	ready           chan struct{} // close when ready
 	readyChanClosed atomic.Bool   // protecting `ready` field from double-close (on unwind). Consumers don't need check this field.
 
@@ -178,8 +178,8 @@ func (c *Coherent) selectOrCreateRoot(viewID ViewID) *CoherentRoot {
 
 	r = &CoherentRoot{
 		ready:     make(chan struct{}),
-		cache:     btree.New(DEGREE),
+		cache:     btree.NewG[*Element](DEGREE, Less),
-		codeCache: btree.New(DEGREE),
+		codeCache: btree.NewG[*Element](DEGREE, Less),
 	}
 	c.roots[viewID] = r
 	return r
@@ -202,15 +202,15 @@ func (c *Coherent) advanceRoot(viewID ViewID) (r *CoherentRoot) {
 		c.codeEvict.Init()
 		if r.cache == nil {
 			//log.Info("advance: new", "to", viewID)
-			r.cache = btree.New(DEGREE)
+			r.cache = btree.NewG[*Element](DEGREE, Less)
-			r.codeCache = btree.New(DEGREE)
+			r.codeCache = btree.NewG[*Element](DEGREE, Less)
 		} else {
-			r.cache.Ascend(func(i btree.Item) bool {
+			r.cache.Ascend(func(i *Element) bool {
-				c.stateEvict.PushFront(i.(*Element))
+				c.stateEvict.PushFront(i)
 				return true
 			})
-			r.codeCache.Ascend(func(i btree.Item) bool {
+			r.codeCache.Ascend(func(i *Element) bool {
-				c.codeEvict.PushFront(i.(*Element))
+				c.codeEvict.PushFront(i)
 				return true
 			})
 		}
@@ -309,7 +309,7 @@ func (c *Coherent) View(ctx context.Context, tx kv.Tx) (CacheView, error) {
 	return &CoherentView{viewID: ViewID(tx.ViewID()), tx: tx, cache: c}, nil
 }
 
-func (c *Coherent) getFromCache(k []byte, id ViewID, code bool) (btree.Item, *CoherentRoot, error) {
+func (c *Coherent) getFromCache(k []byte, id ViewID, code bool) (*Element, *CoherentRoot, error) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
 	r, ok := c.roots[id]
@@ -318,14 +318,14 @@ func (c *Coherent) getFromCache(k []byte, id ViewID, code bool) (btree.Item, *Co
 	}
 	isLatest := c.latestViewID == id
 
-	var it btree.Item
+	var it *Element
 	if code {
-		it = r.codeCache.Get(&Element{K: k})
+		it, _ = r.codeCache.Get(&Element{K: k})
 	} else {
-		it = r.cache.Get(&Element{K: k})
+		it, _ = r.cache.Get(&Element{K: k})
 	}
 	if it != nil && isLatest {
-		c.stateEvict.MoveToFront(it.(*Element))
+		c.stateEvict.MoveToFront(it)
 	}
 
 	return it, r, nil
@@ -339,7 +339,7 @@ func (c *Coherent) Get(k []byte, tx kv.Tx, id ViewID) ([]byte, error) {
 	if it != nil {
 		//fmt.Printf("from cache:  %#x,%x\n", k, it.(*Element).V)
 		c.hits.Inc()
-		return it.(*Element).V, nil
+		return it.V, nil
 	}
 	c.miss.Inc()
 
@@ -364,7 +364,7 @@ func (c *Coherent) GetCode(k []byte, tx kv.Tx, id ViewID) ([]byte, error) {
 	if it != nil {
 		//fmt.Printf("from cache:  %#x,%x\n", k, it.(*Element).V)
 		c.codeHits.Inc()
-		return it.(*Element).V, nil
+		return it.V, nil
 	}
 	c.codeMiss.Inc()
 
@@ -395,13 +395,13 @@ func (c *Coherent) removeOldestCode(r *CoherentRoot) {
 }
 func (c *Coherent) add(k, v []byte, r *CoherentRoot, id ViewID) *Element {
 	it := &Element{K: k, V: v}
-	replaced := r.cache.ReplaceOrInsert(it)
+	replaced, _ := r.cache.ReplaceOrInsert(it)
 	if c.latestViewID != id {
 		//fmt.Printf("add to non-last viewID: %d<%d\n", c.latestViewID, id)
 		return it
 	}
 	if replaced != nil {
-		c.stateEvict.Remove(replaced.(*Element))
+		c.stateEvict.Remove(replaced)
 	}
 	c.stateEvict.PushFront(it)
 	evict := c.stateEvict.Len() > c.cfg.KeysLimit
@@ -413,13 +413,13 @@ func (c *Coherent) add(k, v []byte, r *CoherentRoot, id ViewID) *Element {
 }
 func (c *Coherent) addCode(k, v []byte, r *CoherentRoot, id ViewID) *Element {
 	it := &Element{K: k, V: v}
-	replaced := r.codeCache.ReplaceOrInsert(it)
+	replaced, _ := r.codeCache.ReplaceOrInsert(it)
 	if c.latestViewID != id {
 		//fmt.Printf("add to non-last viewID: %d<%d\n", c.latestViewID, id)
 		return it
 	}
 	if replaced != nil {
-		c.codeEvict.Remove(replaced.(*Element))
+		c.codeEvict.Remove(replaced)
 	}
 	c.codeEvict.PushFront(it)
 	evict := c.codeEvict.Len() > c.cfg.CodeKeysLimit
@@ -475,8 +475,8 @@ func AssertCheckValues(ctx context.Context, tx kv.Tx, cache Cache) (int, error)
 	if !ok {
 		return 0, nil
 	}
-	root.cache.Ascend(func(i btree.Item) bool {
+	root.cache.Ascend(func(i *Element) bool {
-		k, v := i.(*Element).K, i.(*Element).V
+		k, v := i.K, i.V
 		var dbV []byte
 		dbV, err = tx.GetOne(kv.PlainState, k)
 		if err != nil {
@@ -536,9 +536,7 @@ type Element struct {
 	K, V []byte
 }
 
-func (e *Element) Less(than btree.Item) bool {
+func Less(a, b *Element) bool { return bytes.Compare(a.K, b.K) < 0 }
-	return bytes.Compare(e.K, than.(*Element).K) < 0
-}
 
 type ThreadSafeEvictionList struct {
 	l    *List

txpool/pool.go

@@ -266,13 +266,11 @@ func newSender(nonce uint64, balance uint256.Int) *sender {
 
 var emptySender = newSender(0, *uint256.NewInt(0))
 
-type sortByNonce struct{ *metaTx }
+func SortByNonceLess(a, b *metaTx) bool {
-
+	if a.Tx.SenderID != b.Tx.SenderID {
-func (i sortByNonce) Less(than btree.Item) bool {
+		return a.Tx.SenderID < b.Tx.SenderID
-	if i.metaTx.Tx.SenderID != than.(sortByNonce).metaTx.Tx.SenderID {
-		return i.metaTx.Tx.SenderID < than.(sortByNonce).metaTx.Tx.SenderID
 	}
-	return i.metaTx.Tx.Nonce < than.(sortByNonce).metaTx.Tx.Nonce
+	return a.Tx.Nonce < b.Tx.Nonce
 }
 
 func calcProtocolBaseFee(baseFee uint64) uint64 {
@@ -333,8 +331,8 @@ func New(newTxs chan types.Hashes, coreDB kv.RoDB, cfg Config, cache kvcache.Cac
 	}
 
 	byNonce := &BySenderAndNonce{
-		tree:             btree.New(32),
+		tree:             btree.NewG[*metaTx](32, SortByNonceLess),
-		search:           sortByNonce{&metaTx{Tx: &types.TxSlot{}}},
+		search:           &metaTx{Tx: &types.TxSlot{}},
 		senderIDTxnCount: map[uint64]int{},
 	}
 	tracedSenders := make(map[string]struct{})
@@ -1898,18 +1896,17 @@ func (sc *sendersBatch) onNewBlock(stateChanges *remote.StateChangeBatch, unwind
 //  - All senders stored inside 1 large BTree - because iterate over 1 BTree is faster than over map[senderId]BTree
 //  - sortByNonce used as non-pointer wrapper - because iterate over BTree of pointers is 2x slower
 type BySenderAndNonce struct {
-	tree             *btree.BTree
+	tree             *btree.BTreeG[*metaTx]
-	search           sortByNonce
+	search           *metaTx
 	senderIDTxnCount map[uint64]int // count of sender's txns in the pool - may differ from nonce
 }
 
 func (b *BySenderAndNonce) nonce(senderID uint64) (nonce uint64, ok bool) {
 	s := b.search
-	s.metaTx.Tx.SenderID = senderID
+	s.Tx.SenderID = senderID
-	s.metaTx.Tx.Nonce = math.MaxUint64
+	s.Tx.Nonce = math.MaxUint64
 
-	b.tree.DescendLessOrEqual(s, func(i btree.Item) bool {
+	b.tree.DescendLessOrEqual(s, func(mt *metaTx) bool {
-		mt := i.(sortByNonce).metaTx
 		if mt.Tx.SenderID == senderID {
 			nonce = mt.Tx.Nonce
 			ok = true
@@ -1919,17 +1916,15 @@ func (b *BySenderAndNonce) nonce(senderID uint64) (nonce uint64, ok bool) {
 	return nonce, ok
 }
 func (b *BySenderAndNonce) ascendAll(f func(*metaTx) bool) {
-	b.tree.Ascend(func(i btree.Item) bool {
+	b.tree.Ascend(func(mt *metaTx) bool {
-		mt := i.(sortByNonce).metaTx
 		return f(mt)
 	})
 }
 func (b *BySenderAndNonce) ascend(senderID uint64, f func(*metaTx) bool) {
 	s := b.search
-	s.metaTx.Tx.SenderID = senderID
+	s.Tx.SenderID = senderID
-	s.metaTx.Tx.Nonce = 0
+	s.Tx.Nonce = 0
-	b.tree.AscendGreaterOrEqual(s, func(i btree.Item) bool {
+	b.tree.AscendGreaterOrEqual(s, func(mt *metaTx) bool {
-		mt := i.(sortByNonce).metaTx
 		if mt.Tx.SenderID != senderID {
 			return false
 		}
@@ -1938,10 +1933,9 @@ func (b *BySenderAndNonce) ascend(senderID uint64, f func(*metaTx) bool) {
 }
 func (b *BySenderAndNonce) descend(senderID uint64, f func(*metaTx) bool) {
 	s := b.search
-	s.metaTx.Tx.SenderID = senderID
+	s.Tx.SenderID = senderID
-	s.metaTx.Tx.Nonce = math.MaxUint64
+	s.Tx.Nonce = math.MaxUint64
-	b.tree.DescendLessOrEqual(s, func(i btree.Item) bool {
+	b.tree.DescendLessOrEqual(s, func(mt *metaTx) bool {
-		mt := i.(sortByNonce).metaTx
 		if mt.Tx.SenderID != senderID {
 			return false
 		}
@@ -1961,21 +1955,20 @@ func (b *BySenderAndNonce) hasTxs(senderID uint64) bool {
 }
 func (b *BySenderAndNonce) get(senderID, txNonce uint64) *metaTx {
 	s := b.search
-	s.metaTx.Tx.SenderID = senderID
+	s.Tx.SenderID = senderID
-	s.metaTx.Tx.Nonce = txNonce
+	s.Tx.Nonce = txNonce
-	if found := b.tree.Get(s); found != nil {
+	if found, ok := b.tree.Get(s); ok {
-		return found.(sortByNonce).metaTx
+		return found
 	}
 	return nil
 }
 
 //nolint
 func (b *BySenderAndNonce) has(mt *metaTx) bool {
-	found := b.tree.Get(sortByNonce{mt})
+	return b.tree.Has(mt)
-	return found != nil
 }
 func (b *BySenderAndNonce) delete(mt *metaTx) {
-	if b.tree.Delete(sortByNonce{mt}) != nil {
+	if _, ok := b.tree.Delete(mt); ok {
 		senderID := mt.Tx.SenderID
 		count := b.senderIDTxnCount[senderID]
 		if count > 1 {
@@ -1986,9 +1979,9 @@ func (b *BySenderAndNonce) delete(mt *metaTx) {
 	}
 }
 func (b *BySenderAndNonce) replaceOrInsert(mt *metaTx) *metaTx {
-	it := b.tree.ReplaceOrInsert(sortByNonce{mt})
+	it, ok := b.tree.ReplaceOrInsert(mt)
-	if it != nil {
+	if ok {
-		return it.(sortByNonce).metaTx
+		return it
 	}
 	b.senderIDTxnCount[mt.Tx.SenderID]++
 	return nil