Optimize CPU utilization on Bundler go-routine (#96)

Author: hazim-j

internal/config/values.go

@@ -56,8 +56,8 @@ func GetValues() *Values {
 	viper.SetDefault("erc4337_bundler_supported_entry_points", "0x0F46c65C17AA6b4102046935F33301f0510B163A")
 	viper.SetDefault("erc4337_bundler_max_verification_gas", 1500000)
 	viper.SetDefault("erc4337_bundler_max_ops_for_unstaked_sender", 4)
-	viper.SetDefault("erc4337_bundler_debug_mode", false)
 	viper.SetDefault("erc4337_bundler_blocks_in_the_future", 25)
+	viper.SetDefault("erc4337_bundler_debug_mode", false)
 	viper.SetDefault("erc4337_bundler_gin_mode", gin.ReleaseMode)
 
 	// Read in from .env file if available
@@ -84,6 +84,7 @@ func GetValues() *Values {
 	_ = viper.BindEnv("erc4337_bundler_max_ops_for_unstaked_sender")
 	_ = viper.BindEnv("erc4337_bundler_eth_builder_url")
 	_ = viper.BindEnv("erc4337_bundler_blocks_in_the_future")
+	_ = viper.BindEnv("erc4337_bundler_debug_mode")
 	_ = viper.BindEnv("erc4337_bundler_gin_mode")
 
 	// Validate required variables

pkg/bundler/bundler.go

@@ -24,6 +24,8 @@ type Bundler struct {
 	logger               logr.Logger
 	isRunning            bool
 	stop                 chan bool
+	watch                chan bool
+	onStop               func()
 }
 
 // New initializes a new EIP-4337 bundler which can be extended with modules for validating batches and
@@ -37,6 +39,8 @@ func New(mempool *mempool.Mempool, chainID *big.Int, supportedEntryPoints []comm
 		logger:               logger.NewZeroLogr().WithName("bundler"),
 		isRunning:            false,
 		stop:                 make(chan bool),
+		watch:                make(chan bool),
+		onStop:               func() {},
 	}
 }
 
@@ -63,7 +67,7 @@ func (i *Bundler) Run() error {
 			select {
 			case <-i.stop:
 				return
-			default:
+			case <-i.watch:
 				for _, ep := range i.supportedEntryPoints {
 					start := time.Now()
 					l := logger.
@@ -115,6 +119,7 @@ func (i *Bundler) Run() error {
 	}(i)
 
 	i.isRunning = true
+	i.onStop = i.mempool.OnAdd(i.watch)
 	return nil
 }
 
@@ -126,4 +131,5 @@ func (i *Bundler) Stop() {
 
 	i.isRunning = false
 	i.stop <- true
+	i.onStop()
 }

pkg/mempool/instance.go

@@ -5,26 +5,40 @@ package mempool
 import (
 	badger "github.com/dgraph-io/badger/v3"
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/google/uuid"
 	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 )
 
+type watch struct {
+	key string
+	sig chan bool
+}
+
 // Mempool provides read and write access to a pool of pending UserOperations which have passed all Client
 // checks.
 type Mempool struct {
-	db    *badger.DB
-	queue *userOpQueues
+	db      *badger.DB
+	queue   *userOpQueues
+	watches []watch
 }
 
 // New creates an instance of a mempool that uses an embedded DB to persist and load UserOperations from disk
 // incase of a reset.
 func New(db *badger.DB) (*Mempool, error) {
 	queue := newUserOpQueue()
+	watches := []watch{}
 	err := loadFromDisk(db, queue)
 	if err != nil {
 		return nil, err
 	}
 
-	return &Mempool{db, queue}, nil
+	return &Mempool{db, queue, watches}, nil
+}
+
+func (m *Mempool) pushSignals() {
+	for _, watch := range m.watches {
+		watch.sig <- true
+	}
 }
 
 // GetOps returns all the UserOperations associated with an EntryPoint and Sender address.
@@ -49,6 +63,7 @@ func (m *Mempool) AddOp(entryPoint common.Address, op *userop.UserOperation) err
 	}
 
 	m.queue.AddOp(entryPoint, op)
+	m.pushSignals()
 	return nil
 }
 
@@ -77,6 +92,29 @@ func (m *Mempool) RemoveOps(entryPoint common.Address, ops ...*userop.UserOperat
 	return nil
 }
 
+// OnAdd allows entities to register a channel that will receive a signal every time a UserOperation is added.
+// It returns a function to stop receiving.
+func (m *Mempool) OnAdd(sig chan bool) func() {
+	key := uuid.New().String()
+	w := watch{key, sig}
+	m.watches = append(m.watches, w)
+
+	// initial push if queue is non-empty
+	if m.queue.opCount > 0 {
+		w.sig <- true
+	}
+
+	return func() {
+		f := []watch{}
+		for _, w := range m.watches {
+			if w.key != key {
+				f = append(f, w)
+			}
+		}
+		m.watches = f
+	}
+}
+
 // Dump will return a list of UserOperations from the mempool by EntryPoint in the order it arrived.
 func (m *Mempool) Dump(entryPoint common.Address) ([]*userop.UserOperation, error) {
 	return m.queue.All(entryPoint), nil

pkg/mempool/queue.go

@@ -1,6 +1,8 @@
 package mempool
 
 import (
+	"sync/atomic"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 	"github.com/wangjia184/sortedset"
@@ -24,6 +26,7 @@ func (s *set) getSenderSortedSet(sender common.Address) *sortedset.SortedSet {
 
 type userOpQueues struct {
 	maxBatchSize     int
+	opCount          uint64
 	setsByEntryPoint map[common.Address]*set
 }
 
@@ -47,6 +50,7 @@ func (q *userOpQueues) AddOp(entryPoint common.Address, op *userop.UserOperation
 	eps.all.AddOrUpdate(key, sortedset.SCORE(op.MaxPriorityFeePerGas.Int64()), op)
 	eps.arrival.AddOrUpdate(key, sortedset.SCORE(eps.all.GetCount()), op)
 	sss.AddOrUpdate(key, sortedset.SCORE(op.Nonce.Int64()), op)
+	atomic.AddUint64(&q.opCount, 1)
 }
 
 func (q *userOpQueues) GetOps(entryPoint common.Address, sender common.Address) []*userop.UserOperation {
@@ -101,6 +105,7 @@ func (q *userOpQueues) RemoveOps(entryPoint common.Address, ops ...*userop.UserO
 		eps.arrival.Remove(key)
 		sss.Remove(key)
 	}
+	atomic.AddUint64(&q.opCount, ^uint64(0))
 }
 
 func newUserOpQueue() *userOpQueues {