account/utxo_keeper.go

   1 package account
   2
   3 import (
   4         "encoding/json"
   5         "sort"
   6         "sync"
   7         "sync/atomic"
   8         "time"
   9
  10         log "github.com/sirupsen/logrus"
  11         dbm "github.com/tendermint/tmlibs/db"
  12
  13         "github.com/bytom/errors"
  14         "github.com/bytom/protocol/bc"
  15 )
  16
  17 // pre-define error types
  18 var (
  19         ErrInsufficient = errors.New("reservation found insufficient funds")
  20         ErrImmature     = errors.New("reservation found immature funds")
  21         ErrReserved     = errors.New("reservation found outputs already reserved")
  22         ErrMatchUTXO    = errors.New("can't find utxo with given hash")
  23         ErrReservation  = errors.New("couldn't find reservation")
  24 )
  25
  26 // UTXO describes an individual account utxo.
  27 type UTXO struct {
  28         OutputID            bc.Hash
  29         SourceID            bc.Hash
  30         AssetID             bc.AssetID
  31         Amount              uint64
  32         SourcePos           uint64
  33         ControlProgram      []byte
  34         AccountID           string
  35         Address             string
  36         ControlProgramIndex uint64
  37         ValidHeight         uint64
  38         Change              bool
  39 }
  40
  41 // reservation describes a reservation of a set of UTXOs
  42 type reservation struct {
  43         id     uint64
  44         utxos  []*UTXO
  45         change uint64
  46         expiry time.Time
  47 }
  48
  49 type utxoKeeper struct {
  50         // `sync/atomic` expects the first word in an allocated struct to be 64-bit
  51         // aligned on both ARM and x86-32. See https://goo.gl/zW7dgq for more details.
  52         nextIndex     uint64
  53         db            dbm.DB
  54         mtx           sync.RWMutex
  55         currentHeight func() uint64
  56
  57         unconfirmed  map[bc.Hash]*UTXO
  58         reserved     map[bc.Hash]uint64
  59         reservations map[uint64]*reservation
  60 }
  61
  62 func newUtxoKeeper(f func() uint64, walletdb dbm.DB) *utxoKeeper {
  63         uk := &utxoKeeper{
  64                 db:            walletdb,
  65                 currentHeight: f,
  66                 unconfirmed:   make(map[bc.Hash]*UTXO),
  67                 reserved:      make(map[bc.Hash]uint64),
  68                 reservations:  make(map[uint64]*reservation),
  69         }
  70         go uk.expireWorker()
  71         return uk
  72 }
  73
  74 func (uk *utxoKeeper) AddUnconfirmedUtxo(utxos []*UTXO) {
  75         uk.mtx.Lock()
  76         defer uk.mtx.Unlock()
  77
  78         for _, utxo := range utxos {
  79                 uk.unconfirmed[utxo.OutputID] = utxo
  80         }
  81 }
  82
  83 // Cancel canceling the reservation with the provided ID.
  84 func (uk *utxoKeeper) Cancel(rid uint64) {
  85         uk.mtx.Lock()
  86         uk.cancel(rid)
  87         uk.mtx.Unlock()
  88 }
  89
  90 // ListUnconfirmed return all the unconfirmed utxos
  91 func (uk *utxoKeeper) ListUnconfirmed() []*UTXO {
  92         uk.mtx.Lock()
  93         defer uk.mtx.Unlock()
  94
  95         utxos := []*UTXO{}
  96         for _, utxo := range uk.unconfirmed {
  97                 utxos = append(utxos, utxo)
  98         }
  99         return utxos
 100 }
 101
 102 func (uk *utxoKeeper) RemoveUnconfirmedUtxo(hashes []*bc.Hash) {
 103         uk.mtx.Lock()
 104         defer uk.mtx.Unlock()
 105
 106         for _, hash := range hashes {
 107                 delete(uk.unconfirmed, *hash)
 108         }
 109 }
 110
 111 func (uk *utxoKeeper) Reserve(accountID string, assetID *bc.AssetID, amount uint64, useUnconfirmed bool, exp time.Time) (*reservation, error) {
 112         uk.mtx.Lock()
 113         defer uk.mtx.Unlock()
 114
 115         utxos, immatureAmount := uk.findUtxos(accountID, assetID, useUnconfirmed)
 116         optUtxos, optAmount, reservedAmount := uk.optUTXOs(utxos, amount)
 117         if optAmount+reservedAmount+immatureAmount < amount {
 118                 return nil, ErrInsufficient
 119         }
 120
 121         if optAmount+reservedAmount < amount {
 122                 return nil, ErrImmature
 123         }
 124
 125         if optAmount < amount {
 126                 return nil, ErrReserved
 127         }
 128
 129         result := &reservation{
 130                 id:     atomic.AddUint64(&uk.nextIndex, 1),
 131                 utxos:  optUtxos,
 132                 change: optAmount - amount,
 133                 expiry: exp,
 134         }
 135
 136         uk.reservations[result.id] = result
 137         for _, u := range optUtxos {
 138                 uk.reserved[u.OutputID] = result.id
 139         }
 140         return result, nil
 141 }
 142
 143 func (uk *utxoKeeper) ReserveParticular(outHash bc.Hash, useUnconfirmed bool, exp time.Time) (*reservation, error) {
 144         uk.mtx.Lock()
 145         defer uk.mtx.Unlock()
 146
 147         if _, ok := uk.reserved[outHash]; ok {
 148                 return nil, ErrReserved
 149         }
 150
 151         u, err := uk.findUtxo(outHash, useUnconfirmed)
 152         if err != nil {
 153                 return nil, err
 154         }
 155
 156         if u.ValidHeight > uk.currentHeight() {
 157                 return nil, ErrImmature
 158         }
 159
 160         result := &reservation{
 161                 id:     atomic.AddUint64(&uk.nextIndex, 1),
 162                 utxos:  []*UTXO{u},
 163                 expiry: exp,
 164         }
 165         uk.reservations[result.id] = result
 166         uk.reserved[u.OutputID] = result.id
 167         return result, nil
 168 }
 169
 170 func (uk *utxoKeeper) cancel(rid uint64) {
 171         res, ok := uk.reservations[rid]
 172         if !ok {
 173                 return
 174         }
 175
 176         delete(uk.reservations, rid)
 177         for _, utxo := range res.utxos {
 178                 delete(uk.reserved, utxo.OutputID)
 179         }
 180 }
 181
 182 func (uk *utxoKeeper) expireWorker() {
 183         ticker := time.NewTicker(1000 * time.Millisecond)
 184         for now := range ticker.C {
 185                 uk.expireReservation(now)
 186         }
 187 }
 188 func (uk *utxoKeeper) expireReservation(t time.Time) {
 189         uk.mtx.Lock()
 190         defer uk.mtx.Unlock()
 191
 192         for rid, res := range uk.reservations {
 193                 if res.expiry.Before(t) {
 194                         uk.cancel(rid)
 195                 }
 196         }
 197 }
 198
 199 func (uk *utxoKeeper) findUtxos(accountID string, assetID *bc.AssetID, useUnconfirmed bool) ([]*UTXO, uint64) {
 200         immatureAmount := uint64(0)
 201         currentHeight := uk.currentHeight()
 202         utxos := []*UTXO{}
 203         appendUtxo := func(u *UTXO) {
 204                 if u.AccountID != accountID || u.AssetID != *assetID {
 205                         return
 206                 }
 207                 if u.ValidHeight > currentHeight {
 208                         immatureAmount += u.Amount
 209                 } else {
 210                         utxos = append(utxos, u)
 211                 }
 212         }
 213
 214         utxoIter := uk.db.IteratorPrefix([]byte(UTXOPreFix))
 215         defer utxoIter.Release()
 216         for utxoIter.Next() {
 217                 u := &UTXO{}
 218                 if err := json.Unmarshal(utxoIter.Value(), u); err != nil {
 219                         log.WithField("err", err).Error("utxoKeeper findUtxos fail on unmarshal utxo")
 220                         continue
 221                 }
 222                 appendUtxo(u)
 223         }
 224         if !useUnconfirmed {
 225                 return utxos, immatureAmount
 226         }
 227
 228         for _, u := range uk.unconfirmed {
 229                 appendUtxo(u)
 230         }
 231         return utxos, immatureAmount
 232 }
 233
 234 func (uk *utxoKeeper) findUtxo(outHash bc.Hash, useUnconfirmed bool) (*UTXO, error) {
 235         if u, ok := uk.unconfirmed[outHash]; useUnconfirmed && ok {
 236                 return u, nil
 237         }
 238
 239         u := &UTXO{}
 240         if data := uk.db.Get(StandardUTXOKey(outHash)); data != nil {
 241                 return u, json.Unmarshal(data, u)
 242         }
 243         if data := uk.db.Get(ContractUTXOKey(outHash)); data != nil {
 244                 return u, json.Unmarshal(data, u)
 245         }
 246         return nil, ErrMatchUTXO
 247 }
 248
 249 func (uk *utxoKeeper) optUTXOs(utxos []*UTXO, amount uint64) ([]*UTXO, uint64, uint64) {
 250         sort.Slice(utxos, func(i, j int) bool {
 251                 return utxos[i].Amount < utxos[j].Amount
 252         })
 253
 254         var optAmount, reservedAmount uint64
 255         optUtxos := []*UTXO{}
 256         for _, u := range utxos {
 257                 if _, ok := uk.reserved[u.OutputID]; ok {
 258                         reservedAmount += u.Amount
 259                         continue
 260                 }
 261
 262                 optAmount += u.Amount
 263                 optUtxos = append(optUtxos, u)
 264                 if optAmount >= amount {
 265                         break
 266                 }
 267         }
 268         return optUtxos, optAmount, reservedAmount
 269 }