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
  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 }
  35
  36 func NewTxPool() *TxPool {
  37         return &TxPool{
  38                 lastUpdated: time.Now().Unix(),
  39                 pool:        make(map[bc.Hash]*TxDesc),
  40                 errCache:    lru.New(maxCachedErrTxs),
  41         }
  42 }
  43
  44 func (mp *TxPool) AddTransaction(tx *legacy.Tx, weight, height, fee uint64) *TxDesc {
  45         txD := &TxDesc{
  46                 Tx:       tx,
  47                 Added:    time.Now(),
  48                 Weight:   weight,
  49                 Height:   height,
  50                 Fee:      fee,
  51                 FeePerKB: fee * 1000 / tx.TxHeader.SerializedSize,
  52         }
  53
  54         mp.mtx.Lock()
  55         defer mp.mtx.Unlock()
  56
  57         mp.pool[tx.Tx.ID] = txD
  58         atomic.StoreInt64(&mp.lastUpdated, time.Now().Unix())
  59         return txD
  60 }
  61
  62 func (mp *TxPool) AddErrCache(txHash *bc.Hash) {
  63         mp.mtx.Lock()
  64         defer mp.mtx.Unlock()
  65
  66         mp.errCache.Add(txHash, nil)
  67 }
  68
  69 func (mp *TxPool) RemoveTransaction(txHash *bc.Hash) {
  70         mp.mtx.Lock()
  71         defer mp.mtx.Unlock()
  72
  73         if _, ok := mp.pool[*txHash]; ok {
  74                 delete(mp.pool, *txHash)
  75                 atomic.StoreInt64(&mp.lastUpdated, time.Now().Unix())
  76         }
  77 }
  78
  79 func (mp *TxPool) GetTransaction(txHash *bc.Hash) (*TxDesc, error) {
  80         mp.mtx.RLock()
  81         defer mp.mtx.RUnlock()
  82
  83         if txD, ok := mp.pool[*txHash]; ok {
  84                 return txD, nil
  85         }
  86
  87         return nil, ErrTransactionNotExist
  88 }
  89
  90 func (mp *TxPool) GetTransactions() []*TxDesc {
  91         mp.mtx.RLock()
  92         defer mp.mtx.RUnlock()
  93
  94         txDs := make([]*TxDesc, len(mp.pool))
  95         i := 0
  96         for _, desc := range mp.pool {
  97                 txDs[i] = desc
  98                 i++
  99         }
 100         return txDs
 101 }
 102
 103 func (mp *TxPool) IsTransactionInPool(txHash *bc.Hash) bool {
 104         mp.mtx.RLock()
 105         defer mp.mtx.RUnlock()
 106
 107         if _, ok := mp.pool[*txHash]; ok {
 108                 return true
 109         }
 110         return false
 111 }
 112
 113 func (mp *TxPool) IsTransactionInErrCache(txHash *bc.Hash) bool {
 114         mp.mtx.RLock()
 115         defer mp.mtx.RUnlock()
 116
 117         _, ok := mp.errCache.Get(txHash)
 118         return ok
 119 }
 120
 121 func (mp *TxPool) HaveTransaction(txHash *bc.Hash) bool {
 122         return mp.IsTransactionInPool(txHash) || mp.IsTransactionInErrCache(txHash)
 123 }
 124
 125 func (mp *TxPool) Count() int {
 126         mp.mtx.RLock()
 127         defer mp.mtx.RUnlock()
 128
 129         count := len(mp.pool)
 130         return count
 131 }