test/mock/UTXO.go

   1 package mock
   2
   3 import (
   4         "bytes"
   5         "container/list"
   6         "encoding/json"
   7         "sort"
   8         "sync"
   9         "sync/atomic"
  10         "time"
  11
  12         log "github.com/sirupsen/logrus"
  13
  14         acc "github.com/vapor/account"
  15         "github.com/vapor/protocol/bc"
  16 )
  17
  18 const (
  19         desireUtxoCount = 5
  20         logModule       = "account"
  21 )
  22
  23 type Reservation struct {
  24         ID     uint64
  25         UTXOs  []*acc.UTXO
  26         Change uint64
  27         Expiry time.Time
  28 }
  29
  30 type UTXOKeeper struct {
  31         // `sync/atomic` expects the first word in an allocated struct to be 64-bit
  32         // aligned on both ARM and x86-32. See https://goo.gl/zW7dgq for more details.
  33         NextIndex     uint64
  34         Store         acc.AccountStorer
  35         mtx           sync.RWMutex
  36         CurrentHeight func() uint64
  37
  38         Unconfirmed  map[bc.Hash]*acc.UTXO
  39         Reserved     map[bc.Hash]uint64
  40         Reservations map[uint64]*Reservation
  41 }
  42
  43 func NewUtxoKeeper(f func() uint64, store acc.AccountStorer) *UTXOKeeper {
  44         uk := &UTXOKeeper{
  45                 Store:         store,
  46                 CurrentHeight: f,
  47                 Unconfirmed:   make(map[bc.Hash]*acc.UTXO),
  48                 Reserved:      make(map[bc.Hash]uint64),
  49                 Reservations:  make(map[uint64]*Reservation),
  50         }
  51         go uk.expireWorker()
  52         return uk
  53 }
  54
  55 func (uk *UTXOKeeper) expireWorker() {
  56         ticker := time.NewTicker(1000 * time.Millisecond)
  57         defer ticker.Stop()
  58
  59         for now := range ticker.C {
  60                 uk.expireReservation(now)
  61         }
  62 }
  63
  64 func (uk *UTXOKeeper) expireReservation(t time.Time) {
  65         uk.mtx.Lock()
  66         defer uk.mtx.Unlock()
  67
  68         for rid, res := range uk.Reservations {
  69                 if res.Expiry.Before(t) {
  70                         uk.cancel(rid)
  71                 }
  72         }
  73 }
  74
  75 func (uk *UTXOKeeper) cancel(rid uint64) {
  76         res, ok := uk.Reservations[rid]
  77         if !ok {
  78                 return
  79         }
  80
  81         delete(uk.Reservations, rid)
  82         for _, utxo := range res.UTXOs {
  83                 delete(uk.Reserved, utxo.OutputID)
  84         }
  85 }
  86
  87 func (uk *UTXOKeeper) Reserve(accountID string, assetID *bc.AssetID, amount uint64, useUnconfirmed bool, vote []byte, exp time.Time) (*Reservation, error) {
  88         uk.mtx.Lock()
  89         defer uk.mtx.Unlock()
  90
  91         utxos, immatureAmount := uk.findUtxos(accountID, assetID, useUnconfirmed, vote)
  92         optUtxos, optAmount, reservedAmount := uk.optUTXOs(utxos, amount)
  93         if optAmount+reservedAmount+immatureAmount < amount {
  94                 return nil, acc.ErrInsufficient
  95         }
  96
  97         if optAmount+reservedAmount < amount {
  98                 if vote != nil {
  99                         return nil, acc.ErrVoteLock
 100                 }
 101                 return nil, acc.ErrImmature
 102         }
 103
 104         if optAmount < amount {
 105                 return nil, acc.ErrReserved
 106         }
 107
 108         result := &Reservation{
 109                 ID:     atomic.AddUint64(&uk.NextIndex, 1),
 110                 UTXOs:  optUtxos,
 111                 Change: optAmount - amount,
 112                 Expiry: exp,
 113         }
 114
 115         uk.Reservations[result.ID] = result
 116         for _, u := range optUtxos {
 117                 uk.Reserved[u.OutputID] = result.ID
 118         }
 119         return result, nil
 120 }
 121
 122 func (uk *UTXOKeeper) ReserveParticular(outHash bc.Hash, useUnconfirmed bool, exp time.Time) (*Reservation, error) {
 123         uk.mtx.Lock()
 124         defer uk.mtx.Unlock()
 125
 126         if _, ok := uk.Reserved[outHash]; ok {
 127                 return nil, acc.ErrReserved
 128         }
 129
 130         u, err := uk.FindUtxo(outHash, useUnconfirmed)
 131         if err != nil {
 132                 return nil, err
 133         }
 134
 135         if u.ValidHeight > uk.CurrentHeight() {
 136                 return nil, acc.ErrImmature
 137         }
 138
 139         result := &Reservation{
 140                 ID:     atomic.AddUint64(&uk.NextIndex, 1),
 141                 UTXOs:  []*acc.UTXO{u},
 142                 Expiry: exp,
 143         }
 144         uk.Reservations[result.ID] = result
 145         uk.Reserved[u.OutputID] = result.ID
 146         return result, nil
 147 }
 148
 149 func (uk *UTXOKeeper) FindUtxo(outHash bc.Hash, useUnconfirmed bool) (*acc.UTXO, error) {
 150         if u, ok := uk.Unconfirmed[outHash]; useUnconfirmed && ok {
 151                 return u, nil
 152         }
 153
 154         u := &acc.UTXO{}
 155         if data := uk.Store.GetStandardUTXO(outHash); data != nil {
 156                 return u, json.Unmarshal(data, u)
 157         }
 158         if data := uk.Store.GetContractUTXO(outHash); data != nil {
 159                 return u, json.Unmarshal(data, u)
 160         }
 161         return nil, acc.ErrMatchUTXO
 162 }
 163
 164 func (uk *UTXOKeeper) findUtxos(accountID string, assetID *bc.AssetID, useUnconfirmed bool, vote []byte) ([]*acc.UTXO, uint64) {
 165         immatureAmount := uint64(0)
 166         currentHeight := uk.CurrentHeight()
 167         utxos := []*acc.UTXO{}
 168         appendUtxo := func(u *acc.UTXO) {
 169                 if u.AccountID != accountID || u.AssetID != *assetID || !bytes.Equal(u.Vote, vote) {
 170                         return
 171                 }
 172                 if u.ValidHeight > currentHeight {
 173                         immatureAmount += u.Amount
 174                 } else {
 175                         utxos = append(utxos, u)
 176                 }
 177         }
 178
 179         rawUTXOs := uk.Store.GetUTXOs()
 180         for _, rawUTXO := range rawUTXOs {
 181                 utxo := new(acc.UTXO)
 182                 if err := json.Unmarshal(rawUTXO, utxo); err != nil {
 183                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("utxoKeeper findUtxos fail on unmarshal utxo")
 184                         continue
 185                 }
 186                 appendUtxo(utxo)
 187         }
 188
 189         if !useUnconfirmed {
 190                 return utxos, immatureAmount
 191         }
 192
 193         for _, u := range uk.Unconfirmed {
 194                 appendUtxo(u)
 195         }
 196         return utxos, immatureAmount
 197 }
 198
 199 func (uk *UTXOKeeper) optUTXOs(utxos []*acc.UTXO, amount uint64) ([]*acc.UTXO, uint64, uint64) {
 200         //sort the utxo by amount, bigger amount in front
 201         var optAmount, reservedAmount uint64
 202         sort.Slice(utxos, func(i, j int) bool {
 203                 return utxos[i].Amount > utxos[j].Amount
 204         })
 205
 206         //push all the available utxos into list
 207         utxoList := list.New()
 208         for _, u := range utxos {
 209                 if _, ok := uk.Reserved[u.OutputID]; ok {
 210                         reservedAmount += u.Amount
 211                         continue
 212                 }
 213                 utxoList.PushBack(u)
 214         }
 215
 216         optList := list.New()
 217         for node := utxoList.Front(); node != nil; node = node.Next() {
 218                 //append utxo if we haven't reached the required amount
 219                 if optAmount < amount {
 220                         optList.PushBack(node.Value)
 221                         optAmount += node.Value.(*acc.UTXO).Amount
 222                         continue
 223                 }
 224
 225                 largestNode := optList.Front()
 226                 replaceList := list.New()
 227                 replaceAmount := optAmount - largestNode.Value.(*acc.UTXO).Amount
 228
 229                 for ; node != nil && replaceList.Len() <= desireUtxoCount-optList.Len(); node = node.Next() {
 230                         replaceList.PushBack(node.Value)
 231                         if replaceAmount += node.Value.(*acc.UTXO).Amount; replaceAmount >= amount {
 232                                 optList.Remove(largestNode)
 233                                 optList.PushBackList(replaceList)
 234                                 optAmount = replaceAmount
 235                                 break
 236                         }
 237                 }
 238
 239                 //largestNode remaining the same means that there is nothing to be replaced
 240                 if largestNode == optList.Front() {
 241                         break
 242                 }
 243         }
 244
 245         optUtxos := []*acc.UTXO{}
 246         for e := optList.Front(); e != nil; e = e.Next() {
 247                 optUtxos = append(optUtxos, e.Value.(*acc.UTXO))
 248         }
 249         return optUtxos, optAmount, reservedAmount
 250 }
 251
 252 func (uk *UTXOKeeper) FindUtxos(accountID string, assetID *bc.AssetID, useUnconfirmed bool, vote []byte) ([]*acc.UTXO, uint64) {
 253         immatureAmount := uint64(0)
 254         currentHeight := uk.CurrentHeight()
 255         utxos := []*acc.UTXO{}
 256         appendUtxo := func(u *acc.UTXO) {
 257                 if u.AccountID != accountID || u.AssetID != *assetID || !bytes.Equal(u.Vote, vote) {
 258                         return
 259                 }
 260                 if u.ValidHeight > currentHeight {
 261                         immatureAmount += u.Amount
 262                 } else {
 263                         utxos = append(utxos, u)
 264                 }
 265         }
 266
 267         rawUTXOs := uk.Store.GetUTXOs()
 268         for _, rawUTXO := range rawUTXOs {
 269                 utxo := new(acc.UTXO)
 270                 if err := json.Unmarshal(rawUTXO, utxo); err != nil {
 271                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("utxoKeeper findUtxos fail on unmarshal utxo")
 272                         continue
 273                 }
 274                 appendUtxo(utxo)
 275         }
 276
 277         if !useUnconfirmed {
 278                 return utxos, immatureAmount
 279         }
 280
 281         for _, u := range uk.Unconfirmed {
 282                 appendUtxo(u)
 283         }
 284         return utxos, immatureAmount
 285 }