Update Mempool to support multiple ops per Sender (#72)

Author: hazim-j

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/go-logr/logr v1.2.3
 	github.com/go-logr/zerologr v1.2.2
 	github.com/go-playground/validator/v10 v10.11.1
+	github.com/google/go-cmp v0.5.9
 	github.com/mitchellh/mapstructure v1.5.0
 	github.com/rs/zerolog v1.28.0
 	github.com/spf13/cobra v1.6.0

go.sum

@@ -186,7 +186,8 @@ github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXiIaQGbYVAs8BPL6v8lEs=
 github.com/google/martian/v3 v3.0.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=

internal/testutils/constants.go

@@ -4,12 +4,14 @@ import (
 	"math/big"
 	"time"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/stackup-wallet/stackup-bundler/pkg/entrypoint"
 )
 
 var (
 	OneETH                 = big.NewInt(1000000000000000000)
 	DefaultUnstakeDelaySec = uint32(86400)
+	ValidAddress           = common.HexToAddress("0x7357b8a705328FC283dF72D7Ac546895B596DC12")
 	StakedDepositInfo      = &entrypoint.IStakeManagerDepositInfo{
 		Deposit:         big.NewInt(OneETH.Int64()),
 		Staked:          true,

internal/testutils/dbmock.go

@@ -0,0 +1,16 @@
+package testutils
+
+import (
+	"log"
+
+	badger "github.com/dgraph-io/badger/v3"
+)
+
+func DBMock() *badger.DB {
+	db, err := badger.Open(badger.DefaultOptions("").WithInMemory(true).WithLoggingLevel(badger.ERROR))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return db
+}

internal/testutils/opmock.go

@@ -1,7 +1,10 @@
 package testutils
 
 import (
+	"math/big"
+
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/google/go-cmp/cmp"
 	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 )
 
@@ -27,3 +30,19 @@ func MockValidInitUserOp() *userop.UserOperation {
 	op, _ := userop.New(MockUserOpData)
 	return op
 }
+
+func IsOpsEqual(op1 *userop.UserOperation, op2 *userop.UserOperation) bool {
+	return cmp.Equal(
+		op1,
+		op2,
+		cmp.Comparer(func(a *big.Int, b *big.Int) bool { return a.Cmp(b) == 0 }),
+	)
+}
+
+func GetOpsDiff(op1 *userop.UserOperation, op2 *userop.UserOperation) string {
+	return cmp.Diff(
+		op1,
+		op2,
+		cmp.Comparer(func(a *big.Int, b *big.Int) bool { return a.Cmp(b) == 0 }),
+	)
+}

pkg/bundler/bundler.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stackup-wallet/stackup-bundler/pkg/mempool"
 	"github.com/stackup-wallet/stackup-bundler/pkg/modules"
 	"github.com/stackup-wallet/stackup-bundler/pkg/modules/noop"
+	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 )
 
 // Bundler controls the end to end process of creating a batch of UserOperations from the mempool and sending
@@ -72,8 +73,9 @@ func (i *Bundler) Run() error {
 					continue
 				}
 
-				senders := append(getSenders(ctx.Batch), getSenders(ctx.PendingRemoval)...)
-				if err := i.mempool.RemoveOps(ep, senders...); err != nil {
+				rmOps := append([]*userop.UserOperation{}, ctx.Batch...)
+				rmOps = append(rmOps, ctx.PendingRemoval...)
+				if err := i.mempool.RemoveOps(ep, rmOps...); err != nil {
 					l.Error(err, "bundler run error")
 					continue
 				}

pkg/bundler/senders.go

@@ -85,33 +85,33 @@ func (i *Client) SendUserOperation(op map[string]any, ep string) (string, error)
 	//	1. the nonce remains the same
 	//	2. the new maxPriorityFeePerGas is higher
 	//	3. the new maxFeePerGas is increased equally
-	memOp, err := i.mempool.GetOp(epAddr, userOp.Sender)
-	if err != nil {
-		l.Error(err, "eth_sendUserOperation error")
-		return "", err
-	}
-	if memOp != nil {
-		if memOp.Nonce.Cmp(memOp.Nonce) != 0 {
-			err := errors.New("sender: Has userOp in mempool with a different nonce")
-			l.Error(err, "eth_sendUserOperation error")
-			return "", err
-		}
-
-		if memOp.MaxPriorityFeePerGas.Cmp(memOp.MaxPriorityFeePerGas) <= 0 {
-			err := errors.New("sender: Has userOp in mempool with same or higher priority fee")
-			l.Error(err, "eth_sendUserOperation error")
-			return "", err
-		}
-
-		diff := big.NewInt(0)
-		mf := big.NewInt(0)
-		diff.Sub(memOp.MaxPriorityFeePerGas, memOp.MaxPriorityFeePerGas)
-		if memOp.MaxFeePerGas.Cmp(mf.Add(memOp.MaxFeePerGas, diff)) != 0 {
-			err := errors.New("sender: Replaced userOp must have an equally higher max fee")
-			l.Error(err, "eth_sendUserOperation error")
-			return "", err
-		}
-	}
+	// memOp, err := i.mempool.GetOp(epAddr, userOp.Sender)
+	// if err != nil {
+	// 	l.Error(err, "eth_sendUserOperation error")
+	// 	return "", err
+	// }
+	// if memOp != nil {
+	// 	if memOp.Nonce.Cmp(memOp.Nonce) != 0 {
+	// 		err := errors.New("sender: Has userOp in mempool with a different nonce")
+	// 		l.Error(err, "eth_sendUserOperation error")
+	// 		return "", err
+	// 	}
+
+	// 	if memOp.MaxPriorityFeePerGas.Cmp(memOp.MaxPriorityFeePerGas) <= 0 {
+	// 		err := errors.New("sender: Has userOp in mempool with same or higher priority fee")
+	// 		l.Error(err, "eth_sendUserOperation error")
+	// 		return "", err
+	// 	}
+
+	// 	diff := big.NewInt(0)
+	// 	mf := big.NewInt(0)
+	// 	diff.Sub(memOp.MaxPriorityFeePerGas, memOp.MaxPriorityFeePerGas)
+	// 	if memOp.MaxFeePerGas.Cmp(mf.Add(memOp.MaxFeePerGas, diff)) != 0 {
+	// 		err := errors.New("sender: Replaced userOp must have an equally higher max fee")
+	// 		l.Error(err, "eth_sendUserOperation error")
+	// 		return "", err
+	// 	}
+	// }
 
 	// Run through client module stack.
 	ctx := modules.NewUserOpHandlerContext(userOp, epAddr, i.chainID)

pkg/client/client.go

@@ -2,6 +2,7 @@ package mempool
 
 import (
 	"encoding/json"
+	"math/big"
 	"strings"
 
 	badger "github.com/dgraph-io/badger/v3"
@@ -12,16 +13,15 @@ import (
 const keySeparator = ":"
 const keyPrefix = "mempool"
 
-func getDBKey(entryPoint common.Address, sender common.Address) []byte {
-	return []byte(keyPrefix + keySeparator + entryPoint.String() + keySeparator + sender.String())
+func getUniqueKey(entryPoint common.Address, sender common.Address, nonce *big.Int) []byte {
+	return []byte(
+		keyPrefix + keySeparator + entryPoint.String() + keySeparator + sender.String() + keySeparator + nonce.String(),
+	)
 }
 
-func getEntryPointAndSenderFromDBKey(key []byte) (common.Address, common.Address) {
+func getEntryPointFromDBKey(key []byte) common.Address {
 	slc := strings.Split(string(key), keySeparator)
-	ep := common.HexToAddress(slc[1])
-	sender := common.HexToAddress(slc[2])
-
-	return ep, sender
+	return common.HexToAddress(slc[1])
 }
 
 func getUserOpFromDBValue(value []byte) (*userop.UserOperation, error) {
@@ -48,7 +48,7 @@ func loadFromDisk(db *badger.DB, q *userOpQueues) error {
 
 		for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() {
 			item := it.Item()
-			ep, _ := getEntryPointAndSenderFromDBKey(item.Key())
+			ep := getEntryPointFromDBKey(item.Key())
 
 			err := item.Value(func(v []byte) error {
 				op, err := getUserOpFromDBValue(v)

pkg/mempool/db.go

@@ -27,21 +27,22 @@ func New(db *badger.DB) (*Mempool, error) {
 	return &Mempool{db, queue}, nil
 }
 
-// GetOp checks if a UserOperation is in the mempool and returns it.
-func (m *Mempool) GetOp(entryPoint common.Address, sender common.Address) (*userop.UserOperation, error) {
-	op := m.queue.GetOp(entryPoint, sender)
-	return op, nil
+// GetOps returns all the UserOperations associated with an EntryPoint and Sender address.
+func (m *Mempool) GetOps(entryPoint common.Address, sender common.Address) ([]*userop.UserOperation, error) {
+	ops := m.queue.GetOps(entryPoint, sender)
+	return ops, nil
 }
 
-// AddOp adds a UserOperation to the mempool.
+// AddOp adds a UserOperation to the mempool or replace an existing one with the same EntryPoint, Sender, and
+// Nonce values.
 func (m *Mempool) AddOp(entryPoint common.Address, op *userop.UserOperation) error {
 	data, err := op.MarshalJSON()
 	if err != nil {
 		return err
 	}
 
 	err = m.db.Update(func(txn *badger.Txn) error {
-		return txn.Set(getDBKey(entryPoint, op.Sender), data)
+		return txn.Set(getUniqueKey(entryPoint, op.Sender, op.Nonce), data)
 	})
 	if err != nil {
 		return err
@@ -51,16 +52,16 @@ func (m *Mempool) AddOp(entryPoint common.Address, op *userop.UserOperation) err
 	return nil
 }
 
-// BundleOps builds a bundle of ops from the mempool to be sent to the EntryPoint.
+// BundleOps builds a bundle of UserOperations from the mempool to be sent to the EntryPoint.
 func (m *Mempool) BundleOps(entryPoint common.Address) ([]*userop.UserOperation, error) {
 	return m.queue.Next(entryPoint), nil
 }
 
-// RemoveOps removes a list of UserOperations from the mempool by sender address.
-func (m *Mempool) RemoveOps(entryPoint common.Address, senders ...common.Address) error {
+// RemoveOps removes a list of UserOperations from the mempool by EntryPoint, Sender, and Nonce values.
+func (m *Mempool) RemoveOps(entryPoint common.Address, ops ...*userop.UserOperation) error {
 	err := m.db.Update(func(txn *badger.Txn) error {
-		for _, s := range senders {
-			err := txn.Delete(getDBKey(entryPoint, s))
+		for _, op := range ops {
+			err := txn.Delete(getUniqueKey(entryPoint, op.Sender, op.Nonce))
 			if err != nil {
 				return err
 			}
@@ -72,6 +73,6 @@ func (m *Mempool) RemoveOps(entryPoint common.Address, senders ...common.Address
 		return err
 	}
 
-	m.queue.RemoveOps(entryPoint, senders)
+	m.queue.RemoveOps(entryPoint, ops...)
 	return nil
 }