Fixed implementation of herocache transactions

Author: durkmurder

module/mempool/herocache/transactions.go

@@ -1,8 +1,6 @@
 package herocache
 
 import (
-	"fmt"
-
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/model/flow"
@@ -14,7 +12,7 @@ import (
 )
 
 type Transactions struct {
-	c       *stdmap.Backend
+	*stdmap.Backend[flow.Identifier, *flow.TransactionBody]
 	byPayer map[flow.Address]map[flow.Identifier]struct{}
 }
 
@@ -26,31 +24,26 @@ func NewTransactions(limit uint32, logger zerolog.Logger, collector module.HeroC
 	}
 
 	tracer := &ejectionTracer{transactions: t}
-	t.c = stdmap.NewBackend(stdmap.WithBackData(
-		herocache.NewCache(limit,
-			herocache.DefaultOversizeFactor,
-			heropool.LRUEjection,
-			logger.With().Str("mempool", "transactions").Logger(),
-			collector,
-			herocache.WithTracer(tracer))))
+	t.Backend = stdmap.NewBackend(
+		stdmap.WithMutableBackData[flow.Identifier, *flow.TransactionBody](
+			herocache.NewCache[*flow.TransactionBody](limit,
+				herocache.DefaultOversizeFactor,
+				heropool.LRUEjection,
+				logger.With().Str("mempool", "transactions").Logger(),
+				collector,
+				herocache.WithTracer[*flow.TransactionBody](tracer))))
 
 	return t
 }
 
-// Has checks whether the transaction with the given hash is currently in
-// the memory pool.
-func (t *Transactions) Has(id flow.Identifier) bool {
-	return t.c.Has(id)
-}
-
 // Add adds a transaction to the mempool.
 func (t *Transactions) Add(tx *flow.TransactionBody) bool {
 	added := false
-	err := t.c.Run(func(backdata mempool.BackData) error {
+	err := t.Run(func(backdata mempool.BackData[flow.Identifier, *flow.TransactionBody]) error {
 		// Warning! reference pointer must be dereferenced before adding to HeroCache.
 		// This is crucial for its heap object optimizations.
 		txID := tx.ID()
-		added = backdata.Add(txID, *tx)
+		added = backdata.Add(txID, tx)
 		if !added {
 			return nil
 		}
@@ -68,37 +61,9 @@ func (t *Transactions) Add(tx *flow.TransactionBody) bool {
 	return added
 }
 
-// ByID returns the transaction with the given ID from the mempool.
-func (t *Transactions) ByID(txID flow.Identifier) (*flow.TransactionBody, bool) {
-	entity, exists := t.c.ByID(txID)
-	if !exists {
-		return nil, false
-	}
-	tx, ok := entity.(flow.TransactionBody)
-	if !ok {
-		panic(fmt.Sprintf("invalid entity in transaction pool (%T)", entity))
-	}
-	return &tx, true
-}
-
-// All returns all transactions from the mempool. Since it is using the HeroCache, All guarantees returning
-// all transactions in the same order as they are added.
-func (t *Transactions) All() []*flow.TransactionBody {
-	entities := t.c.All()
-	txs := make([]*flow.TransactionBody, 0, len(entities))
-	for _, entity := range entities {
-		tx, ok := entity.(flow.TransactionBody)
-		if !ok {
-			panic(fmt.Sprintf("invalid entity in transaction pool (%T)", entity))
-		}
-		txs = append(txs, &tx)
-	}
-	return txs
-}
-
 // Clear removes all transactions stored in this mempool.
 func (t *Transactions) Clear() {
-	err := t.c.Run(func(backdata mempool.BackData) error {
+	err := t.Run(func(backdata mempool.BackData[flow.Identifier, *flow.TransactionBody]) error {
 		backdata.Clear()
 		t.byPayer = make(map[flow.Address]map[flow.Identifier]struct{})
 		return nil
@@ -108,15 +73,10 @@ func (t *Transactions) Clear() {
 	}
 }
 
-// Size returns total number of stored transactions.
-func (t *Transactions) Size() uint {
-	return t.c.Size()
-}
-
 // Remove removes transaction from mempool.
 func (t *Transactions) Remove(id flow.Identifier) bool {
 	removed := false
-	err := t.c.Run(func(backdata mempool.BackData) error {
+	err := t.Run(func(backdata mempool.BackData[flow.Identifier, *flow.TransactionBody]) error {
 		var entity flow.Entity
 		entity, removed = backdata.Remove(id)
 		if !removed {
@@ -136,15 +96,14 @@ func (t *Transactions) Remove(id flow.Identifier) bool {
 // by the given payer.
 func (t *Transactions) ByPayer(payer flow.Address) []*flow.TransactionBody {
 	var result []*flow.TransactionBody
-	err := t.c.Run(func(backdata mempool.BackData) error {
+	err := t.Run(func(backdata mempool.BackData[flow.Identifier, *flow.TransactionBody]) error {
 		ids := t.byPayer[payer]
 		for id := range ids {
-			entity, exists := backdata.ByID(id)
+			txBody, exists := backdata.Get(id)
 			if !exists {
 				continue
 			}
-			txBody := entity.(flow.TransactionBody)
-			result = append(result, &txBody)
+			result = append(result, txBody)
 		}
 		return nil
 	})
@@ -170,19 +129,16 @@ type ejectionTracer struct {
 	transactions *Transactions
 }
 
-var _ herocache.Tracer = (*ejectionTracer)(nil)
+var _ herocache.Tracer[*flow.TransactionBody] = (*ejectionTracer)(nil)
 
-func (t *ejectionTracer) EntityEjectionDueToEmergency(ejectedEntity flow.Entity) {
-	t.cleanupIndex(ejectedEntity)
-}
-
-func (t *ejectionTracer) EntityEjectionDueToFullCapacity(ejectedEntity flow.Entity) {
-	t.cleanupIndex(ejectedEntity)
+// EntityEjectionDueToEmergency calls removeFromIndex on the transactions to clean up the index. This is safe since
+// backdata is locked by the caller when this function is called.
+func (t *ejectionTracer) EntityEjectionDueToEmergency(txBody *flow.TransactionBody) {
+	t.transactions.removeFromIndex(txBody.ID(), txBody.Payer)
 }
 
-// cleanupIndex calls removeFromIndex on the transactions to clean up the index. This is safe since
+// EntityEjectionDueToFullCapacity calls removeFromIndex on the transactions to clean up the index. This is safe since
 // backdata is locked by the caller when this function is called.
-func (t *ejectionTracer) cleanupIndex(ejectedEntity flow.Entity) {
-	txBody := ejectedEntity.(flow.TransactionBody)
+func (t *ejectionTracer) EntityEjectionDueToFullCapacity(txBody *flow.TransactionBody) {
 	t.transactions.removeFromIndex(txBody.ID(), txBody.Payer)
 }