protocol/mempool.go

   1 package protocol
   2
   3 import (
   4         "errors"
   5         "sync"
   6         "sync/atomic"
   7         "time"
   8
   9         "github.com/bytom/protocol/bc"
  10         "github.com/bytom/protocol/bc/legacy"
  11         "github.com/golang/groupcache/lru"
  12 )
  13
  14 var (
  15         maxCachedErrTxs        = 1000
  16         maxNewTxChSize         = 1000
  17         ErrTransactionNotExist = errors.New("transaction are not existed in the mempool")
  18 )
  19
  20 type TxDesc struct {
  21         Tx       *legacy.Tx
  22         Added    time.Time
  23         Height   uint64
  24         Weight   uint64
  25         Fee      uint64
  26         FeePerKB uint64
  27 }
  28
  29 type TxPool struct {
  30         lastUpdated int64
  31         mtx         sync.RWMutex
  32         pool        map[bc.Hash]*TxDesc
  33         errCache    *lru.Cache
  34         newTxCh     chan *legacy.Tx
  35 }
  36
  37 func NewTxPool() *TxPool {
  38         return &TxPool{
  39                 lastUpdated: time.Now().Unix(),
  40                 pool:        make(map[bc.Hash]*TxDesc),
  41                 errCache:    lru.New(maxCachedErrTxs),
  42                 newTxCh:     make(chan *legacy.Tx, maxNewTxChSize),
  43         }
  44 }
  45
  46 func (mp *TxPool) GetNewTxCh() chan *legacy.Tx {
  47         return mp.newTxCh
  48 }
  49
  50 func (mp *TxPool) AddTransaction(tx *legacy.Tx, height, fee uint64) *TxDesc {
  51         txD := &TxDesc{
  52                 Tx:       tx,
  53                 Added:    time.Now(),
  54                 Weight:   tx.TxData.SerializedSize,
  55                 Height:   height,
  56                 Fee:      fee,
  57                 FeePerKB: fee * 1000 / tx.TxHeader.SerializedSize,
  58         }
  59
  60         mp.mtx.Lock()
  61         defer mp.mtx.Unlock()
  62
  63         mp.pool[tx.Tx.ID] = txD
  64         atomic.StoreInt64(&mp.lastUpdated, time.Now().Unix())
  65
  66         mp.newTxCh <- tx
  67         return txD
  68 }
  69
  70 func (mp *TxPool) AddErrCache(txHash *bc.Hash, err error) {
  71         mp.mtx.Lock()
  72         defer mp.mtx.Unlock()
  73
  74         mp.errCache.Add(txHash, err)
  75 }
  76
  77 func (mp *TxPool) GetErrCache(txHash *bc.Hash) error {
  78         mp.mtx.Lock()
  79         defer mp.mtx.Unlock()
  80
  81         v, ok := mp.errCache.Get(txHash)
  82         if ok {
  83                 return nil
  84         }
  85         return v.(error)
  86 }
  87
  88 func (mp *TxPool) RemoveTransaction(txHash *bc.Hash) {
  89         mp.mtx.Lock()
  90         defer mp.mtx.Unlock()
  91
  92         if _, ok := mp.pool[*txHash]; ok {
  93                 delete(mp.pool, *txHash)
  94                 atomic.StoreInt64(&mp.lastUpdated, time.Now().Unix())
  95         }
  96 }
  97
  98 func (mp *TxPool) GetTransaction(txHash *bc.Hash) (*TxDesc, error) {
  99         mp.mtx.RLock()
 100         defer mp.mtx.RUnlock()
 101
 102         if txD, ok := mp.pool[*txHash]; ok {
 103                 return txD, nil
 104         }
 105
 106         return nil, ErrTransactionNotExist
 107 }
 108
 109 func (mp *TxPool) GetTransactions() []*TxDesc {
 110         mp.mtx.RLock()
 111         defer mp.mtx.RUnlock()
 112
 113         txDs := make([]*TxDesc, len(mp.pool))
 114         i := 0
 115         for _, desc := range mp.pool {
 116                 txDs[i] = desc
 117                 i++
 118         }
 119         return txDs
 120 }
 121
 122 func (mp *TxPool) IsTransactionInPool(txHash *bc.Hash) bool {
 123         mp.mtx.RLock()
 124         defer mp.mtx.RUnlock()
 125
 126         if _, ok := mp.pool[*txHash]; ok {
 127                 return true
 128         }
 129         return false
 130 }
 131
 132 func (mp *TxPool) IsTransactionInErrCache(txHash *bc.Hash) bool {
 133         mp.mtx.RLock()
 134         defer mp.mtx.RUnlock()
 135
 136         _, ok := mp.errCache.Get(txHash)
 137         return ok
 138 }
 139
 140 func (mp *TxPool) HaveTransaction(txHash *bc.Hash) bool {
 141         return mp.IsTransactionInPool(txHash) || mp.IsTransactionInErrCache(txHash)
 142 }
 143
 144 func (mp *TxPool) Count() int {
 145         mp.mtx.RLock()
 146         defer mp.mtx.RUnlock()
 147
 148         count := len(mp.pool)
 149         return count
 150 }