github.com/ava-labs/avalanchego@v1.11.11/vms/txs/mempool/mempool.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package mempool
     5  
     6  import (
     7  	"errors"
     8  	"fmt"
     9  	"sync"
    10  
    11  	"github.com/ava-labs/avalanchego/cache"
    12  	"github.com/ava-labs/avalanchego/ids"
    13  	"github.com/ava-labs/avalanchego/utils/linked"
    14  	"github.com/ava-labs/avalanchego/utils/set"
    15  	"github.com/ava-labs/avalanchego/utils/setmap"
    16  	"github.com/ava-labs/avalanchego/utils/units"
    17  )
    18  
    19  const (
    20  	// MaxTxSize is the maximum number of bytes a transaction can use to be
    21  	// allowed into the mempool.
    22  	MaxTxSize = 64 * units.KiB
    23  
    24  	// droppedTxIDsCacheSize is the maximum number of dropped txIDs to cache
    25  	droppedTxIDsCacheSize = 64
    26  
    27  	// maxMempoolSize is the maximum number of bytes allowed in the mempool
    28  	maxMempoolSize = 64 * units.MiB
    29  )
    30  
    31  var (
    32  	ErrDuplicateTx          = errors.New("duplicate tx")
    33  	ErrTxTooLarge           = errors.New("tx too large")
    34  	ErrMempoolFull          = errors.New("mempool is full")
    35  	ErrConflictsWithOtherTx = errors.New("tx conflicts with other tx")
    36  )
    37  
    38  type Tx interface {
    39  	InputIDs() set.Set[ids.ID]
    40  	ID() ids.ID
    41  	Size() int
    42  }
    43  
    44  type Metrics interface {
    45  	Update(numTxs, bytesAvailable int)
    46  }
    47  
    48  type Mempool[T Tx] interface {
    49  	Add(tx T) error
    50  	Get(txID ids.ID) (T, bool)
    51  	// Remove [txs] and any conflicts of [txs] from the mempool.
    52  	Remove(txs ...T)
    53  
    54  	// Peek returns the oldest tx in the mempool.
    55  	Peek() (tx T, exists bool)
    56  
    57  	// Iterate iterates over the txs until f returns false
    58  	Iterate(f func(tx T) bool)
    59  
    60  	// Note: dropped txs are added to droppedTxIDs but are not evicted from
    61  	// unissued decision/staker txs. This allows previously dropped txs to be
    62  	// possibly reissued.
    63  	MarkDropped(txID ids.ID, reason error)
    64  	GetDropReason(txID ids.ID) error
    65  
    66  	// Len returns the number of txs in the mempool.
    67  	Len() int
    68  }
    69  
    70  type mempool[T Tx] struct {
    71  	lock           sync.RWMutex
    72  	unissuedTxs    *linked.Hashmap[ids.ID, T]
    73  	consumedUTXOs  *setmap.SetMap[ids.ID, ids.ID] // TxID -> Consumed UTXOs
    74  	bytesAvailable int
    75  	droppedTxIDs   *cache.LRU[ids.ID, error] // TxID -> Verification error
    76  
    77  	metrics Metrics
    78  }
    79  
    80  func New[T Tx](
    81  	metrics Metrics,
    82  ) *mempool[T] {
    83  	m := &mempool[T]{
    84  		unissuedTxs:    linked.NewHashmap[ids.ID, T](),
    85  		consumedUTXOs:  setmap.New[ids.ID, ids.ID](),
    86  		bytesAvailable: maxMempoolSize,
    87  		droppedTxIDs:   &cache.LRU[ids.ID, error]{Size: droppedTxIDsCacheSize},
    88  		metrics:        metrics,
    89  	}
    90  	m.updateMetrics()
    91  
    92  	return m
    93  }
    94  
    95  func (m *mempool[T]) updateMetrics() {
    96  	m.metrics.Update(m.unissuedTxs.Len(), m.bytesAvailable)
    97  }
    98  
    99  func (m *mempool[T]) Add(tx T) error {
   100  	txID := tx.ID()
   101  
   102  	m.lock.Lock()
   103  	defer m.lock.Unlock()
   104  
   105  	if _, ok := m.unissuedTxs.Get(txID); ok {
   106  		return fmt.Errorf("%w: %s", ErrDuplicateTx, txID)
   107  	}
   108  
   109  	txSize := tx.Size()
   110  	if txSize > MaxTxSize {
   111  		return fmt.Errorf("%w: %s size (%d) > max size (%d)",
   112  			ErrTxTooLarge,
   113  			txID,
   114  			txSize,
   115  			MaxTxSize,
   116  		)
   117  	}
   118  	if txSize > m.bytesAvailable {
   119  		return fmt.Errorf("%w: %s size (%d) > available space (%d)",
   120  			ErrMempoolFull,
   121  			txID,
   122  			txSize,
   123  			m.bytesAvailable,
   124  		)
   125  	}
   126  
   127  	inputs := tx.InputIDs()
   128  	if m.consumedUTXOs.HasOverlap(inputs) {
   129  		return fmt.Errorf("%w: %s", ErrConflictsWithOtherTx, txID)
   130  	}
   131  
   132  	m.bytesAvailable -= txSize
   133  	m.unissuedTxs.Put(txID, tx)
   134  	m.updateMetrics()
   135  
   136  	// Mark these UTXOs as consumed in the mempool
   137  	m.consumedUTXOs.Put(txID, inputs)
   138  
   139  	// An added tx must not be marked as dropped.
   140  	m.droppedTxIDs.Evict(txID)
   141  	return nil
   142  }
   143  
   144  func (m *mempool[T]) Get(txID ids.ID) (T, bool) {
   145  	m.lock.RLock()
   146  	defer m.lock.RUnlock()
   147  
   148  	return m.unissuedTxs.Get(txID)
   149  }
   150  
   151  func (m *mempool[T]) Remove(txs ...T) {
   152  	m.lock.Lock()
   153  	defer m.lock.Unlock()
   154  
   155  	for _, tx := range txs {
   156  		txID := tx.ID()
   157  		// If the transaction is in the mempool, remove it.
   158  		if _, ok := m.consumedUTXOs.DeleteKey(txID); ok {
   159  			m.unissuedTxs.Delete(txID)
   160  			m.bytesAvailable += tx.Size()
   161  			continue
   162  		}
   163  
   164  		// If the transaction isn't in the mempool, remove any conflicts it has.
   165  		inputs := tx.InputIDs()
   166  		for _, removed := range m.consumedUTXOs.DeleteOverlapping(inputs) {
   167  			tx, _ := m.unissuedTxs.Get(removed.Key)
   168  			m.unissuedTxs.Delete(removed.Key)
   169  			m.bytesAvailable += tx.Size()
   170  		}
   171  	}
   172  	m.updateMetrics()
   173  }
   174  
   175  func (m *mempool[T]) Peek() (T, bool) {
   176  	m.lock.RLock()
   177  	defer m.lock.RUnlock()
   178  
   179  	_, tx, exists := m.unissuedTxs.Oldest()
   180  	return tx, exists
   181  }
   182  
   183  func (m *mempool[T]) Iterate(f func(T) bool) {
   184  	m.lock.RLock()
   185  	defer m.lock.RUnlock()
   186  
   187  	it := m.unissuedTxs.NewIterator()
   188  	for it.Next() {
   189  		if !f(it.Value()) {
   190  			return
   191  		}
   192  	}
   193  }
   194  
   195  func (m *mempool[_]) MarkDropped(txID ids.ID, reason error) {
   196  	if errors.Is(reason, ErrMempoolFull) {
   197  		return
   198  	}
   199  
   200  	m.lock.RLock()
   201  	defer m.lock.RUnlock()
   202  
   203  	if _, ok := m.unissuedTxs.Get(txID); ok {
   204  		return
   205  	}
   206  
   207  	m.droppedTxIDs.Put(txID, reason)
   208  }
   209  
   210  func (m *mempool[_]) GetDropReason(txID ids.ID) error {
   211  	err, _ := m.droppedTxIDs.Get(txID)
   212  	return err
   213  }
   214  
   215  func (m *mempool[_]) Len() int {
   216  	m.lock.RLock()
   217  	defer m.lock.RUnlock()
   218  
   219  	return m.unissuedTxs.Len()
   220  }