Fixed mempool implementaiton. Fixed tests for collection builder

Author: durkmurder

module/builder/collection/builder_test.go

@@ -1088,6 +1088,9 @@ func (suite *BuilderSuite) TestBuildOn_MaxCollectionSizeRateLimiting() {
 	// start with an empty mempool
 	suite.ClearPool()
 
+	cfg := updatable_configs.DefaultBySealingLagRateLimiterConfigs()
+	suite.Require().NoError(cfg.SetMinSealingLag(50)) // set min sealing lag to 50 blocks so we can hit rate limiting
+	suite.Require().NoError(cfg.SetMaxSealingLag(50)) // set max sealing lag to 50 blocks so we can hit rate limiting
 	suite.builder, _ = builder.NewBuilder(
 		suite.db,
 		trace.NewNoopTracer(),
@@ -1100,7 +1103,7 @@ func (suite *BuilderSuite) TestBuildOn_MaxCollectionSizeRateLimiting() {
 		suite.pool,
 		unittest.Logger(),
 		suite.epochCounter,
-		updatable_configs.DefaultBySealingLagRateLimiterConfigs(),
+		cfg,
 		builder.WithMaxCollectionSize(100),
 	)
 
@@ -1125,27 +1128,36 @@ func (suite *BuilderSuite) TestBuildOn_MaxCollectionSizeRateLimiting() {
 	// build a long chain of blocks that were finalized but not sealed
 	// this will lead to a big sealing lag.
 	for i := 0; i < 100; i++ {
-		block := unittest.BlockWithParentFixture(head)
-		block.SetPayload(unittest.PayloadFixture(unittest.WithProtocolStateID(protocolStateID)))
-		err = suite.protoState.ExtendCertified(context.Background(), block, unittest.CertifyBlock(block.Header))
+		block := unittest.BlockWithParentAndPayload(head, unittest.PayloadFixture(unittest.WithProtocolStateID(protocolStateID)))
+		err = suite.protoState.ExtendCertified(context.Background(), unittest.NewCertifiedBlock(block))
 		suite.Require().NoError(err)
 		err = suite.protoState.Finalize(context.Background(), block.ID())
 		suite.Require().NoError(err)
-		head = block.Header
+		head = block.ToHeader()
 	}
 
 	rateLimiterCfg := updatable_configs.DefaultBySealingLagRateLimiterConfigs()
 
 	// rate-limiting should be applied, resulting in minimum collection size.
 	parentID := suite.genesis.ID()
+	setter := func(h *flow.HeaderBodyBuilder) error {
+		h.WithChainID(flow.Emulator).
+			WithParentID(parentID).
+			WithView(1337).
+			WithParentView(1336).
+			WithParentVoterIndices(unittest.SignerIndicesFixture(4)).
+			WithParentVoterSigData(unittest.QCSigDataFixture()).
+			WithProposerID(unittest.IdentifierFixture())
+		return nil
+	}
 	for i := 0; i < 10; i++ {
-		header, err := suite.builder.BuildOn(parentID, noopSetter, noopSigner)
+		header, err := suite.builder.BuildOn(parentID, setter, signer)
 		suite.Require().NoError(err)
-		parentID = header.ID()
+		parentID = header.Header.ID()
 
-		// each collection should be half-full with 5 transactions
+		// each collection should be equal to the minimum collection size
 		var built model.Block
-		err = suite.db.View(procedure.RetrieveClusterBlock(header.ID(), &built))
+		err = suite.db.View(procedure.RetrieveClusterBlock(header.Header.ID(), &built))
 		suite.Assert().NoError(err)
 		suite.Assert().Len(built.Payload.Collection.Transactions, int(rateLimiterCfg.MinCollectionSize()))
 	}
@@ -1382,11 +1394,11 @@ func (suite *BuilderSuite) TestBuildOn_SystemTxAlwaysIncluded() {
 
 	// rate-limiting should not be applied, since the payer is marked as unlimited
 	parentID := suite.genesis.ID()
-	header, err := suite.builder.BuildOn(parentID, noopSetter, noopSigner)
+	header, err := suite.builder.BuildOn(parentID, setter, signer)
 	suite.Require().NoError(err)
 
 	var built model.Block
-	err = suite.db.View(procedure.RetrieveClusterBlock(header.ID(), &built))
+	err = suite.db.View(procedure.RetrieveClusterBlock(header.Header.ID(), &built))
 	suite.Assert().NoError(err)
 	suite.Assert().Len(built.Payload.Collection.Transactions, 2)
 	for _, tx := range built.Payload.Collection.Transactions {

module/mempool/herocache/transactions.go

@@ -16,6 +16,8 @@ type Transactions struct {
 	byPayer map[flow.Address]map[flow.Identifier]struct{}
 }
 
+var _ mempool.Transactions = (*Transactions)(nil)
+
 // NewTransactions implements a transactions mempool based on hero cache.
 func NewTransactions(limit uint32, logger zerolog.Logger, collector module.HeroCacheMetrics) *Transactions {
 	byPayer := make(map[flow.Address]map[flow.Identifier]struct{})
@@ -37,12 +39,11 @@ func NewTransactions(limit uint32, logger zerolog.Logger, collector module.HeroC
 }
 
 // Add adds a transaction to the mempool.
-func (t *Transactions) Add(tx *flow.TransactionBody) bool {
+func (t *Transactions) Add(txID flow.Identifier, tx *flow.TransactionBody) bool {
 	added := false
 	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)
 		if !added {
 			return nil
@@ -77,12 +78,11 @@ func (t *Transactions) Clear() {
 func (t *Transactions) Remove(id flow.Identifier) bool {
 	removed := false
 	err := t.Run(func(backdata mempool.BackData[flow.Identifier, *flow.TransactionBody]) error {
-		var entity flow.Entity
-		entity, removed = backdata.Remove(id)
+		var txBody *flow.TransactionBody
+		txBody, removed = backdata.Remove(id)
 		if !removed {
 			return nil
 		}
-		txBody := entity.(flow.TransactionBody)
 		t.removeFromIndex(id, txBody.Payer)
 		return nil
 	})

module/mempool/mempool.go

@@ -22,6 +22,8 @@ type Mempool[K comparable, V any] interface {
 	Adjust(key K, f func(V) V) (V, bool)
 	// Size will return the size of the mempool.
 	Size() uint
+	// Values returns all stored values from the mempool.
+	Values() []V
 	// All returns all stored key-value pairs as a map from the mempool.
 	All() map[K]V
 	// Clear removes all key-value pairs from the mempool.

module/mempool/transactions.go

@@ -6,35 +6,8 @@ import (
 
 // Transactions represents a concurrency-safe memory pool for transactions.
 type Transactions interface {
-
-	// Has checks whether the transaction with the given hash is currently in
-	// the memory pool.
-	Has(txID flow.Identifier) bool
-
-	// Add will add the given transaction body to the memory pool. It will
-	// return false if it was already in the mempool.
-	Add(txId flow.Identifier, tx *flow.TransactionBody) bool
-
-	// Remove will remove the given transaction from the memory pool; it will
-	// return true if the transaction was known and removed.
-	Remove(txID flow.Identifier) bool
-
-	// Get retrieve the transaction with the given ID from the memory
-	// pool. It will return false if it was not found in the mempool.
-	Get(txID flow.Identifier) (*flow.TransactionBody, bool)
-
+	Mempool[flow.Identifier, *flow.TransactionBody]
 	// ByPayer retrieves all transactions from the memory pool that are sent
 	// by the given payer.
 	ByPayer(payer flow.Address) []*flow.TransactionBody
-
-	// Size will return the current size of the memory pool.
-	Size() uint
-
-	// Values will retrieve all transactions that are currently in the memory pool
-	// as a slice.
-	// Values guarantees returning all transactions in the same order as they are added.
-	Values() []*flow.TransactionBody
-
-	// Clear removes all transactions from the mempool.
-	Clear()
 }